Why Habu

Triton #10927 computes a gathered GEMM. On Blackwell, the tensor-descriptor path gets it badly wrong when K = 511: 96% of the result differs from the reference. The operation itself fits on one line. Most of the kernel is tile bookkeeping and data movement. I am building Habu to generate that part and check it before it reaches the GPU.

In the issue, a gathered GEMM from a vLLM mixture-of-experts path goes wrong on Blackwell when K = 511 in bf16. 96.42% of the output differs from the reference. Change K to 512 and it works. The descriptor is accepted. The kernel launches. It returns numbers. They are simply the wrong numbers.

# the two paths compute the identical GEMM
pointer + masked load :  max diff 7.6e-05   ← correct
tensor descriptor / TMA:  max diff 130.45    ← 96.42% of outputs wrong

The failure was reproduced upstream on an NVIDIA B200, and I reproduced it on my DGX Spark (GB10). It is deterministic. That makes it easier to debug, but it also means a system can run the bad kernel over and over without giving any other sign that something is wrong.

1Most of the kernel is not the GEMM

Here is the kernel from the issue. Ignore the exact failure for a moment and look at what the author had to write. The gathered GEMM is acc += tl.dot(a, b). Nearly everything around it arranges tiles or moves data.

import torch, triton
import triton.language as tl

def _alloc_fn(size, alignment, stream):
    return torch.empty(size, dtype=torch.int8, device="cuda")

triton.set_allocator(_alloc_fn)              # host callback the TMA path needs for scratch

