TOP dense parity for Qwen3 on Blackwell: flashinfer attention, batch-invariant kernel seam, non-reentrant recompute

[adrenaline21]

TOP dense parity for Qwen3 on Blackwell: flashinfer attention, batch-invariant kernel seam, non-reentrant recompute

Summary

Makes dense Qwen3 true-on-policy (TOP) training bitwise-correct and NaN-free, and brings it to Blackwell. TOP requires the training-recomputed logprobs to be bitwise-identical to the rollout logprobs (train_rollout_logprob_abs_diff == 0); this PR closes the remaining gaps to hold that on Blackwell and to train past step 0 without NaN.

Commits

1. chore: gitignore .rxignore — ignore the rx devbox sync artifact.
2. feat(top): batch-invariant kernel seam with analytic backward — kernels.py.
3. feat(top): dense parity — O(L) attention backward + non-reentrant recompute.
4. refactor(top): rename flashinfer ragged wrapper + configurable workspace.
5. test(top): guard standard Megatron RoPE == SGLang RoPE bitwise.
6. test(top): grad tests for hand-written RMSNorm + attention backwards.
7. fix(top): pin flashinfer prefill backend to sglang's (contractual parity).
8. test(top): grad test for tp-invariant row-linear backward.

What changed

Batch-invariant kernel seam (kernels.py)

kernels.py wraps SGLang's forward-only batch-invariant Triton kernels (rms_norm, tp-invariant row-linear) in autograd.Functions with analytic backwards: the forward is the exact inference kernel (so parity comes from running the same kernel as inference, not a reimplementation that can drift), and the backward is defined (the bare kernel has no grad_fn). norm.py and matmul.py route through this seam. Row-parallel matmul now always uses the tp-invariant kernel (a single deterministic reduction), so parity holds regardless of TP degree.

RoPE: dropped the custom sglang-rope path

Previously a custom sglang_apply_rotary_pos_emb_with_freqs forced Megatron to use SGLang's rotary formula. This PR removes it and relies on Megatron's own apply_rotary_pos_emb (unfused via --no-rope-fusion, non-interleaved, mscale=1), which is bitwise-identical to SGLang's neox rotary. Guarded by a new CPU unit test (test_standard_megatron_rope_matches_sglang_rope_bitwise) asserting torch.equal between the two on fp32 and bf16 inputs, so the removal cannot silently change the policy. Removes rope.py and its call sites.

Blackwell attention: flashinfer + O(L) backward

FA3/flash-attn varlen is Hopper-only (sm90), so TOP attention couldn't run on Blackwell. The attention backend is selected by an explicit flag TOP_ATTN_BACKEND, which miles' policy sets to match SGLang's --sglang-attention-backend so training and rollout run the same kernel: flashinfer on Blackwell (and usable on Hopper for parity tests), fa3 on Hopper (faster, native fused backward). It is not GPU-arch-gated; the standalone fallback is fa3. flashinfer's ragged prefill mirrors SGLang's own design — a shared workspace singleton + plan-per-batch (_ragged_prefill_wrapper in sglang_attention.py, workspace size via MEGATRON_TRUE_ON_POLICY_FLASHINFER_WORKSPACE_SIZE, default 1 GiB). The wrapper's backend= is set to fa2 — the flashinfer kernel SGLang's deterministic prefill uses (prefill_backend="fa2", paged wrapper). Matching it makes training and rollout run the identical flashinfer kernel; the choice of ragged-vs-paged wrapper is irrelevant (verified bitwise-equal under fa2). This corrects an earlier setting of cutlass (SGLang's non-deterministic ragged backend), which silently diverged by abs_diff ≈ 0.012. Guarded by a real output-parity test (ragged+_fmha_backend vs paged+fa2, torch.equal), not just a rule assertion.

Backward strategy: flashinfer's run is forward-only (no grad_fn), so training needs an autograd.Function wrapper. The backward dispatches aten's fused mem-efficient attention kernel directly, per packed segment, with is_causal=True and no dense [L,L] mask — so it's O(max segment length) memory, not O(seq²), essential for long single-sample segments and long context. It deliberately does NOT use torch.autograd.grad (a recompute-then-autograd.grad nests a second autograd traversal inside Function.backward, which is reentrant-unsafe); dispatching the fused kernel directly is the safe, analytic path — same rule as the norm and row-linear backwards. The custom O(L) backward is therefore Blackwell-only; on Hopper FA3 supplies its own fused backward.

