[Kh4L]

Original work: #3426 by @xiuhu17. Re-cut at his request; his original authorship and sign-offs are preserved.

Adds NVFP4 SGLang rollouts and Megatron weight refit on top of #3614. The shared iterator expands weight/scale groups, preserves padded NVFP4 storage, refreshes graph-captured scales in place and rejects semantic refit failures. MoE and FP4 GEMM runner settings can explicitly select CuteDSL.

Startup requires a locally converted checkpoint with the expected MoE-only layout marker. Incompatible external ModelOpt NVFP4 layouts are rejected before serving starts.

Stack: #3613 has merged; #3614 remains the parent. The full diff against main includes MXFP8. Review the NVFP4 increment.

Head 035a6003f219f324ab2dc136e381ef95849f1b0f restacks on refreshed #3614 and main d48b2e0f3d7fc4891b418ae603c8e8fa71482b9c. Quantization modules are unchanged. Dependency files and submodule pins exactly match main, including Transformer Engine 2.18; this PR no longer introduces a separate Transformer Engine version bump.

Validation on the refreshed head:

• 118 focused CPU tests pass, covering shared MXFP8/NVFP4 checkpoint handling, paired scales, layout rejection, backend forwarding, semantic update failures and cache-policy behavior.
• Repository-pinned Ruff 0.9.9, import sorting, formatting and diff checks pass across the PR's Python files.
• CPU tests use normal module imports on macOS, with project conftest/config disabled. This is supplementary coverage, not full CI or a fresh GPU run.

Historical GPU evidence: HSG job 6593575 completed successfully on 3 nodes / 12 GB200 GPUs at 736472843ba46eca4c16f127009452bdb148f090, with four optimizer steps and four NVFP4 broadcast refits. This predates the current main refresh and does not validate its dependency updates.

Current-head maintainer CI and review remain required. A maintainer must run /ok to test 035a6003f219f324ab2dc136e381ef95849f1b0f.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Kh4L]

Self-review (/review-pr-team), per the contributor self-review policy

Scoped to this PR's own increment (35d910009..4f56b148a — the nvfp4 feature, the TE 2.17 bump, the lock regen, and my two rebase fixes). The three layers below have their own posted self-reviews and were not re-reviewed.

The headline: the lint gate is red and the two tests guarding the pair-completeness logic have never executed. pytest is used at test_megatron_sglang_weight_iterator.py:240,251 but never imported — ruff fails with F821 twice, and both tests error with NameError instead of running. Fixed in 6e6599998. Beyond that, the offline/online consistency I most expected to be wrong is actually correct, and the dependency bump is justified and clean — but there are five correctness issues in the new code and a checkpoint-compatibility gap that stops stock ModelOpt NVFP4 checkpoints from being recognized at all.

Fixed in 6e6599998

The missing import pytest. Verified rather than assumed: same harness, pre-fix file → 2 failed, 5 passed (NameError: name 'pytest' is not defined); post-fix → 7 passed. ruff check now clean across all 15 changed files. The two dead tests were the only coverage for _iter_nvfp4_hf_named_tensors' incomplete-pair and duplicate-role guards, so replacing either raise with a bare return was indistinguishable from the shipped state.

Provenance note, since I'd rather say it than have it found: this is my own miss. My commit 4f56b148a says these fixes were "surfaced by running the repo's own pre-commit against this layer" — I ran it against the production files and not the tests, which is exactly how an F821 in a test file survived a commit whose stated purpose was fixing an F821.

Left inline for your judgement

Seven items, in rough severity order. All are marked traced-vs-executed inline.