@triton.jit
def gemm_gather_kernel(
    a_ptr, b_ptr, c_ptr, M, N, K,
    stride_am, stride_ak, stride_bn, stride_bk, stride_cm, stride_cn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    USE_TD: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    offs_k = tl.arange(0, BLOCK_K)

    if USE_TD:                               # the tensor-descriptor (TMA) path
        a_desc = tl.make_tensor_descriptor(
            base=a_ptr, shape=(M, K), strides=(stride_am, stride_ak),
            block_shape=(1, BLOCK_K),
        )
        b_desc = tl.make_tensor_descriptor(
            base=b_ptr, shape=(N, K), strides=(stride_bn, stride_bk),
            block_shape=(BLOCK_N, BLOCK_K),
        )
        gather_idx = offs_m.to(tl.int32)     # gather rows of A by index
    else:                                    # the pointer + masked-load path
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        if USE_TD:
            a = a_desc.gather(gather_idx, k * BLOCK_K)
            b = b_desc.load([pid_n * BLOCK_N, k * BLOCK_K]).T
        else:
            k_mask = offs_k < K - k * BLOCK_K
            a = tl.load(a_ptrs, mask=k_mask[None, :], other=0.0)
            b = tl.load(b_ptrs, mask=k_mask[:, None].broadcast_to(BLOCK_K, BLOCK_N).T, other=0.0)
            b = tl.trans(b)
        acc += tl.dot(a, b)
        if not USE_TD:
            a_ptrs += BLOCK_K * stride_ak
            b_ptrs += BLOCK_K * stride_bk

    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
    tl.store(c_ptrs, acc)

The source spells out base pointers, six strides, two tensor shapes, three block dimensions, two descriptors, the gather index, the K-tile loop, mask arithmetic, a transpose, and the accumulator type. Triton exposes this control because control is how you get performance. The cost is that the source also carries a large collection of hardware assumptions, and any mistake in the source, the lowering, or one of those assumptions can change the answer.

That is the part I want to get rid of. The operation is simple enough to fit on one line, yet its implementation contains enough shape, stride, tile, masking, and descriptor state to hide a silent numerical failure. The descriptor API accepts the base, shape, strides, and block shape without checking that the strides satisfy TMA's 16-byte alignment requirement.

2The part I actually want to write

Strip out the descriptors, masks, transpose, and K-tile loop, and the kernel reduces to this:

The equation says everything needed to define the result. stride_ak, the transpose, the masks, and the tile loop tell a particular GPU how to produce it. They do not change what the operation means. My starting point for Habu was simple: could the equation be the program instead of a comment above more than fifty lines of implementation?

3Habu starts with the equation

This is the same indexed GEMM in Habu:

128 EXTENT: #M                     \ rows of A and O
 64 EXTENT: #N                     \ rows of B, columns of O
 32 EXTENT: #K                     \ contraction dimension

TENSOR:  A  ( #M #K )             \ A@  ( ix<#M> ix<#K> -- r )
TENSOR:  B  ( #N #K )             \ B@  ( ix<#N> ix<#K> -- r )
TENSOR:  O  ( #M #N )             \ O!  ( r ix<#M> ix<#N> -- )
ITENSOR: IX ( #M #M )             \ IX@ ( ix<#M> -- ix<#M> )

MODEL: GGEMM  O[m n] = Σ k  A[ IX[m] k ] · B[n k] ;

The MODEL: line is the kernel definition. O[m n] names the two free output dimensions. A[IX[m] k] gathers a row from A. B[n k] supplies the other operand. Σ k sums over the contraction dimension, and · multiplies the two factors. From that equation Habu derives the checked kernel, a dataflow record, and the shape obligations. There is no tile loop or descriptor in the source because the compiler has not chosen them yet.

The declarations above the line are what make the short form useful rather than decorative. An index into #M and an index into #K are different types. Swap them and the definition does not compile. The compiler keeps those extent types when it later chooses layouts, movement, and tiles. One line is only an improvement if the compiler still has enough information to catch mistakes and build a serious kernel.

4What has to happen after that one line

A short source program is useless if it produces a slow or unreliable kernel. With fixed tensor shapes and a target GPU, Habu can choose the tile sizes, data layout, pipeline depth, tensor-core instruction, fusion, data movement, masking, and numerical precision at compile time. When a dimension is known only at runtime, it can compile a small set of checked kernels and select one after seeing the actual shape. Those choices appear in the generated kernel and the compilation report, not in the model definition, unless someone explicitly pins one for an experiment.

How Habu chooses depends on what limits the operation. Fusion helps when writing an intermediate tensor to global memory costs more than the extra work inside one kernel. A compute-bound GEMM needs the right tensor-core tile, pipeline depth, and accumulated precision. Habu can derive candidates from the operation, the shapes, and the target GPU; reject candidates that break bounds, alignment, layout, precision, or resource limits; then benchmark the rest. For #10927, any candidate that uses TMA must pass the 16-byte stride check before Habu emits the descriptor. This part is still being built.

TMA is a lowering, not part of the language

A[IX[m] k] says which elements are read. It does not say which instruction moves them. Depending on the target, the best answer might be predicated per-lane loads, a pointer table, cp.async, or TMA. Before Habu chooses TMA, it checks the target's base and stride rules. For K = 511, it can pad the storage, choose a predicated path, or report a compile-time error. It cannot knowingly emit the invalid descriptor.

Concatenation gives the compiler the complete chain

This is where the concatenative language matters. A Habu program already is a sequence of typed words. Put two operations next to each other and their combined sequence is explicit; there is no Python trace to rebuild before looking for fusion. If a value is produced by one word and consumed by the next, the compiler can keep it in registers or shared memory rather than write it to global memory and read it back.

\ y = relu(a·x + y), one kernel, no global-memory round trips
: FUSED ( -- )
   x g LOAD-V4  a SCALE-V4  y g LOAD-V4  ADD-V4  RELU-V4  y g STORE-V4 ;

Every word declares a typed stack effect. A sequence that leaves an extra value, consumes a value that is not there, or connects incompatible layouts fails to compile. Automatic differentiation for the backward pass uses the same sequence: walk the forward words in reverse, substitute the derivative word for each operation, and record which forward values need to be saved or recomputed. Each differentiable word still needs a derivative rule; the useful part is that the transformation runs over the same typed program instead of a separate runtime graph.

The same representation continues above a single kernel

MODEL: also accepts a signature followed by a sequence of operation words. The shapes are not repeated in the signature. They come from the tensor declarations and the composition checker unifies their extents:

128 EXTENT: #B
  1 EXTENT: #ONE
  8 EXTENT: #D
 16 EXTENT: #H

TENSOR: X  ( #B   #D )
TENSOR: W1 ( #D   #H )
TENSOR: B1 ( #ONE #H )
TENSOR: W2 ( #H   #D )
TENSOR: B2 ( #ONE #D )
TENSOR: Y  ( #B   #D )

MODEL: FFN-SKIP ( x w1 b1 w2 b2 -- y )
   LINEAR GELU LINEAR  x RESIDUAL-ADD  RMSNORM ;

The backend partitions the sequence into fused regions, selects tiles for each region, and lowers each region to a GPU kernel. Before I keep a generated kernel, Habu records four kinds of evidence:

The K = 511 case belongs in both places. It is a golden test for every legal lowering of the gathered GEMM, and a forced TMA plan with the bad stride is a negative compile test. The rule prevents this known-invalid descriptor from being emitted. The golden catches mistakes in plans that passed the rule. One does not replace the other.

I am not building Habu merely to save more than fifty lines in one GEMM. I want most machine-learning kernels to stop being source code that somebody has to carry from one GPU generation to the next. The model program would contain the operations, shapes, and precision requirements. Tiling, fusion, data movement, tensor-core selection, and the backward pass would be regenerated for H100, B200, or whatever comes after them. Hand-written kernels would still exist for cases the compiler cannot handle, but they would be the exception rather than the foundation of the stack.

Triton's useful abstraction is the tile. Habu asks whether the programmer can stop one level earlier, at the operation itself. The hard part is not the one-line syntax. It is producing the kernel an expert would have written, across irregular shapes and several GPU generations, without dropping the checks that motivated the project in the first place. That is the research problem.