[Common][PyTorch] Fuse NVFP4 quantization launches for MoE experts

[liujshi]

Summary

Two complementary fusions that reduce NVFP4 quantization kernel launches by ~96% for MoE expert layers while keeping outputs bit-identical to the per-tensor loop.

1. Row-scaled activations (fused whole-buffer quantization)

Row-scaled NVFP4 quantization is row-local (per-row amax, 1x16 block scales and 4/6 candidate selection never cross rows), so the whole expert input buffer can be quantized in a single kernel launch and split into zero-copy per-expert row views.

• Add NVFP4TensorStorage.slice_rows: zero-copy row-range view of a rowwise-quantized NVFP4 tensor, with optional scale padding to the 128-row cuBLAS block-scaling layout.
• grouped_linear: fused input quantization for row-scaled NVFP4 (env: NVTE_ROW_SCALED_FUSED_QUANTIZE, default on).

2. Batched 4over6 weight quantization (per-tensor-scaled tensors)

A new batched kernel quantizes all same-shaped tensors in two launches (one batched amax pass + one batched quantize pass, with grid.z tensor segments) instead of two launches per tensor. Each tensor keeps its own scalar amax and global scale, so results are bit-identical to the per-tensor loop. This also fills the grouped quantization gap for the non-RHT path noted in split_quantize_nvfp4_impl_helper.

• batched_amax_kernel (per-segment atomicMax; max is exact and order-independent) and quantize_4over6_batched_kernel (reuses the single-tensor pipeline per segment).
• New C API nvte_nvfp4_quantize_4over6_multi and torch extension entry nvfp4_quantize_4over6_multi, with optional caller-provided output workspaces to preserve weight-cache semantics.
• grouped_linear: batched weight quantization (env: NVTE_NVFP4_BATCHED_WEIGHT_QUANTIZE, default on).

Benchmarks

Setup: Qwen3-30B-A3B, 48 layers, seq_len=8192, 4xB300, mock data, nsys over 2 steady-state training steps.

Metric	Before	After	Delta
quantize_4over6 kernel time	2053 ms	832 ms	-59.5%
quantization launches (2 steps)	71564	3072	-95.7%
quantization total (amax + swizzle + bwd dequant)	3061 ms	1529 ms	-50.0%
total compute time	9.197 s	7.606 s	-17.3%
GEMM time (sanity check)	1598 ms	1603 ms	unchanged

Correctness

• Kernel-level bitwise equality vs the per-tensor loop for 64 weight tensors (packed fp4 data, e4m3 block scales, scalar amax).
• Workspace-write path bitwise equality (32/32 tensors).
• Identical first-iteration loss and grad norm in mock training (5-layer and 48-layer); later-iteration deviations match run-to-run noise.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add NVFP4TensorStorage.slice_rows (zero-copy row-range view, optional 128-row scale padding).
• grouped_linear: fused whole-buffer input quantization for row-scaled NVFP4 activations (NVTE_ROW_SCALED_FUSED_QUANTIZE=1 default).
• Add batched_amax_kernel and quantize_4over6_batched_kernel in common/cast/nvfp4/quantize_4over6_nvfp4.cuh (grid.z segmented, per-segment scalar amax).
• Add C API nvte_nvfp4_quantize_4over6_multi (common/cast/cast.cu, cast.h).
• Add torch extension entry tex.nvfp4_quantize_4over6_multi (pytorch/csrc/extensions/nvfp4_multi_quantize.cpp) with optional output workspaces.
• grouped_linear: batched 4over6 weight quantization for same-shaped expert weights (NVTE_NVFP4_BATCHED_WEIGHT_QUANTIZE=1 default).

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works (bitwise comparison tests vs the per-tensor loop: packed data / scales / amax, plus mock-training loss equivalence)
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR reduces NVFP4 quantization launches for MoE expert layers.

• Adds whole-buffer row-scaled activation quantization with per-expert row views.
• Adds batched per-tensor 4over6 weight quantization across same-shaped experts.
• Exposes the batched operation through the common C API and PyTorch extension.
• Integrates both optimizations into GroupedLinear behind default-enabled environment flags.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/cast.cu	Adds validation, temporary workspace management, and dispatch for the batched NVFP4 4over6 C API.
transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh	Adds segmented batched-amax and 4over6 quantization kernels for same-shaped tensors.
transformer_engine/pytorch/csrc/extensions/nvfp4_multi_quantize.cpp	Bridges lists of PyTorch tensors and optional cached outputs to the new batched C API.
transformer_engine/pytorch/module/grouped_linear.py	Selects fused activation and weight quantization paths for eligible NVFP4 GroupedLinear calls.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Adds row-range NVFP4 views and optional scale-buffer padding for expert splits.

Sequence Diagram

sequenceDiagram
  participant GL as GroupedLinear
  participant Q as NVFP4 Quantizer
  participant Ext as PyTorch Extension
  participant Core as CUDA Core
  participant GEMM as Grouped GEMM
  GL->>Q: Quantize full activation buffer
  Q-->>GL: Row-scaled NVFP4 tensor
  GL->>GL: Create per-expert row views
  GL->>Ext: Batch same-shaped expert weights
  Ext->>Core: nvte_nvfp4_quantize_4over6_multi
  Core->>Core: Batched amax pass
  Core->>Core: Batched 4over6 quantization
  Core-->>GL: Per-expert NVFP4 weights
  GL->>GEMM: Expert activation and weight tensors
  GEMM-->>GL: Grouped output

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into nvfp4-fused-qua..." | Re-trigger Greptile}