Writing a Triton kernel; CUTLASS at a glance

Take the two lines every transformer block runs twice per layer:

s = x.float() + residual.float()
v = s.square().mean(dim=-1, keepdim=True)
out = (s * torch.rsqrt(v + eps)).to(weight.dtype) * weight
residual_out = s.to(x.dtype)  # preserve the pre-norm residual

Kernel count and physical HBM traffic need a trace. For bf16 inputs/weights with fp32 intermediates, the logical materialization accounting follows the explicit reference in §8.2:

At 2048 rows, 253,976 logical bytes per row correspond to about 155 microseconds at peak 3.35 TB/s, versus about 25 microseconds for the fused payload if each row reloads weight. A warm shared weight can reduce HBM traffic. These are conditional payload estimates: the final norm lacks residual semantics, and no measured multi-layer saving follows.

SGLang's answer in CUDA is python/sglang/kernels/jit/csrc/elementwise/rmsnorm.cuh — 371 lines holding four kernel variants, a warp-level path for $d \le 256$, a CTA path, a double-load path for pre-Blackwell, a wide path for Blackwell, hand-rolled two-level reductions, explicit __shared__ staging, and packed 16-byte vector types. It is excellent code. It is also exactly the kind of code you will not write for the eleventh fused operation you need this quarter. Triton exists for the other ten.

CUDA makes you name a thread. Everything else follows: you compute which elements that thread owns from blockIdx and threadIdx, you choose a vector width so the loads coalesce, you reduce inside a warp with shuffles, you stage the per-warp partials through __shared__, you place two __syncthreads() around that staging, and you get one of them wrong the first time.

Triton makes you name a block of elements. You write tl.arange(0, BLOCK_N) and operate on the whole vector; the compiler decides how many threads back it, how to vectorise the loads, which lanes carry which elements, and how to lower tl.sum into a shuffle tree plus shared-memory staging plus the barriers. You still choose the parallel decomposition — what one program owns, and how many programs there are — because that is a modelling decision no compiler can make for you. You stop choosing everything below it.

Here is the same reduction, both ways, from the two files this chapter reads.

The CUDA column is real code. Here is its reduction, verbatim:

