Add fused NVFP4 fake-QAT QDQ kernels

[zianglih]

What does this PR do?

@HumansAnd

Adds fused CuTe DSL NVFP4 fake-QAT QDQ, strict numerical tests, and a colocated
GLM-5.2 W4A16 RL end-to-end recipe in Miles, paired with the minimal Megatron integration in
Megatron-LM #91.

ci-megatron-pr: #91

• Ownership split: Miles owns the kernel, FP32-amax/STE adapter, strict Blackwell tests, and
  E2E recipe. The manual benchmark was measured and removed in a separate cleanup commit; its source and raw results remain linked below. The paired Megatron PR adds only the env-gated lazy call in
  TEGroupedLinear._get_weight_tensors().
• Fused data path: computes the FP32 per-tensor amax with PyTorch, keeps each E4M3 block scale
  and packed E2M1 value in registers, and writes only the dequantized BF16/FP16 result. This removes
  the quantized-data and block-scale global-memory round trips in the former
  quantize(...).dequantize(...) fake-QAT path.
• General list API: maps an arbitrary number of equal-contract discrete rank-2 weights. There is
  no assumption that the weight count is 4, 8, or another fixed multiple; eight weights are only
  the requested benchmark proxy.
• Static dispatch: dtype, standard/4over6 selection, E4M3 max, error metric, and exact/FP16
  candidate-error behavior are compile-time CuTe specializations cached outside the steady-state
  path. The grid-stride kernels do not branch on those modes at runtime.
• Numerical contract: BF16 and FP16; 1D 1x16 weight scaling; caller-provided FP32 per-tensor
  amax; NVTE_USE_FAST_MATH=0; standard NVFP4 plus the complete 4over6 MAE/MSE, E4M3 256/448,
  exact-error/FP16-error matrix. The public config field is error_use_fast_math, matching
  NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH, while its TE 2.17 meaning is explicitly documented as
  FP16-rounded candidate-error math rather than ordinary instruction fast math.
• E2E diagnostics: enables FP32 SGLang LM-head output and explicit Torch router matmul; disables the known-fragile MLA signed-KL assertion while retaining all RL numerical metrics. Avoids unsupported NVFP4 event checksums with explicit data-dump flags.
• Recipe premise: the enabled integration targets TE 2.17+ discrete rank-2 weights. It
  deliberately carries no pre-2.17 version fallback or runtime support validator; recipe
  compatibility is configured by the caller.
• W4A16 RL recipe: copies test_glm5_2_744b_a40b_5layer_nvfp4.py into
  test_glm5_2_744b_a40b_5layer_nvfp4_w4a16.py. It selects
  --sglang-moe-runner-backend flashinfer_cutedsl,
  SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1, and SGLANG_FLASHINFER_MOE_FUSED_FINALIZE=0.
  The actor enables the fake-QAT hook and uses a pure BF16 TE recipe for forward/backward GEMMs;
  --no-gradient-accumulation-fusion lets STE gradients reach the original expert parameters.
• Recipe settings: keeps weight 4over6 MSE with NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1
  (FP16 candidate-error math) and sets NVTE_NVFP4_4OVER6_E4M3_USE_256=none (E4M3 max 448).
  Removes NVTE_BACKWARD_OVERRIDE, the unused FLASHINFER_NVFP4_4OVER6* settings, and
  SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK. There are no first/last-layer BF16
  exclusions; shared experts remain BF16. Actor TP8/EP8 and four 2-GPU rollout engines share
  one 8-GPU node with --colocate; the recipe runs two GRPO rollouts with routing replay.
• CI scope: enables the new E2E in stage-c-8-gpu-b200 (est_time=3600, model-scripts).
  The existing test_nvfp4_quantizer.py registration remains unchanged; Blackwell routing for
  that file is handled separately.
• Deliberate non-goals: stochastic rounding, RHT, 2D quantization, transpose/columnwise output,
  row scaling, ordinary quantization fast math, pre-TE-2.17 compatibility, TE's experimental
  single-grouped-parameter path, and a grouped/persistent rank-3 QDQ kernel.

The vectorized loading, FP4 conversion, and 4over6 structure are adapted from the bit-exact work in
FlashInfer #3448. The per-tensor arithmetic
order follows the current Transformer Engine contract and the work in
TE #2972 and
TE #3068. The strict oracle follows TE's
test_nvfp4_quantize_exact.py,
while the four data families mirror FlashInfer's
test_fp4_quantize.py.
The timing method follows
TE commit 83e23087.

This preserves Megatron-LM #88 unchanged
as the complete single-repository implementation and re-expresses it as a cleaner paired split.
It remains a standalone follow-up to Megatron-LM #75.

Current rebase and C1 validation

