Fuse NVFP4 fake-QAT quantize-dequantize with CuTe DSL

[ziang-and]

What does this PR do ?

@HumansAnd

Replaces the fake-QAT NVFP4 quantize/dequantize global-memory round trip from #75
with a standalone CuTe DSL kernel that quantizes and immediately dequantizes in
registers, while preserving Transformer Engine output bit-for-bit. Grouped expert
weights use one packed rank-3 API and one grouped launch rather than a Python QDQ
loop.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

Scope and numerical contract

• Input/API: one contiguous, 16-byte-aligned [G, M, N] BF16 or FP16 tensor,
  1 <= G <= 2048, homogeneous M/N, N % 16 == 0, and one FP32 amax per
  tensor. G=1 uses the same API and kernel as larger groups.
• Scaling: weight-side 1x16 block scaling; a single PyTorch
  torch.linalg.vector_norm(..., dim=(1, 2), ord=inf, dtype=torch.float32) call
  produces the [G] per-tensor amax vector consumed directly by the kernel.
• Supported modes: standard NVFP4 and the complete requested 4over6 matrix:
  MAE/MSE, E4M3 max 256/448, and exact/FP16-rounded candidate error.
• Fast-math meaning: NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH=1 selects current
  TE's FP16 candidate-error contract. It is not the stale general fast-math
  toggle. Ordinary NVTE_USE_FAST_MATH must be 0.
• Non-goals: SR, RHT, 2D quantization, row-scaled activations, transpose,
  activation/gradient roles, mixed shapes in one group, and non-SM10x GPUs.

Implementation

• Ports FlashInfer's vectorized CuTe DSL NVFP4/4over6 structure, but follows TE's
  per-tensor expression association and final dequantization expression exactly.
• Uses a row-major rank-3 CuTe ABI and a two-dimensional grid:
  blockIdx.x = group, blockIdx.y = persistent CTA. Int64 group-base addressing
  is separated from validated Int32 within-weight offsets.
• Divides a tuned total-grid target across runtime G, caps it by natural work,
  and keeps dtype, standard/4over6, error mode, E4M3 max, error precision, thread
  count, and occupancy bound compile-time static. Shape and G do not enter the
  compilation-cache key.
• Writes BF16/FP16 output directly; no FP4 payload or E4M3 scale buffer reaches
  global memory. The STE has identity backward and preserves main_grad.
• On the target Miles TE 2.17 image, TEGroupedLinear requests one packed
  parameter with NVTE_GROUPED_LINEAR_SINGLE_PARAM=1, QDQs its rank-3 backing
  storage once, and exposes zero-copy member views to TE. Older discrete-weight
  paths call the same API with G=1.
• Newer TE's native grouped-tensor GEMM has a different one-wrapper and CUDA
  m_splits contract. That path is rejected explicitly rather than silently
  claiming compatibility; supporting it would require a broader MCore/TE adapter.

Numerical references

• FlashInfer ff22228d2fa144e9ac6a0d841f2e9ba767ba0f0a:
  flashinfer/quantization/kernels/nvfp4_quantize.py and
  tests/utils/test_fp4_quantize.py.
• Transformer Engine ea1a165decdd09c1ef98b273d965796c89c47473:
  tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py and current native/reference
  NVFP4 implementations.
• Historical context: FlashInfer Multimodal: Limit transformer version in Dockerfile NVIDIA/Megatron-LM#3448 and Transformer Engine Refit EP support NVIDIA/Megatron-LM#2972, Revert "Support multimodule communication (#2031)" NVIDIA/Megatron-LM#3068.

Native TE QDQ is the oracle. Every numerical case compares FP32 amax integer bits
and BF16/FP16 output integer bits, then also runs rtol=0, atol=0. Tolerance is
never relaxed. Per-token FlashInfer agreement is not used as the per-tensor oracle.

Performance history

The retained history uses the current 8 x 6144 x 4096 proxy wherever stated.
Each matrix is stored [4096, 6144], so 1x16 blocks lie along 6144. The geomean
covers BF16 and FP16 across standard NVFP4 and the four 4over6 FP16-error modes;
exact-error remains in correctness coverage but is excluded from performance.
Older obsolete versions were not rerun.

