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.

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 a new NVFP4PerTokenBlockScaling recipe that replaces the per-tensor outer amax with a per-row (M,) / per-col (K,) vector, giving each token its own global scale. It introduces per-token CUDA cast kernels (K1 vector-amax + K2 encode/swizzle), a fused CUTLASS EVT GEMM (nvfp4_cutlass_per_token_gemm) that rescales with those vectors in the epilogue, and an env-var-based activation path (NVTE_NVFP4_PER_TOKEN=1) so frameworks that only construct a plain NVFP4BlockScaling (e.g. Megatron-Core) can opt in without code changes. As noted in the PR description, this targets accuracy evaluation rather than optimized production deployment.

• Adds per-token cast kernels, a fused-EVT CUTLASS GEMM (single + grouped), and a NVFP4PerTokenBlockScaling recipe class (re-exported from transformer_engine.pytorch.fp8) covering forward + backward (dgrad NN / wgrad NT layouts).
• Extends NVFP4Quantizer in Python and C++ with per_token / per_token_weight_2d flags; general_gemm / general_grouped_gemm auto-dispatch to the per-token CUTLASS GEMM based on the _per_token tensor flag.
• Documents all known limitations (unfused norm required, no wgrad-accumulation fusion, no comm-overlap/CUDA graphs) and adds an end-to-end Megatron-Core MoE example.

Confidence Score: 4/5

Safe to merge for evaluation/research use as described in the PR; the feature is explicitly gated as non-production and the known unsupported paths (fuse_wgrad_accumulation, comm-overlap, CUDA graphs) all raise clear runtime errors rather than silently corrupting results.

The implementation is large but well-structured, with comprehensive guards in both Python and C++. Fresh findings are style/robustness observations. The return-type divergence in _nvfp4_per_token_grouped_gemm (tensor vs list for single_output=True) is harmless today because all call sites discard the return value and rely on in-place writes, but could trap a future caller. The env-var recipe-state divergence path produces misleading logs in an unusual (mid-run env-var change) scenario. Both are minor concerns for a feature intentionally scoped to accuracy evaluation.

transformer_engine/pytorch/cpp_extensions/gemm.py (return-type consistency in grouped per-token path) and transformer_engine/common/recipe/init.py (env-var re-evaluation in nvfp4_per_token classmethod).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/cpp_extensions/gemm.py	Adds _nvfp4_per_token_gemm and _nvfp4_per_token_grouped_gemm helpers plus dispatch logic in general_gemm / general_grouped_gemm; well-guarded but the grouped helper's return type differs from the non-per-token path when single_output=True (returns tensor vs list), though current callers ignore the return value so this is benign.
transformer_engine/pytorch/csrc/quantizer.cpp	Adds per_token / per_token_weight_2d branches to NVFP4Quantizer: create_tensor, convert_and_update_tensor, and quantize_impl. The per_token_weight_2d amax-set-before-quantize ordering is correct for the reachable (rowwise != nullptr) code path; the latent null-ptr edge case remains (previously flagged).
transformer_engine/common/recipe/init.py	Adds NVFP4PerTokenBlockScaling subclass and env-var activation path (NVTE_NVFP4_PER_TOKEN=1) on the base class; nvfp4_per_token() classmethod + _force_per_token_settings() keep the two activation paths consistent.
transformer_engine/pytorch/quantization.py	Correctly computes per_token / per_token_weight_2d flags per tensor_type and mode; per-token feature overrides (rht, SR, 2d, 4over6) are properly gated in the NVFP4Quantizer constructor call.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Propagates per_token flag through NVFP4Tensor lifecycle (copy, reshape, view, FSDP2 metadata, reduce_ex); per_token and per_token_weight_2d field docstrings are string literals placed after the annotation (not before), so they are discarded by Python and won't appear in help() output.
transformer_engine/pytorch/custom_recipes/quantization_nvfp4_per_token_group.py	Public grouped-quantize entry with validation; rejects split_sections[i] <= 0 which prevents use with zero-token experts (previously flagged). Shape checks and 64-group cap are correct.
transformer_engine/pytorch/custom_recipes/quantization_nvfp4_per_token.py	Pure-PyTorch reference quantizer plus production composite/K1-only/K2-only Python wrappers; arithmetic matches the kernel (row amax floor, S_enc computation, FP8 saturating cast) and validation matches kernel constraints.
transformer_engine/pytorch/custom_recipes/gemm_nvfp4_per_token.py	Stand-alone production GEMM wrapper with clear TN/NN/NT layout dispatch table and good shape-mismatch error messages.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Adds _per_token class annotation, constructor parameter, copy_from_storage mismatch check, and __init_kwargs propagation; straightforward and complete.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["RecipeState.get_quantizer()"] --> B{nvfp4_per_token?}
    B -- "False\n(default NVFP4BlockScaling)" --> C["per-tensor path\n(cuBLASLt 1A/2A)"]
    B -- "True\n(NVFP4PerTokenBlockScaling\nOR NVTE_NVFP4_PER_TOKEN=1)" --> D{per_token_weight_2d?}
    D -- "False (default)" --> E["NVFP4Quantizer\n(per_token=True)"]
    D -- "True\n(NVTE_NVFP4_PER_TOKEN_WEIGHT_2D=1)" --> F["NVFP4Quantizer\n(per_token=True,\nper_token_weight_2d=True)"]
    E --> G["quantize_impl:\nnvte_nvfp4_per_token_quantize\n(K1 row/col amax + K2 FP4 encode)"]
    F --> H["quantize_impl:\nnvte_compute_amax → nvte_quantize_v2\n(2D cast) → broadcast amax vector"]
    G --> I["NVFP4Tensor\n_amax_rowwise shape: (M,)\n_amax_columnwise shape: (K,)"]
    H --> I
    I --> J["general_gemm / general_grouped_gemm:\n_is_nvfp4_per_token_tensor check"]
    J -- "True" --> K["_nvfp4_per_token_gemm\nor\n_nvfp4_per_token_grouped_gemm\n(CUTLASS EVT fused epilogue)"]
    J -- "False" --> L["tex.generic_gemm\n(cuBLASLt)"]
    K --> M["Output: BF16\nD = alpha_row x alpha_col · (A @ B^T)"]
    L --> N["Output: BF16/FP8\n(standard path)"]

Loading

_{Reviews (2): Last reviewed commit: "Add native CUTLASS grouped per-token NVF..." | Re-trigger Greptile}

[greptile-apps]

Latent null-pointer dereference in per-token weight-2D amax restore

After nvte_compute_amax_with_config writes the global amax into amax_ptr[0], the code restores the output tensor's amax field with out.set_amax(rowwise_amax_ptr, ...). If rowwise_amax_ptr is nullptr (i.e., the quantizer was constructed with rowwise=False), this sets the output's amax descriptor to a null pointer. The immediately following nvte_quantize_v2 then tries to read amax[0] to derive S_enc and will crash.

Currently this path is unreachable because per_token_weight_2d is only set for weight quantizers, and all weight quantizers in the recipe are constructed with rowwise=True, columnwise=True. However, the guard in step 3 (if (rowwise_amax_ptr != nullptr && w2d_rows > 1)) shows the author anticipated both pointers could be null, while the critical out.set_amax call on this line does not. Using amax_ptr (the non-null pointer already validated by the NVTE_CHECK above) would be safe in all configurations: out.set_amax(amax_ptr, DType::kFloat32, std::vector<size_t>{1}).

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