• Final Miles head / cleanup: 142e9047c805d2486f48304309b413f6195e744c removes only tests/manual/benchmark_fused_nvfp4_qdq.py in a separate commit, after the E2E passed and both PR bodies were refreshed. The E2E was measured on 1d1c1c11b87b2fc392b3acd92818524fe4f70394; the benchmark-file deletion was not rerun on GPU. Runtime code, the E2E recipe, tests/fast-gpu/test_nvfp4_quantizer.py, and the paired Megatron hook are unchanged. The manual benchmark at its historical tested revision and complete raw C2 results remain referenced below. The strict quantizer test remains in the final tree.
• Validated sources: Miles 1d1c1c11b87b2fc392b3acd92818524fe4f70394 on 4dff6f262576bf751fefa17c3c4db985fc935062 (main); Megatron 12ea91ff82f9e38fd47d63fe0a06c729654d27d4 on 8c1e05747eb612b382df2632783df5c83a853646 (miles-main). Range-diff preserved the original semantic patches; the rebase-only Miles head was pushed before adding the E2E. The paired Megatron hook remains four lines.
• Image: radixark/miles:dev-202609081227, amd64 digest sha256:5b060f9c8304394fb8cb789d5b30ae54e9057e0487ed81c5a85cb7e17c3a0269. Selected from the live tag list immediately before allocation as the latest eligible CUDA 13 development image. Image CUDA 13.0.3; no SGLang, FlashInfer, or TE package override.
• Hardware/runtime: C1 node b200-70, 8x NVIDIA B200 (SM100, 183359 MiB each), driver 580.126.09; PyTorch 2.13.0+cu130, runtime CUDA 13.0, TE 2.17.0, FlashInfer 0.6.18, SGLang 0.5.20.dev54+ga8e5c63 (a8e5c632fe40555f720d4f2c69771ea8cf24f3c4), CuTe DSL 4.6.2.
• Inputs: Pinaster/GLM-5.2_5layer at 1c749139f70e158e4420ba67f342bef1de2e650d; zhuzilin/dapo-math-17k at 2e65612930298bde4c5d58fd97b3f23a483aaff9. Five layers (three dense, two routed MoE), hidden size 6144, 256 experts. Shared experts retain BF16; no first/last-layer exclusions.
• Workload: 8 physical GPUs, colocated actor TP8/EP8 and four 2-GPU SGLang DP2/EP2 engines. Two GRPO rollouts, 8 prompts x 8 samples = 64 samples/rollout, response cap 100, global batch 64, temperature 1, DAPO math/deepscaler rewards, LR 1e-6, TIS [0.5,2], KL/entropy coefficients zero, routing replay enabled. Training uses BF16 TE GEMMs with routed-weight fake QDQ; serving uses CuTe DSL NVFP4 W4A16 with A2A none and unfused finalization.
• Precision flags: --sglang-enable-fp32-lm-head produces FP32 LM-head outputs with BF16 operands; --moe-router-use-torch-mm uses Torch router matmul with FP32 router dtype. The latter was already enabled by Miles defaults and is now explicit in the test.
• Checker scope: --ci-disable-kl-checker skips the existing one-sided MLA signed-PPO-KL assertion, following Miles #2855. PPO KL, trainer/rollout KL, absolute logprob difference, and other diagnostics remain computed and logged. Logprob and tensor-weight-equality checkers are also disabled. Explicit rollout/train/trajectory dumps replace --dump-details, avoiding its event logger and unsupported NVFP4 weight-checksum request. This is integration coverage, not actor/rollout numerical parity.
• Local checks: changed-file pre-commit hooks, Python syntax, git diff --check, CI registration discovery, and preparation/training argument/environment/BF16-recipe capture passed. Parser spellings were verified in the actual image. The quantizer's integer-view and zero-tolerance assertions were not changed.

Exact test environment (the E2E sets these for preparation/training/serving as applicable):

export OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG=1
export SGLANG_FLASHINFER_CUTEDSL_NVFP4_W4A16=1
export SGLANG_FLASHINFER_MOE_FUSED_FINALIZE=0
export NVTE_NVFP4_DISABLE_2D_QUANTIZATION=1
export NVTE_NVFP4_DISABLE_RHT=1
export NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1
export NVTE_USE_FAST_MATH=0
export NVTE_NVFP4_4OVER6=all
export NVTE_NVFP4_4OVER6_E4M3_USE_256=none
export NVTE_NVFP4_4OVER6_ERR_MODE=MSE
export NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1
export SGLANG_DSA_FUSE_TOPK=1
export SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD=0
export SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK=large
export INDEXER_ROPE_NEOX_STYLE=0
export NVSHMEM_DISABLE_NCCL=1

Reproduction from a fresh bare devbox using the image above:

mkdir -p /hai-workspace/nvfp4-qdq-split
cd /hai-workspace/nvfp4-qdq-split
git clone https://github.com/zianglih/miles.git miles
git -C miles checkout --detach 1d1c1c11b87b2fc392b3acd92818524fe4f70394
git -C /root/Megatron-LM fetch https://github.com/zianglih/Megatron-LM.git agent/nvfp4-qdq-miles-integration
git -C /root/Megatron-LM checkout --detach 12ea91ff82f9e38fd47d63fe0a06c729654d27d4
cd miles
export PYTHONPATH=/hai-workspace/nvfp4-qdq-split/miles:/root/TransformerEngine:/root/Megatron-LM
export PYTHONUNBUFFERED=1 RAY_DEDUP_LOGS=0
python -u tests/e2e/megatron/test_glm5_2_744b_a40b_5layer_nvfp4_w4a16.py