Commit	Change	Geomean result
d8c78ebf3	Eight discrete expert weights; exact former for w in weight_tensors loop	1.518x vs naive TE QDQ
4e005bb60	Experimental pointer-table grouped launch	1.142x vs same-run scalar fused loop; later reverted
0a213eb34	Revert pointer-table design	1.518x; content-identical evidence reused
fa8262134	Hoist global decode scale	1.567x vs naive TE QDQ
30814f602	Experimental fused 4over6 candidate evaluation	1.622x vs naive TE QDQ; reverted after the kernel diagnostic regressed
4480b2fff	Revert candidate-evaluation experiment	1.567x; content-identical evidence reused
f1cba72d3	Clean packed [G,M,N] API, vectorized amax, and one grouped persistent launch	2.141x vs naive; 1.345x vs eight fused singleton calls
f25fbfa40	Tune grouped runtime grids to 96 standard / 64 4over6 target CTAs per SM	2.344x vs naive; 1.431x vs eight fused singleton calls

Absolute native-TE timings vary across runs more than the fused kernel timings.
Every speedup above is an interleaved same-run ratio; cross-commit geomean changes
are not attributed solely to code. The narrower grid-tuning evidence is below.

Current best full mode table

• Image: radixark/miles:dev-202608230022
• Digest: sha256:44795dc9abafaa584a81742803541f287b5c7c6ff938c1ad991d59c706b1362f
• Commit: f25fbfa408e4912945097bef60561e7f6f6d60d6
• Hardware: one NVIDIA B300 SXM6 AC, SM10.3, GPU 2 on an 8-GPU C2 node;
  BF16 then FP16 sequentially in one process
• Software: PyTorch 2.11.0+cu130, CUDA 13.0, TE 2.17.0, Cutlass DSL 4.6.2
• Timing: 1.0-second blocked autorange; three interleaved naive/grouped/naive
  repeats per row; input/view/quantizer construction and warmup excluded
• Primary grouped timing includes vectorized FP32 amax, STE, validation, output
  allocation, one QDQ launch, and packed main_grad assignment. The zero-copy
  TE 2.17 post-QDQ unbind adapter is not included.
• Raw-log SHA256:
  78304b03bc22f295a91350e45c483daa81a361e285070bae5e728bf3bab5ec36

dtype	stored shape [G,rows,block-axis]	mode	naive TE 8-slice loop median [A/B/A raw] (us)	grouped fused one-call median [raw] (us)	fused singleton 8-call diagnostic median [raw] (us)	grouped vs naive	grouped vs singleton
bfloat16	[8,4096,6144]	nvfp4	675.237 [674.959, 677.612, 678.795, 675.515, 672.407, 673.794]	245.727 [245.742, 245.707, 245.727]	382.832 [384.402, 382.832, 381.341]	2.748x	1.558x
bfloat16	[8,4096,6144]	4over6-mae-e4m3-448-fp16-error	670.079 [676.108, 674.198, 670.338, 665.662, 668.290, 669.821]	294.270 [294.275, 294.248, 294.270]	411.102 [411.087, 411.102, 411.154]	2.277x	1.397x
bfloat16	[8,4096,6144]	4over6-mae-e4m3-256-fp16-error	669.379 [672.268, 672.971, 665.585, 669.977, 668.782, 666.783]	294.264 [294.271, 294.264, 294.248]	411.106 [411.170, 411.102, 411.106]	2.275x	1.397x
bfloat16	[8,4096,6144]	4over6-mse-e4m3-448-fp16-error	673.003 [669.919, 674.634, 671.372, 670.526, 679.971, 675.079]	303.179 [303.179, 303.126, 303.234]	426.259 [426.259, 426.490, 426.222]	2.220x	1.406x
bfloat16	[8,4096,6144]	4over6-mse-e4m3-256-fp16-error	673.105 [673.447, 668.276, 677.700, 674.806, 668.006, 672.762]	303.215 [303.215, 303.221, 303.200]	426.209 [426.557, 426.209, 426.087]	2.220x	1.406x
float16	[8,4096,6144]	nvfp4	657.559 [658.334, 655.053, 657.002, 662.231, 656.236, 658.117]	241.774 [241.779, 241.774, 241.773]	380.689 [380.226, 380.689, 384.022]	2.720x	1.575x
float16	[8,4096,6144]	4over6-mae-e4m3-448-fp16-error	666.478 [668.166, 672.196, 666.377, 664.374, 665.410, 666.578]	290.279 [290.196, 290.280, 290.279]	410.319 [410.406, 410.302, 410.319]	2.296x	1.414x
float16	[8,4096,6144]	4over6-mae-e4m3-256-fp16-error	667.688 [668.155, 669.047, 670.236, 665.090, 667.222, 663.661]	290.247 [290.253, 290.247, 290.202]	410.344 [410.344, 410.364, 410.250]	2.300x	1.414x
float16	[8,4096,6144]	4over6-mse-e4m3-448-fp16-error	666.357 [668.179, 663.657, 663.525, 667.665, 665.049, 673.444]	298.737 [298.737, 298.736, 298.785]	412.011 [412.341, 412.011, 411.927]	2.231x	1.379x
float16	[8,4096,6144]	4over6-mse-e4m3-256-fp16-error	666.399 [674.430, 669.154, 663.611, 663.500, 665.304, 667.494]	298.707 [298.701, 298.707, 298.728]	412.261 [412.013, 412.261, 412.356]	2.231x	1.380x

