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.

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 (NVFP4PerTokenBlockScaling) for model pre-training. Instead of a single per-tensor global amax, the recipe computes a per-row outer amax (M,) for rowwise data and a per-column outer amax (K,) for columnwise data, enabling finer-grained scaling per token.

• New cast kernels (quantize_nvfp4_per_token.cu, quantize_nvfp4_per_token_group.cu): K1 (vector amax) + K2 (FP4 encode/swizzle) producing per-token NVFP4 tensors for both single and grouped (MoE) inputs.
• Fused CUTLASS GEMM (nvfp4_cutlass_gemm.cu, nvfp4_cutlass_grouped_gemm.cu): nvfp4_cutlass_per_token_gemm consumes per-row/per-col amax vectors in the epilogue (EVT), covering TN/NN/NT layouts for fwd/dgrad/wgrad.
• Recipe and Python integration: NVFP4PerTokenBlockScaling dataclass plus NVTE_NVFP4_PER_TOKEN=1 env-var opt-in on the base recipe, quantizer flag propagation through NVFP4TensorStorage._per_token, and automatic general_gemm routing to the fused EVT kernel.

Confidence Score: 5/5

The new per-token recipe is additive (opt-in via explicit recipe class or env var) and does not modify the default NVFP4BlockScaling code path, so existing users are unaffected.

The three previously flagged issues (null-ptr restore in weight-2D path, silent alpha=1 assumption in per-token GEMM, 0-token split rejection) remain open but are either gated behind an unreachable code path today or are documented limitations. No new blocking defects were found. The dispatch table, amax-vector allocation, FSDP2 metadata propagation, and grouped-GEMM empty-expert handling are all correct.

