tile-ai / tile-ai/tilelang

[BUG] `T.gemm` K-alignment diagnostic points at a padding value that then fails a second, different check

Open
#2,907 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.