Strict C1 quantizer validation

Measured on Miles 67d7fe7b2601470005cbe6ac17112738490c8bec and the same paired Megatron head, on two visible B200 GPUs. Subsequent commits change only the E2E diagnostic/precision flags; the tested quantizer, adapter, and Megatron hook are unchanged.

CUDA_VISIBLE_DEVICES=0,1 PYTHONWARNINGS=ignore \
PYTHONPATH=/hai-workspace/nvfp4-qdq-split/miles:/root/TransformerEngine:/root/Megatron-LM \
pytest -q -o addopts= tests/fast-gpu/test_nvfp4_quantizer.py

........................................................................ [  4%]
........................................................................ [  9%]
........................................................................ [ 13%]
........................................................................ [ 18%]
........................................................................ [ 22%]
........................................................................ [ 27%]
........................................................................ [ 31%]
........................................................................ [ 36%]
........................................................................ [ 41%]
........................................................................ [ 45%]
........................................................................ [ 50%]
........................................................................ [ 54%]
........................................................................ [ 59%]
........................................................................ [ 63%]
........................................................................ [ 68%]
........................................................................ [ 73%]
........................................................................ [ 77%]
........................................................................ [ 82%]
........................................................................ [ 86%]
........................................................................ [ 91%]
........................................................................ [ 95%]
................................................................         [100%]
1576 passed in 10.40s

Exit 0; log SHA256 03ced2512a924b9b390678ecc0f2dbc41a83fd28e89d6d6b90cad4671e5a5845.

E2E attempts before the final configuration

• Attempt 1 — Miles 67d7fe7b, exit 1: conversions, W4A16 engine initialization and initial weight update completed, then event logging requested a checksum that SGLang does not implement for ModelOptNvFp4FusedMoEMethod. Commit 30c101240 replaces --dump-details with the three explicit data-dump flags above. No rollout numerics were produced in this attempt.
• Attempt 2 — Miles 30c101240, exit 1: both RL iterations logged metrics, but step 1 failed the one-sided train/ppo_kl < 1e-8 condition with 0.00023829405836295336. Step 0 passed by sign with -0.00007193960482254624; that does not imply near-zero magnitude. Only initial and post-step-0 weight synchronization completed. Both steps had zero rewards, advantages, losses and gradient norms. The final configuration disables that assertion and enables FP32 serving LM-head outputs; this is not a controlled ablation of either change.

Attempt 2 step	ppo_kl	train_rollout_logprob_abs_diff	train_rollout_kl	grad_norm
0	-7.193960482254624e-05	0.011775476858019829	0.00013989763101562858	0.0
1	0.00023829405836295336	0.01186961680650711	0.00014199175348039716	0.0

NotImplementedError: weight checker has no ComparableWeight for ModelOptNvFp4FusedMoEMethod
Attempt 1: exit 1, before rollout generation
Attempt 2: assertion train/ppo_kl=0.00023829405836295336 < 1e-8 failed
2026-09-09 05:51:08,953 ERR cli.py:75 -- Job 'raysubmit_t6VEmuWC6pUPAgNL' failed
Attempt 2: exit 1
Attempt 2 log SHA256=e6c64391e527ab248607ee729d2113262727b423d451be8902e24efd7946a9fd

Final E2E completion proof

• Validated Miles commit: 1d1c1c11b (the later manual-benchmark deletion is a separate commit).
• Raw log: artifacts/c1-w4a16-e2e-attempt3.log.
• Raw log SHA256: 92ff1f8d4d6c592d51d1ba04cc91f5adc54da87b360fecb3747067deeab191d7.
• Recorded E2E Python process exit: 0 from artifacts/c1-w4a16-e2e-attempt3.exitcode.
• Rollout IDs 0 and 1 each produced 64 samples; train-step IDs 0 and 1 both logged metrics. All 8 actor ranks returned successfully twice (16 returns).
• Exactly three rank-0 weight-update completions: startup, after rollout 0, after rollout 1. Served rollout weight versions advanced from 1 to 2; final update completed before Ray success.
• R3: 864 logged observations (576 forward, 288 backward), 884736 token observations, zero mismatches. Counts include repeated rank/layer/recompute checks, not unique sampled tokens.
• R3 retains the 1% mismatch-fraction limit and default requirement of at least one common valid expert pick; zero mismatches does not mean all top-8 IDs are equal. Padding is excluded from mismatch counts; the threshold denominator is the logged token count. The printed threshold is rounded.
• KL, logprob, and full-weight-equality assertions were disabled as requested. Metrics remain recorded; completion is integration evidence, not numerical parity or learning.
• Eight Ray Stack (most recent call first) dumps at lines 7932–7953 occur after final synchronization and explicit server disposal, show Ray worker main_loop frames, and precede normal WebSocket closure and Ray success. They are retained as shutdown diagnostics.

