Triton MoE improvements and chunked GRPO improvements

[michaelroyzen]

Summary

Performance-tuning pass over the fused MoE grouped-GEMM op (liger_kernel.ops.fused_moe), plus a tensor-core fix in the chunked PPO loss backward.

Headline numbers (bf16, Qwen3-30B-A3B expert shape E=128 / H=2048 / I=768 / K=8, vs. the current main implementation):

• H100: up to 1.95x forward / 1.91x full fwd+bwd at T=32768 in the tuning environment (triton 3.7.1); an independent re-run on torch 2.10.0 / triton 3.6.0 measured 1.25–1.33x forward and 1.47–1.48x full
• B300 (sm103): up to 2.77x backward, 2.35x full fwd+bwd (T=32768) — and a fix for silently wrong dx gradients caused by a Triton 3.7.1 tcgen05 miscompile
• Peak sustained throughput on H100: 523 TFLOPS inference / 501 TFLOPS forward / 413 TFLOPS full fwd+bwd — roughly half of the H100's dense bf16 peak for a gathered/scattered grouped GEMM with SwiGLU fused in
• Default peak memory unchanged; opt-in LIGER_FUSED_MOE_MEMORY_EFFICIENT=1 saves −288 MB at T=8192 / −1151 MB at T=32768 for a ~10-15% slower backward on H100

Full tables in the Performance section below.

Fused MoE (ops/fused_moe.py, ops/fused_moe_kernels.py)

