[samsja]

Hopper (SM90) training kernels for DeepSeek-V4.1 Flash. They port what cuDNN frontend 1.30 ships for the DeepSeek-V4 family on Blackwell only. All three are Python-only (Triton / Gluon), arch = ["9.0"].

• dsa_indexer_topk: Lightning Indexer forward fused with top-k and the two-level candidate-block filter. It is a drop-in for prime-rl's dsv41_index_topk: same arguments, same picks up to ties. It never builds the (queries, entries) score matrix. A dense FP8 pass keeps per-group maxima, the best groups are rescored, and candidate consumer layers score only their candidates. Forward only, because the indexer is frozen.
• moe_experts: expert MLP of a rank's local experts: grouped GEMMs with the clamped SwiGLU fused into a warp-specialized epilogue, forward and backward. Only x, gate and up are saved. Weight gradients use torch._grouped_mm.
• moe_experts(..., fp8=True): DeepSeek-V3 blockwise FP8 (1x128 activation/gradient scales, 128x128 weight blocks, FP8 forward, dgrad and wgrad). It runs DeepGEMM's grouped FP8 GEMMs with the quantization fused into the surrounding passes:
  • one pass over x (and dout) writes the row-quantized copy and the transposed column-quantized copy that wgrad reads;
  • the SwiGLU forward writes FP8 h, and the SwiGLU backward writes FP8 dgate|dup;
  • there is no unpack pass.
• dsa_sparse_attn_bwd: SM90 backward for V4.1's sparse attention (64 heads × 512, shared latent K = V, attention sink), a drop-in for cuDNN's DSA backward.
  • A Triton preprocess computes delta, the sink-aware LSE and dsink, and compacts each query's non-empty slots.
  • A CuTe DSL wgmma main kernel loops over queries, with dKV accumulated by fp32 vector atomics.
  • It is 1.20–1.23x faster than cuDNN at V4.1 shapes, and its error against an fp32 reference matches cuDNN's to 3 digits.
• mhc_projection: mHC projection (RMS + gate logits + collapse of the streams by the previous sublayer's pre) in one read of the streams per direction, plus fused gate activations (sigmoid, Sinkhorn). dW is deterministic.

Each has a torch.library custom-op surface with fake impls and works under torch.compile(fullgraph=True). prime-rl picks each kernel up behind prime_kernels.is_available and keeps its own path otherwise.

Results (H200)

kernel	today's prime-rl path	this PR
indexer, 16k tokens, ratio 1 / ratio 2	6.06 / 3.54 ms	4.19 / 3.04 ms
indexer, CP 131k (rank 7/8): plain / candidate source / consumer	48.3 / 90.7 / 35.6 ms	24.1 / 27.6 / 9.05 ms
MoE experts fwd+bwd, EP=96 load (98k rows)	34.3 ms	32.0 ms
mHC fwd+bwd per sublayer, 16k tokens (GPU time)	1.52 ms	0.75–0.80 ms
MoE experts FP8 fwd+bwd, EP=96 load (same run: bf16 kernel 31.7 ms)	31.8 ms (DeepGEMM + separate quantization)	22.5 ms

• Indexer picks: match prime-rl's on 100% of tokens in the benchmark. Peak extra memory per call drops from 1.3–4.6 GB to 0.5–1.0 GB.
• End to end: on a 6-layer CP=8 131k proxy, the indexer alone takes the step from 9.5 s to 8.7 s. On 12-node 16k training the three kernels together gain about 1%, because that run is bound by communication.

FP8 weight cache, split backward, fused mHC stream gradient

• moe_experts_fp8_quantize_weights / moe_experts_fp8_forward_quantized: quantize expert weights once per optimizer step and reuse them across micro-batches (+8.7% on a PP2×4 V4.1 proxy, less memory).
• moe_experts_fp8_backward_accumulate: accumulate FP8 weight gradients straight into a caller's fp32 buffer such as FSDP's (+9.1% on the same proxy).
• Split backward ops (data gradient first, weight gradients after), so a caller can send the data gradient before the weight-gradient GEMMs run. prime-rl uses this to hide DeepEP's backward combine (−2.5 ms per layer at EP48, bitwise-identical gradients).
• mhc_projection_streams / hyper_connection_streams: return the streams as a view so the post-update's gradient is added inside the projection backward (grad_x_addend), bit for bit the sum autograd computes, instead of a separate full-size add (−0.47 ms per V4.1 layer). tests/mhc_projection gains one test; the module's 16 tests pass.

Sparse-attention backward and FP8 dual cast (DeepSeek-V4.1 pipeline stage)

• dsa_sparse_attn_backward(..., out_tail=): for a forward output whose per-head last channels the caller rewrote in place (prime-rl's in-place RoPE on the attention output), out_tail holds the original channels so delta = rowsum(dO * O) stays exact. tests/dsa_sparse_attn_bwd gains one test.
• dKV MMAs of the second warpgroup read P^T / dS^T from registers instead of shared memory: main kernel 4.41 → 4.33 ms at the V4.1 stage shape (8k tokens, 64 heads × 512, 640 slots). dq and dsink are bitwise; dkv is bitwise per contribution, so only the existing atomic order differs. KERNEL_OPTS["rs"] = 0 restores the old path.
• New fp8_dual_cast (CuTe): per-token and transposed per-token e4m3 casts in one pass, with an optional inverse RoPE. prime-rl uses it for the query-gradient path and every Float8BlockwiseLinear grad_output: cast kernels per V4.1 layer 0.93 → 0.77 ms. FP8 bytes and scales are bitwise-identical to the separate casts.

These three were validated on H200 inside prime-rl's stage bench (bitwise checks as above). The module tests were not rerun after rebasing them onto this branch.

Testing

tests/{dsa_indexer_topk,moe_experts,mhc_projection}: forward and backward against fp32/fp64 torch references on small odd shapes and V4.1 shapes, plus meta/compile checks. 52 passed on H200 (FP8 outputs and grads also checked against an exact torch emulation of the recipe and against prime-rl's DeepGEMM path). Each tests/<name>/bench_<name>.py reproduces its numbers.

Known gaps

• Only validated on SM90.
• moe_experts uses Gluon, which is experimental in Triton 3.7.
• mhc_projection autotunes on first use per shape.

🤖 Generated with Claude Code