Exact log excerpts (ANSI colors removed; original physical line numbers):

L4947: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:03:37.793 actor_cell0_rank0] memory_utils.py:41 - [Rank 0] Memory-Usage after update_weights: {'gpu': '0', 'total_GB': 178.35, 'free_GB': 126.08, 'used_GB': 52.27, 'allocated_GB': 16.54, 'reserved_GB': 31.27}
L6288: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:04:24.151 actor_cell0_rank0] structured_log.py:28 - actor cls=MegatronTrainRayActorWithConcurrencyGroups fn=train phase=end ok=true elapsed_s=40.6
L6614: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:04:32.494 actor_cell0_rank0] memory_utils.py:41 - [Rank 0] Memory-Usage after update_weights: {'gpu': '0', 'total_GB': 178.35, 'free_GB': 124.16, 'used_GB': 54.2, 'allocated_GB': 16.93, 'reserved_GB': 32.0}
L7494: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:04:48.311 actor_cell0_rank0] replay_base.py:196 - Replay check (rank 0, stage replay_forward): mismatch 0/1024 tokens, threshold 10
L7515: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:04:48.479 actor_cell0_rank0] replay_base.py:196 - Replay check (rank 0, stage replay_backward): mismatch 0/1024 tokens, threshold 10
L7537: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:04:54.150 actor_cell0_rank0] structured_log.py:28 - actor cls=MegatronTrainRayActorWithConcurrencyGroups fn=train phase=end ok=true elapsed_s=18.5
L7875: (MegatronTrainRayActorWithConcurrencyGroups pid=99758) [2026-09-09 06:05:02.218 actor_cell0_rank0] memory_utils.py:41 - [Rank 0] Memory-Usage after update_weights: {'gpu': '0', 'total_GB': 178.35, 'free_GB': 124.16, 'used_GB': 54.2, 'allocated_GB': 16.93, 'reserved_GB': 32.28}
L7958: 2026-09-09 06:05:16,920	SUCC cli.py:67 -- Job 'raysubmit_E2VZfEw5FiY5Fckg' succeeded

E2E Python exitcode file: 0

Recorded RL numerics

• Source log: c1-w4a16-e2e-attempt3.log (SHA256 below).
• Source SHA256: 92ff1f8d4d6c592d51d1ba04cc91f5adc54da87b360fecb3747067deeab191d7.
• Completion evidence: completion_evidence_present; recorded Python exit: 0.
• Values below are logged aggregates. Raw payloads retain every printed digit.
• KL assertion disabled by the recorded RL ci_disable_kl_checker=True; numerical metrics remain active.

Metric	Rollout / step 0	Rollout / step 1
Training samples	64	64
Raw reward	0.0	0.0
Mean response length	100.0	100.0
Truncated fraction	1.0	1.0
Served weight version (min)	1	2
Served weight version (max)	1	2
Mean serving logprob	-10.295594215393066	-10.237783432006836
Mean trainer logprob	-10.295638084411621	-10.237601280212402
Mean reference logprob	-10.295459747314453	-10.238136291503906
Mean advantage	0.0	0.0
Mean return	0.0	0.0
Trainer/rollout mean absolute logprob difference	0.011665616184473038	0.011925606988370419
Trainer/rollout sampled k3 KL	0.00011017757060471922	0.00011546722089406103
PPO KL	8.531162893632427e-05	0.0004593973862938583
Reference KL loss	0.00019067288667429239	0.00023488527222070843
Entropy loss metric	0.0	0.0
Policy loss	0.0	0.0
Total loss	0.0	0.0
Gradient norm	0.0	0.0
Policy clip fraction	0.0	0.0
Effective sample size ratio	0.9998125433921814	0.9998195171356201
OIS ratio	1.0000089406967163	0.9996318817138672
TIS ratio (before clamp)	1.0000665187835693	1.0002976655960083
TIS absolute distance from 1	0.011662938632071018	0.011929665692150593
TIS clip fraction	0.0	0.0
Existing MLA PPO KL condition (< 1e-8)	disabled	disabled

Checker scope:

• When enabled, training_utils/ci_utils.py:13-16 checks train/ppo_kl < 1e-8 on optimizer step 0 of each rollout for MLA. This test has one optimizer step per rollout. The condition is one-sided; a negative value satisfies it regardless of magnitude. A disabled condition is never reported as passed.
• Step dictionaries remain logged when the KL assertion is disabled. Their presence alone does not establish final train completion or synchronization.
• Recorded rank-0 weight-update completions: 3. A full two-rollout run requires three (startup, after rollout 0, after rollout 1).

Measured limitations:

