Skip to content

Fuse NVFP4 fake-QAT quantize-dequantize with CuTe DSL - #88

Draft
zianglih wants to merge 3 commits into
radixark:miles-mainfrom
zianglih:agent/nvfp4-qdq-cute-dev-sync
Draft

Fuse NVFP4 fake-QAT quantize-dequantize with CuTe DSL#88
zianglih wants to merge 3 commits into
radixark:miles-mainfrom
zianglih:agent/nvfp4-qdq-cute-dev-sync

Conversation

@zianglih

@zianglih zianglih commented Aug 24, 2026 •
edited
Loading

Copy link
Copy Markdown
  • I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

@HumansAnd

Adds a fused CuTe DSL NVFP4 quantize-dequantize path for fake QAT, reimplemented on the
current miles-main grouped-linear integration after the upstream-dev sync.

  • Integration style: mirrors upstream INT4 QAT's dedicated extension helper: the normal
    TEGroupedLinear._get_weight_tensors() path obtains TE's native discrete weight list and
    delegates to one env-gated helper. There is no constructor cache, packed-parameter adapter,
    or grouped rank-3 kernel.
  • General list API: maps an arbitrary number of discrete rank-2 weights. The production API
    has no assumption that the list length is 4, 8, or another fixed multiple; eight weights are
    only the requested performance proxy.
  • Fused data path: computes the FP32 per-tensor amax with PyTorch, keeps E4M3 block scales
    and packed E2M1 values 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 TE
    quantize(...).dequantize(...) fake-QAT path.
  • Packed pair arithmetic: the FP16-error MSE specialization squares each FP32 error pair with
    one mul.f32x2. Adjacent values multiplied by one common scale also share mul.f32x2 across
    standard/4over6 quantization, FP16-error comparison, and final dequantization. TE's term order
    for sequential error accumulation remains unchanged, and no runtime mode branch is added.
  • 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.
  • Safety contract: NVFP4 QAT requires TE >= 2.17 and discrete rank-2 parameters. It rejects
    simultaneous INT4 QAT, gradient-accumulation fusion, delayed wgrad, FSDP-patched weights, and
    moe_single_grouped_weight=True, matching the safety rationale of current upstream INT4 QAT.
  • 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 legacy-named
    NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH switch follows TE 2.17's current FP16
    candidate-error contract, not its older instruction-fast-math meaning.
  • 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 CuTe arithmetic and test design follow the bit-exact work in
FlashInfer #3448,
Transformer Engine #2972, and
Transformer Engine #3068. The strict
oracle follows TE's
test_nvfp4_quantize_exact.py;
the four data families mirror FlashInfer's
test_fp4_quantize.py.
The timing method follows
TE commit 83e23087.

Issue tracking

Supersedes draft #87 after miles-main was rewritten to the synced upstream-dev base. The earlier
draft remains open and draft for history. This remains a standalone follow-up to #75.

Linked issue: N/A (fork follow-up)

Validation

  • Commit: 61b2c1170dd41d5094e81d918d8d3a95fc56cde1
  • Base: 235952df607b3820716e5e67728a5ab470ca33ae (miles-main)
  • Image: radixark/miles:dev-202608241212
    (sha256:c60775b5b11e6aac47c1f076fbef7a29506de3442ac8bbfef1efeef29e1c90b2)
  • Hardware: C2, NVIDIA B300 SXM6 AC, compute capability 10.3; one GPU for the Cartesian
    matrix, two GPUs for device restoration, four GPUs for the CI-shaped torchrun launch
  • Runtime: PyTorch 2.11.0+cu130, CUDA 13.0, Transformer Engine 2.17.0, Cutlass DSL
    4.6.2

The Cartesian test covers 2 dtypes x 12 shapes x 4 data patterns 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.

cd /hai-workspace/megatron-nvfp4-qdq-dev-sync

PYTHONWARNINGS=ignore PYTHONPATH=. \
pytest -q -o addopts='' \
  tests/unit_tests/extension/test_transformer_engine_nvfp4_qat.py

