Triton MoE improvements and chunked GRPO improvements

benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py

@@ -0,0 +1,191 @@
+"""Head-to-head GRPO loss benchmark: torch chunked vs triton unchunked vs triton chunked.
+
+Compares the three Liger GRPO loss implementations from the hidden-state
+boundary (lm_head projection included, since the chunked variants fuse it):
+
+  chunked_torch:  LigerFusedLinearGRPOLoss (fused linear, torch/cuBLAS chunking)
+  triton:         triton_grpo_loss on materialized (B, L+1, V) logits
+  chunked_triton: chunked_triton_grpo_loss (fused linear, Triton kernels)
+
+Measures forward+backward wall time (CUDA events) and peak memory above the
+resident inputs, per micro-batch. Config mirrors GRPO training on Qwen3.5-MoE:
+dapo loss, sequence-level importance sampling, beta=0, temperature 1.0,
+eps 0.2/0.2, hidden 2048, vocab 248320, bf16.
+
+Run from the repo root:
+  PYTHONPATH=src python benchmark/scripts/benchmark_chunked_grpo_loss_head_to_head.py
+"""
+
+import argparse
+
+import torch
+
+from liger_kernel.chunked_loss import LigerFusedLinearGRPOLoss
+from liger_kernel.transformers.chunked_grpo_loss import chunked_triton_grpo_loss
+from liger_kernel.transformers.grpo_loss import triton_grpo_loss
+
+HIDDEN_SIZE = 2048
+VOCAB_SIZE = 248320
+LOSS_KWARGS = dict(
+    temperature=1.0,
+    beta=0.0,
+    eps_low=0.2,
+    eps_high=0.2,
+    loss_type="dapo",
+    importance_sampling_level="sequence",
+)
+
+
+def make_inputs(batch, seq_len, device, seed=0):
+    gen = torch.Generator(device=device).manual_seed(seed)
+    hidden = torch.randn(batch, seq_len + 1, HIDDEN_SIZE, device=device, generator=gen).to(torch.bfloat16).mul_(0.02)
+    weight = torch.randn(VOCAB_SIZE, HIDDEN_SIZE, device=device, generator=gen).to(torch.bfloat16).mul_(0.02)
+    completion_ids = torch.randint(0, VOCAB_SIZE, (batch, seq_len), device=device, generator=gen)
+    lengths = torch.randint(seq_len // 2, seq_len + 1, (batch,), device=device, generator=gen)
+    mask = (torch.arange(seq_len, device=device).unsqueeze(0) < lengths.unsqueeze(1)).float()
+    advantages = torch.randn(batch, device=device, dtype=torch.float32, generator=gen)
+    return {
+        "hidden": hidden,
+        "weight": weight,
+        "completion_ids": completion_ids,
+        "mask": mask,
+        "advantages": advantages,
+        "num_items_in_batch": mask.sum(),
+    }
+
+
+def run_variant(variant, inputs, chunked_torch_module):
+    hidden, weight = inputs["hidden"], inputs["weight"]
+    common = dict(num_items_in_batch=inputs["num_items_in_batch"])
+    if variant == "chunked_torch":
+        loss, _ = chunked_torch_module(
+            hidden[:, :-1, :],
+            weight,
+            inputs["completion_ids"],
+            inputs["mask"],
+            inputs["advantages"],
+            **common,
+        )
+    elif variant == "triton":
+        logits = hidden @ weight.t()
+        loss, _ = triton_grpo_loss(
+            logits,
+            None,
+            None,
+            inputs["completion_ids"],
+            inputs["advantages"],
+            inputs["mask"],
+            inplace=True,
+            reduce=True,
+            **LOSS_KWARGS,
+            **common,
+        )
+    elif variant == "chunked_triton":
+        loss, _ = chunked_triton_grpo_loss(
+            hidden[:, :-1, :].contiguous(),
+            weight,
+            None,
+            None,
+            inputs["completion_ids"],
+            inputs["advantages"],
+            inputs["mask"],
+            reduce=True,
+            **LOSS_KWARGS,
+            **common,
+        )
+    else:
+        raise ValueError(variant)
+    return loss
+
+
+def bench(variant, inputs, chunked_torch_module, warmup, iters):
+    hidden, weight = inputs["hidden"], inputs["weight"]
+
+    def once():
+        hidden.grad = None
+        weight.grad = None
+        hidden.requires_grad_(True)
+        weight.requires_grad_(True)
+        torch.cuda.synchronize()
+        torch.cuda.reset_peak_memory_stats()
+        baseline = torch.cuda.memory_allocated()
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        loss = run_variant(variant, inputs, chunked_torch_module)
+        loss.backward()
+        end.record()
+        torch.cuda.synchronize()
+        peak = (torch.cuda.max_memory_allocated() - baseline) / 1024**3
+        hidden.requires_grad_(False)
+        weight.requires_grad_(False)
+        return start.elapsed_time(end), peak
+
+    try:
+        for _ in range(warmup):
+            once()
+        times, peaks = zip(*[once() for _ in range(iters)])
+    except torch.cuda.OutOfMemoryError:
+        torch.cuda.empty_cache()
+        return "OOM"
+    except RuntimeError as err:
+        return f"FAIL: {err}"
+    t = torch.tensor(times)
+    return t.mean().item(), t.std().item(), max(peaks)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--seq-lens", type=int, nargs="+", default=[1024, 4096, 16384, 32768, 65535, 65536])
+    parser.add_argument("--warmup", type=int, default=2)
+    parser.add_argument("--iters", type=int, default=5)
+    args = parser.parse_args()
+
+    device = torch.device("cuda:0")
+    torch.cuda.set_device(device)
+    chunked_torch_module = LigerFusedLinearGRPOLoss(
+        beta=0.0,
+        compiled=False,
+        use_ref_model=False,
+        epsilon_low=LOSS_KWARGS["eps_low"],
+        epsilon_high=LOSS_KWARGS["eps_high"],
+        loss_type=LOSS_KWARGS["loss_type"],
+        importance_sampling_level=LOSS_KWARGS["importance_sampling_level"],
+        temperature=LOSS_KWARGS["temperature"],
+    )
+    print(f"Device: {torch.cuda.get_device_name(device)}")
+    print(f"Hidden {HIDDEN_SIZE}, vocab {VOCAB_SIZE}, batch {args.batch_size}, bf16")
+    print(f"Config: {LOSS_KWARGS}\n")
+
+    # quick loss parity sanity check
+    inputs = make_inputs(args.batch_size, 1024, device)
+    losses = {
+        v: run_variant(v, inputs, chunked_torch_module).item() for v in ["chunked_torch", "triton", "chunked_triton"]
+    }
+    print(f"Loss parity @1024: {losses}\n")
+    del inputs
+
+    variants = ["chunked_torch", "triton", "chunked_triton"]
+    header = f"{'seq_len':>8} {'logits_GiB':>11}" + "".join(f" {v + '_ms':>22} {v + '_peak_GiB':>18}" for v in variants)
+    print(header)
+    print("-" * len(header))
+    for seq_len in args.seq_lens:
+        inputs = make_inputs(args.batch_size, seq_len, device)
+        logits_gib = args.batch_size * (seq_len + 1) * VOCAB_SIZE * 2 / 1024**3
+        row = f"{seq_len:>8} {logits_gib:>11.1f}"
+        for variant in variants:
+            result = bench(variant, inputs, chunked_torch_module, args.warmup, args.iters)
+            if isinstance(result, str):
+                label = result if len(result) < 20 else "LAUNCH FAIL"
+                row += f" {label:>22} {'-':>18}"
+            else:
+                mean_ms, std_ms, peak = result
+                row += f" {mean_ms:>14.1f} ±{std_ms:>5.1f} {peak:>18.2f}"
+        print(row)
+        del inputs
+        torch.cuda.empty_cache()
+
+
+if __name__ == "__main__":
+    main()

benchmark/scripts/benchmark_fused_moe.py

@@ -25,6 +25,7 @@
 from utils import run_speed_benchmark
 
 from liger_kernel.ops import LigerFusedMoEFunction
+from liger_kernel.ops.fused_moe import _pick_block_m_token
 from liger_kernel.utils import get_total_gpu_memory
 from liger_kernel.utils import infer_device
 
@@ -133,12 +134,11 @@ def bench_memory_fused_moe(input: SingleBenchmarkRunInput) -> SingleBenchmarkRun
 
 
 def _warmup_liger(T, E, H, intermediate_dim, K, dtype, sweep_dim):
-    """Run one full fwd+bwd to exhaust Triton autotune for (H, intermediate_dim).
+    """Run one full fwd+bwd to exhaust Triton autotune for one autotune key.
 
-    Triton autotune key is (H_dim, I_dim), so a single call is sufficient to
-    cache the best config for all subsequent calls with the same H and intermediate_dim.
-    For the num_experts sweep we also call this once per E value to warm up
-    CUDA caches for each expert count before do_bench starts timing.
+    The GEMM autotune key is (H_dim, I_dim, BLOCK_M[, USE_TMA]) where BLOCK_M is
+    picked adaptively from tokens-per-expert, so the caller warms one
+    representative point per distinct BLOCK_M bucket of the sweep.
     """
     warmup_input = SingleBenchmarkRunInput(
         x=T if sweep_dim == "T" else E,
@@ -221,34 +221,6 @@ def _probe():
     peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
     kernel_bpt = peak_bytes // probe_T
 
-    # Pre-warm Liger's Triton autotune before benchmarks start.
-    #
-    # Autotune key is (H_dim, I_dim) — one warmup per (H, intermediate_dim) pair is sufficient
-    # to cache the best config for the entire sweep.
-    #
-    # For num_tokens sweep: one pass with the model's base T is enough.
-    # For num_experts sweep: one pass per E value in EXPERT_SWEEP_VALUES to also
-    #   warm up CUDA caches for each expert count, since weight tensor sizes differ.
-    print(f"Pre-warming Liger autotune (H={H}, intermediate_dim={intermediate_dim})...")
-
-    if args.sweep_dim == "num_tokens":
-        _warmup_liger(probe_T, E, H, intermediate_dim, K, dtype, sweep_dim="T")
-    else:  # num_experts
-        for e_val in EXPERT_SWEEP_VALUES:
-            print(f"  warmup E={e_val}...")
-            _warmup_liger(probe_T, e_val, H, intermediate_dim, K, dtype, sweep_dim="E")
-
-    if device == "cuda":
-        torch.cuda.synchronize()
-    elif device == "npu":
-        torch.npu.synchronize()
-    elif device == "xpu":
-        torch.xpu.synchronize()
-    else:
-        torch.cpu.synchronize()
-
-    print("Autotune warmup complete.\n")
-
     if args.sweep_dim == "num_tokens":
         # Derive a memory-safe upper bound for T from the probe measurement.
         # Target 40% GPU memory utilisation to leave headroom for framework overhead.
@@ -284,6 +256,37 @@ def _probe():
         ]
         x_name, x_label = "E", "num_experts"
 
+    # Pre-warm Liger's Triton autotune before benchmarks start.
+    #
+    # The GEMM autotune key includes the adaptive BLOCK_M (a function of tokens per
+    # expert), so warm one representative x-value per distinct BLOCK_M bucket.
+    # For the num_experts sweep this also warms CUDA caches per expert count.
+    print(f"Pre-warming Liger autotune (H={H}, intermediate_dim={intermediate_dim})...")
+
+    if args.sweep_dim == "num_tokens":
+        warmed = set()
+        for t_val in x_values:
+            bucket = _pick_block_m_token(t_val * K, E)
+            if bucket not in warmed:
+                print(f"  warmup T={t_val} (BLOCK_M={bucket})...")
+                _warmup_liger(t_val, E, H, intermediate_dim, K, dtype, sweep_dim="T")
+                warmed.add(bucket)
+    else:  # num_experts
+        for e_val in EXPERT_SWEEP_VALUES:
+            print(f"  warmup E={e_val}...")
+            _warmup_liger(probe_T, e_val, H, intermediate_dim, K, dtype, sweep_dim="E")
+
+    if device == "cuda":
+        torch.cuda.synchronize()
+    elif device == "npu":
+        torch.npu.synchronize()
+    elif device == "xpu":
+        torch.xpu.synchronize()
+    else:
+        torch.cpu.synchronize()
+
+    print("Autotune warmup complete.\n")
+
     common_configs = {
         "kernel_name": "fused_moe",
         "x_name": x_name,

src/liger_kernel/chunked_loss/fused_linear_ppo.py

@@ -4,8 +4,12 @@
 import torch
 import torch._dynamo.config
 
-_SELECTIVE_LOGPROB_VOCAB_CHUNK_SIZE = 4096
-_SELECTIVE_LOGPROB_SEQ_CHUNK_SIZE = 2048
+# Chunk temporaries are seq_chunk x vocab_chunk fp32 (~128 MB at these sizes) —
+# negligible next to the fp32 grad_weight buffer, and larger chunks amortize the
+# per-chunk elementwise/launch overhead: 4096x8192 measured 21% faster than the
+# previous 2048x4096 at identical peak memory (B300, V=248320, 65K tokens).
+_SELECTIVE_LOGPROB_VOCAB_CHUNK_SIZE = 8192
+_SELECTIVE_LOGPROB_SEQ_CHUNK_SIZE = 4096
 
 
 def _maybe_mark_dynamic_dim1(tensor):
@@ -72,6 +76,16 @@ def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logpr
     """Dual-chunked (sequence × vocab) backward for selective logprob.
 
     Recomputes logits per chunk for memory efficiency.
+
+    The two grad GEMMs run with operands in ``hidden.dtype`` (bf16/fp16 in
+    practice) rather than fp32: fp32×fp32 matmuls dispatch to SIMT CUDA-core
+    kernels (~57 TFLOPS on B300, ~20x below the bf16 tensor-core rate) and were
+    ~78% of this backward's runtime. Precision is preserved where it matters:
+    cuBLAS accumulates each chunk GEMM in fp32 internally, and the cross-chunk
+    accumulation buffers (grad_hidden/grad_weight) stay fp32. Only the per-chunk
+    GEMM inputs/outputs round to the compute dtype — the same rounding a
+    non-chunked autograd backward through a bf16 lm_head applies everywhere.
+    For fp32 inputs the casts are no-ops and behavior is unchanged.
     """
     inv_t = 1.0 / temperature
     n_rows, _ = hidden.shape
@@ -109,8 +123,11 @@ def _selective_logprob_backward(hidden, weight, targets, bias, log_z, grad_logpr
             grad_logits[row_idx, local_idx] += grad_chunk * in_chunk
             grad_logits.mul_(inv_t)
 
-            grad_hidden[seq_start:seq_end].add_(grad_logits @ weight_chunk.float())
-            grad_weight[vocab_start:vocab_end].add_(grad_logits.t() @ hidden_chunk.float())
+            # Tensor-core GEMMs in the input dtype; fp32 accumulation across
+            # chunks via the fp32 grad buffers (see docstring).
+            grad_logits_lp = grad_logits.to(hidden.dtype)
+            grad_hidden[seq_start:seq_end].add_(grad_logits_lp @ weight_chunk.to(hidden.dtype))
+            grad_weight[vocab_start:vocab_end].add_(grad_logits_lp.t() @ hidden_chunk)
             if has_bias:
                 grad_bias[vocab_start:vocab_end].add_(grad_logits.sum(dim=0))