9 · The router and the Triton split-K kernels

This chapter covers the last code region, lines 475–583. It holds two Triton GEMM kernels and the split-K "epilogue" they share. They run the router at m = 128 and 256 and the pooler at m = 128. The chapter starts with the reason the router needs special treatment at all: its weights are FP32.

9.1 Why the router is different

The router scores 256 experts, and the top scores pick which experts process each token. A tiny error can flip a choice, so:

• its weights and outputs are FP32;
• its tolerance is 0.01 % relative error, 100 times stricter than the other layers.

Tensor cores multiply BF16, not FP32. The rules allow the standard workaround, which prepare() applies (lines 713–717):

hi = bf16(r)             # the nearest BF16 value: about 8 significant bits
lo = bf16(r - hi)        # the nearest BF16 value to what hi missed: about 8 more bits
x @ r.T  ≈  x @ hi.T + x @ lo.T        (both products accumulated in FP32)

Lab 4 measures it:

output error vs the FP64 reference (router tolerance 1e-4):
  FP32 weights : 1.67e-06  pass
  BF16 hi only : 1.58e-03  FAIL
  BF16 hi + lo : 2.57e-06  pass

SGLang uses FP32 math (F.linear) up to 8 tokens and two separate BF16 GEMMs plus an add above that. We always use hi + lo, fused into one kernel:

• the GEMV at m ≤ 32 (DUAL = true, chapter 6);
• the Triton kernel below at m = 128/256.

One pass reads both matrices, accumulates into one set of FP32 sums, and writes once. That is why the router is our biggest per-shape win: 1.7–4.5×.

9.2 Triton in five minutes

Triton is a Python-embedded language for GPU kernels. You write the code for one program instance, which is one CTA working on one tile, using whole blocks of values instead of individual threads. Triton decides how the block is spread over the threads.

9.3 _proj_tma_kernel, line by line

559
560
561
562
563
564
565

@triton.jit
def _proj_tma_kernel(
    x_desc, w_desc, w2_desc, out_ptr, ws_ptr, sem_ptr,
    M, N, K,
    BN: tl.constexpr, BM: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr,
    DUAL: tl.constexpr, OUT_F32: tl.constexpr, NS: tl.constexpr, ATOM: tl.constexpr = False,
):