__global__ __launch_bounds__(kDim / 16) void rmsnorm_cta_double(const RMSNormParams __grid_constant__ params) {
  using namespace device;
  using Float2 = packed_t<Float>;
  using Storage = AlignedVector<Float2, 4>;

  constexpr auto kNumThreads = kDim / 16;
  constexpr auto kNumWarps = kNumThreads / kWarpThreads;

  const auto& [input, weight_ptr, output, input_stride, output_stride, num_tokens, eps] = params;
  const auto gmem = tile::Memory<Storage>::cta(kNumThreads);
  __shared__ float smem[32];
// ...
  float sum_of_squares = 0.0f;
#pragma unroll
  for (auto j = 0u; j < 4u; ++j) {
    const auto [x, y] = cast<fp32x2_t>(input_first[j]);
    sum_of_squares += x * x + y * y;
  }
#pragma unroll
  for (auto j = 0u; j < 4u; ++j) {
    const auto [x, y] = cast<fp32x2_t>(input_second[j]);
    sum_of_squares += x * x + y * y;
  }

  sum_of_squares = warp::reduce_sum(sum_of_squares);
  float norm_factor;
  if constexpr (kNumWarps == 1) {
    norm_factor = math::rsqrt(sum_of_squares / kDim + eps);
  } else {
    const auto warp_id = threadIdx.x / kWarpThreads;
    smem[warp_id] = sum_of_squares;
    __syncthreads();
    if (warp_id == 0) {
      const auto tx = threadIdx.x;
      const auto local_sum = tx < kNumWarps ? smem[tx] : 0.0f;
      sum_of_squares = warp::reduce_sum(local_sum);
      smem[tx] = math::rsqrt(sum_of_squares / kDim + eps);
    }
    __syncthreads();
    norm_factor = smem[warp_id];
  }

Every construct there was introduced in §0.3: warp shuffles, shared memory, barriers, packed vector loads, __launch_bounds__. Nothing is gratuitous. Note kDim / 16 — the thread count is a compile-time function of the hidden size, which is why this kernel is a template instantiated per shape, and why python/sglang/kernels/ops/layernorm/norm.py:L129-L144 raises "jit rmsnorm: unsupported hidden_size=..." for shapes nobody instantiated.

Now the Triton version of the same operation. This is not illustrative code — it is in-tree, tested, and shipping:

@triton.jit
def _gemma_rmsnorm_kernel(
    x_ptr,
    w_ptr,
    out_ptr,
    n_cols,
    stride_row,
    stride_col,
    eps,
    BLOCK_N: tl.constexpr,
):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_N)
    mask = cols < n_cols
    x = tl.load(x_ptr + row * stride_row + cols * stride_col, mask=mask, other=0.0).to(
        tl.float32
    )
    var = tl.sum(x * x, axis=0) / n_cols
    rstd = 1.0 / tl.sqrt(var + eps)
    w = tl.load(w_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    out = x * rstd * (1.0 + w)
    tl.store(
        out_ptr + row * n_cols + cols,
        out.to(out_ptr.dtype.element_ty),
        mask=mask,
    )

Nine constructs, in order:

The launch

def _num_warps(block_n: int) -> int:
    if block_n >= 4096:
        return 16
    if block_n >= 1024:
        return 8
    return 4


def gemma_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
    """Gemma RMSNorm = normalize(x) * (1 + weight), fp32 math, single pass."""
    orig_shape = x.shape
    n = orig_shape[-1]
    x2 = x.reshape(-1, n)
    m = x2.shape[0]
    out = torch.empty((m, n), dtype=x.dtype, device=x.device)
    block_n = triton.next_power_of_2(n)
    _gemma_rmsnorm_kernel[(m,)](
        x2,
        weight,
        out,
        n,
        x2.stride(0),
        x2.stride(1),
        eps,
        BLOCK_N=block_n,
        num_warps=_num_warps(block_n),
    )
    return out.reshape(orig_shape)

kernel[(m,)](args...) is the whole launch API. The subscript is the grid; the call is the arguments. num_warps is a reserved keyword Triton consumes rather than passing to your kernel — it is how you tell the compiler how many threads to spend on one program, and it is the one occupancy knob you keep from CUDA.

The masking idiom, with numbers

tl.arange requires a power-of-two length, so the launcher rounds the hidden size up with triton.next_power_of_2(n) and the kernel masks the excess. This is the thing everyone gets wrong first, in two ways: forgetting other=0.0 (masked lanes then hold garbage and poison the sum), and dividing by BLOCK_N instead of n_cols (the mean is then wrong by a constant factor that looks like a bad epsilon). Notice line 35 divides by n_cols, the true width.

The fused-add variant, _gemma_fused_add_rmsnorm_kernel at python/sglang/kernels/ops/layernorm/minimax_m3_rmsnorm.py:L46-L85, is the same kernel with a second input pointer, a second stride pair, and one extra store. That store is the whole point of fusing: it emits the pre-norm sum for the next block's residual add, so the sum is never re-read from HBM by a separate kernel.

You now have a kernel. Is it any good? The procedure is three steps and it does not require a profiler.

1. Count bytes at the relevant memory level. Fused residual RMSNorm reads x, residual and weight and writes sum and output: 40,960 logical bytes at d=4096 bf16. If the shared weight is warm in L2, the input/output HBM component is 32,768 bytes per row. State which assumption applies rather than silently omitting weight.
2. Count FLOPs. One add for the residual, one multiply and one add for the sum of squares, one multiply by rstd, one by the weight: $F \approx 5d = 20480$ FLOP per token.
3. Divide, and place it on the roofline. $I = F/Q = 20480/32768 = 0.625$ FLOP/byte.

This norm uses scalar fp32 arithmetic and reductions, so peak bf16 tensor-core FLOPs is the wrong compute ceiling. Low arithmetic intensity motivates minimizing traffic, but occupancy, vectorization, reduction layout, launch latency and registers also matter. Fusion is one important optimization, not the only possible one.

This conditional warm-weight payload bound scales to about 20 microseconds for 2048 rows. For a valid HBM accounting model, payload divided by measured time gives effective bandwidth; if most accesses hit L2, that ratio is not physical HBM bandwidth. Also inspect device counters, small-batch latency and end-to-end fraction before deciding optimization is complete.

The bound also tells you when to stop. If you measure 24 µs against a 20 µs bound you are at 83% of peak bandwidth and there is at most 17% left in the entire kernel; go find a different bottleneck. That is a far more useful answer than "it got 8% faster when I changed BLOCK_N".

BLOCK_N, num_warps and num_stages trade off against each other in ways that depend on the shape, the dtype and the architecture, and nobody can predict the winner. All three strategies for dealing with that are in these trees.

1. Online autotuning

# The best (BLOCK_N, num_warps, num_stages) is shape- and arch-sensitive, so
# autotune over a grid keyed on the attention shape. Benchmarked once per key
# (a one-time stall on the first prefill of each new shape), then cached.
_AUTOTUNE_CONFIGS = [
    triton.Config({"BLOCK_N": bn}, num_warps=w, num_stages=ns)
    for bn in (32, 64, 128)
    for w in (1, 2, 4)
    for ns in (1, 2)
]


@triton.autotune(
    configs=_AUTOTUNE_CONFIGS,
    key=["topk", "H", "DIM"],
    prune_configs_by={"early_config_prune": _prune_configs},

Three things to take from this. The config list is a Cartesian product — $3 \times 3 \times 2 = 18$ configs — each of which must be separately compiled and separately timed the first time a new key appears. The key is the tuple of runtime arguments whose values change the right answer: ["topk", "H", "DIM"]. Change topk and Triton re-benchmarks; change the batch size, which is not in the key, and it reuses the cached choice. Picking the key is a modelling decision: too narrow and you tune for a shape that behaves differently, too wide and you stall on every new request shape. prune_configs_by lets you drop configs that cannot work — a config needing more shared memory than the SM's 228 KB, for instance — before they are compiled.

vLLM's Mamba kernels use the explicit form, nine hand-written configs:

@triton.autotune(
    configs=[
        triton.Config(
            {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64},
            num_stages=3,
            num_warps=8,
        ),
        triton.Config(
            {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32},
            num_stages=4,
            num_warps=4,
        ),
# ... seven more, down to BLOCK_SIZE_M=64, BLOCK_SIZE_N=64, num_stages=4, num_warps=2
    ],
    key=["chunk_size", "K", "IS_CAUSAL"],
)

The shape of that list is itself instructive: big M×N tiles get more warps and fewer pipeline stages, small tiles get two warps and five stages. num_stages is the depth of the software pipeline over the K-loop — how many BLOCK_K slabs are in flight in shared memory at once. More stages hides more latency and costs more shared memory, which is why the small-tile configs can afford five.

2. Offline tuning, checked into the repo

Autotuning at serving time is a latency spike on the critical path. vLLM's MoE kernel dodges it entirely: it ships 330 JSON files of pre-tuned configs under vllm/model_executor/layers/fused_moe/configs/, one per (expert count, intermediate size, device, dtype), each mapping a batch size to a config.

{
    "1": {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 64,
        "BLOCK_SIZE_K": 128,
        "GROUP_SIZE_M": 64,
        "num_warps": 8,
        "num_stages": 5
    },
    "2": {
        "BLOCK_SIZE_M": 64,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 256,
        "GROUP_SIZE_M": 64,
        "num_warps": 4,
        "num_stages": 3
    },

Read the keys as batch sizes. At $M = 1$ the tuner chose a 64×64 tile with 5 stages; by $M = 2$ it wants 64×128×256 with 3. The tile that is right for one token is not the tile that is right for two, which is the entire reason this file exists. get_moe_configs (vllm/model_executor/layers/fused_moe/fused_moe.py:L1110-L1122) looks the file up by name, and when it misses:

    # If no optimized configuration is available, we will use the default
    # configuration
    logger.warning_once(
        "Using default MoE config. Performance might be sub-optimal! "
        "Config file not found at %s",
        ", ".join(config_file_paths),
    )

If you see that line in a production log, you are running an untuned MoE. The mechanism is credited in-source to SGLang PR #2628 (vllm/model_executor/layers/fused_moe/fused_moe.py:L1109) — a rare, explicit cross-pollination between the two projects.

3. Cache the tuning, and shout when it happens late

# Enable Triton autotuning result caching to disk by default.
# Without this, Triton re-runs autotuning on every process restart,
# adding significant latency to the first inference request.
# This writes autotuning results to TRITON_CACHE_DIR.
# It can still be overridden by setting TRITON_CACHE_AUTOTUNING=0
# in the environment.
os.environ.setdefault("TRITON_CACHE_AUTOTUNING", "1")

And after warmup, vLLM installs hooks so that any compile or autotune during serving is an incident: vllm/utils/jit_monitor.py:L119-L141 raises or logs "Triton kernel JIT compilation during inference: <name>. This causes a latency spike; consider extending warmup to cover this shape/config.". Run with --jit-monitor-mode=error to make it fatal. For benchmarking, the opposite knob exists: VLLM_TRITON_FORCE_FIRST_CONFIG=1 monkey-patches the autotuner to take the first valid config, "to eliminate autotuning variability when measuring kernel performance" (vllm/envs.py:L1177-L1184).

Everything above appears at scale in vLLM's Triton MoE GEMM, vllm/model_executor/layers/fused_moe/fused_moe.py:L298-L610. It is a grouped GEMM: each block of rows belongs to one expert, and the block must fetch that expert's weight matrix. Read it in five moves.

Move 1 — a one-dimensional grid, decoded into a 2-D tile index, deliberately out of order.

    # -----------------------------------------------------------
    # Map program ids `pid` to the block of C it should compute.
    # This is done in a grouped ordering to promote L2 data reuse.
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

The naive mapping pid_m = pid // num_pid_n sweeps a whole row of output tiles before moving down, so consecutive concurrent blocks share A rows but touch every column of B. The grouped ordering walks a GROUP_SIZE_M-tall column of tiles instead, so the blocks resident at any moment share both a small set of A rows and a small set of B columns, and both stay in L2. GROUP_SIZE_M is in the tuned JSON above — 64 at batch 1, 1 at batch 4. This is a pure Triton-level optimisation with no CUDA analogue in the source; it is scheduling, and Triton lets you express scheduling because you own the grid.

Move 2 — the gather, and the early exit. The row indices are not pid_m * BLOCK + arange; they are loaded from the alignment tensor that moe_align_block_size (vllm/model_executor/layers/fused_moe/moe_align_block_size.py:L11-L103) produced, so that each block's rows all belong to one expert:

    offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
        return
    if not naive_block_assignment:
        offs_token_id = pid_m * BLOCK_SIZE_M + offs
        offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
# ...
    offs_token = offs_token.to(tl.int64)

    token_mask = offs_token < num_valid_tokens

Three idioms in ten lines. The grid is sized for the worst case and programs return early when the actual token count is smaller — a device-side bound that avoids relaunching with a different grid. offs_token is a block of indices read from memory, which is how you write a gather in Triton: use a block value as an offset. And token_mask — the padding tokens the aligner inserted carry an out-of-range index, so the same mask that guards the loads also silently drops them from the store at the end. Below, at L423-L440, a block whose expert is -1 (not on this rank) writes zeros and returns.

Move 3 — the K-loop with a 2-D mask.

        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        # Load the next block of A and B, generate a mask by checking the
        # K dimension.
# ...
        else:
            a = tl.load(
                a_ptrs,
                mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                other=0.0,
            )
            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)

token_mask[:, None] & (offs_k[None, :] < ...) is the RMSNorm mask idiom in two dimensions: [:, None] and [None, :] broadcast two 1-D block values into a [BLOCK_SIZE_M, BLOCK_SIZE_K] boolean. The row mask handles padded tokens, the column mask handles a K that is not a multiple of BLOCK_SIZE_K. The accumulator is fp32 regardless of input dtype — the same discipline as the .to(tl.float32) in the norm kernel, for the same reason.

Move 4 — the epilogue, all of it before a single store.

    if use_int8_w8a16:
        accumulator = accumulator * b_scale
    elif (use_fp8_w8a8 or use_int8_w8a8) and not (group_k > 0 and group_n > 0):
        accumulator = accumulator * a_scale * b_scale
# ...
    if HAS_BIAS:
        accumulator += bias[None, :]

    # Router (MoE) weight multiplication:
    # This multiplication MUST be performed in float32 before any precision
    # conversion to ensure numerical stability, which is especially critical
    # on ROCm platforms.
    if MUL_ROUTED_WEIGHT:
        moe_weight = tl.load(
            topk_weights_ptr + offs_token,
            mask=token_mask,
            other=0,
        )
        accumulator *= moe_weight[:, None]

    # Final precision conversion:
    # Cast once at the end to the desired compute/output dtype.
    accumulator = accumulator.to(compute_type)

    # -----------------------------------------------------------
    # Write back the block of the output
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
    tl.store(c_ptrs, accumulator, mask=c_mask)

Dequantisation, bias and the router weight are all applied to the fp32 accumulator in registers, and the tile touches HBM exactly once. This is epilogue fusion, hand-written. Hold on to it — it is the concept CUTLASS turns into a type system in the next section.

Move 5 — the launch. All the if use_fp8_w8a8 and if HAS_BIAS branches you just read are tl.constexpr, so each combination compiles to a separate kernel with the dead branches gone. The grid is a lambda over the chosen config, because the grid depends on the block sizes the tuner picked:

    grid = lambda META: (
        triton.cdiv(EM, META["BLOCK_SIZE_M"])
        * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]),
    )

