[UPSTREAM][TILELANG] WGMMA descriptor assert 'Not a canonical GMMA_MN layout' for narrow [64,16] B operand with 256 threads

[lcy-seso]

Summary

Upstream tilelang regression, staged here for refinement before forwarding to tile-ai/tilelang. On sm_90, T.gemm(A_s, B_s, C, transpose_A=True) with a narrow [64, 16] shared-memory B operand and 256 threads (2 warpgroups) fails while building the WGMMA B-operand descriptor with tilelang c7fabc4:

AssertionError: Not a canonical GMMA_MN layout: stride<1,0>=2 != W=1
  tilelang/cuda/intrinsics/macro/wgmma_macro_generator.py:166 (compute_gmma_descriptor)

The same kernel passes with 128 threads, for both fp16 and bf16. Instrumentation shows the rejected operand layout is a plain row-major (64, 16) -> (_i * 16 + _j) tile.

The assertion family arrived with the descriptor rework around b1e1e25 (tilelang#2452, 2026-07-01); these shapes compiled and ran on 65dbc98.

Minimal repro (verified on H200, sm_90a, image c7fabc4-torch2.10)

import tilelang
import tilelang.language as T
import torch

@tilelang.jit
def atb_narrow_kernel(BT: int, DK: int, DVP: int, dtype: str, accum_dtype: str, threads: int):
    @T.prim_func
    def main(
        q: T.Tensor((BT, DK), dtype),
        do: T.Tensor((BT, DVP), dtype),
        out: T.Tensor((DK, DVP), accum_dtype),
    ):
        with T.Kernel(1, threads=threads):
            q_s = T.alloc_shared([BT, DK], dtype)
            do_s = T.alloc_shared([BT, DVP], dtype)
            acc = T.alloc_fragment([DK, DVP], accum_dtype)
            T.copy(q, q_s, disable_tma=True)
            T.copy(do, do_s, disable_tma=True)
            T.fill(acc, 0.0)
            T.gemm(q_s, do_s, acc, transpose_A=True,
                   policy=T.GemmWarpPolicy.FullRow)
            T.copy(acc, out)
    return main

# threads=128 -> OK; threads=256 -> AssertionError (both fp16 and bf16)
kernel = atb_narrow_kernel(64, 64, 16, "bfloat16", "float32", 256)

Affected TileOps smoke tests (2)

• tests/ops/test_gla_chunkwise_bwd.py::test_gla_bwd[2-64-2-64-64-64-dtype1-False] and [...dtype2-False] — the GLA bwd sequential kernel uses threads_seq=256 with num_v_partitions=4, so its do_s slice is [chunk=64, dim_v/4=16].

Notes for upstream

Either narrow-N operands should fall back to a supported path (smaller wgmma atom / mma), or layout inference should give the operand a canonical layout for 2-warpgroup tiling instead of asserting.

Repro script: drafts/tilelang-c7fabc4-repros/repro4_gmma_mn_assert.py in TileOpsGov.

[lcy-seso]

Reclassification: this is NOT an upstream regression — the assert exposes pre-existing silent wrong-code

Bisect baseline (both versions built from source, run in the current CI image on H200):

tilelang	narrow [64,16] B + 256 threads
65dbc98 (2026-06-18, previous image)	compiles, produces wrong results (fp16 and bf16, mismatch vs torch reference beyond 1e-2)
dfc7e95 (2026-06-24)	compiles, produces wrong results
c7fabc4 (2026-08-07)	hard assert Not a canonical GMMA_MN layout

So the WGMMA path for a narrow B operand with 2 warpgroups was silently miscomputing before the upgrade window; the July descriptor assert (b1e1e25 / #2452 family) turned it into a compile-time error. The upstream report should be reframed as: "narrow-N B operand with 2 warpgroups silently miscomputes (still reproducible on 0.1.11-era commits); the new assert catches it but the shape needs either a working fallback or a clear unsupported-shape diagnostic."

Action needed on the TileOps side, independent of upstream: GLABwdKernel's sequential kernel uses exactly this shape (do_s = [chunk=64, dim_v/num_v_partitions=16], threads_seq=256, T.gemm(..., transpose_A=True, policy=FullRow) at tileops/kernels/gla/gla_bwd.py:127). The GLA bwd tests passed on the old image, which now looks suspicious — either the test tolerances mask the error or the kernel's layout differs from the repro in some relevant detail. The kernel's numerical correctness on 65dbc98 should be re-verified against a tightened reference before trusting existing GLA bwd results.