Radix cache: flashinfer is not in SGLang's RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND, so the rollout side must run with radix cache disabled (--disable-radix-cache). TOP's recompute-prefill is radix-free anyway, so this costs nothing here — but it's a required flag, noted so it isn't dropped.

Non-reentrant activation recompute

• Observed: Megatron's reentrant CheckpointFunction yields NaN gradients in TOP training (reproduced tp=1 and tp=4); non-reentrant (use_reentrant=False) is clean.
• What the change is: it generalizes the existing non-reentrant recompute path (authored by @maocheng23 in the TOP init, originally scoped to the Ulysses-CP fallback because that was the only case his runs exercised) to all TOP runs — a field rename, a default flip (disabled → non_reentrant), and turning it on in the contract. Not a new mechanism.
• Why safe for TOP: the forward is deterministic (dropout 0, batch-invariant kernels), so on-demand recompute reproduces activations exactly and reentrant's RNG-tracking is unnecessary.
• Not established: the exact reason reentrant fails. Reentrant recompute re-enters the autograd engine via a nested backward and composes poorly with TOP's stack of custom autograd.Functions, but the precise failure point isn't isolated (a nested-autograd.grad hypothesis was already refuted). @maocheng23 owns this path and should confirm the mechanism / that generalizing his CP fallback is sound.

Required config flags

--no-rope-fusion (rope parity), --disable-radix-cache (flashinfer rollout), batch-invariant mode, and the TOP contract (--sglang-true-on-policy-contract qwen3_dense_true_on_policy_v1).

Testing

Unit (CPU): rope bitwise-equivalence guard; flashinfer backend selector is fa2; existing contract-wiring, matmul, and norm tests.

Unit (GPU, real parity): test_flashinfer_output_matches_deterministic_paged_reference — runs our training flashinfer path (ragged + _fmha_backend) and SGLang's deterministic reference (paged + fa2) on identical q/k/v and asserts torch.equal. This is the test that would have caught the cutlass regression (the prior rule-assertion test blessed it instead).

Unit (GPU, CUDA-gated): gradchecks for the three hand-written backwards — RMSNorm vjp, O(L) flashinfer attention vjp, tp-invariant row-linear vjp — each against a differentiable reference (the backwards have no inference counterpart, so nothing else validates them). All four (the three gradchecks + the CPU backend guard) pass on H200.

E2E — parity across both parallelism axes:

• fa3 / Hopper: 4×H200, tp=4/cp=1 (39 steps) and tp=2/cp=4 Ulysses-a2a (Qwen3-4B) — Δ=0, ess_ratio = 1.0, grad flows. fa3 is used by both training and rollout, so parity is contractual; CP holds bitwise (a2a-gathered layout runs the same kernel as non-CP). (Ring/online-merge CP untested — out of scope.)
• flashinfer / Blackwell: 4×GB300, tp=2/cp=2 — Δ=0 with the flashinfer backend set to fa2 (matching SGLang's deterministic prefill; see below). CP included. (An earlier cutlass setting gave Δ≈0.012 — a self-inflicted regression, now fixed and covered by a unit test.)
• 0.6B, GSM8K, real learning: Δ=0 across 20 steps of an actual RL run with ess_ratio = 1.0 and nonzero grad_norm — TOP holds during sustained training, not just a smoke step.

Honest scoping: the abs_diff gate is computed pre-backward, so E2E proves forward parity; backward correctness rests on the gradchecks + runs not NaN-ing. Reward convergence is a separate follow-up — GSM8K/0.6B starts near-ceiling (~0.81); a 4B/DAPO run (starts ~0.56, real headroom) is the setup for a reward curve.

Not in scope / follow-ups

• FP8 / MoE TOP parity (the tp-invariant row-linear has no FP8 path yet) — separate work.
• The reentrant-failure mechanism (see above) — reviewer confirmation wanted.
• Reward-convergence run at length.
• Does not touch @maocheng23's gradient-debug tooling (debug.py + hooks) — left intact.