Everything in this kernel maps back to the twenty-six-line norm: program_id, arange, masked load/store, constexpr specialisation, an fp32 accumulator, and a grid you compute yourself. The only genuinely new primitive is tl.dot; the emitted MMA instruction family depends on dtype, shape, layout, hardware and compiler version. Inspect generated PTX/SASS rather than assuming WGMMA.

§0.3 left this pair for here. SGLang's split-KV merge exists twice — a compiled op in the external sgl_kernel wheel and a Triton kernel in-tree — and the wrapper picks:

# Automatically fallback to the Triton kernel in some cases
# (e.g., for AMD GPUs, when the head dimension is not a multiple
# of 4 or 8, and in FP8 precision)
def _supported_dtypes(o: torch.Tensor) -> bool:
    return o.dtype in [torch.float32, torch.half, torch.bfloat16]


def _supported_headdim(o: torch.Tensor) -> bool:
    headdim = o.shape[2]  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
    if o.dtype == torch.float32:
        return headdim % 4 == 0
    return headdim % 8 == 0
# ...
    if (
        _is_cuda
        and _supported_dtypes(prefix_output)
        and _supported_headdim(prefix_output)
    ):
        return merge_state_v2(
            prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse
        )
    else:
        # Fallback to Triton kernel
        return merge_state_triton(
            prefix_output, prefix_lse, suffix_output, suffix_lse, output, output_lse
        )

