Add fused NVFP4 fake-QAT QDQ kernels

[zianglih]

What does this PR do?

@HumansAnd

Moves the fused CuTe DSL NVFP4 fake-QAT QDQ implementation, numerical tests, and benchmark into
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
  manual benchmark. 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.
• 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.
• CI scope: this PR leaves the existing test_nvfp4_quantizer.py CI registration 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 draft 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.

Validation

• Miles head: ca0c3425f2c5b5f8ec97fc4af6e8fd976d32f704
• C2 validation/benchmark commit: ffa7dc4ac3741c1e4f2c896aa78273730c418fe4
  (the head differs only by restoring the pre-existing CI registration metadata)
• Miles base: 2799fe386320c156334bf763ad4d7ca0f85dca4e (main)
• Paired Megatron commit: 65940a7197427997bbfd411b3944849f68c1aaf0
• 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.

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

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 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
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
current paired implementation and latest 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

• C2 validation covers B300/SM103; 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.
• The PR remains draft while paired review and repository CI are pending.

[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.