1. A stranded gate/up half is detected only after part of the refit has already shipped. The completeness check sits at the end of the generator, but buckets are consumed lazily and transmitted as they're produced — so the engines are left holding a mix of new and stale weights with no rollback. The name-axis skip is correctly symmetric; this is the dtype/rank rejections, which have no pair-level counterpart.
2. should_quantize_nvfp4 raises where all three of its sibling checks return False. The same expert tensor is fatal or harmless depending only on whether the user happened to list it in modules_to_not_convert.
3. The converted checkpoint advertises FP8 KV-cache quantization it does not contain. SGLang reads kv_cache_scheme → attaches ModelOptFp8KVCacheMethod to every RadixAttention → finds no k_scale/v_scale → silently substitutes 1.0. Inert at kv_cache_dtype: auto, uncalibrated and unwarned otherwise. mxfp8 emits no KV metadata at all, so this is new here.
4. Stock ModelOpt NVFP4 checkpoints are never recognized. is_existing_nvfp4_checkpoint reads only config.json, but real releases put the block in hf_quant_config.json — the file SGLang itself declares. Confirmed against nvidia/Qwen3-235B-A22B-NVFP4 and nvidia/DeepSeek-R1-0528-NVFP4: both have no quantization_config in config.json. The converter then tries to re-quantize an already-NVFP4 checkpoint and dies, after copying its metadata into the cache dir.
5. exclude_modules entries are globs, matched with in. Real ModelOpt checkpoints ship "model.layers.0.self_attn*"; plain substring containment never matches, so those entries are silent no-ops for skipping and make validate_checkpoint_high_precision_layout pass vacuously.
6. _contiguous_pair_view doesn't require canonical strides. A size-1 leading dim can carry any stride[0] and still report contiguous, which defeats both the stride proxy and the adjacency check — and unlike every other precondition failure it returns a wrong result rather than falling back to torch.cat.
7. weight_scale_2 is re-derived from a local copy of TE's formula rather than read back from the quantizer. The copy is faithful today (I diffed it against TE's reference recipe, see below). But the block scales come from TE while the global scale comes from our reimplementation, and the two must agree exactly or every dequantized weight is off by a constant factor. Any upstream change — pow_2_scales, a different row_scaled_nvfp4 default, a tweak to the clamp — silently desynchronizes them, and no test would catch it.

The thing a smoke test has to prove

SGLang re-runs process_weights_after_loading on every quantized module after each refit (model_runner.end_weight_update → postprocess_weight). Its counterpart restore_weight dispatches to restore_weights_before_loading, and ModelOpt defines no such hook — so NVFP4 MoE re-finalizes over already-finalized state. Some transforms guard for this deliberately (the W13 deinterleave carries a sticky _w13_deinterleaved flag); the CuteDSL-v2 branch applies interleave_w13_halves unconditionally. Whether that corrupts depends on whether the refit fully overwrites both params first, and nothing here establishes it. The previously validated MoE hot-reload work ran on flashinfer_trtllm, which takes the other branch — so the CuteDSL path is unexercised. This needs ≥3 refits, not 1; a single-refit run cannot tell "correct" from "corrupts on re-finalize".

Verified good — recording these so nobody re-litigates them