The predicate selects the compiled CUDA path for eligible aligned inputs and Triton otherwise. Selection is not proof of universal speed superiority. Both have supported shape/dtype/backend limits; Triton does not automatically support every shape or platform. Benchmark the eligible alternatives while preserving the operation's numerical contract.

SGLang has generalised this into a registry. python/sglang/kernels/spec.py:L29-L51 enumerates provenance as an enum — TORCH, TORCH_COMPILE, TRITON, JIT, AOT, CUTE_DSL, FLASHINFER, DEEPGEMM, AITER, TORCH_NPU and more — and python/sglang/kernels/README.md:L39-L64 describes the contract: registration records metadata only, and no Triton compile or JIT build happens until a kernel is actually called. Every operator carries a required forward_native, "the pure-torch correctness reference every other implementation is checked against" (python/sglang/kernels/README.md:L75-L76). If you write a Triton kernel for either engine, that is the shape of the contract you are joining.

CUTLASS/CuTe expose detailed tile, pipeline and epilogue abstractions. These are not exclusive capabilities: Triton's official persistent-matmul tutorial includes persistent scheduling, TMA, warp specialization and fused epilogues on supported hardware. Compare the exact version, kernel, layout control and authoring complexity; no fixed "last 20%" performance advantage follows from choosing one language.

