mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-20 06:30:07 +08:00
Diao Han
2c31e071c6
* [BugFix][Hopper][Blackwell] Read the WGMMA/UMMA K-panel stride from the layout
A shared operand sliced along MN silently computed wrong results whenever
block_K spanned more than one swizzle atom:
T.gemm(a_s, b_s[0:64, :], c0, transpose_B=True) # wrong past the first K atom
When block_K exceeds one swizzle atom the shared layout splits K into panels
stored panel-major -- a (128, 128) bf16 tile decodes as
(128,(64,2)):(64,(1,8192)) -- so stepping to the next K panel moves by the
*buffer's* MN extent. The atom offset formulas instead reconstructed that step
as `mn_extent * swizzle_atom_elems` using the extent of the operand, which is
smaller for a slice. The slice origin was already layout-derived and correct, so
only K atoms past the first were affected: the two halves of a split buffer read
each other's rows, in bounds and without any diagnostic.
block_K == 64 was correct only by accident -- at 16-bit dtypes that is exactly
one 128B atom, so `ki // k_atom_size` is always zero and the bad multiplier was
multiplied away. MN-major operands never reach the term at all.
Take the panel stride from the decoded layout instead of reconstructing it:
cute.restrict already returns the sublayout carrying it, and slicing does not
change a stride. Expose it as `k_panel_elems` on WGMMADescriptorParams /
TCGEN05DescriptorParams, None when the layout has a single K panel or the
operand is MN-major -- exactly the cases where the term is multiplied by zero,
so those keep the previous expression and behaviour.
The same defect was duplicated across all seven atom emitters: dense Hopper
(ss/rs), sparse Hopper (ss/rs) and Blackwell (ss/ts/block-scaled).
Verified on H200 (sm_90a):
- Whole-buffer path is unchanged: across 120 K-major layouts (fp16/bf16/fp8_e4m3/
fp32 x 5 MN extents x 6 K extents), k_panel_elems is either None or exactly
equal to the old reconstruction, for both descriptor builders. The fix is a
strict generalization of the path that already worked.
- New test_tilelang_kernel_gemm_sliced_operand.py cross-checks a sliced operand
against one-buffer-per-slice over operand x majorness x block_K; it fails on
the four K-major block_K>64 combinations before this change and passes after.
- A layout-level assertion covers the Blackwell descriptor, which has no numeric
coverage here for lack of the hardware.
Closes #2964
* [BugFix][Hopper][Blackwell] Refuse a multi-panel K mode with no constant stride
Review follow-up on the K-panel stride decode.
The decode returned None whenever the K mode was not a canonical 2-tuple, and
every consumer then fell back to `mn_extent * swizzle_atom_elems` -- the exact
reconstruction this series removes. The fallback is only sound for a single-panel
layout, where `ki // k_atom_size` is always zero, and nothing checked that. So a
swizzled multi-panel layout that failed to decode would silently re-enter the
wrong-code path: the same shape as the original defect, where a formula's
precondition was assumed but never enforced.
Move the decode into a shared `decode_k_panel_elems` used by both descriptor
builders, and make it assert instead of returning None when a swizzled operand
spans more than one panel without exposing a constant panel stride. Panel count
is rounded up, so a K extent that is not a whole number of atoms cannot slip
through as single-panel.
Also tighten the accept condition: reading `stride[-1]` as the panel step is only
valid once the leading sub-mode is confirmed to be the contiguous atom. On a
panels-first `(2, atom):(panel_stride, 1)` ordering the old condition would have
returned 1 -- a silently wrong step rather than a refusal.
Consumers now read the step through `k_panel_stride(mn_extent)` on the params
dataclass, so the contract and its single-panel default live in one place
instead of being restated at eleven call sites.
New test_tilelang_cuda_wgmma_operand_layout.py covers the decode directly
(hardware-independent): canonical multi-panel strides, the single-panel and
MN-major None cases, and each refusal -- panels-first, flat multi-atom, a
three-sub-mode K, and a partial-atom extent.
Re-verified on H200: the whole-buffer invariant is unchanged by the tightening
(120 K-major layouts x 2 builders, 60 single-panel, 0 disagreements), every
multi-panel K slice is still accepted with the buffer-derived stride, and the
full suite set passes (174 passed, 40 skipped).