[BUG][Fuzzer][ice-on-valid-code] `T.gemm_sp` with `uint8` operands `KeyError`s at compile instead of compiling like `int8`

Required prerequisites

• I have read the documentation https://tilelang.com.
• I have searched the Issue Tracker and Discussions that this hasn't already been reported. (+1 or comment there if it has.)

What version of TileLang are you using?

0.1.13 (latest release, reproduced this session on L40S).

System information

TileLang 0.1.13 / CUDA 13.0 / PyTorch 2.13.0 / NVIDIA L40S (sm_89). The crash is in Python during layout inference, before any target codegen, so it is architecture-independent.

Problem description

A sparse GEMM whose A/B operands are uint8 aborts compilation with KeyError: "dtype('uint8')". The identical kernel with int8 operands compiles. uint8 is treated as a supported sparse-GEMM operand dtype everywhere the emitter computes sparsity metadata (GROUP_CONFIG, get_e_factor, get_e_replicate_factor all accept it), but the dtype_abbrv lookup table used to name the MMA intrinsic omits it, so construction of the sparse emitter raises before the kernel can be built. Not a regression — see Provenance.

Trigger vs. control (same kernel, operand dtype swapped)

operand dtype	in dtype_abbrv?	in GROUP_CONFIG/get_e_factor?	result
int8	yes	yes	compiles
uint8	no	yes	KeyError: "dtype('uint8')"

uint16/uint32 are absent from dtype_abbrv as well, but unlike uint8 they are not listed as valid operand dtypes in GROUP_CONFIG, so uint8 is the one dtype the tables disagree on.

Reproducible example code

import tilelang
import tilelang.language as T

