Add NVFP4 per-token quantization recipe

[cael-ling]

Description

This PR adds an NVFP4 per-token quantization recipe for model pre-training. The default NVFP4BlockScaling recipe computes a single per-tensor outer amax (s_global) per tensor. The per-token variant instead computes a per-row outer amax (length M) for rowwise data and a per-col outer amax (length K) for columnwise data, giving each token/row its own global scale.

Changes

• Per-token cast kernels: vector-amax + encode/swizzle producing NVFP4 tensors whose _amax_rowwise / _amax_columnwise are per-row/per-col vectors.
• CUTLASS GEMM (nvfp4_cutlass_per_token_gemm) that rescales with the per-row/per-col outer-amax vectors inside the epilogue;
• Forward + backward coverage (dgrad NN / wgrad NT layouts).
• NVFP4PerTokenBlockScaling recipe (re-exported from transformer_engine.pytorch.fp8), plus an equivalent NVTE_NVFP4_PER_TOKEN=1 env-var switch on a plain NVFP4BlockScaling so frameworks that only build a default recipe (e.g. Megatron-Core) can opt in with no code change.
• Opt-in RHT / SR (per_token_rht / per_token_sr) — off by default on the per-token path.
• Opt-in 2D weight quantization (per_token_weight_2d): transposition-invariant 16×16 cast emitted in per-token layout.
• Docs: API reference entry, NVTE_NVFP4_* env-var docs, and a "Per-token NVFP4" feature section with Megatron-Core launch instructions.
• Example: examples/pytorch/nvfp4_per_token_megatron — single-GPU MoE example comparing per-token vs per-tensor vs BF16 with identical model/data/seed.

Ongoing work

The per-token recipe currently targets accuracy evaluation, not optimized production deployment:

• Requires NVTE_NORM_FWD_USE_CUDNN=1 (unfused norm forward); the fused norm+amax path rejects per-token quantizers.
• fuse_wgrad_accumulation=True is unsupported → launch with --no-gradient-accumulation-fusion in Mcore.
• Forward/backward output quantization, communication/bulk overlap, and CUDA graphs are not yet supported/validated.
• Kernels are functional but not perf-tuned; use for numerical comparison, not perf benchmarking.
• 2026-06-15: Native CUTLASS grouped per-token NVFP4 GEMM (single ptr-array launch) — done (b0b72cf)
• 2026-06-15: Support fuse_wgrad_accumulation in NVFP4 per-token GEMM — done (c3ce651)

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

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
• New and existing unit tests pass locally with my changes

[qijiaxing]

maybe use DIVUP here to handle the remainder case?

[cael-ling]

This fast path has a hard precondition that M and K are exact multiples of CHUNK_DIM (128): validate() does NVTE_CHECK(M % CHUNK_DIM_Y == 0) / NVTE_CHECK(K % CHUNK_DIM_X == 0), and is_supported() returns false unless both hold — so any non-multiple shape is rejected / routed to the generic per-token fallback before it ever reaches this launcher.

[qijiaxing]

typo: Row direction never sees RHT -> Row direction never uses RHT

[cael-ling]

Good catch

[qijiaxing]

For these quantization kernel, TMA only require SM 9.0+ only. Is there any other constraints that limit to sm 10.0+?

[cael-ling]

The CUDA_ARCH >= 1000 guard is intentional but not because of a hardware op in this kernel. Two reasons:

1. The shared TE PTX wrappers it calls — cp_async_bulk_tensor_2d_global_to_shared and mbarrier_wait_parity_acquire_cta_shared_cta in util/ptx.cuh — are themselves guarded to >= 1000 and emit NVTE_DEVICE_ERROR below that. They were authored/validated only for the Blackwell path.
2. The whole NVFP4 quantize path is host-gated to SM100 anyway (NVTE_ERROR("NVFP4 requires SM100 ...")), since NVFP4 is a Blackwell datatype and the downstream FP4 GEMM that consumes these scales only exists on SM100. So the amax kernel is never launched off <SM100; the per-arch guard just yields a clean error instead of an undefined symbol.

[greptile-apps]

Greptile Summary