CUDA_VISIBLE_DEVICES=0 PYTHONWARNINGS=ignore PYTHONPATH=. \
pytest -q -o addopts='' \
  tests/unit_tests/fusions/test_fused_nvfp4_qdq.py

CUDA_VISIBLE_DEVICES=0,1 PYTHONWARNINGS=ignore PYTHONPATH=. \
pytest -q -o addopts='' \
  tests/unit_tests/fusions/test_fused_nvfp4_qdq.py \
  -k uses_and_restores_non_current_device

CUDA_VISIBLE_DEVICES=0,1,3,4 PYTHONWARNINGS=ignore PYTHONPATH=. \
python -m torch.distributed.run --nproc-per-node 4 -m pytest \
  -q -o addopts='' -m launch_on_gb200 \
  tests/unit_tests/fusions/test_fused_nvfp4_qdq.py
..........                                                               [100%]
10 passed in 0.60s

........................................................................ [  7%]
........................................................................ [ 15%]
........................................................................ [ 23%]
........................................................................ [ 30%]
........................................................................ [ 38%]
........................................................................ [ 46%]
........................................................................ [ 53%]
........................................................................ [ 61%]
........................................................................ [ 69%]
........................................................................ [ 76%]
........................................................................ [ 84%]
........................................................................ [ 92%]
........................................................................ [ 99%]
.s                                                                       [100%]
937 passed, 1 skipped in 11.66s

.                                                                        [100%]
1 passed, 937 deselected in 1.79s

# Four-worker GB200-CI-shaped launch; each worker selects LOCAL_RANK.
937 passed, 1 skipped in 6.92s
937 passed, 1 skipped in 7.04s
937 passed, 1 skipped in 6.82s
937 passed, 1 skipped in 11.44s

The single-GPU run intentionally skips only the cross-device state-restoration case; the next
command runs that exact case with two visible GPUs and passes. The four-worker run intentionally
skips it in each worker to avoid cross-worker device contention.

Formatter/linter command:

BASE_REF=miles-main CHECK_ONLY=true SKIP_DOCS=false bash tools/autoformat.sh
Black: 6 files would be left unchanged.
isort: passed.
pylint: 10.00/10.
ruff: All checks passed!

autoformat.sh completed successfully. Its internally non-gating mypy invocation reported
missing optional CUDA-package stubs and CuTe DSL tuple-inference limitations in the isolated
formatter environment; Black, isort, pylint, and Ruff all passed.

Performance

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 1x16 blocks run
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 production maybe_fake_quantize_nvfp4_weight_tensors(...) call,
    including env/config validation, one FP32 PyTorch amax reduction per weight, STE wrapping,
    output allocation, TVM-FFI argument marshalling, kernel launches, and the full 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.
CUDA_VISIBLE_DEVICES=2 NVTE_USE_FAST_MATH=0 PYTHONPATH=. \
python tests/unit_tests/fusions/benchmark_fused_nvfp4_qdq.py \
  --num-weights 8 --shape 6144x4096 --dtype both \
  --min-run-time 1 --repeats 3 \
  --image radixark/miles:dev-202608241212@sha256:c60775b5b11e6aac47c1f076fbef7a29506de3442ac8bbfef1efeef29e1c90b2 \
  --commit 61b2c1170dd41d5094e81d918d8d3a95fc56cde1
