xai-org/x-algorithm · error · ValueError

Only Swizzle<2,5,2> supported for 128B_BASE32B

Error message

Only Swizzle<2,5,2> supported for 128B_BASE32B

What it means

_layout_type maps Swizzle<2,5,2> to the special SWIZZLE_128B_BASE32B SMEM layout mode. For swizzles with num_base M==5, only the exact triple (B,S)==(2,2) is supported; anything else raises ValueError.

Source

Thrown at phoenix/xrex/cutedsl/ranker_fa4/mma_sm100_desc.py:223

    SWIZZLE_64B = 4
    SWIZZLE_32B = 6


def _layout_type(swizzle: cute.Swizzle) -> LayoutType:
    B, M, S = swizzle.num_bits, swizzle.num_base, swizzle.num_shift

    if M == 4:
        if S != 3:
            raise ValueError("Unexpected swizzle shift – want S==3 for M==4")
        return {
            0: LayoutType.SWIZZLE_NONE,
            1: LayoutType.SWIZZLE_32B,
            2: LayoutType.SWIZZLE_64B,
            3: LayoutType.SWIZZLE_128B,
        }[B]
    if M == 5:
        if (B, S) != (2, 2):
            raise ValueError("Only Swizzle<2,5,2> supported for 128B_BASE32B")
        return LayoutType.SWIZZLE_128B_BASE32B

    raise ValueError("Unsupported swizzle triple for UMMA smem descriptor")


def make_smem_desc_base(layout: cute.Layout, swizzle: cute.Swizzle, major: Major) -> int:
    layout_type = _layout_type(swizzle)

    VERSION = 1
    LBO_MODE = 0
    BASE_OFFSET = 0

    swizzle_atom_mn_size = {
        LayoutType.SWIZZLE_NONE: 1,
        LayoutType.SWIZZLE_32B: 2,
        LayoutType.SWIZZLE_64B: 4,
        LayoutType.SWIZZLE_128B: 8,
        LayoutType.SWIZZLE_128B_BASE32B: 8,

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Use exactly Swizzle<2,5,2> when you need SWIZZLE_128B_BASE32B
  2. If you need a different swizzle width, switch to an M==4 atom (Swizzle<B,4,3>) which supports 32/64/128B modes
  3. Source the layout from a standard swizzle atom copy (e.g. copy(swizzle_atom)) rather than constructing Swizzle manually

Example fix

# before
sw = cute.Swizzle(3, 5, 2)
# after
sw = cute.Swizzle(2, 5, 2)  # only legal M==5 triple
Defensive patterns

Strategy: validation

Validate before calling

B, M, S = sw.num_bits, sw.num_base, sw.num_shift
if M == 5:
    assert (B, S) == (2, 2), f'M==5 requires Swizzle<2,5,2>, got Swizzle<{B},{M},{S}>'

Type guard

def is_valid_umma_swizzle(sw) -> bool:
    if sw.num_base == 4:
        return sw.num_shift == 3 and 0 <= sw.num_bits <= 3
    if sw.num_base == 5:
        return (sw.num_bits, sw.num_shift) == (2, 2)
    return False

Prevention

When it happens

Trigger: Calling make_smem_desc_base with a Swizzle whose num_base is 5 but with (num_bits, num_shift) != (2, 2), e.g. Swizzle<3,5,2> or Swizzle<2,5,3>.

Common situations: Enabling the 128B_BASE32B layout for fp8/MX-format operands and getting the swizzle atom parameters wrong, or adapting an SM90 swizzle where base-5 atoms have other legal bit widths.

Related errors


AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/7d3092ce1987d857. Report an issue: GitHub.