• Logged mean advantages are zero for rollout IDs [0, 1].
• Logged gradient norms are zero for step IDs [0, 1]. This run does not demonstrate a nonzero policy-gradient update or learning.
• The five-layer checkpoint and 100-token response cap are an integration smoke test; report observed rewards and truncation without treating them as task accuracy.
• Weight-version advancement records completed synchronization. It is not proof that transmitted parameter bits changed; optimizer weight decay and parameter updates are not measured by these aggregate logs.
• Trainer/serving logprob mismatch metrics are observed values, not parity guarantees. The logprob and full-weight-equality checkers are disabled; KL-assertion state is listed above, and R3 retains its original limits.
• PPO KL compares the train forward with the old trainer policy; trainer/rollout k3 KL compares the serving and old trainer-scored logprobs. These are distinct metrics.
• Reported token-level aggregates use existing loss masks; loss code also masks or replaces nonfinite values, so these aggregates alone are not complete tensor-finiteness checks.
• train/entropy_loss=0 is not a measurement of zero entropy: entropy computation is disabled with entropy coefficient 0 and observe-training-entropy false in this test. The reference KL metric is computed but contributes zero to total loss because its coefficient is 0.

Complete raw logged numeric payloads

Source line(s) 5324:

perf 0: {'rollout/num_training_samples': 64, 'rollout/episode_raw_reward': 0.0, 'rollout/episode_response_length/mean': 100.0, 'rollout/episode_response_length/median': 100.0, 'rollout/episode_response_length/max': 100, 'rollout/episode_response_length/min': 100, 'rollout/episode_total_response_length/mean': 100.0, 'rollout/response_len/mean': 100.0, 'rollout/response_len/median': 100.0, 'rollout/response_len/max': 100, 'rollout/response_len/min': 100, 'rollout/zero_std/count_0': 8, 'rollout/zero_std/all_zero_percentage': 0.0, 'rollout/zero_std/all_one_percentage': 0.0, 'rollout/prefix_cache_hit_rate': 0.0, 'rollout/avg_cached_tokens_per_sample': 0.0, 'rollout/repetition_frac': 0.0, 'rollout/truncated_ratio': 1.0, 'rollout/weight_version/mean': 1.0, 'rollout/weight_version/median': 1.0, 'rollout/weight_version/max': 1, 'rollout/weight_version/min': 1, 'rollout/weight_version/mixed_version_ratio': 0.0, 'perf/rollout_time': 3.254490375518799, 'perf/tokens_per_gpu_per_sec': 245.8142159576895, 'perf/longest_sample_tokens_per_sec': 30.726776994711187, 'perf/effective_tokens_per_gpu_per_sec': 245.8142159576895, 'perf/longest_effective_sample_tokens_per_sec': 30.726776994711187}

Source line(s) 5746:

rollout 0: {'rollout/response_lengths': 100.0, 'rollout/rewards': 0.0, 'rollout/truncated': 1.0, 'rollout/rollout_log_probs': -10.295594215393066, 'rollout/raw_reward': 0.0, 'rollout/total_lengths': 256.0, 'rollout/ref_log_probs': -10.295459747314453, 'rollout/log_probs': -10.295638084411621, 'rollout/advantages': 0.0, 'rollout/returns': 0.0}

Source line(s) 6275, 6276:

step 0: {'train/loss': 0.0, 'train/pg_loss': 0.0, 'train/entropy_loss': 0.0, 'train/pg_clipfrac': 0.0, 'train/ppo_kl': 8.531162893632427e-05, 'train/ess_ratio': 0.9998125433921814, 'train/train_rollout_logprob_abs_diff': 0.011665616184473038, 'train/train_rollout_kl': 0.00011017757060471922, 'train/kl_loss': 0.00019067288667429239, 'train/ois': 1.0000089406967163, 'train/tis': 1.0000665187835693, 'train/tis_clipfrac': 0.0, 'train/tis_abs': 0.011662938632071018, 'train/grad_norm': 0.0, 'train/lr-pg_0': 1e-06, 'train/lr-pg_1': 1e-06, 'train/lr-pg_2': 1e-06, 'train/step': 0}

Source line(s) 6286:

perf 0: {'perf/sleep_time': 0.9600825309753418, 'perf/update_weights_implementation_time': 4.316767692565918, 'perf/finalize_and_resume_engines_time': 1.080359935760498, 'perf/update_weights_time': 6.077657461166382, 'perf/wake_up_time': 1.0262548923492432, 'perf/data_preprocess_time': 0.005175113677978516, 'perf/train_wait_time': 50.6668016910553, 'perf/ref_log_probs_time': 10.940051317214966, 'perf/log_probs_time': 0.8811109066009521, 'perf/actor_train_time': 27.357726097106934, 'perf/train_time': 39.401143312454224, 'perf/log_probs_tflops': 13.212280530410382, 'perf/ref_log_probs_tflops': 1.0641160757717176, 'perf/actor_train_tflops': 1.2765846585817398, 'perf/actor_train_tok_per_s': 598.8801825796698, 'perf/mfu_peak_tflops': 2250.0, 'perf/actor_train_mfu': 0.0005673709593696621, 'perf/step_time': 90.06794500350952, 'perf/wait_time_ratio': 0.5625397769326373}

