tile-ai / tile-ai/tilelang

[BUG][Fuzzer][ice-on-invalid-code] `T.gemm` on a narrow dtype (`int8`/`fp8`/`bfloat16`) with any sub-atom `block_K` crashes nvcc instead of cleanly rejecting

Open
#3,032 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.4k
Forks
745
Avg merge
1d 1h
Merged PRs (30d)
104

Description

### Required prerequisites

- [x] I have read the documentation .
- [x] I have searched the [Issue Tracker](https://github.com/tile-ai/tilelang/issues) that this hasn't already been reported. (comment there if it has.)

### What version of TileLang are you using?

0.1.13 (latest release)

### System information

NVIDIA L40S (sm_89), CUDA 12.8 (nvcc), PyTorch 2.8.0. The defect is in the shared `mma.sync` macro generator and is not arch-gated within the sm_75–sm_89 tensor-core range (see Additional context).

### Problem description

A `T.gemm` whose reduction tile `block_K` is smaller than the operand dtype's MMA K-atom fails to compile: nvcc aborts with a `static_assert` `"tl::mma_sync: unsupported configuration"`. This affects the whole class of sub-atom tiles on the narrow dtypes whose atom exceeds a common `block_K` — verified this session on `int8` (atom 32, crashes at `block_K∈{8,16}`), `float8_e4m3` (atom 32, crashes at `block_K=16`), and `bfloat16` (atom 16, crashes at `block_K=8`). The identical kernel with `block_K` at or above the atom compiles and is correct in every case (see §17). `float16` at a sub-atom `block_K` does NOT crash — it silently miscomputes instead (distinct symptom, same upstream clamp; see §17).

The emitted call is `tl::mma_sync` with `K` below the dtype's atom. For `int8`/`fp8`/`bf16` no `MmaDispatcher` specialization matches that under-atom `K`, so the `static_assert` in the unspecialized `MmaDispatcher::exec` at [`mma.h#L167`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/tl_templates/cuda/instruction/mma.h#L167) fires. (`float16` differs only because mma.h *does* ship a real `m16n8k8` fp16 dispatcher at [`mma.h#L198`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/src/tl_templates/cuda/instruction/mma.h#L198), so the under-selected fp16 call binds to a valid-but-wrong-K instruction and miscomputes rather than crashing.)

A sub-atom `block_K` has no legal lowering for these dtypes (one int8/fp8 MMA consumes 32 K-elements, one bf16 MMA consumes 16). The expected outcome is a clean compile-time rejection, not an nvcc back-end `static_assert`. Regression introduced by a change to `_initialize_k_dim` — see Provenance.

nvcc error (int8, block_K=16)

```
tl::mma_sync(...);
mma.h(282): error: static assertion failed with "tl::mma_sync: unsupported configuration"
static_assert(!std::is_void_v, ...)
mma.h(167): error: static assertion failed with "tl::mma_sync: unsupported configuration"
static_assert(always_false_v>, ...)
[with AType=kInt8, BType=kInt8, CType=kInt32, M=16, N=8, K=16, ...]
2 errors detected in the compilation of "tvm_kernels.cu".
```

### Reproducible example code

```python
import tilelang, tilelang.language as T, torch

def gemm(M, N, K, dt, acc):
@T.prim_func
def main(A: T.Tensor((M, K), dt), B: T.Tensor((N, K), dt), C: T.Tensor((M, N), acc)):
with T.Kernel(1, threads=32):
sa = T.alloc_shared((M, K), dt); sb = T.alloc_shared((N, K), dt)
c = T.alloc_fragment((M, N), acc)
T.copy(A, sa); T.copy(B, sb); T.clear(c)
T.gemm(sa, sb, c, transpose_B=True) # int8 tensor-core requires B transposed
T.copy(c, C)
return main

def run(tag, M, N, K, dt, acc, tdt):
try:
k = tilelang.compile(gemm(M, N, K, dt, acc), out_idx=[2])
a = torch.randint(-3, 3, (M, K), device="cuda").to(tdt)
b = torch.randint(-3, 3, (N, K), device="cuda").to(tdt)
out = k(a, b); exp = a.float() @ b.float().T
print(f"{tag}: compiled, maxerr={(out.cpu().float()-exp.cpu()).abs().max().item()}")
except Exception as e:
print(f"{tag}: CRASH {type(e).__name__}")

run("int8 block_K=16 (bug) ", 16, 16, 16, "int8", "int32", torch.int8) # -> CRASH (nvcc static_assert, K=16)
run("int8 block_K=8 (bug) ", 16, 16, 8, "int8", "int32", torch.int8) # -> CRASH (nvcc static_assert, K=8, same root)
run("int8 block_K=32 (control)", 16, 16, 32, "int8", "int32", torch.int8) # -> compiled, maxerr=0.0
run("fp16 block_K=16 (control)", 16, 16, 16, "float16", "float32", torch.float16) # -> compiled, maxerr~0
```

### Traceback

The failure surfaces as an nvcc back-end `static_assert` (shown collapsed above), not a TileLang-level diagnostic — the front-end accepts the kernel and emits an int8 MMA call that has no valid instruction.

### Expected behavior

A clean compile-time rejection at the TileLang level (e.g. "`block_K` must be a multiple of the int8 MMA K-atom (32)"). TileLang already has this exact guard on the same path — [`gemm_mma.py#L104`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/op/gemm/gemm_mma.py#L104): `assert block_K >= micro_size_k` — but it is defeated because `micro_size_k` is shrunk to `block_K` first (see Root cause). Restoring the true atom makes this existing assert fire cleanly (verified below).

### Additional context

**Root cause.** `_initialize_k_dim` picks an MMA K smaller than the operand dtype's real atom when `block_K` is small. It computes [`self.k_dim = min(256 // a_dtype.bits, self.chunk)`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/intrinsics/macro/mma_macro_generator.py#L131), where `chunk` is `block_K`. For `int8` (8 bits) with `block_K=16`: `min(32, 16) = 16`. [`_initialize_mma_prefix`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/intrinsics/macro/mma_macro_generator.py#L167) then maps `k_dim==16` to `m16n8k16` — the `float16` atom — and `_initialize_micro_size` sets `micro_size_k = k_dim = 16`, so the guard at `gemm_mma.py#L104` sees `block_K(16) >= micro_size_k(16)` and passes. The `min(…)` clamp both selects a nonexistent int8 instruction and neutralizes the check that would have caught it.

Why this reads as reject (no legal int8 target) and a fix sketch, verified before→after

The `int8` tensor-core MMA has no atom with `K<32`; `block_K=16` provides fewer K-elements than one int8 MMA consumes, so there is no legal lowering — the natural outcome is rejection. The Hopper WGMMA sibling on the same file never clamps: [`self.k_dim = 256 // DataType(a_dtype).bits`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/intrinsics/macro/mma_macro_generator.py#L952) (int8 → 32 regardless of `block_K`).

One direction that matches the WGMMA sibling: pin `k_dim` to the dtype atom rather than clamping it below the atom, so `micro_size_k` stays 32 for int8 and the existing `block_K >= micro_size_k` assert rejects `block_K=16` with a clear message. Whether the clamp is instead meant to enable a smaller-`block_K` int4 packed path (its introducing PR was int4-focused) is a maintainer call.

Verified this session (0.1.13, sm_89): with `_initialize_k_dim` pinned to `256 // bits` (no clamp), `int8`/`block_K=16` now raises the clean `AssertionError: block_K (16) must be >= micro_size_k (32)` instead of the nvcc `static_assert`, and `int8`/`block_K=32` still compiles with `maxerr=0.0`.

**Suggested fix.** Pin `k_dim` to the dtype atom rather than clamping it below the atom (matching the WGMMA sibling's `256 // bits`), so `micro_size_k` stays 32 for int8 and the existing `block_K >= micro_size_k` assert at `gemm_mma.py#L104` rejects `block_K=16` with a clear message. Verified this session: with `_initialize_k_dim` pinned to `256 // bits` (no clamp), `int8`/`block_K=16` raises the clean `AssertionError: block_K (16) must be >= micro_size_k (32)` and `int8`/`block_K=32` still compiles with `maxerr=0.0`. Whether the clamp is instead meant to enable a smaller-`block_K` int4 packed path is a maintainer call.

**Provenance.** Introduced by [#2073](https://github.com/tile-ai/tilelang/pull/2073) ("[CUDA] Improve int4 GEMM lowering and packed codegen support", merged 2026-04-21), commit `174c4f5d`, which changed `self.k_dim = 256 // a_dtype.bits` → `self.k_dim = min(256 // a_dtype.bits, self.chunk)`. Before that change, `int8` `k_dim` was pinned to 32 and `block_K=16` was rejected by the `block_K >= micro_size_k` guard (a correct clean reject) — so this is a regression of the reject behavior, not a never-worked case. Still present on `main` (the clamp is at `tilelang/cuda/intrinsics/macro/mma_macro_generator.py:131`; the WGMMA sibling at `:952` still uses `256 // bits`).

**Dedup.** I searched the open and closed tracker and found no existing report of this crash. Distinct from [#2381](https://github.com/tile-ai/tilelang/issues/2381) (`block_K` *not a multiple* of the atom → silent K-tail miscompile): that is a supported tile size with a silent wrong result; this is a *sub-atom* tile with a compile-time crash and a different root (atom under-selection vs a floored K-loop). The `float16` sub-atom *silent miscompile* twin (see §17) is a distinct symptom of the same clamp and is not filed separately.

**Reach.** The trigger is a narrow-dtype `T.gemm` (documented dtypes) at any `block_K` below that dtype's MMA K-atom. `block_K=16` is the common `float16` default and a natural first choice a user carries over to `int8`/`fp8`; `block_K=8` breaks even `bfloat16` (atom 16). No shipped example or test exercises a sub-atom `block_K` — that is why CI is green (see example-run below).

**Generalization (§17).**

_Two-level root._
- **SOURCE-level (fragile implementation):** the clamp `self.k_dim = min(256 // a_dtype.bits, self.chunk)` in `TensorCoreIntrinEmitter._initialize_k_dim` ([`mma_macro_generator.py#L131`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/intrinsics/macro/mma_macro_generator.py#L131)) lets `k_dim` fall below the operand dtype's real K-atom when `block_K < atom`. `_initialize_micro_size` then sets `micro_size_k = k_dim` ([`#L221`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/intrinsics/macro/mma_macro_generator.py#L221)), which neutralizes the `block_K >= micro_size_k` guard at [`gemm_mma.py#L104`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/op/gemm/gemm_mma.py#L104). The `min(…)` both selects a nonexistent instruction and disarms the check that would catch it.
- **OPERATOR-level (what correlates):** any operand dtype routed through this SIMT `TensorCoreIntrinEmitter` whose K-atom exceeds the requested `block_K`. Whether the under-selection *crashes* or *silently miscomputes* depends only on whether `mma.h` happens to ship a `MmaDispatcher` at the under-selected `K` for that dtype.

**4-axis findings** (each cell run in its own process, fresh `TILELANG_CACHE_DIR`, 0.1.13 / sm_89; "emit-K" = the `K` in the generated `mma_sync<…>` call):

| axis | cell tested | observed result | same-root? |
|---|---|---|---|
| related-type | `float8_e4m3`, `block_K=16` | **CRASH** — emits `mma_sync`, same `mma.h(282)`/`mma.h(167)` static_assert | **YES** — broadens the class |
| related-type | `bfloat16`, `block_K=8` | **CRASH** — emits `mma_sync`, same `mma.h(282)`/`mma.h(167)` static_assert (bf16 atom is 16, so `min(16,8)=8`) | **YES** — broadens the class |
| related-type | `int8`, `block_K=8` | **CRASH** — emits `mma_sync` (`min(32,8)=8`), same static_assert | **YES** — same dtype, other sub-atom |
| related-type | `float16`, `block_K=8` | compiled, **maxerr=31.0** (silent wrong result) — same clamp under-selects to K=8, but `mma.h#L198` ships a real `m16n8k8` fp16 dispatcher so it binds a valid-but-wrong-K instruction | **DISTINCT symptom**, same upstream clamp — not claimed here |
| related-type | `int4`, `block_K=32` | compiled, emits `mma_sync<…,K=32>` (int4 atom is 64, but `min(64,32)=32` maps to the valid `m16n8k32` shape → no static_assert); fails later at a runtime packed-ABI check | **DISTINCT** — clamp lands on an existing instruction, not this bug |
| related-source | `int8`/`fp8`/`bf16` at/above atom (controls) | compiled, `maxerr=0.0` in every case | confirms the boundary is exactly `block_K < atom` |
| related-operator | Hopper WGMMA sibling `_initialize_k_dim` at [`#L952`](https://github.com/tile-ai/tilelang/blob/8001cc4ccf6149382d2019654a19f59c1d4d0482/tilelang/cuda/intrinsics/macro/mma_macro_generator.py#L952) | `self.k_dim = 256 // DataType(a_dtype).bits` — **no clamp** (int8 → 32 regardless of `block_K`) | correct-behavior contrast; the fix should match it |

Axes skipped: **similar-logic** beyond the WGMMA sibling — the `min(atom, chunk)` idiom appears only in this one emitter's `_initialize_k_dim`; the WGMMA copy is the sole sibling and it does not clamp (already tested above).

**Class boundary.** The crash class is `{int8, float8_e4m3, bfloat16}` (and any other dtype with no `MmaDispatcher` at its under-selected `K`) at `block_K < atom`. `float16` sub-atom is a *silent-miscompile* twin (distinct symptom, flagged distinct, not filed separately here — same fix closes it since pinning `k_dim` to the atom makes the `block_K >= micro_size_k` guard reject it too). `int4` sub-atom does not reach this failure because its clamp lands on a real `m16n8k32` shape.

**Example-run (PART 1, MANDATORY).** The prior Reach cited `int8` GEMM shipping in `examples/bitnet-1.58b/kernel_benchmark/tl_int8xint8.py`. I located it in the v0.1.13 tree and inspected/ran it:
- It **cannot run on the 0.1.13 `tilelang` install** — it imports `bitblas` and the old `tvm.tl` API (`from tvm import tl as TL`, `import tvm.tl.language`), neither of which is present (`ModuleNotFoundError: No module named 'bitblas'` / `'tvm'`). It is a legacy BitBLAS-era example, not a modern-API kernel.
- Even structurally it would **NOT exercise this bug**: its int8 config uses `chunk = 64` (i.e. `block_K=64`), a multiple of the int8 atom of 32, so it is a supported tile.
- The cited test `testing/python/kernel/test_tilelang_kernel_gemm_sm75.py` likewise builds its int8 GEMM with a `block_K` of 64 (a multiple of 32) and is sm_75-gated, so it neither exercises nor could run the sub-atom case on this host.

Corrected conclusion: **no shipped example or test exercises a sub-atom `block_K`**, so the example does not demonstrate the bug — the Reach rests on the tested cells above (which reproduce cleanly), not on any shipped example. The presence of `int8`/`fp8`/`bf16` `T.gemm` as documented, supported dtypes is what makes the sub-atom tile a reachable user mistake.

**Impact.** Trigger is a narrow-dtype `T.gemm` (`int8`/`fp8`/`bfloat16`) at a sub-atom `block_K` — already-invalid tiles, since one MMA consumes a full atom of K-elements. When it fires the failure is a loud compile-time crash (an nvcc back-end `static_assert` on every affected build) — deterministic, corrupts no output, and cannot reach a running workload silently; the cost is that a smaller tile size fails to build with a confusing back-end assert instead of a TileLang-level message. Fixing it (pin `k_dim` to the atom, matching the WGMMA sibling) moves the failure to a clear front-end rejection, closes the sub-atom `block_K` boundary class for all three crash dtypes, and — as a bonus — turns the `float16` sub-atom *silent miscompile* into the same clean rejection.

Contributor guide

Open the contributing guide

Research direction

Start in tilelang/cuda/intrinsics/macro/mma_macro_generator.py at _initialize_k_dim and compare it with the WGMMA sibling near line 952. Trace how _initialize_micro_size reaches the guard at tilelang/cuda/op/gemm/gemm_mma.py:104, then reproduce the int8 sub-atom and block_K=32 control cases from the issue. Done means invalid sub-atom configurations receive a clear TileLang assertion while valid configurations still compile correctly.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, python
Domain
backend, performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.