Add SM100 CuTe DSL GRPO fast path

[Celaena24]

Summary

• adds an SM100 CuTe DSL fused-linear selective-logprob autograd path for chunked GRPO;
• rebases the change onto the final merged implementation from linkedin/Liger-Kernel#1339;
• reuses the merged persistent SM100 GEMM and in-place CuTe DSL cross-entropy primitive so full logits are never materialized;
• uses a bounded 1024-row logits workspace, direct log-probability output, lazy gradient ownership, and FP32 multi-chunk weight-gradient accumulation;
• preserves the existing chunked implementation for non-SM100 GPUs, FP32, and temperatures other than 1.0.

Why it is faster and uses less memory

• The existing path splits a 128,256-token vocabulary into about 32 smaller blocks. Each block needs another matrix multiplication and probability calculation in both forward and backward. The new path uses one large, optimized SM100 matrix multiplication for each group of up to 1,024 input tokens, followed by one in-place cross-entropy operation. This means fewer GPU launches, less Python-loop overhead, and better use of B200 tensor cores.
• Matrix multiplication uses BF16 or FP16 inputs for speed while accumulating dot products internally in FP32. If several input chunks contribute to the same weight gradient, those chunk results are also accumulated in FP32 before conversion back to the model dtype.
• The logits workspace is reused and capped at 1,024 rows, so logits for the full batch or sequence are never kept in memory at once.
• The first computed weight and bias gradients become the output accumulators directly. The old path allocated additional zero-filled buffers first; avoiding that duplicate weight-gradient buffer saves about 1 GiB for the benchmarked 128,256 x 4,096 matrix.

Together, these changes explain the roughly 9-10x lower latency and about 2 GiB lower peak memory in the B200 benchmark.

B200 results

NVIDIA B200 (SM100), BF16, T=1024, H=4096, V=128256. Values include policy and reference log-probabilities, GRPO loss math, and backward.

Batch	Existing Liger	CuTe DSL	Speedup	Existing peak	CuTe DSL peak	Memory reduction
1	58.42 ms	6.69 ms	8.74x	7126 MiB	5114 MiB	1.39x
2	114.52 ms	12.09 ms	9.47x	7158 MiB	5389 MiB	1.33x
4	226.77 ms	23.07 ms	9.83x	7222 MiB	5453 MiB	1.32x
8	450.80 ms	43.79 ms	10.29x	7350 MiB	5581 MiB	1.32x

Compared with the first merged-compatible implementation, the final optimization lowers token-level full latency by 9.0-14.3% and sequence-level full latency by 11.2-14.7%, with peak memory staying within 2.4%.

PyTorch 2.13.0.1 hits a Dynamo shape-guard bug in the benchmark's compiled loss-only section, so both baseline and CuTe DSL measurements used TORCHDYNAMO_DISABLE=1 for an equal kernel-path comparison.

Native scope

• NVIDIA SM100, exact compute capability (10, 0)
• BF16 and FP16
• temperature 1.0
• optional linear bias
• aligned and internally padded hidden dimensions
• partial input/weight/bias gradient requirements

Unsupported inputs delegate to the existing _ChunkedSelectiveLogProbFunction.

Validation

• 1094 passed, 631 skipped — CuTe DSL, chunked GRPO, and low-level GRPO suites
• explicit BF16/FP16, bias, hidden-padding, partial-gradient, and FP32 multi-chunk accumulation coverage
• make checkstyle
• independent final review completed; its multi-chunk accumulation finding is fixed and regression-tested

Performance plots

Latency and peak memory

Improvement ratios

[kolehma8]

I think we should follow the same signature from the current fused linear ppo for this.

[Celaena24]

Do you mean this one @kolehma8 ? - #1334