Source line(s) 6968:

perf 1: {'rollout/num_training_samples': 64, 'rollout/episode_raw_reward': 0.0, 'rollout/episode_response_length/mean': 100.0, 'rollout/episode_response_length/median': 100.0, 'rollout/episode_response_length/max': 100, 'rollout/episode_response_length/min': 100, 'rollout/episode_total_response_length/mean': 100.0, 'rollout/response_len/mean': 100.0, 'rollout/response_len/median': 100.0, 'rollout/response_len/max': 100, 'rollout/response_len/min': 100, 'rollout/zero_std/count_0': 8, 'rollout/zero_std/all_zero_percentage': 0.0, 'rollout/zero_std/all_one_percentage': 0.0, 'rollout/prefix_cache_hit_rate': 0.0, 'rollout/avg_cached_tokens_per_sample': 0.0, 'rollout/repetition_frac': 0.0, 'rollout/truncated_ratio': 1.0, 'rollout/weight_version/mean': 2.0, 'rollout/weight_version/median': 2.0, 'rollout/weight_version/max': 2, 'rollout/weight_version/min': 2, 'rollout/weight_version/mixed_version_ratio': 0.0, 'perf/rollout_time': 0.39023685455322266, 'perf/tokens_per_gpu_per_sec': 2050.037024093765, 'perf/longest_sample_tokens_per_sec': 256.2546280117206, 'perf/effective_tokens_per_gpu_per_sec': 2050.037024093765, 'perf/longest_effective_sample_tokens_per_sec': 256.2546280117206}

Source line(s) 7212:

rollout 1: {'rollout/response_lengths': 100.0, 'rollout/rewards': 0.0, 'rollout/truncated': 1.0, 'rollout/rollout_log_probs': -10.237783432006836, 'rollout/raw_reward': 0.0, 'rollout/total_lengths': 250.625, 'rollout/ref_log_probs': -10.238136291503906, 'rollout/log_probs': -10.237601280212402, 'rollout/advantages': 0.0, 'rollout/returns': 0.0}

Source line(s) 7523, 7524:

step 1: {'train/loss': 0.0, 'train/pg_loss': 0.0, 'train/entropy_loss': 0.0, 'train/pg_clipfrac': 0.0, 'train/ppo_kl': 0.0004593973862938583, 'train/ess_ratio': 0.9998195171356201, 'train/train_rollout_logprob_abs_diff': 0.011925606988370419, 'train/train_rollout_kl': 0.00011546722089406103, 'train/kl_loss': 0.00023488527222070843, 'train/ois': 0.9996318817138672, 'train/tis': 1.0002976655960083, 'train/tis_clipfrac': 0.0, 'train/tis_abs': 0.011929665692150593, 'train/grad_norm': 0.0, 'train/lr-pg_0': 1e-06, 'train/lr-pg_1': 1e-06, 'train/lr-pg_2': 1e-06, 'train/step': 1}

Source line(s) 7535:

perf 1: {'perf/sleep_time': 2.589731216430664, 'perf/update_weights_implementation_time': 4.132070779800415, 'perf/finalize_and_resume_engines_time': 1.0400664806365967, 'perf/update_weights_time': 6.035803318023682, 'perf/wake_up_time': 1.133188009262085, 'perf/data_preprocess_time': 0.005443572998046875, 'perf/train_wait_time': 12.778081178665161, 'perf/ref_log_probs_time': 2.229121208190918, 'perf/log_probs_time': 0.856961727142334, 'perf/actor_train_time': 13.910450220108032, 'perf/train_time': 17.219899654388428, 'perf/log_probs_tflops': 13.299065175762617, 'perf/ref_log_probs_tflops': 5.1126828906936215, 'perf/actor_train_tflops': 2.4578909414288153, 'perf/actor_train_tok_per_s': 1153.0899249266304, 'perf/mfu_peak_tflops': 2250.0, 'perf/actor_train_mfu': 0.0010923959739683623, 'perf/step_time': 29.99798083305359, 'perf/wait_time_ratio': 0.4259647090842027}

Historical C2 numerical validation (not rerun after rebase)

The following raw results were measured on Miles ffa7dc4ac3741c1e4f2c896aa78273730c418fe4
and Megatron 65940a7197427997bbfd411b3944849f68c1aaf0. They do not validate the current heads,
new image, or new E2E recipe.

• Historical image: radixark/miles:dev-202609010119
  (sha256:3ee2af22223ee8f0d0b0bf19340fc06bf8da6ddb26a36c76136f8f5e5c462728 on amd64/B300)
• Hardware/runtime: C2, 8x NVIDIA B300 SXM6 AC (SM103); PyTorch 2.13.0+cu130, CUDA 13.0,
  Transformer Engine 2.17.0, Cutlass DSL 4.6.2