• 559–562The inputs:
  • three TMA descriptors: x, the weight, and an optional second weight (the router's lo);
  • the output, the scratch buffer and the counters;
  • the sizes M, N and K.
• 563–564Compile-time constants:
  • the tile shape BN (weight rows) × BM (tokens) × BK (K per step);
  • the number of K splits;
  • DUAL (hi + lo) and OUT_F32;
  • NS, the pipeline depth;
  • ATOM, the epilogue variant.

566
567
568
569
570
571
572
573
574

    pid_n = tl.program_id(0)
    pid_k = tl.program_id(1)
    pid_m = tl.program_id(2)
    grid_n = tl.num_programs(0)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_m = pid_m * BM + tl.arange(0, BM)
    k_per = K // SPLIT
    k0 = pid_k * k_per
    acc = tl.zeros((BN, BM), dtype=tl.float32)

• 566–568The grid is three-dimensional: weight-row tile, K slice and token tile.
• 569The number of row tiles, used to number the output tiles.
• 570–571This tile's weight-row indices, which are the output columns, and its token indices.
• 572–573This program's K range: one SPLIT-th of K.
• 574A [BN, BM] FP32 accumulator, held in the program's registers.

575
576
577
578
579
580
581
582

    for k in tl.range(k0, k0 + k_per, BK, num_stages=NS):
        xt = tl.trans(x_desc.load([pid_m * BM, k]))
        w = w_desc.load([pid_n * BN, k])
        acc = tl.dot(w, xt, acc)
        if DUAL:
            w2 = w2_desc.load([pid_n * BN, k])
            acc = tl.dot(w2, xt, acc)
    _epilogue(acc, out_ptr, ws_ptr, sem_ptr, M, N, pid_n, pid_k, pid_m, grid_n, offs_n, offs_m, BN, BM, SPLIT, OUT_F32, ATOM)

• 575tl.range(…, num_stages=NS) asks Triton to software-pipeline the loop. Triton allocates NS shared-memory buffers and issues the copies for the next NS−1 steps while computing the current one: chapter 6.5's idea, generated automatically.
• 576A TMA copy of a [BM, BK] block of x, transposed to [BK, BM].
• 577A TMA copy of a [BN, BK] block of weights.
• 578acc += w @ xᵀ on the tensor cores. As in our GEMV, the weight rows form the MMA's M side and the tokens its N side ("swap A and B"). On Hopper, Triton turns this into the wgmma tensor-core instruction.
• 579–581For the router, also load the lo block and accumulate it into the same accumulator, with the same x block already in shared memory.
• 582Hand the finished partial tile to the epilogue.

524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556

@triton.jit
def _proj_kernel(
    x_ptr, w_ptr, w2_ptr, out_ptr, ws_ptr, sem_ptr,
    M, N, K,
    BN: tl.constexpr, BM: tl.constexpr, BK: tl.constexpr, SPLIT: tl.constexpr,
    DUAL: tl.constexpr, OUT_F32: tl.constexpr,
):
    pid_n = tl.program_id(0)
    pid_k = tl.program_id(1)
    pid_m = tl.program_id(2)
    grid_n = tl.num_programs(0)
    offs_n = pid_n * BN + tl.arange(0, BN)
    offs_m = pid_m * BM + tl.arange(0, BM)
    offs_k = tl.arange(0, BK)
    k_per = K // SPLIT
    k0 = pid_k * k_per
    w_ptrs = w_ptr + offs_n[:, None] * K + (k0 + offs_k)[None, :]
    x_ptrs = x_ptr + offs_m[None, :] * K + (k0 + offs_k)[:, None]
    m_mask = offs_m[None, :] < M
    acc = tl.zeros((BN, BM), dtype=tl.float32)
    if DUAL:
        w2_ptrs = w2_ptr + offs_n[:, None] * K + (k0 + offs_k)[None, :]
    for _ in range(0, k_per, BK):
        xt = tl.load(x_ptrs, mask=m_mask, other=0.0)
        w = tl.load(w_ptrs)
        acc = tl.dot(w, xt, acc)
        if DUAL:
            w2 = tl.load(w2_ptrs)
            acc = tl.dot(w2, xt, acc)
            w2_ptrs += BK
        w_ptrs += BK
        x_ptrs += BK
    _epilogue(acc, out_ptr, ws_ptr, sem_ptr, M, N, pid_n, pid_k, pid_m, grid_n, offs_n, offs_m, BN, BM, SPLIT, OUT_F32)

• 524–556The same algorithm, written with explicit pointer blocks (lines 540–541) and masked tl.loads instead of TMA descriptors. It is kept as a fallback for configurations with tma = 0; none exists today. Compare lines 546–555 with 575–581: the structure is identical.

9.4 The split-K epilogue

The router has only 256 rows. With 64-row tiles that is 4 row tiles, far too few CTAs, so K is split 16 ways (m = 128) or 8 ways (m = 256). Several programs therefore produce partial sums for the same output tile, and someone has to add them up. _epilogue has two strategies.

475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495

@triton.jit
def _epilogue(acc, out_ptr, ws_ptr, sem_ptr, M, N, pid_n, pid_k, pid_m, grid_n, offs_n, offs_m,
              BN: tl.constexpr, BM: tl.constexpr, SPLIT: tl.constexpr, OUT_F32: tl.constexpr, ATOM: tl.constexpr = False):
    m_mask = offs_m[None, :] < M
    if ATOM:
        # split-K partials are atomically added into a zeroed fp32 [M, N] buffer; the last CTA of the tile
        # copies it out and re-zeroes it (no serial S-way read)
        tile = pid_m * grid_n + pid_n
        a_ptrs = ws_ptr + offs_m[None, :] * N + offs_n[:, None]
        tl.atomic_add(a_ptrs, acc, mask=m_mask, sem="relaxed")
        tl.debug_barrier()
        cnt = tl.atomic_add(sem_ptr + tile, 1, sem="acq_rel")
        if cnt == SPLIT - 1:
            res = tl.load(a_ptrs, mask=m_mask, other=0.0, cache_modifier=".cg")
            o_ptrs = out_ptr + offs_m[None, :] * N + offs_n[:, None]
            if OUT_F32:
                tl.store(o_ptrs, res, mask=m_mask)
            else:
                tl.store(o_ptrs, res.to(tl.bfloat16), mask=m_mask)
            tl.store(a_ptrs, tl.zeros((BN, BM), dtype=tl.float32), mask=m_mask)
            tl.atomic_xchg(sem_ptr + tile, 0)

• 475–477The inputs: the finished partial tile acc [BN, BM], the pointers, this program's ids, and the index vectors.
• 478A mask for tokens beyond M.
• 479ATOM is a compile-time constant, so only one branch is compiled.
• 482–483One counter per output tile. The scratch here is racc, an FP32 buffer laid out like the output [M, N].
• 484Every program atomically adds its whole tile into the scratch. Each element is a single "relaxed" atomic float add, performed inside the L2 cache; relaxed means no ordering promises, which is the cheapest kind.
• 485tl.debug_barrier() is a __syncthreads(): every thread of this program has issued its adds.
• 486One atomic increment of the tile's counter, with acq_rel ordering. "Release" publishes this program's adds before the increment can be seen; "acquire" lets the last program see everyone else's published adds. It returns the old value.
• 487–493The program that sees SPLIT−1 is the last. It reads the finished sum from L2 (.cg bypasses L1) and stores it, as FP32 for the router.
• 494–495It re-zeroes the scratch tile and resets the counter, so the next call starts clean without a separate clearing kernel.

496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521

    else:
        if SPLIT == 1:
            res = acc
        else:
            tile = pid_m * grid_n + pid_n
            base = ws_ptr + (tile * SPLIT + pid_k) * (BN * BM)
            tl.store(base + tl.arange(0, BN)[:, None] * BM + tl.arange(0, BM)[None, :], acc)
            tl.debug_barrier()
            cnt = tl.atomic_add(sem_ptr + tile, 1, sem="acq_rel")
            res = acc
        if SPLIT == 1:
            do_out = True
        else:
            do_out = cnt == SPLIT - 1
        if do_out:
            if SPLIT > 1:
                res = tl.zeros((BN, BM), dtype=tl.float32)
                for s in tl.static_range(SPLIT):
                    b2 = ws_ptr + (tile * SPLIT + s) * (BN * BM)
                    res += tl.load(b2 + tl.arange(0, BN)[:, None] * BM + tl.arange(0, BM)[None, :], cache_modifier=".cg")
                tl.atomic_xchg(sem_ptr + tile, 0)
            o_ptrs = out_ptr + offs_m[None, :] * N + offs_n[:, None]
            if OUT_F32:
                tl.store(o_ptrs, res, mask=m_mask)
            else:
                tl.store(o_ptrs, res.to(tl.bfloat16), mask=m_mask)

• 497–498Without a split, the tile is already the answer.
• 499–505Otherwise:
  • store the partial tile into its own slot of ws;
  • wait for the program's threads (barrier);
  • increment the tile's counter.
• 506–509Only the last program goes on.
• 510–516It adds all SPLIT partial tiles in a fixed order (tl.static_range unrolls the loop at compile time), then resets the counter.
• 517–521It stores the result: BF16 for the pooler.

Why two variants?

• In the serial variant, the last program reads all SPLIT tiles one after another. With SPLIT = 16 that is 16 × 32 KB through a single SM, about 3–4 µs of tail. Round 4's switch to atomics took the router at m = 128 from 12.4 to 9.3–10.2 µs (9.86 in the check run).
• With atomics, the additions happen in L2 as each program finishes, and the last program reads just one tile.
• The pooler splits only 2 ways, so the serial version's tail is short, and it avoids atomic traffic on a 64 × 128 tile.

Lab 5 simulates both protocols, including what goes wrong if the reset lines are removed.

9.5 The three Triton configurations

The router's floor is about 2.2–2.7 µs, and we take about 9.9–10.5 µs. A kernel this small is dominated by fixed costs: start-up, a pipeline that only runs 4–8 steps, and the split-K handshake. Shrinking those is one of the open problems (chapter 10).

cd ~/projects/okc/muse_bf16_proj/learn/labs
python3 lab4_router.py

Exercises.

1. In lab 4, add the third term lo2 to the output and see how far the error drops.
2. In lab 5, set REZERO = False and watch the atomic accumulator keep the previous call's sum.
3. Chapter 10's experiment C turns ATOM off for the router at m = 128 on the real H200. Predict the time first: round 4 measured about 12.4 µs for the serial version.

Check yourself