image=radixark/miles:dev-202608241212@sha256:c60775b5b11e6aac47c1f076fbef7a29506de3442ac8bbfef1efeef29e1c90b2
megatron_commit=61b2c1170dd41d5094e81d918d8d3a95fc56cde1
gpu=NVIDIA B300 SXM6 AC
compute_capability=(10, 3)
torch=2.11.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 mode naive TE 8-weight loop median [raw A/B/A] (us) fused QAT 8-weight loop median [raw] (us) precomputed-amax median [raw] (us) speedup
BF16 6144x4096 NVFP4 664.135 [664.736, 663.534, 667.028, 661.265, 669.943, 662.558] 412.044 [406.457, 413.457, 412.044] 156.813 [157.256, 155.723, 156.813] 1.612x
BF16 6144x4096 4over6 MAE E4M3-448 FP16-error 661.998 [666.312, 666.869, 664.484, 659.512, 657.276, 657.830] 411.017 [413.612, 410.785, 411.017] 213.055 [213.055, 213.056, 213.048] 1.611x
BF16 6144x4096 4over6 MAE E4M3-256 FP16-error 657.782 [657.328, 658.236, 656.429, 662.755, 666.895, 657.203] 406.754 [406.754, 407.097, 403.862] 213.064 [213.064, 213.059, 213.066] 1.617x
BF16 6144x4096 4over6 MSE E4M3-448 FP16-error 663.994 [664.984, 670.722, 667.309, 663.004, 661.595, 661.737] 415.201 [415.677, 415.161, 415.201] 216.156 [216.342, 216.156, 216.146] 1.599x
BF16 6144x4096 4over6 MSE E4M3-256 FP16-error 666.663 [668.181, 668.712, 668.162, 663.273, 665.006, 665.163] 415.198 [415.436, 415.198, 415.151] 216.315 [216.315, 216.201, 216.336] 1.606x
FP16 6144x4096 NVFP4 642.919 [645.737, 644.035, 641.056, 641.804, 638.106, 645.481] 407.686 [412.815, 407.686, 404.214] 157.402 [157.402, 158.133, 156.344] 1.577x
FP16 6144x4096 4over6 MAE E4M3-448 FP16-error 660.497 [661.669, 664.938, 659.325, 652.906, 663.591, 654.602] 407.738 [408.462, 407.738, 407.414] 213.018 [213.010, 213.029, 213.018] 1.620x
FP16 6144x4096 4over6 MAE E4M3-256 FP16-error 658.420 [662.064, 662.525, 658.447, 658.394, 655.616, 657.986] 412.831 [412.831, 414.789, 407.744] 213.021 [213.022, 213.012, 213.021] 1.595x
FP16 6144x4096 4over6 MSE E4M3-448 FP16-error 659.033 [664.879, 659.269, 656.537, 658.670, 661.884, 658.798] 409.730 [409.560, 410.687, 409.730] 213.625 [213.601, 213.625, 213.629] 1.608x
FP16 6144x4096 4over6 MSE E4M3-256 FP16-error 657.652 [660.544, 658.854, 655.317, 658.832, 655.627, 656.472] 409.798 [409.566, 409.798, 410.076] 213.596 [213.650, 213.596, 213.580] 1.605x

Geomean history

Exact-error 4over6 is excluded from every row. New optimization commits will append rows while
retaining earlier exact-commit results.

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; confirmation run 1.642x 1.651x 1.633x
61b2c1170 packed adjacent FP32 multiplies by a common scale 1.605x 1.609x 1.601x

The first exact-commit run of 1e63642d7 measured 1.626x overall (1.635x BF16,
1.617x FP16). The confirmation table above is the latest full run; both runs reproduce the
MSE precomputed-amax reduction from approximately 229.3 us at baseline to 227.8-228.0 us
BF16 and 226.6-226.8 us FP16. The end-to-end geomean range is retained here rather than
presenting its sub-percent change as noise-free.

The first exact-commit run of 61b2c1170 measured 1.545x overall (1.527x BF16,
1.563x FP16), with the standard TE loop in a distinct faster host/launch regime. The repeat
above returned to the longer-running regime seen in earlier tables. Both exact runs reproduce the
kernel result: MSE precomputed-amax falls from 227.8-228.0 us to 216.1-216.4 us BF16 and from
226.6-226.8 us to 213.6 us FP16. End-to-end geomeans are retained as measured rather than
normalizing away the host-heavy standard-path variability.

Packed-MSE mechanism check

Nsight Compute 2025.3.1.0 captured one BF16 [4096,6144], 4over6 MSE,
E4M3-448 FP16-error launch with precomputed amax. This is a diagnostic mechanism check, not a
replacement for the full-loop benchmark above.