transformer_engine/pytorch/csrc/quantizer.cpp — the weight-2D amax restore after nvte_compute_amax_with_config uses rowwise_amax_ptr without a null guard; transformer_engine/pytorch/cpp_extensions/gemm.py — the grouped GEMM bias+accumulate guard uses assert.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/cpp_extensions/gemm.py	Adds per-token GEMM dispatch via _nvfp4_per_token_gemm and _nvfp4_per_token_grouped_gemm; return-type annotation for single_output=True is incorrect and a bias+accumulate mutual-exclusion guard uses assert (disableable with -O).
transformer_engine/pytorch/csrc/quantizer.cpp	Per-token and per-token weight-2D quantize_impl paths added; previously-flagged latent null-pointer restore (rowwise_amax_ptr after nvte_compute_amax_with_config) remains; new per-token and split-quantize dispatches are otherwise correct.
transformer_engine/common/recipe/init.py	Adds NVFP4PerTokenBlockScaling subclass and nvfp4_per_token() classmethod; mutex knobs are correctly forced in __post_init__ and guarded again in get_quantizer(), so env-var activation after construction is safe.
transformer_engine/pytorch/quantization.py	Per-token quantizer factory correctly gates all mutex features (RHT, SR, 2D-quant, 4over6, row-scaled) with and not per_token guards, making env-var opt-in safe even when recipe was constructed with default settings.
transformer_engine/pytorch/custom_recipes/quantization_nvfp4_per_token_group.py	Grouped per-token quantize wrapper; split_sections[i] <= 0 raises unconditionally (previously flagged), making this unusable with zero-token experts from dynamic routing without pre-filtering.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Propagates per_token flag through FSDP2 shard/unshard metadata, __reduce_ex__, _ViewFunc, and _ReshapeFunc consistently; FSDP2 fast-fail guard for per-token mode is correctly placed.
transformer_engine/pytorch/csrc/extensions/nvfp4_cutlass_gemm.cpp	Per-token CUTLASS GEMM C++ binding with thorough shape/dtype/contiguity checks; SF-swizzle handled internally or skipped when pre-swizzled; accumulate path correctly restricted to fp32 output.
transformer_engine/pytorch/csrc/extensions/nvfp4_per_token.cpp	K1/K2/composite quantize C++ bindings; shape/dtype/alignment checks are thorough and match the kernel constraints; grouped bulk-allocation path correctly handles per-token amax vector sizing.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Recipe as NVFP4PerTokenBlockScaling
    participant RecipeState as NVFP4BlockScalingRecipeState
    participant Quantizer as NVFP4Quantizer (Python)
    participant QuantizerCPP as NVFP4Quantizer (C++)
    participant KernelK1 as K1: vector amax kernel
    participant KernelK2 as K2: FP4 encode/swizzle kernel
    participant GEMM as nvfp4_cutlass_per_token_gemm (EVT)

    Recipe->>RecipeState: get_quantizer(tensor_type, mode)
    RecipeState->>RecipeState: "per_token = nvfp4_per_token() and tensor_type in scope"
    RecipeState->>Quantizer: "NVFP4Quantizer(per_token=True, ...)"
    Quantizer->>QuantizerCPP: quantize_impl(input, out)

    alt per_token_weight_2d (weight only)
        QuantizerCPP->>KernelK1: nvte_compute_amax_with_config (scalar global amax)
        QuantizerCPP->>KernelK2: nvte_quantize_v2 (2D 16x16 cast)
        QuantizerCPP->>QuantizerCPP: broadcast scalar amax to (M,) rowwise + (K,) colwise vectors
    else standard per_token
        QuantizerCPP->>KernelK1: nvte_nvfp4_per_token_quantize (K1+K2 composite)
        KernelK1-->>QuantizerCPP: amax_rowwise (M,) + amax_columnwise (K,) + FP4 data
    end

    Note over QuantizerCPP: out._per_token = True
    QuantizerCPP-->>Quantizer: NVFP4Tensor with per-token amaxes

    Note over RecipeState,GEMM: general_gemm dispatch
    RecipeState->>GEMM: "_is_nvfp4_per_token_tensor(A/B) = True"
    GEMM->>GEMM: select rowwise/columnwise view per TN/NN/NT layout
    GEMM->>GEMM: tex.nvfp4_cutlass_per_token_gemm(ka_data, kb_data, ka_sf, kb_sf, ka_amax, kb_amax, out)
    Note over GEMM: EVT epilogue: D[i,j] = bf16(alpha_a[i] * alpha_b[j] * (A@B^T)[i,j])
    GEMM-->>RecipeState: bf16 / fp32 output tensor

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 RecipeState as NVFP4BlockScalingRecipeState
    participant Quantizer as NVFP4Quantizer (Python)
    participant QuantizerCPP as NVFP4Quantizer (C++)
    participant KernelK1 as K1: vector amax kernel
    participant KernelK2 as K2: FP4 encode/swizzle kernel
    participant GEMM as nvfp4_cutlass_per_token_gemm (EVT)

    Recipe->>RecipeState: get_quantizer(tensor_type, mode)
    RecipeState->>RecipeState: "per_token = nvfp4_per_token() and tensor_type in scope"
    RecipeState->>Quantizer: "NVFP4Quantizer(per_token=True, ...)"
    Quantizer->>QuantizerCPP: quantize_impl(input, out)

    alt per_token_weight_2d (weight only)
        QuantizerCPP->>KernelK1: nvte_compute_amax_with_config (scalar global amax)
        QuantizerCPP->>KernelK2: nvte_quantize_v2 (2D 16x16 cast)
        QuantizerCPP->>QuantizerCPP: broadcast scalar amax to (M,) rowwise + (K,) colwise vectors
    else standard per_token
        QuantizerCPP->>KernelK1: nvte_nvfp4_per_token_quantize (K1+K2 composite)
        KernelK1-->>QuantizerCPP: amax_rowwise (M,) + amax_columnwise (K,) + FP4 data
    end

    Note over QuantizerCPP: out._per_token = True
    QuantizerCPP-->>Quantizer: NVFP4Tensor with per-token amaxes

    Note over RecipeState,GEMM: general_gemm dispatch
    RecipeState->>GEMM: "_is_nvfp4_per_token_tensor(A/B) = True"
    GEMM->>GEMM: select rowwise/columnwise view per TN/NN/NT layout
    GEMM->>GEMM: tex.nvfp4_cutlass_per_token_gemm(ka_data, kb_data, ka_sf, kb_sf, ka_amax, kb_amax, out)
    Note over GEMM: EVT epilogue: D[i,j] = bf16(alpha_a[i] * alpha_b[j] * (A@B^T)[i,j])
    GEMM-->>RecipeState: bf16 / fp32 output tensor

Loading

_{Reviews (7): Last reviewed commit: "Apply pre-commit formatting for NVFP4 pe..." | 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!