• Offline and online quantization selection agree. This is the property most likely to silently break an RL run, so I executed it rather than reasoning. The raw predicates do diverge — with extra_high_precision_layers_hf: ["down_proj"] the offline converter quantizes 0 tensors while the online iterator quantizes 8, because offline skips at MoE-container granularity and online at pair granularity. But _sync_checkpoint_ignore merges the checkpoint's container-level ignore list back into modules_to_not_convert, and the synchronizer re-reads that config fresh from the driver-side generation object every refit, so the mutation propagates. Replaying the merge with the real functions: 4/4 cases agree. No test pins this reconciliation, though — the two shipped tests each cover one side in isolation.
• The TE bump is necessary, not incidental. nvfp4_use_4over6, nvfp4_e4m3_max and nvfp4_4over6_err_mode are all new in v2.17 (v2.16 adds only row_scaled_nvfp4), so the >=2.17 requirement in the docstring is exactly right. Tag v2.17 resolves to 2e559f06…, matching the lock; VERSION.txt says 2.17.0, matching the hand-declared metadata; v2.17.0 does not exist, so @v2.17 is the correct spelling. TE's requires-dist is unchanged between 2.15 and 2.17, so the 12-line lock diff is correct rather than suspicious. No second, now-stale TE pin exists anywhere in the repo. v2.16 is the only breaking release crossed and the repo uses none of what it removed. Moving from the release_v2.15 branch to the v2.17 tag is a real improvement — release_v2.17 has already drifted past its own release.
• The Transformer Engine contract holds. Two things here could have been quietly catastrophic, so I checked both against upstream rather than trusting the shape asserts. _rowwise_scale_inv is a plain row-major block-scale matrix, not a swizzled/tiled one (with_gemm_swizzled_scales defaults False), so [:num_rows, :num_cols // 16] reads the region it intends to. And _amax_rowwise is a per-tensor scalar here, so amax_rowwise.reshape(-1)[0] is right: TE's reference recipe only reshapes the global amax to a per-row [m, 1, 1] under row_scaled_nvfp4, and this PR passes row_scaled_nvfp4=False. nvfp4_global_encode_scale_te also reproduces TE's formula exactly — same e4m3_max * 6.0 / amax, same finfo.max clamp, same zero→one guard, same reciprocal for the decode scale.
• Cross-shard deferred-pair handling is correct and genuinely tested, and the emitted companion tensors have the right ranks and dtypes.

Two build notes worth a line in the PR description

• TE v2.17.1 was available before this lock was regenerated. Published 2026-08-07; it fixes checkpoint "extra state" being unpickled for stateless recipes and adds an explicit env-var guard. A framework that loads checkpoints routinely probably wants it. Caveat: the fix introduces an opt-in env var, so it may not be purely drop-in.
• NeMo-Automodel declares transformer-engine[pytorch]<=2.11.0, and NeMo-RL requests nemo-automodel[moe]. The global override-dependencies already overshot that ceiling at 2.15; this widens it from +4 to +6 minors. Pre-existing, not caused here, but crossing 2.16 makes it worth stating.

Tests and docs

No test exercises any real quantization math — every one monkeypatches the leaf quantizer. Three mutations survive the full suite undetected: nvfp4_global_decode_scale_te → constant 1.0; nvfp4_weight_e4m3_max → hardcoded 448 (ignoring NVTE_NVFP4_4OVER6* entirely); _contiguous_pair_view → always None. All three are pure CPU functions that need neither a GPU nor Transformer Engine, so this gap is closable today.

Relatedly, NVTE_NVFP4_4OVER6* is undeclared cross-process configuration that scales every emitted weight_scale_2: offline conversion runs on the driver, online refit in the Megatron worker, and nothing checks the two saw the same environment.

There is no documentation for SGLang rollout weight precision anywhere — not for nvfp4, and not for mxfp8 below it. docs/fp8.md is not the natural home (different backend, different mechanism); docs/guides/quantization-aware-rl.md is, and it already tabulates verified configurations. Which raises the point worth making explicitly: NeMo-RL already has a real-quant NVFP4 rollout lane — ModelOpt + vLLM, policy.generation.real_quant: true — with published convergence results on the very models this targets, including a 300-step Qwen3-30B-A3B W4A4 run and a Nano3 W4A16 run. That lane is vLLM-only so an SGLang implementation is legitimately needed, but this PR adds a second, independent NVFP4 packing implementation without saying how the two relate or why the existing one couldn't be reused. Its real_quant_ignore list is also the existing precedent for the skip-list design here. Separately, docs/guides/refit.md:26 now describes constraints that quantized SGLang violates.

Environment caveats

No GPU, and sglang/flashinfer/transformer_engine are not installed, so nothing involving their runtime behaviour was executed by anyone — SGLang and TE findings are traced against fetched upstream sources at the pinned revisions, and quoted inline. pytest cannot run the repo suite here (tests/unit/conftest.py imports ray), so test executions used a stubbed module loader; the counts above are real runs of the real test bodies. Nothing in this PR has run on hardware, which the PR body already says.

Generated by Claude Code

[Kh4L]

/ok to test 74c4836

[Kh4L]

/ok to test cdd2097

[Kh4L]

/ok to test 679ab0c

[Kh4L]

Current head 7d429c388bdaade042f6aa0e108abd289fddc6b1 is restacked on refreshed #3614 and adds the NVFP4 declaration to the failing runner-backend fixture. Compared with HSG-validated 736472843ba46eca4c16f127009452bdb148f090, only five SGLang unit-test fixtures changed, so job 6593575 remains applicable to the runtime tree. Fresh CI is triggered on the exact head.

[Kh4L]

/ok to test d98ccf0

[Kh4L]

/ok to test 31fd89a