Derived geomeans: 2.344x vs naive TE QDQ and 1.431x vs eight fused
singleton calls. BF16 and FP16 geomeans versus naive are 2.340x and 2.349x.

cd /root/Megatron-LM-grouped3d
CUDA_VISIBLE_DEVICES=2 PYTHONPATH=. \
python tests/unit_tests/fusions/benchmark_fused_nvfp4_qdq.py \
  --groups 8 --shape 6144x4096 --dtype both \
  --min-run-time 1 --repeats 3 \
  --image radixark/miles:dev-202608230022 \
  --commit f25fbfa408e4912945097bef60561e7f6f6d60d6

Persistent-grid tuning

Full target runs varied only the runtime total-grid targets. These are geometric
means of grouped-call microseconds; lower is better. The retained 96/64 pair is
the smallest observed best plateau before each mode regresses.

standard / 4over6 target CTAs per SM	standard geomean (us)	standard vs 24	4over6 geomean (us)	4over6 vs 8
24 / 8	249.197	1.000x	309.218	1.000x
48 / 16	244.092	1.021x	303.129	1.020x
96 / 32	243.617	1.023x	298.245	1.037x
192 / 64	250.881	0.993x	296.616	1.043x
96 / 64	243.642	1.023x	296.623	1.043x
96 / 128	243.667	1.023x	298.188	1.037x

Validation

The strict Cartesian is 2 dtypes x 3 group counts (G=1,3,8) x 12 shapes x
4 data patterns (random, boundary, zeros, maxes) x 9 numerical modes =
2,592 TE-oracle cases. Group-distinct values and amaxes exercise non-power-of-two
G; an additional random TE-oracle case covers the G=2048 upper bound.
Additional tests cover STE/main-grad, config parsing, invalid layouts/bounds,
non-current-device restoration, discrete G=1 fallback, and an actual TE 2.17
packed GroupedLinear forward/backward.

cd /root/Megatron-LM-grouped3d
CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. \
pytest -q tests/unit_tests/fusions/test_fused_nvfp4_qdq.py

CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=. \
pytest -q tests/unit_tests/fusions/test_fused_nvfp4_qdq.py \
  -k uses_and_restores_non_current_device

2675 passed, 1 skipped, 26 warnings in 24.51s
1 passed, 2675 deselected, 25 warnings in 2.01s

Test-log SHA256: full
5903489ee223f3db240a6ff6ec8b14c0c72fd4cd45cb202e65075aee110e3c02;
two-device
7fe6c09244cb893c2ceea6529c9a409f5c0a56ff4307e2c5a3e27b89fad0d4e7.

Static checks at f25fbfa40:

ruff check <four changed Python files>
All checks passed!
python3 -m py_compile <four changed Python files>
git diff --check

Limitations and pending work

• Only C2 B300 and the pinned Miles TE 2.17 member-view grouped-weight API have
  been validated.
• Current upstream TE's native grouped-tensor GEMM is deliberately not integrated;
  doing so also requires preserving CUDA m_splits through MCore.
• The performance table measures the clean rank-3 QAT API and excludes TE 2.17's
  zero-copy output-view adapter.
• Classic main_grad propagation is covered; MCore FSDP and a full training job
  have not been run.
• CI has not run yet; the PR remains draft while follow-up optimization continues.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Loading

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

Feel free to message or comment on the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

1. Attach the Expert Review label when your PR is ready for review.
2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge conflicts are resolved, and CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

1. Add Final Review label
2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches after merge, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge this PR.