[BUG] `T.gemm` K-alignment diagnostic points at a padding value that then fails a second, different check
- Dominant language
- Python
- Stars
- 7.4k
- Forks
- 742
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 104
Description
**Version:** v0.1.11 (`VERSION` = 0.1.11), commit `607a9144cf375e1176ad487dd5b0e3d87c2cb79a`
**Platform:** Linux, CUDA 13.0, NVIDIA RTX 4000 Ada Generation (sm_89), Python 3.10, torch 2.12.1+cu130
**Related:** #2715 (same K=20 abort; this report is about the guidance the error gives, not the abort itself)
## Summary
`T.gemm` enforces K alignment in two places with two different moduli, and the first
failure's message implies a padding value that the second check then rejects. A user who
follows the first diagnostic gets a second, differently-worded failure.
| K | Result |
|---|---|
| 20 | `Check failed: (stride == 4 \|\| stride % 8 == 0) is false: stride=20` |
| 24 | `AssertionError: block_K (24) must be a multiple of micro_size_k (16)` |
| 32 | works, max abs err 3.8e-06 |
| 48 | works, max abs err 5.7e-06 |
| 64 | works, max abs err 3.8e-06 |
24 is the natural reading of the first message — the smallest multiple of 8 above 20 — and
it is wrong. The effective constraint is a multiple of 16, but nothing in the first error
says so.
## Reproduction
```python
import tilelang
import tilelang.language as T
M = N = 128
K = 20 # vary: 20, 24, 32
@tilelang.jit(out_idx=[2])
def build():
@T.prim_func
def k(A: T.Tensor((M, K), "float16"),
Bm: T.Tensor((K, N), "float16"),
C: T.Tensor((M, N), "float")):
with T.Kernel(1, threads=128) as bx:
As = T.alloc_shared([M, K], "float16")
Bs = T.alloc_shared([K, N], "float16")
Cf = T.alloc_fragment([M, N], "float")
T.clear(Cf)
T.copy(A, As)
T.copy(Bm, Bs)
T.gemm(As, Bs, Cf)
T.copy(Cf, C)
return k
build()
```
For K=24 and K=32, allocate the shared buffers at the padded K, `T.clear` them, and copy
the real `[.., 0:20]` region in — the failure at 24 happens at build time regardless of
how the data is staged.
## Why this matters beyond the individual checks
The two constraints live in different layers — the `stride % 8` rule in the copy/stride
path and `micro_size_k = 16` in the MMA tile selection — and neither diagnostic mentions
the other. From the outside they read as one constraint reported inconsistently. For
anyone padding irregular shapes to satisfy `T.gemm`, the correct target is the **union**
of the constraints, and that union is not discoverable from the error messages.
## Suggested fix
Either is a clear improvement:
1. **Make the first message state the effective requirement.** If the MMA path will demand
a multiple of `micro_size_k`, say so at the point of first failure:
`K=20 is not supported: K must be a multiple of 16 (micro_size_k) for the fp16 MMA path`.
2. **Check the strongest constraint first**, so users hit the binding one immediately
rather than being sent to an intermediate value that fails later.
A one-line note in the docs on residual-tile handling — that padding K must target
`micro_size_k`, not the stride rule — would also help.
Contributor guide
Research direction
Start by reproducing the issue with the provided Python T.gemm example at K=20, 24, and 32, then trace the two diagnostics reached by the T.gemm build. Check how the stride % 8 rule and micro_size_k requirement are reported; done means the first guidance leads directly to a valid padding target or clearly states the effective constraint.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100