[Feature Request] SM120 block-scaled GEMM support for register-resident attention

[sepcnt]

Required prerequisites

• I have searched the Issue Tracker that this hasn't already been reported. (comment there if it has.)

Motivation

Replacing the MMA macros in #2253 with T.gemm_blockscaled is blocked by:

• Operand scope: fragment A / shared B is rejected. SageAttention3 keeps Q and quantized P in registers to avoid reloads and shared-memory round trips.
• Warp geometry: for M=N=K=128 and 256 threads, FullRow with sf_layout="blockscaled_chunk_kmajor" rejects the per-warp 1x8 MMA atom grid.

Solution

Support register-resident A and scales, and the 8x1 consumer warp partition, through the existing SM120 block-scaled GEMM API. Preserve the register-only path without implicit shared-memory staging.

Alternatives

Keep the low-level MMA macros. Shared/shared GEMM works with Square + compact scales or FullRow + rowmajor scales, but does not preserve the current attention pipeline.

Additional context

Tested on main eab74a4a, with unchanged block-scaled GEMM implementation; PyTorch 2.11.0+cu130, NVRTC, sm_120a.

Both shared/shared controls passed GPU correctness checks. The replacement fails compilation, so no performance regression is claimed. Migration should retain correctness and stay within 5% of CUDA SageAttention3 raw-core latency on the same GPU.

Minimal PoC

Lowering only; no GPU execution or external inputs required.

import tilelang
import tilelang.language as T


def repro(fragment_a):
    @T.prim_func
    def main(O: T.Tensor((128, 128), T.float32)):
        with T.Kernel(1, threads=256):
            A = (T.alloc_fragment((128, 128), T.float4_e2m1fn)
                 if fragment_a else T.alloc_shared((128, 128), T.float4_e2m1fn))
            B = T.alloc_shared((128, 128), T.float4_e2m1fn)
            S = T.alloc_shared((128, 2), T.uint32)
            C = T.alloc_fragment((128, 128), T.float32)
            T.clear(A)
            T.clear(B)
            T.clear(S)
            T.gemm_blockscaled(
                A, B, C, S, S, transpose_B=True, clear_accum=True,
                policy=T.GemmWarpPolicy.FullRow, k_start=0,
                sf_a_granularity_k=16, sf_b_granularity_k=16,
                sf_layout="rowmajor" if fragment_a else "blockscaled_chunk_kmajor",
            )
            T.copy(C, O)
    return main


target = tilelang.tvm.target.Target({"kind": "cuda", "arch": "sm_120a"})
for fragment_a in (True, False):
    try:
        with target:
            tilelang.lower(repro(fragment_a), target=target, enable_device_compile=False)
    except Exception as error:
        print(f"fragment_a={fragment_a}: {error}")

Observed failures:

• fragment_a=True: rejects A scope=local.fragment.
• fragment_a=False: SM120 compact scale packages require a positive even MMA atom grid per warp, got 1x8.