CUTLASS is a C++ template library, not a compiler. Its organising idea is a hierarchy of tiles, each level a template parameter you instantiate rather than a runtime argument you pass.

The composition is explicit in vLLM's own code. A GEMM is a mainloop plus an epilogue, both built from the same tile shape:

  using CollectiveEpilogue =
      typename cutlass::epilogue::collective::CollectiveBuilder<
          cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
          ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
          ElementAcc, float, ElementC, StrideC, AlignmentCD, ElementD, StrideD,
          AlignmentCD, EpilogueSchedule, EVTCompute>::CollectiveOp;

  static constexpr size_t CEStorageSize =
      sizeof(typename CollectiveEpilogue::SharedStorage);
  using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
      static_cast<int>(CEStorageSize)>;

  // clang-format off
  using CollectiveMainloop =
      typename cutlass::gemm::collective::CollectiveBuilder<
          cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
          ElementAB, cutlass::layout::RowMajor, AlignmentAB,
          ElementAB, cutlass::layout::ColumnMajor, AlignmentAB,
          ElementAcc, TileShape, ClusterShape,
          Stages,
          KernelSchedule>::CollectiveOp;
  // clang-format on

  using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
      cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
      cutlass::gemm::PersistentScheduler>>;

StageCountAutoCarveout<sizeof(CollectiveEpilogue::SharedStorage)> is the whole philosophy in one line: the pipeline depth of the mainloop is computed at compile time from how much of the SM's 228 KB of shared memory the epilogue already claimed. Triton's num_stages is a number you guess and then autotune; CUTLASS's is a number the type system solves for. That is the trade — less exploration, more compile time, and a much steeper authoring curve.