The existing Blackwell quantization test is extended with
2 dtypes x 12 shapes x 4 data families x 9 numerical modes = 864 strict TE-oracle cases.
It compares BF16/FP16 integer views to retain signed-zero sensitivity and also calls
torch.testing.assert_close(..., rtol=0, atol=0). Tolerance is never relaxed. The same file also
tests env mapping, invalid contracts, STE/main-grad behavior, two-device state restoration, the
general arbitrary-count adapter, and a real three-weight TEGroupedLinear forward/backward pass.

cd /hai-workspace/nvfp4-qdq-split/miles
CUDA_VISIBLE_DEVICES=0,1 PYTHONWARNINGS=ignore \
PYTHONPATH=/hai-workspace/nvfp4-qdq-split/miles:/hai-workspace/nvfp4-qdq-split/Megatron-LM \
pytest -q -o addopts= tests/fast-gpu/test_nvfp4_quantizer.py

........................................................................ [  4%]
........................................................................ [  9%]
........................................................................ [ 13%]
........................................................................ [ 18%]
........................................................................ [ 22%]
........................................................................ [ 27%]
........................................................................ [ 31%]
........................................................................ [ 36%]
........................................................................ [ 41%]
........................................................................ [ 45%]
........................................................................ [ 50%]
........................................................................ [ 54%]
........................................................................ [ 59%]
........................................................................ [ 63%]
........................................................................ [ 68%]
........................................................................ [ 73%]
........................................................................ [ 77%]
........................................................................ [ 82%]
........................................................................ [ 86%]
........................................................................ [ 91%]
........................................................................ [ 95%]
................................................................         [100%]
1576 passed in 20.68s

Log SHA-256: cbe0279a94cbd7adcdce5cd1f8d20dece5a8036fd4ec2ae717ac50097f34731f.

Historical repository checks:

isort: passed
Black: 4 files left unchanged
Ruff: All checks passed
compileall: passed
git diff --check: passed
CI registration: unchanged in this PR; Blackwell routing is handled separately

Historical C2 performance (not rerun after rebase)

The target is the complete loop over eight independently stored TE-style parameters. Each weight
is contiguous [4096, 6144]; the reported logical shape is 6144 x 4096, so every 1x16 NVFP4
block lies along the 6144 input-feature axis. --num-weights is configurable and the production
helper accepts arbitrary list lengths.

• Naive baseline: for every weight, pad M to TE's required multiple, then
  quantizer.quantize(x_padded).dequantize(dtype=x.dtype)[:m, :n].contiguous().
• Fused primary: the paired Megatron outer env gate/local import and Miles
  maybe_fake_quantize_nvfp4_weight_tensors(...) call, including env/config resolution, one FP32
  PyTorch amax reduction per weight, STE wrapping, output allocation, TVM-FFI marshalling, kernel
  launches, and the complete weight loop.
• Precomputed-amax diagnostic: fused kernel/API time after amax values are prepared; not used in
  the speedup or geomean.
• Timing: torch.utils.benchmark.Timer.blocked_autorange(min_run_time=1); three repeats;
  interleaved naive/fused/naive order. The naive median uses all six A/B/A samples and the fused
  median uses three samples.
• Modes in geomean: standard NVFP4 and the four 4over6 FP16-error modes for each dtype.
  Exact-error remains fully covered by zero-tolerance tests but is deliberately excluded from the
  benchmark target and geomean.

cd /hai-workspace/nvfp4-qdq-split/miles
git checkout --detach ffa7dc4ac3741c1e4f2c896aa78273730c418fe4
CUDA_VISIBLE_DEVICES=2 PYTHONPATH=/hai-workspace/nvfp4-qdq-split/miles \
python tests/manual/benchmark_fused_nvfp4_qdq.py \
  --num-weights 8 --shape 6144x4096 --dtype both \
  --min-run-time 1 --repeats 3 \
  --image radixark/miles:dev-202609010119@sha256:3ee2af22223ee8f0d0b0bf19340fc06bf8da6ddb26a36c76136f8f5e5c462728 \
  --miles-commit ffa7dc4ac3741c1e4f2c896aa78273730c418fe4 \
  --megatron-commit 65940a7197427997bbfd411b3944849f68c1aaf0

image=radixark/miles:dev-202609010119@sha256:3ee2af22223ee8f0d0b0bf19340fc06bf8da6ddb26a36c76136f8f5e5c462728
miles_commit=ffa7dc4ac3741c1e4f2c896aa78273730c418fe4
megatron_commit=65940a7197427997bbfd411b3944849f68c1aaf0
gpu=NVIDIA B300 SXM6 AC
compute_capability=(10, 3)
torch=2.13.0+cu130 cuda=13.0
transformer_engine=2.17.0
cutlass_dsl=4.6.2
min_run_time_s=1.0
repeats=3
num_weights=8
num_gemms=8
weight_storage=discrete_parameters
gradient_accumulation_fusion=false
moe_single_grouped_weight=false
fused_path=maybe_fake_quantize_nvfp4_weight_tensors
shape_contract=logical_block_axis_x_rows
tensor_layout=contiguous_[rows,block_axis]
in_features=6144
out_features=4096
stored_weight_shape=[4096,6144]
primary_order=naive/fused/naive per repeat
NVTE_USE_FAST_MATH=0

