xai-org/x-algorithm · error · ValueError

Unexpected swizzle shift – want S==3 for M==4

Error message

Unexpected swizzle shift – want S==3 for M==4

What it means

_layout_type converts a cute.Swizzle into the SMEM descriptor LayoutType. For swizzles with num_base M==4 (the common 128B swizzle family), hardware requires num_shift S==3; any other shift raises ValueError.

Source

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

        Major.K if op.a_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN,
        Major.K if op.b_major_mode == cute.nvgpu.tcgen05.mma.OperandMajorMode.K else Major.MN,
    )


class LayoutType(IntEnum):
    SWIZZLE_NONE = 0
    SWIZZLE_128B_BASE32B = 1
    SWIZZLE_128B = 2
    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

View on GitHub (pinned to 24c60942c5)

Solutions

  1. Use a canonical swizzle for M==4: Swizzle<B,4,3> with B in {0,1,2,3} (none/32B/64B/128B)
  2. Build the layout via cute algorithms (e.g. tile_to_shape with a standard swizzle atom) instead of hand-specifying swizzle params
  3. Double-check the Swizzle constructor argument order (bits, base, shift) — a swapped base/shift is the usual cause

Example fix

# before
sw = cute.Swizzle(3, 4, 2)  # bad shift
# after
sw = cute.Swizzle(3, 4, 3)  # SWIZZLE_128B
Defensive patterns

Strategy: validation

Validate before calling

B, M, S = sw.num_bits, sw.num_base, sw.num_shift
if M == 4:
    assert S == 3, f'Swizzle<,{M},> requires shift 3, got {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 4 but num_shift is not 3, e.g. Swizzle<3,4,2> or Swizzle<3,4,4>, typically from a hand-built smem layout.

Common situations: Constructing custom shared-memory layouts for tcgen05 MMA operands and copying swizzle parameters from SM90 layouts, or using cute.make_swizzle with wrong base/shift arguments.

Related errors


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