metric 30d5fd632 baseline 1e63642d7 packed squares delta
grid x block 1184 x 128 1184 x 128 unchanged
registers/thread 62 60 -2
executed SASS warp instructions 22,970,496 22,184,064 -3.424%
branch instructions 367,744 367,744 unchanged
local spilling requests 0 0 unchanged
static FMUL / FMUL2 rows 102 / 0 70 / 16 16 scalar pairs packed
NCU replay duration 49,888 ns 49,344 ns -1.090%

The only opcode-count changes are 32 fewer static FMUL rows and 16 new FMUL2 rows. At
runtime, 1,572,864 scalar multiplies become 786,432 packed multiplies, exactly accounting for the
786,432-instruction reduction. Report SHA256s are
4236ce067e74895d3cb0edeb0d92bebc78a210591043e72c51a666edbf620374 (baseline) and
baf3190683d89d433a2166a8b6f1a8f281360fdfd05f4fe207d634c60f32bdac (packed squares).

Packed-scale mechanism check

The next one-launch BF16 NCU comparison uses the same shape and precomputed amax. Standard and
MAE compare to 30d5fd632; MSE compares to 1e63642d7, which already contains packed error
squares.

mode executed instructions before -> 61b2c1170 registers/thread before -> after allocated registers before -> after NCU replay duration before -> after
standard 11,698,944 -> 10,912,512 (-6.722%) 32 -> 32 32 -> 32 30,368 -> 29,472 ns (-2.950%)
MAE 21,397,632 -> 19,824,768 (-7.351%) 62 -> 56 64 -> 56 48,352 -> 47,168 ns (-2.449%)
MSE 22,184,064 -> 20,660,352 (-6.868%) 60 -> 47 64 -> 48 49,344 -> 47,328 ns (-4.086%)

Standard replaces 32 static scalar FMUL instructions with 16 FMUL2; MAE and MSE each
replace 64 with 32. Branch counts are unchanged and all reports have zero local spilling. The MSE
specialization crosses from the 64-register to the 48-register allocation bucket. Candidate
report SHA256s are 0979b9c9764913eb911b6693d0307930288ae09c8e7913bf103e0a011b67d2cb
(standard), 8e70b863a79673f7e778b754e8b6a30b3b527b99bb5d24afda6fc790319f9429
(MAE), and 1a169563b920c1ec62202ca7322ec4bb1b6c94e10bf2bb096f71444d7a2d41b3
(MSE).

Limitations and pending work

  • Enabled NVFP4 QAT is supported only with Transformer Engine >= 2.17. The synced repository's
    development lock still names an older TE source revision, so generic older environments fail
    clearly when the flag is enabled; the validated Miles runtime supplies TE 2.17.0.
  • The kernel is restricted to contiguous rank-2 BF16/FP16 tensors on SM10x, with K divisible by
    16 and at most INT32_MAX elements per tensor.
  • Recipes must disable gradient-accumulation fusion and delayed expert-wgrad computation. Megatron
    FSDP-patched weights and moe_single_grouped_weight=True are rejected.
  • The production path intentionally retains the simple per-weight loop used by current INT4 QAT.
    It does not depend on NVTE_GROUPED_LINEAR_SINGLE_PARAM and does not include the abandoned
    grouped rank-3 kernel direction.
  • Only B300/SM103 has been measured in this PR. SM100 is accepted by the architecture gate but is
    not performance-validated here.

Contribution process

Pre-checks

  • 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

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

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

  1. When your PR is ready, click Ready for Review.
  2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
    • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

@zianglih
Add fused NVFP4 fake QAT kernel
30d5fd6 
Signed-off-by: Ziang Li <ziangli@umich.edu>
@zianglih
Pack NVFP4 4over6 MSE error squares
1e63642 
Signed-off-by: Ziang Li <ziangli@umich.edu>
@zianglih
Pack NVFP4 scale multiplies
61b2c11 
Signed-off-by: Ziang Li <ziangli@umich.edu>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@zianglih