dtype	logical shape (block-axis x rows)	mode	naive TE 8-weight loop median [A/B/A raw] (us)	fused QAT 8-weight loop median [raw] (us)	precomputed-amax 8-weight loop median [raw] (us)	end-to-end speedup
bfloat16	6144x4096	nvfp4	783.213 [784.843, 784.276, 781.282, 782.204, 783.747, 782.680]	479.023 [481.504, 479.023, 478.790]	198.966 [199.267, 198.966, 198.214]	1.635x
bfloat16	6144x4096	4over6-mae-e4m3-448-fp16-error	783.807 [787.407, 782.779, 780.833, 787.074, 783.746, 783.869]	479.510 [477.882, 481.158, 479.510]	213.060 [213.060, 213.067, 213.049]	1.635x
bfloat16	6144x4096	4over6-mae-e4m3-256-fp16-error	789.933 [780.698, 783.321, 790.619, 789.907, 789.959, 796.159]	482.917 [483.986, 482.790, 482.917]	213.055 [213.055, 213.067, 213.050]	1.636x
bfloat16	6144x4096	4over6-mse-e4m3-448-fp16-error	791.755 [792.208, 791.338, 792.172, 791.312, 788.658, 792.480]	483.064 [485.371, 483.064, 482.382]	215.964 [215.964, 215.915, 215.979]	1.639x
bfloat16	6144x4096	4over6-mse-e4m3-256-fp16-error	790.953 [789.688, 792.217, 793.973, 792.474, 789.118, 788.348]	485.259 [485.430, 483.784, 485.259]	215.975 [215.936, 215.975, 215.989]	1.630x
float16	6144x4096	nvfp4	767.243 [768.792, 767.251, 765.439, 768.793, 762.112, 767.234]	478.398 [486.605, 477.924, 478.398]	200.673 [200.597, 200.726, 200.673]	1.604x
float16	6144x4096	4over6-mae-e4m3-448-fp16-error	779.465 [779.556, 777.795, 782.546, 779.374, 776.773, 781.146]	481.640 [482.056, 480.007, 481.640]	213.010 [213.027, 213.010, 213.009]	1.618x
float16	6144x4096	4over6-mae-e4m3-256-fp16-error	780.493 [784.028, 778.290, 781.596, 775.166, 779.390, 782.957]	483.382 [486.014, 481.998, 483.382]	213.010 [213.009, 213.014, 213.010]	1.615x
float16	6144x4096	4over6-mse-e4m3-448-fp16-error	785.673 [785.508, 785.754, 784.338, 787.427, 787.589, 785.592]	486.955 [484.638, 486.955, 487.258]	213.556 [213.556, 213.581, 213.546]	1.613x
float16	6144x4096	4over6-mse-e4m3-256-fp16-error	783.492 [784.610, 783.674, 783.311, 772.600, 785.571, 781.779]	488.394 [488.744, 488.394, 487.303]	213.523 [213.560, 213.518, 213.523]	1.604x

geomean_speedup=1.623x
geomean_speedup_bfloat16=1.635x
geomean_speedup_float16=1.611x

Benchmark log SHA-256:
ed4e0d77d6d8bcf947a9fddf978bbf1cb078c6e2de71c791667ec682d71bbfe3.

Geomean history

Exact-error 4over6 is excluded from every row. The first three rows are retained historical runs
from the single-repository PR #88. They used an older image/PyTorch build, so they document the
optimization progression rather than a strict cross-image A/B comparison. The last row is the
historically measured paired implementation in the full table above.

commit	change	overall	BF16	FP16
30d5fd632	dev-sync baseline; upstream-style production helper included	1.616x	1.624x	1.608x
1e63642d7	packed FP32x2 4over6 MSE FP16-error squares	1.642x	1.651x	1.633x
61b2c1170	packed adjacent FP32 multiplies by a common scale	1.605x	1.609x	1.601x
ffa7dc4ac	split ownership plus paired Megatron outer gate	1.623x	1.635x	1.611x

Limitations and untested scope

• Historical C2 validation covers B300/SM103; the current C1 quantizer and W4A16 E2E validation covers B200/SM100.
  No Hopper or pre-Blackwell execution is claimed.
• The paired recipe was tested with TE 2.17.0. Earlier TE behavior is intentionally unsupported
  and no compatibility fallback is included.
• The benchmark is a focused online weight-QDQ proxy, not a full model throughput measurement.
• Exact-error 4over6 is tested bit-for-bit but intentionally excluded from performance targets and
  geomeans.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.