sgl-project/sglang · error · TypeError

Unsupported tcgen05 MMA op kind: {type(op).__name__}

Error message

Unsupported tcgen05 MMA op kind: {type(op).__name__}

What it means

_tcgen05_mma_kind maps a CUTLASS tcgen05 MMA op object to a string kind ('mxf8f6f4', 'mxf4', 'mxf4nvf4') used by the PTX gemm helpers. Any other op type (e.g. a plain fp16/bf16 MmaOp) is unsupported and raises TypeError.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/blackwell_helpers.py:30

def _tcgen05_mma_kind(op: cute.nvgpu.tcgen05.mma.MmaOp) -> str:
    if isinstance(op, tcgen05.mma.MmaF16BF16Op):
        return "f16"
    if isinstance(op, tcgen05.mma.MmaTF32Op):
        return "tf32"
    if isinstance(op, tcgen05.mma.MmaI8Op):
        return "i8"
    # cutlass-dsl >=4.5.2 builds plain FP8 MMAs as MmaF8F6F4Op (make_trivial_tiled_mma's
    # _F8F6F4_TYPES branch); <4.4.x returned the now-legacy MmaFP8Op. Both map to kind::f8f6f4.
    if isinstance(op, (tcgen05.mma.MmaFP8Op, tcgen05.mma.MmaF8F6F4Op)):
        return "f8f6f4"
    if isinstance(op, tcgen05.mma.MmaMXF8Op):
        return "mxf8f6f4"
    if isinstance(op, tcgen05.mma.MmaMXF4Op):
        return "mxf4"
    if isinstance(op, tcgen05.mma.MmaMXF4NVF4Op):
        return "mxf4nvf4"
    raise TypeError(f"Unsupported tcgen05 MMA op kind: {type(op).__name__}")


@cute.jit
def gemm_w_idx(
    tiled_mma: cute.TiledMma,
    acc: cute.Tensor,
    tCrA: cute.Tensor,
    tCrB: cute.Tensor,
    A_idx: Optional[Int32] = None,
    B_idx: Optional[Int32] = None,
    zero_init: bool | Boolean = False,
    swap_AB: bool = False,
    num_unroll_groups: int = 1,
) -> None:
    if const_expr(swap_AB):
        return gemm_w_idx(
            tiled_mma, acc, tCrB, tCrA, B_idx, A_idx, zero_init=zero_init, swap_AB=False
        )

View on GitHub (pinned to 0132848349)

Solutions

1. Build the TiledMma with one of the supported MX op classes (tcgen05.mma.MmaMXF8Op / MmaMXF4Op / MmaMXF4NVF4Op)
2. Use the existing non-PTX gemm path for unsupported dtypes
3. Add a new kind mapping in _tcgen05_mma_kind if you control the kernel tree

Example fix

# before
op = tcgen05.mma.MmaF16Op()  # unsupported
tiled_mma = cute.make_tiled_mma(op)
# after
op = tcgen05.mma.MmaMXF8Op(...)
tiled_mma = cute.make_tiled_mma(op)

from flash_attn.cute.blackwell_helpers import _tcgen05_mma_kind\n_tcgen05_mma_kind(tiled_mma.op)  # raises early if unsupported

def is_supported_mma_op(op) -> bool:\n    import cutlass.cute.tcgen05.mma as m\n    return isinstance(op, (m.MmaMXF8Op, m.MmaMXF4Op, m.MmaMXF4NVF4Op))

When it happens

Trigger: Passing a cute.TiledMma built with an MMA op class not in {MmaMXF8Op, MmaMXF4Op, MmaMXF4NVF4Op} to gemm_ptx/gemm_ptx_loop/gemm_ptx_partial/mma in the Blackwell CUTE helpers.

Common situations: Extending FA4 CUTE kernels on Blackwell with non-MX (bf16/fp16/fp8-legacy) MMA configurations; constructing TiledMma with default ops.

Related errors

• Only Float16 or BFloat16 is supported
• M must be 64, 128 or 256
• N must be a multiple of 8 in the range 8…256
• cta_n={cta_n} invalid for use_2cta={use_2cta}: bf16 K-major
• Cannot find CUTLASS headers required for JIT compilation. Pl

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/4484caabbd478ff2. Report an issue: GitHub.