The tile shapes are selected by C++ dispatch on the problem geometry rather than by benchmarking. The SM90 FP8 path has seven configurations; at $M > 128$ it takes TileShape = Shape<_128, _128, _128> with a Ping-Pong schedule, at $M \ge 8192$ and $K \ge 6144$ it takes Shape<_256, _128, _128> Cooperative, and for $M \in (16, 64]$ with $N \le 1280$ it takes Shape<_64, _16, _256> with ClusterShape = Shape<_1, _4, _1> and operands swapped (csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm90_fp8_dispatch.cuh:L102-L110, L122-L130, L148-L166). A 16-wide N tile is a decode-shaped GEMM: almost no rows, so what matters is streaming B.

Epilogue fusion as a type

The MoE kernel fused its epilogue by writing the arithmetic before the store. CUTLASS makes the epilogue a compile-time expression tree — an epilogue visitor tree, Sm90EVT — so any composition of broadcasts, elementwise ops and casts is a type. §4.3 owns the scaled-GEMM epilogue and quotes the tree; what matters here is what fusion buys, stated by the vendored loader itself:

//
// This file is a modified excerpt of
// include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
// from https://github.com/NVIDIA/cutlass v3.5.0
// It has been modified to support either row/column or scalar broadcasting
// where the tensor being loaded from is always passed in via a device pointer.
// This lets one compiled kernel handle all cases of per-tensor or
// per-channel/per-token quantization.
//
// This interface also allows the scales to be passed in as tensors that
// consistently reside on the device, which avoids an issue with a previous
// implementation where scalars needed to be on the CPU since they
// were passed in via float values. This created a potential performance hazard
// if scales were initially on the device, and caused torch.compile graphs
// breaks when moving scales to the CPU.
//

One compiled kernel covering per-tensor, per-channel and per-token scaling — and note the last clause: the naive alternative moved a scalar to the CPU and broke the §8.2 graph. Kernel design and compiler behaviour are the same problem.

CuTe: layouts as algebra

Underneath sits CuTe, whose one idea is that a tensor's shape and stride form a Layout object you can compose, permute and coalesce at compile time. vLLM's helpers are a clean miniature:

template <size_t... I, typename Layout>
CUTE_HOST_DEVICE static constexpr auto permute_layout(Layout l) {
  static_assert(rank(l) == sizeof...(I), "Invalid permutation, rank mismatch");
  return cute::make_layout(cute::get<I>(l)...);
}