def make(in_dtype):
    M = N = K = 128
    bM = bN = bK = 64
    e_dtype, e_factor, acc = "uint8", 8, "int32"   # e_dtype = metadata; 8-bit operands -> e_factor 8

    @tilelang.jit(out_idx=[-1])
    def build():
        @T.prim_func
        def gemm_sp(
            A_sparse: T.Tensor((M, K // 2), in_dtype),        # sparse operand, dtype = in_dtype
            E:        T.Tensor((M, K // e_factor), e_dtype),
            B:        T.Tensor((K, N), in_dtype),
            C:        T.Tensor((M, N), acc),
        ):
            with T.Kernel(T.ceildiv(N, bN), T.ceildiv(M, bM), threads=128) as (bx, by):
                As = T.alloc_shared((bM, bK // 2), in_dtype)
                Es = T.alloc_shared((bM, bK // e_factor), e_dtype)
                Bs = T.alloc_shared((bK, bN), in_dtype)
                Cl = T.alloc_fragment((bM, bN), acc)
                T.clear(Cl)
                for k in T.Pipelined(T.ceildiv(K, bK), num_stages=1):
                    T.copy(A_sparse[by * bM, k * bK // 2], As)
                    T.copy(E[by * bM, k * bK // e_factor], Es)
                    T.copy(B[k * bK, bx * bN], Bs)
                    T.gemm_sp(As, Es, Bs, Cl)
                T.copy(Cl, C[by * bM, bx * bN])
        return gemm_sp
    return build

make("int8")()    # -> compiles
make("uint8")()   # -> KeyError: "dtype('uint8')"

Traceback

File ".../tilelang/cuda/intrinsics/macro/mma_sp_macro_generator.py", line 136, in _initialize_abbrev
    self.a_dtype_abbrv = self.dtype_abbrv[a_dtype]
KeyError: "dtype('uint8')"

(v0.1.13, mma_sp_macro_generator.py:136.)

Expected behavior

uint8 operands should compile the sparse GEMM (or, if uint8 is not intended to be a supported operand dtype, be rejected up front with a clear message). The sparse metadata layer already treats uint8 as valid: it is a key in GROUP_CONFIG and drives get_e_factor / get_e_replicate_factor, and NVIDIA's mma.sp m16n8k64 has a .u8 operand form (the same kernel routed through the C++ tl::gemm_sp_ss library path instantiates the CUTLASS SparseMma template with uint8_t operands, i.e. it doesn't reject uint8). The int8 sibling reaching kernel construction over the identical inputs shows the operation is expressible on this path; only the abbreviation lookup is missing.

Additional context

Root cause. The sparse MMA emitter's dtype-abbreviation table does not include uint8, so constructing the emitter for a uint8-operand GEMM crashes before it can name the intrinsic. SparseTensorCoreIntrinEmitter.__init__ first computes the sparsity metadata factors for the operand dtype (succeeds for uint8 — uint8 is a valid key in GROUP_CONFIG), then calls _initialize_abbrev, which does self.dtype_abbrv[a_dtype]; dtype_abbrv lists int8 but not uint8, so the lookup raises KeyError. This is reached through the public T.gemm_sp path when it selects MMA SP lowering (GemmSPMMA → the Python emitter, which is where this abbreviation table is consulted). (On 0.1.13 the earlier T.gemm_sp_v2 spelling was merged into T.gemm_sp.)

Generalization tested this session (v0.1.13, L40S) — adjacent unsigned widths, class = inconsistent guard

Each cell run in its own fresh process, same kernel with the operand dtype swapped:

operand dtype	result	fails at
int8	COMPILE_OK	—
uint8	KeyError: "dtype('uint8')"	dtype_abbrv[a_dtype] (mma_sp_macro_generator.py:136) — passed GROUP_CONFIG first
uint16	KeyError: "dtype('uint16')"	GROUP_CONFIG[a_dtype] (sparse_layout.py:28, in get_e_factor) — earlier
uint32	KeyError: "dtype('uint32')"	GROUP_CONFIG[a_dtype] (sparse_layout.py:28) — earlier

This is the discriminating result: uint8 is the ONLY dtype where the two tables disagree — GROUP_CONFIG/get_e_factor accept it (so it clears layout inference) but dtype_abbrv omits it (so it crashes at the abbreviation step). uint16/uint32 are absent from both tables, so they are rejected consistently one step earlier — a distinct root (unsupported dtype, not a table disagreement), not the same bug. Class = inconsistent/too-narrow guard (dtype_abbrv narrower than GROUP_CONFIG). Title kept specific to uint8.

Suggested fix. Add "uint8": "uint8" to the dtype_abbrv table at mma_sp_macro_generator.py:55 so it agrees with GROUP_CONFIG (and uint16/uint32 if those are meant to be admitted too). If uint8 operands are not intended to be supported, make it a single explicit reject shared by all three tables instead of a silently disagreeing table (not verified end-to-end).

Provenance. Reproduced on v0.1.13. The available release history is squashed, so an exact introducing PR isn't determinable; the defect is the standing disagreement between dtype_abbrv (no uint8) and GROUP_CONFIG/get_e_factor (accept uint8).

Dedup. I searched the open and closed tracker (gemm_sp uint8, dtype_abbrv, SparseTensorCoreIntrinEmitter, uint8 sparse) and found no existing report of this defect. It is distinct from the other sparse-GEMM issues (#2603 WGMMA block_N swizzle, #2605 K-tail drop, #2606 metadata layout, #2634 proxy fence), which are all value/codegen defects on already-compiling kernels rather than a dtype-table crash.

Reach. uint8 is explicitly listed as a valid operand dtype in GROUP_CONFIG and accepted by get_e_factor/get_e_replicate_factor, but the dtype_abbrv table omits it. It is not demonstrated in examples/ or testing/ — grepping the sparse-GEMM example and tests this session, the operand dtype is only ever float16/int8 (the uint8 occurrences in examples/gemm_sp/sparse_utils.py are local index buffers, not GEMM operands), so no test exercises a uint8 operand and CI stays green.

Impact. The trigger is narrow: a sparse GEMM (T.gemm_sp) whose operands are specifically uint8, a dtype no example or test uses. When it fires the failure is a compile-time KeyError raised during Python layout inference, before any codegen — loud, caught immediately, and it blocks that one kernel from building; nothing is silently miscompiled and it cannot reach production as a wrong result. Fixing it closes the standing disagreement between the sparse metadata tables (which accept uint8) and the abbreviation table (which omits it), so the accepted operand set is consistent across the emitter.

Start in tilelang/cuda/intrinsics/macro/mma_sp_macro_generator.py at _initialize_abbrev and compare its dtype_abbrv table with GROUP_CONFIG and the factor helpers in tilelang/cuda/intrinsics/sparse_layout.py. Reproduce the supplied T.gemm_sp example with int8 and uint8 operands. Done means uint8 follows the intended supported-or-clearly-rejected behavior without a raw KeyError, with coverage for the regression if the project’s test layout permits it.