• Sync-free, CUDA-graph-capturable hot path. Routing metadata previously needed one .item() sync to size the m-tile buffers. Tile metadata is now allocated at a host-computable upper bound (TK//BLOCK_M + min(E, TK)), the actual tile count stays on device in expert_tile_offset[E], and GEMM CTAs past it exit early. No host↔device sync anywhere in forward or backward.
• TMA expert-weight loads on Hopper+. Weight tiles load through tl.make_tensor_descriptor over a flattened 2D view when rows are 16-byte aligned (async bulk copies that bypass L1 and free LSU issue slots), with the portable pointer-load path kept for other GPUs and unaligned shapes. The TMA scratch allocator is (re)registered in both forward() and backward() because Triton stores it in a ContextVar that does not propagate to the autograd engine's backward thread.
• Adaptive token-tile size. BLOCK_M was a fixed 64; it is now picked per call from expected tokens-per-expert (clamped to [16, 128]), so low-occupancy shapes (e.g. T=128, K=8, E=128 → 8 tokens/expert) stop running mostly-padded MMAs.
• Smarter autotuning.
  • Keys now include BLOCK_M, USE_TMA, and — for the dW kernels — a tokens-per-expert bucket (TPE_BUCKET): the best dW tile is regime-dependent (output-write-bound at small TK/E, K-loop-bound at large TK/E), which H_dim/I_dim alone cannot see.
  • An early_config_prune hook drops tiles wider than the (padded) problem, cutting first-run compile+tune cost dramatically for small shapes such as unit tests.
  • The config space gains wide tiles for H100's 228 KB smem (BLOCK_N=256, BLOCK_K=128), plus Blackwell-datacenter-only extras — see the dedicated section below.
• L2-friendly launch geometry. All GEMM grids are now 1D with a grouped pid swizzle (matmul-tutorial style) so co-resident CTAs share x rows and expert-weight tiles in L2. The dW kernels are expert-major, and empty experts store zeros in-kernel instead of relying on a separate zeros_like memset (grad buffers are torch.empty now).
• Numerics.
  • dS (router-score grads) accumulates through fp32 atomics instead of bf16 — ceil(I/BLOCK_N) partials per element previously rounded on every add.
  • The dX backward is a single GEMM over K=2I (d_pre_act columns and W1 rows share the same [gate; up] ordering) instead of two tl.dot calls chained through one accumulator. Besides being simpler, this avoids a Triton 3.7.1 miscompile on sm103 under the tcgen05 MMA path that made the two-chained-dot pattern produce silently wrong dx on B300-class parts; the single-GEMM form is used unconditionally (neutral-to-positive on other archs, verified clean on H100).
  • The token gather-sum kernel grid became 2D (T, ceil(H/BLOCK_H)) so small-T launches still fill the GPU.
• New modes.
  • Inference (no input requires grad) skips allocating/storing the (TK, 2I) pre-activation buffer.
  • Opt-in LIGER_FUSED_MOE_MEMORY_EFFICIENT=1 (read at import): backward writes SwiGLU gradients in place over the saved pre-activations and recomputes s_k * silu(gate) * up inside dW2 instead of materializing weighted_act — ≈ −1.2 GB peak at T=32768/I=768/bf16 for a ~10-15% slower backward. A second backward over the same graph then raises a version-counter error by design (guarded via torch.autograd.graph.increment_version).

B300 / Blackwell-datacenter-specific optimizations (sm100/sm103)

All Blackwell-only behavior is confined to autotune config selection; kernel code is shared across architectures, and non-Blackwell search spaces stay exactly as tuned on H100.

• Extended autotune config space, gated at launch. sm100/sm103 extras exploit three hardware deltas vs Hopper: TMEM holds the fp32 accumulator, so wide-N tiles (BLOCK_N=256) with deep pipelines are cheap where Hopper would spill registers (a gathered-GEMM probe hit 1103 TFLOPS at BN=256/BK=64/ns=3 vs 808 at the Hopper-favored BN=128/BK=64/ns=4); tcgen05 MMAs are issued by a single warp, making num_warps=4 competitive with 8; and the 132 MB L2 (H100: 50 MB) rewards GROUP_M=8-16 CTA grouping and deep-K (BK=128) dW tiles. The configs are tagged via _blackwell_config and pruned at first kernel launch when infer_device_arch() isn't Blackwell datacenter — deliberately not at import, so importing liger_kernel never initializes CUDA (fork-based workers). Consumer Blackwell (sm120) is excluded: no TMEM, smaller smem, the wide tiles would only waste tuning time.
• Tokens-per-expert autotune key unlocks the Blackwell dW configs. The TPE_BUCKET key is architecture-neutral, but it matters most on B300: dW1 wants BK=128-deep tiles at large TK/E and BK=32 at small TK/E, and the old shared cache entry let whichever T tuned first pin the config for every later T. With the key, realistic mixed-shape runs (small T tunes first, as any warmup or unit test would) measure 1.22x / 1.26x backward at T=8192 / 32768 on B300 over the extended config space alone.
• sm103 correctness fix. The two-chained-tl.dot miscompile (see the Numerics bullet in the Fused MoE section above) made main's dx gradients silently wrong on B300-class parts under Triton 3.7.1; the single-GEMM dX rewrite fixes it (applied unconditionally on all archs).
• Token gather tile sweep extended to BLOCK_H=1024. HBM3e sustains ~6.4 TB/s on B300 vs ~3.3 TB/s on H100 and rewards longer contiguous vectors; harmless elsewhere (autotuner simply won't pick it).
• The PPO-loss GEMM dtype fix below was also motivated by B300 measurements (fp32 SIMT matmul ~57 TFLOPS vs ~1232 TFLOPS bf16 tensor-core on that part), though it helps every architecture.

Net effect on B300: 2.35x full-step at T=32768 vs main (table below), on top of the gradient correctness fix.

Chunked PPO loss (chunked_loss/fused_linear_ppo.py)

The selective-logprob backward's two grad GEMMs ran in fp32, dispatching to SIMT CUDA-core kernels (~20x below the bf16 tensor-core rate) and dominating the backward's runtime. GEMM operands now round to the input dtype while cuBLAS still accumulates each chunk in fp32 internally and the cross-chunk grad_hidden/grad_weight buffers stay fp32 — the same rounding a non-chunked autograd backward through a bf16 lm_head applies everywhere. For fp32 inputs the casts are no-ops. Chunk sizes double to 8192×4096 (identical peak memory, ~21% faster measured).

Tests & benchmarks

• New unit tests: inference no-grad path, retain-graph double backward (guards against in-place clobbering of saved tensors), and the memory-efficient mode (run in a subprocess since the flag is read at import).
• benchmark_fused_moe.py warmup updated for the new autotune keys (one warmup per distinct BLOCK_M bucket of the sweep).

Performance

All numbers bf16 at the Qwen3-30B-A3B expert shape (E=128, H=2048, I=768, K=8), median of triton.testing.do_bench, after autotune warmup. "before" is the current main implementation, "after" is this PR.

H100 80GB SXM — speedup vs main

Measured in the tuning environment (triton 3.7.1):

T	mode	before (ms)	after (ms)	speedup
128	forward	0.551	0.444	1.24x
128	full	1.818	1.275	1.43x
1024	forward	0.646	0.512	1.26x
1024	full	2.354	1.663	1.42x
8192	forward	2.659	1.456	1.83x
8192	backward	7.210	3.998	1.80x
8192	full	9.920	5.417	1.83x
8192	inference	2.651	1.398	1.90x
32768	forward	9.901	5.066	1.95x
32768	backward	26.076	14.477	1.80x
32768	full	36.249	19.005	1.91x
32768	inference	9.969	5.174	1.93x

Independently re-verified on torch 2.10.0+cu128 / triton 3.6.0 (same H100 class, same shapes): 1.33x forward / 1.48x full at T=8192, 1.25x forward / 1.47x full at T=32768. The gap vs the table above tracks the Triton version (3.6.0 lacks some of 3.7.1's TMA/pipelining improvements the tuned configs exploit), so treat 1.47x+ on full steps as the conservative floor.

Peak memory (full fwd+bwd, default mode) is unchanged vs main: 1953 MB at T=8192, 4356→4357 MB at T=32768 (re-verified). Opt-in LIGER_FUSED_MOE_MEMORY_EFFICIENT=1 saves −288 MB at T=8192 and −1151 MB at T=32768 (1665 / 3205 MB peaks) at ~10-15% slower backward on H100 (~32% on B300, where faster MMAs make the dW2 recompute relatively more expensive).

H100 sustained TFLOPS (this PR)

FLOP model per token-expert pair: forward 6·H·I, backward 12·H·I, full 18·H·I.

T	inference	forward	backward	fwd+bwd
1024	149.0	146.6	131.7	136.6
4096	372.8	354.3	283.8	308.9
8192	452.8	418.3	338.1	364.6
16384	474.3	457.8	375.2	406.0
32768	506.6	476.8	376.6	413.2
65536	522.7	501.5	391.5	374.6

B300 (sm103) — speedup vs main

T	mode	before (ms)	after (ms)	speedup
128	full	1.314	0.774	1.70x
1024	full	1.719	1.023	1.68x
8192	forward	1.323	0.924	1.43x
8192	backward	5.060	2.150	2.35x
8192	full	6.351	3.066	2.07x
32768	forward	4.994	3.152	1.58x
32768	backward	19.829	7.147	2.77x
32768	full	24.345	10.344	2.35x
32768	inference	5.100	3.201	1.59x

(Note: main's dx gradients on B300 are silently wrong due to the Triton 3.7.1 sm103 miscompile described above, so on that hardware this PR is a correctness fix as well as a speedup.)

Robustness spot checks (H100)

• Skewed routing (Zipf skew=1.5, Qwen3-30B shape, T=8192): 1.37x forward, 1.44x full
• Mixtral-8x7B shape (E=8, H=4096, I=14336, K=2, T=4096): 1.36x forward, 1.19x full

Testing Done

Environment: torch 2.10.0+cu128, triton 3.6.0, Python 3.12, transformers 5.13.0.

• Unit tests: test/transformers/test_fused_moe.py — 22/22 passed (routing invariants, fwd/bwd vs reference in fp32 + bf16, single-token / all-to-one-expert / K=E edge cases, and the three new tests).

• PPO loss (exercises the fused_linear_ppo.py change): test/chunked_loss/test_grpo_loss.py — 976 passed, 614 skipped.

• Convergence: bf16 MoE mini-models (qwen3_moe, qwen3_vl_moe, glm4v_moe, qwen3_5_moe, hunyuan_v1_moe) — 5/5 passed. fp32 MoE mini-models — 3 passed, 1 skipped; mini_hunyuan_v1_moe is flaky by one logprob element in full-suite runs, passes in isolation, and reproduces identically on unmodified main (pre-existing, unrelated to this PR).

• Benchmarks: see the Performance section above (speedup vs main, TFLOPS sweep, memory, skewed-routing and Mixtral-shape robustness checks).

• Hardware Type: H100-80G-SXM

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

@Tcc0403 @Mecoli1219

[michaelroyzen]

Reported triton-lang/triton#10821 separately, but the patch in this PR serves as a workaround for all B300 code using Triton in Liger.

[michaelroyzen]

Chunked Triton GRPO loss: head-to-head vs existing GRPO loss implementations

Benchmarks a new chunked Triton GRPO loss (liger_kernel/transformers/chunked_grpo_loss.py + liger_kernel/ops/chunked_grpo_loss.py) against the two existing Liger GRPO loss implementations. The new kernel fuses the lm_head projection into the loss with a flash-attention-style online logsumexp, so the (N, V) logits tensor is never materialized — the memory profile of the chunked fused-linear path at roughly Triton speed.

TL;DR

	speed	peak memory (loss stage)	64k+ completions
LigerFusedLinearGRPOLoss (torch chunked)	1x (baseline)	~4–5 GiB, flat	works
triton_grpo_loss (unchunked)	3.9–6.9x	O(logits): 3→123 GiB	crashes at L ≥ 65,536 (CUDA grid-dim limit)
chunked_triton_grpo_loss (new)	1.8–2.5x	~5–7 GiB, flat	works

The new kernel is 1.8–2.5x faster than the torch chunked path at the same flat memory profile (2.3x+ from 4k context up), and within ~1.6x of the unchunked Triton kernel at 16k+ contexts while using up to ~18x less memory. Short contexts show a larger relative gap to the unchunked kernel because fixed per-launch overheads dominate there.

Setup

GPU	NVIDIA B300 SXM6 (sm_103), 275 GiB
Model dims	hidden 2048, vocab 248,320 (Qwen3.5-MoE-35B-A3B), bf16
Loss config	loss_type="dapo", importance_sampling_level="sequence", beta=0, temperature 1.0, eps 0.2/0.2, num_items_in_batch passed
Batch	4 sequences per micro-batch, variable mask lengths in [L/2, L]
Software	torch 2.10 cu130, Triton 3.6, TORCH_DYNAMO_DISABLE=1

Methodology

All three implementations are measured from the hidden-state boundary with the lm_head GEMM included, since the chunked variants fuse it: the unchunked Triton path is timed as logits = hidden @ W.T followed by triton_grpo_loss(logits, ...) (with inplace=True, its production setting), while both chunked paths consume (hidden, W) directly. Timing is forward + backward via CUDA events (2 warmup, 5 measured iterations, mean ± std); memory is peak allocation above the resident inputs. Gradients flow to both hidden and weight.

Results

Per micro-batch of 4 sequences, forward + backward:

completion len	logits size	torch chunked	Triton unchunked	Triton chunked (new)
1,024	1.9 GiB	52.9 ms / 3.81 GiB	7.7 ms / 2.86 GiB	28.8 ms / 4.77 GiB
4,096	7.6 GiB	149.2 ms / 3.88 GiB	34.2 ms / 8.59 GiB	62.2 ms / 4.86 GiB
16,384	30.3 GiB	580.7 ms / 4.16 GiB	144.7 ms / 31.51 GiB	249.7 ms / 5.24 GiB
32,768	60.6 GiB	1156.6 ms / 4.54 GiB	298.3 ms / 62.07 GiB	486.6 ms / 5.74 GiB
65,535	121.2 GiB	2307.0 ms / 5.29 GiB	599.0 ms / 123.20 GiB	935.1 ms / 6.74 GiB
65,536	121.3 GiB	2306.5 ms / 5.29 GiB	launch failure	979.5 ms / 6.74 GiB

Loss values agree across all three implementations to full print precision on identical inputs (1.1300394535064697 at L=1024).

Notes:

• The unchunked Triton path's peak is dominated by the materialized logits and scales linearly with context; both chunked paths stay flat (the ~2 GiB fp32 grad_weight accumulator plus a reusable per-chunk buffer).
• Unchunked triton_grpo_loss hard-fails at L ≥ 65,536: its kernels launch with grid (B, L), mapping L onto gridDim.y, which CUDA caps at 65,535 (Triton Error [CUDA]: invalid argument). Both chunked paths are unaffected. (Fixable separately by transposing the launch grid.)
• In end-to-end training the unchunked path's logits are produced by the model forward either way, so its practical extra cost vs the chunked paths is holding logits + writing gradients into them during the loss stage; the chunked paths remove the (N, V) tensor from the model forward entirely by consuming hidden states.

Correctness

Covered by test/transformers/test_chunked_grpo_loss.py (38 tests, all passing), which validates at context lengths 128–16,384, vocab 248,320 and 50,257 (odd size exercising the vocab-tail masking path), and 10 loss configurations (grpo / bnpo / dr_grpo / dapo / cispo / sapo, token- and sequence-level IS, KL with bias correction, delta two-sided clipping, vllm_is_ratio, num_items_in_batch):

• Intermediates: per-token logp/lse vs an fp32 ground truth (atol 1e-2; only bf16-GEMM accumulation order separates them) and vs fused_selective_log_softmax.
• Per-token outputs (reduce=False): per_token_loss / per_token_kl / is_clipped vs unchunked triton_grpo_loss.
• End results (reduce=True): loss, metrics, grad_hidden, grad_weight vs a plain-torch reference, the unchunked Triton kernel, and LigerFusedLinearGRPOLoss (gradient cosine > 0.999, norm ratio within 2%).
• Bitwise determinism: repeated fwd+bwd produce bit-identical loss and gradients — fixed launch configs, no autotune, no atomics (grad_weight accumulates through deterministic cuBLAS GEMMs into an fp32 buffer).

Kernel design notes (sm_103 / B300)

• Every K-loop issues exactly one tl.dot per accumulator per iteration, and each vocab tile gets a fresh accumulator — the two-dots-through-one-accumulator pattern that miscompiles on tcgen05 tiles (triton-lang/triton#10821) never occurs, while still using BLOCK_M ≥ 64 (tcgen05 MMA) for tensor-core throughput.
• Backward recomputes logits tiles in a fused kernel that emits grad_logits for one 4,096-token chunk into a reusable buffer in the input dtype; the two large grad GEMMs run in cuBLAS (deterministic) with the cross-chunk grad_weight accumulator kept in fp32.
• Tile config (BM=128, BN=256, BK=64, 8 warps, 3 stages) chosen by manual sweep at V=248,320, H=2048, 65k tokens: 249 ms fwd+bwd vs 377 ms for the initial (64, 128, 128) guess; num_warps=16 regressed ~2x, larger tiles exhaust shared memory.

Reproduce

PYTHONPATH=src python benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py
PYTHONPATH=src:. python -m pytest test/transformers/test_chunked_grpo_loss.py -q --override-ini addopts=