// is the layout f(x) = x
template <typename Layout>
CUTE_HOST_DEVICE static constexpr bool is_identity_layout() {
  if constexpr (std::is_same_v<Layout, void>) {
    return true;
  } else {
    constexpr auto coalesced_layout = coalesce(Layout{});
    if constexpr (rank(coalesced_layout) == 1 &&
                  stride<0>(coalesced_layout) == 1) {
      return true;
    }
    return false;
  }
}
// ...
template <typename T, typename Elements>
CUTE_HOST_DEVICE static constexpr auto create_auto_vectorizing_copy() {
  constexpr auto bits = sizeof_bits_v<T> * Elements{};
  if constexpr (bits % 128 == 0) {
    return AutoVectorizingCopyWithAssumedAlignment<128>{};
  } else if constexpr (bits % 64 == 0) {

if constexpr eliminates branches inside each compiled specialization. Runtime dispatch can still choose among C++ template instantiations, as the earlier wrapper does. Triton JIT also specializes compile-time values. Compare build time, binary size, cache misses and supported specializations, not "templates have no runtime dispatch".

Machete is the worked example of all of this: a mixed-precision W4A16 GEMM built on CUTLASS, "a spiritual successor to the Marlin kernel but optimized for Hopper architectures" (csrc/libtorch_stable/quantization/machete/Readme.md:L1-L3), whose type pairs and tile schedules are code-generated by generate.py so that "we can generate multiple type pairs and different tile shapes using the same kernel template" (csrc/libtorch_stable/quantization/machete/Readme.md:L41-L43). Its static shape-to-schedule heuristic and its prepacked weight layout are §4.3's to explain, and it explains them there.

Follow one call to gemma_rmsnorm(x, weight, eps) with x of shape [257, 6144], bf16.

1. gemma_rmsnorm reshapes to [257, 6144], allocates out, computes block_n = triton.next_power_of_2(6144) = 8192 and num_warps = _num_warps(8192) = 16.
2. The subscript _gemma_rmsnorm_kernel[(257,)] binds the grid and returns a launcher.
3. Specialisation. Triton builds a key from everything that changes the generated code: the tl.constexpr values (BLOCK_N=8192), the pointer dtypes (bf16 in, bf16 out, fp32 weight), num_warps and num_stages, and divisibility properties of the integer arguments — a stride divisible by 16 licenses wider vector loads than one that is not. Calls with different runtime n_cols may share a specialization when padded BLOCK_N and other key properties match. Verify actual JIT keys; hidden size is not always a constexpr.
4. Cache lookup in TRITON_CACHE_DIR. vLLM redirects this into VLLM_CACHE_ROOT when compile caching is on and otherwise Triton uses ~/.triton/cache (docs/usage/security.md:L376); SGLang sets it under its own cache root at python/sglang/srt/environ.py:L1658. A hit skips everything below.
5. Compilation on a miss: Python AST → Triton IR → TritonGPU IR (where the layout assignment that decides which lane holds which element happens) → LLVM IR → PTX → cubin via ptxas. All stages land in the cache directory next to the cubin.
6. Launch of 257 blocks × 512 threads. On a warm cache the Python-side cost is a dict lookup and a driver call; that residual cost is what §8.1 removes.

To see the output of step 5, the reliable route is to read the cache directory. Each compiled kernel gets a hashed subdirectory holding the IR at every stage:

# run your kernel once with a private, empty cache
export TRITON_CACHE_DIR=/tmp/tritoncache
rm -rf $TRITON_CACHE_DIR
python -c "import torch, sglang.kernels.ops.layernorm.minimax_m3_rmsnorm as m; \
  x=torch.randn(257,6144,device='cuda',dtype=torch.bfloat16); \
  w=torch.randn(6144,device='cuda',dtype=torch.float32); m.gemma_rmsnorm(x,w,1e-6)"

# every stage of the pipeline is on disk, one directory per specialisation
find $TRITON_CACHE_DIR -type f | sed 's|.*/||' | sort | uniq -c
#   .ttir   Triton IR            .ttgir  TritonGPU IR (layouts assigned)
#   .llir   LLVM IR              .ptx    PTX          .cubin  the binary
#   .json   metadata: num_warps, shared bytes, register count

# SASS, the actual machine code, needs the CUDA toolkit
nvdisasm -c $TRITON_CACHE_DIR/*/*.cubin | head -60

The .json is the file to read first: it carries the shared-memory bytes and the register count per thread, which is what decides how many of your blocks fit concurrently on an SM. The Python object also exposes the same artefacts — a compiled kernel handle carries an asm mapping with "ttir", "ttgir", "llir" and "ptx" keys — but the exact attribute names move between Triton releases, so the cache layouts and handle attributes are both version-dependent inspection aids.

A one-program-per-row grid exposes only one block at m=1. The quoted CUDA grid-stride loop also uses min(num_tokens, max_occupancy*kNumSM), which is one block at num_tokens=1. It caps/reschedules many rows, not splits one row across SMs. More splitting can add synchronization overhead, so choose from measured end-to-end impact.

Use a contiguous and a noncontiguous input of non-power-of-two width, check the exact 1+weight Gemma convention, and keep reduction epsilon explicit. A GPU test should then add empty and wide inputs, supported aliases/strides, extreme finite values, NaN/Inf policy and dtype-specific error bounds. This CPU oracle does not test CUDA/Triton code generation or performance.

import numpy as np

rng = np.random.default_rng(4)
base = rng.normal(size=(3, 14))
x = base[:, ::2]  # noncontiguous width 7
w = rng.normal(size=7)
eps = 1e-6
assert not x.flags.c_contiguous
def gemma_norm(value, weight):
    return value / np.sqrt(np.mean(value**2, axis=-1, keepdims=True) + eps) * (1+weight)
strided = gemma_norm(x, w)
contiguous = gemma_norm(np.ascontiguousarray(x), w)
np.testing.assert_allclose(strided, contiguous)
padded = np.pad(x, ((0, 0), (0, 1)))
right_variance = np.sum(padded**2, axis=-1, keepdims=True) / x.shape[-1]
np.testing.assert_allclose(right_variance, np.mean(x**2, axis=-1, keepdims=True))
wrong_variance = np.mean(padded**2, axis=-1, keepdims=True)
assert not np.allclose(right_variance, wrong_variance)
print("Strided Gemma reference and true-width reduction contract pass.")

1. Check correctness against a torch reference before anything else. SGLang's test for the kernel this chapter read is the template:

def _gemma_rmsnorm_ref(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
    orig_dtype = x.dtype
    x_f = x.float()
    variance = x_f.pow(2).mean(dim=-1, keepdim=True)
    out = x_f * torch.rsqrt(variance + EPS)
    out = out * (1.0 + weight.float())
    return out.to(orig_dtype)


@pytest.mark.parametrize("shape", [(1, 512), (64, 6144), (257, 6144)])
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@torch.inference_mode()
def test_gemma_rmsnorm_matches_reference(shape, dtype):
    torch.manual_seed(0)
    x = torch.randn(*shape, device=DEVICE, dtype=dtype)
    weight = torch.randn(shape[-1], device=DEVICE, dtype=torch.float32)

    got = gemma_rmsnorm(x, weight, EPS)
    ref = _gemma_rmsnorm_ref(x, weight)

    torch.testing.assert_close(got, ref, atol=2e-2, rtol=2e-2)

Read the parametrisation as a checklist: (1, 512) is the degenerate one-program grid, (64, 6144) exercises the mask, (257, 6144) makes the grid non-uniform across SMs, and a separate test at L52-L65 feeds base[:, ::2] to prove the stride arguments are real. Copy this shape set for anything you write.

2. Time it against the bound.

python ~/sglang/test/registered/kernels/benchmark/layernorm/bench_norm.py
# Account for exact residual/Gemma semantics, dtypes, weight reads and cache state.
# Measure cold/warm data and model-level impact, not only an effective GB/s number.

# Separate checkout: run vLLM's own benchmark by absolute path.
python ~/vllm/benchmarks/kernels/benchmark_moe.py --help

Note that run_benchmark in python/sglang/kernels/jit/benchmark/utils.py:L81-L98 uses triton.testing.do_bench_cudagraph, not do_bench — it captures the launch into a graph so the measurement is kernel time, not launch time. Which of the two you want depends on the question you are asking; §8.1 explains the difference.

3. Separate tuning from execution. Use private TRITON_CACHE_DIR roots and compare cold compile/tune time, warm selected-kernel time and end-to-end startup independently. Forcing the first config changes both tuning and the chosen kernel; their latency difference is not just the tuning bill. Reset mutable outputs between autotune trials and validate every candidate before comparing performance.

• Triton tutorials — vector add, fused softmax, matrix multiplication, layer norm. The fused-softmax and layer-norm lessons are the same shape as the kernel in §3 and take under an hour.
• Tillet, Kung and Cox, "Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations" (MAPL 2019) — the original paper. The block abstraction and the layout-assignment pass are the durable parts; the surface syntax has moved on.
• NVIDIA/CUTLASS, pinned to v4.4.2 by CMakeLists.txt:L486. Start with the CuTe quickstart under media for layouts, then the repository's own examples for the collective-builder pattern you saw in scaled_mm.cuh.
• SGLang RFC #29630 and its finale PR #32072 — the reorganisation that produced python/sglang/kernels/, the KernelBackend enum, and the JIT/AOT split this chapter reads.
• SGLang PR #2628 — the tuned MoE config mechanism, credited in vLLM's own source at vllm/model_executor/layers/fused_moe/fused_moe.py:L1109. Worth reading for the methodology of producing the JSON files.
• Neighbours: §0.3 for every CUDA construct in Figure 1, §0.4 for the roofline this chapter places kernels on, §3.3 for a Triton attention kernel read line by line, §4.3 for Marlin, Machete and the scaled-GEMM epilogue, and §8.2 for the Triton you never see because Inductor wrote it.