This PR adds an NVFP4 per-token quantization recipe for model pre-training. The default NVFP4BlockScaling recipe computes a single per-tensor outer scale (s_global), while this variant computes per-row outer amaxes (length M) for row-wise data and per-col amaxes (length K) for column-wise data.

• Adds per-token cast kernels (K1 vector amax + K2 encode/swizzle) and a matching CUTLASS fused-EVT GEMM (nvfp4_cutlass_per_token_gemm) that reads per-row/col outer-amax vectors directly in its epilogue, covering TN (fwd), NN (dgrad), and NT (wgrad) layouts with grouped-GEMM support.
• Introduces NVFP4PerTokenBlockScaling as an explicit recipe subclass and adds an NVTE_NVFP4_PER_TOKEN=1 env-var switch on the base recipe so Megatron-Core can opt in without code changes.
• Documents known limitations: requires NVTE_NORM_FWD_USE_CUDNN=1, no fuse_wgrad_accumulation, no comm-overlap/output quantization.

Confidence Score: 4/5

Safe to merge after addressing the columnwise default inconsistency in the diagnostic split functions.

The PR is large but well-structured. The only new finding is a P1 inconsistency in default argument values for the diagnostic split functions nvfp4_per_token_amax and nvfp4_per_token_encode (columnwise=True vs composite's columnwise=False), which can cause a confusing ValueError for callers using the split API with defaults. Core production paths (composite quantize, CUTLASS GEMM dispatch, recipe routing) look correct.

transformer_engine/pytorch/custom_recipes/quantization_nvfp4_per_token.py — columnwise default mismatch between split and composite APIs.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/custom_recipes/quantization_nvfp4_per_token.py	New file: pure-Python reference quantizer + production wrappers for the per-token cast. Contains a default columnwise=False in nvfp4_per_token_quantize that diverges from columnwise=True defaults in the diagnostic nvfp4_per_token_amax/nvfp4_per_token_encode helpers.
transformer_engine/pytorch/cpp_extensions/gemm.py	Large addition: _nvfp4_per_token_gemm and _nvfp4_per_token_grouped_gemm dispatch helpers plus routing logic in general_gemm/general_grouped_gemm. Layout dispatch (TN/NN/NT), N-D output shape restoration, accumulate/bias/gelu guards, and empty-expert zeroing are all well-handled.
transformer_engine/pytorch/csrc/quantizer.cpp	Extends NVFP4Quantizer with per_token and per_token_weight_2d fields; with_gemm_swizzled_scales is correctly hardcoded false with a TODO for the future optimize_for_gemm path.
transformer_engine/common/recipe/init.py	Adds NVFP4PerTokenBlockScaling subclass and env-var-driven per_token flags; env vars are read at class-definition time for instance fields but at call time for nvfp4_per_token().
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Adds _per_token bool attribute and mismatch guard in copy_from_storage; straightforward and correct.
transformer_engine/pytorch/csrc/extensions/nvfp4_per_token.cpp	New C++ bindings for K1/K2/composite per-token quantize ops; assemble_per_token_tensors validates dtypes and shapes; no critical issues found.
transformer_engine/pytorch/custom_recipes/gemm_nvfp4_per_token.py	Reference dequant and production GEMM wrappers look correct; beta!=0 guard prevents accidental accumulation.
transformer_engine/pytorch/custom_recipes/quantization_nvfp4_per_token_group.py	Grouped (MoE) per-token quantize helper; validates split_sections[i]>0 appropriately.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Threads per_token/per_token_weight_2d flags through NVFP4Tensor lifecycle (new, FSDP2, reduce_ex, reshape); _with_gemm_swizzled_scales hardcoded False per TODO comment.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Recipe as NVFP4PerTokenBlockScaling
    participant Quant as NVFP4Quantizer (Python)
    participant QCpp as NVFP4Quantizer (C++)
    participant K1 as nvte_nvfp4_per_token_amax (K1 kernel)
    participant K2 as nvte_nvfp4_per_token_encode (K2 kernel)
    participant Gemm as _nvfp4_per_token_gemm
    participant CUTLASS as tex.nvfp4_cutlass_per_token_gemm

    Recipe->>Quant: "nvfp4_per_token()=True, per_token_weight_2d"
    Quant->>QCpp: create_tensor(M,K,rowwise/columnwise)
    QCpp-->>Quant: NVFP4Tensor + row_amax(M,) or col_amax(K,)

    Note over Quant,K2: Composite quantize path (production)
    Quant->>K1: "nvte_nvfp4_per_token_quantize(input, with_swizzle=0)"
    K1-->>Quant: inner_sf (E4M3), row_amax / col_amax vector
    Quant->>K2: nvte_nvfp4_per_token_encode(inner_sf, row_amax)
    K2-->>Quant: NVFP4 data + outer_sf

    Note over Gemm,CUTLASS: TN forward GEMM (A rowwise, B columnwise)
    Gemm->>Gemm: _nvfp4_per_token_select(A, rowwise)
    Gemm->>Gemm: _nvfp4_per_token_select(B, columnwise)
    Gemm->>CUTLASS: nvfp4_cutlass_per_token_gemm(A_data, A_sf, A_row_amax, B_data, B_sf, B_col_amax)
    CUTLASS-->>Gemm: output (M, N)
    Gemm-->>Gemm: reshape to N-D leading dims + optional bias add

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Recipe as NVFP4PerTokenBlockScaling
    participant Quant as NVFP4Quantizer (Python)
    participant QCpp as NVFP4Quantizer (C++)
    participant K1 as nvte_nvfp4_per_token_amax (K1 kernel)
    participant K2 as nvte_nvfp4_per_token_encode (K2 kernel)
    participant Gemm as _nvfp4_per_token_gemm
    participant CUTLASS as tex.nvfp4_cutlass_per_token_gemm

    Recipe->>Quant: "nvfp4_per_token()=True, per_token_weight_2d"
    Quant->>QCpp: create_tensor(M,K,rowwise/columnwise)
    QCpp-->>Quant: NVFP4Tensor + row_amax(M,) or col_amax(K,)

    Note over Quant,K2: Composite quantize path (production)
    Quant->>K1: "nvte_nvfp4_per_token_quantize(input, with_swizzle=0)"
    K1-->>Quant: inner_sf (E4M3), row_amax / col_amax vector
    Quant->>K2: nvte_nvfp4_per_token_encode(inner_sf, row_amax)
    K2-->>Quant: NVFP4 data + outer_sf

    Note over Gemm,CUTLASS: TN forward GEMM (A rowwise, B columnwise)
    Gemm->>Gemm: _nvfp4_per_token_select(A, rowwise)
    Gemm->>Gemm: _nvfp4_per_token_select(B, columnwise)
    Gemm->>CUTLASS: nvfp4_cutlass_per_token_gemm(A_data, A_sf, A_row_amax, B_data, B_sf, B_col_amax)
    CUTLASS-->>Gemm: output (M, N)
    Gemm-->>Gemm: reshape to N-D leading dims + optional bias add

Loading

_{Reviews (4): Last reviewed commit: "Batch per-group metadata H2D in NVFP4 gr..." | Re-trigger Greptile}

[greptile-apps]

alpha scalar silently ignored for per-token GEMM

general_gemm validates and stores alpha in kwargs["alpha"], but the per-token short-circuit path dispatches to _nvfp4_per_token_gemm which has no alpha parameter and never forwards the value. The C++ binding nvfp4_cutlass_per_token_gemm also lacks a global scalar alpha argument — only the per-row/per-col alpha_a/alpha_b vectors are supported. For all current TE module call sites alpha=1.0 is the invariant, so numerical output is unaffected today. If a caller ever passes alpha != 1.0 through general_gemm with per-token tensors, the result will be silently wrong instead of raising an error.

[greptile-apps]

Public grouped-quantize API unconditionally rejects 0-token splits

split_sections[i] <= 0 raises ValueError, but in MoE training with dynamic token routing, experts commonly receive zero tokens in a given micro-batch. The general_grouped_gemm per-token loop already handles this by skipping the launch when m_splits[i] == 0, so the GEMM side is fine. If users call this Python wrapper directly (e.g., from bench scripts or custom MoE quantization pipelines), they must pre-filter empty experts. A comment or guard skipping allocation for empty splits would make the API usable in unbalanced-routing scenarios.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!