diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index bf6b6810b5d..069138a1b99 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -26,6 +26,9 @@ from megatron.core.extensions.transformer_engine_int4_fake_qat import ( maybe_fake_quantize_int4_weight_tensors, ) +from megatron.core.extensions.transformer_engine_nvfp4_fake_qat import ( + maybe_fake_quantize_nvfp4_weight_tensors, +) from megatron.core.model_parallel_config import ModelParallelConfig from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.parallel_state import ( @@ -2357,6 +2360,9 @@ def forward(self, x, m_splits): def _get_weight_tensors(self): """Get the weight tensors of the module.""" weight_tensors = super()._get_weight_tensors() + weight_tensors = maybe_fake_quantize_nvfp4_weight_tensors( + self.config, self.delay_wgrad_compute, weight_tensors + ) return maybe_fake_quantize_int4_weight_tensors( self.config, self.delay_wgrad_compute, weight_tensors ) diff --git a/megatron/core/extensions/transformer_engine_nvfp4_fake_qat.py b/megatron/core/extensions/transformer_engine_nvfp4_fake_qat.py new file mode 100644 index 00000000000..75affd719e6 --- /dev/null +++ b/megatron/core/extensions/transformer_engine_nvfp4_fake_qat.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Transformer Engine integration for fused NVFP4 fake QAT.""" + +import os + +import torch + +from megatron.core.extensions.transformer_engine_int4_fake_qat import INT4_FAKE_QAT_FLAG +from megatron.core.model_parallel_config import ModelParallelConfig +from megatron.core.utils import is_te_min_version + +NVFP4_FAKE_QAT_FLAG = "OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG" +_MIN_TE_VERSION = "2.17.0" + + +def _validate_nvfp4_fake_qat_support( + config: ModelParallelConfig, delay_wgrad_compute: bool, weight_tensors: list[torch.Tensor] +) -> None: + """Reject TE weight layouts and gradient paths that cannot safely use STE tensors.""" + if not is_te_min_version(_MIN_TE_VERSION): + raise RuntimeError( + f"{NVFP4_FAKE_QAT_FLAG}=1 requires Transformer Engine >= {_MIN_TE_VERSION}." + ) + if os.getenv(INT4_FAKE_QAT_FLAG, "0") == "1": + raise RuntimeError( + f"{INT4_FAKE_QAT_FLAG}=1 and {NVFP4_FAKE_QAT_FLAG}=1 are mutually exclusive." + ) + if config.gradient_accumulation_fusion: + raise RuntimeError( + f"{NVFP4_FAKE_QAT_FLAG}=1 is not supported with " + "gradient_accumulation_fusion because TE fused wgrad accumulation mutates " + "Python attributes on the original weight tensors." + ) + if delay_wgrad_compute: + raise RuntimeError( + f"{NVFP4_FAKE_QAT_FLAG}=1 is not supported with delayed wgrad compute because " + "the delayed TE path mutates Python attributes on the original weight tensors." + ) + if getattr(config, "moe_single_grouped_weight", False): + raise RuntimeError( + f"{NVFP4_FAKE_QAT_FLAG}=1 requires TE's discrete grouped-linear parameters; " + "moe_single_grouped_weight is not supported." + ) + if any( + hasattr(weight, "__fsdp_param__") or hasattr(weight, "get_main_grad") + for weight in weight_tensors + ): + raise RuntimeError( + f"{NVFP4_FAKE_QAT_FLAG}=1 is not supported with Megatron FSDP because " + "FSDP patches weight tensors with main-gradient attributes and methods." + ) + if any(getattr(weight, "ndim", None) != 2 for weight in weight_tensors): + raise RuntimeError( + f"{NVFP4_FAKE_QAT_FLAG}=1 requires one rank-2 tensor per grouped-linear GEMM." + ) + + +def maybe_fake_quantize_nvfp4_weight_tensors( + config: ModelParallelConfig, delay_wgrad_compute: bool, weight_tensors: list[torch.Tensor] +) -> list[torch.Tensor]: + """Optionally apply fused NVFP4 fake QAT to discrete TE grouped-linear weights.""" + if os.getenv(NVFP4_FAKE_QAT_FLAG, "0") != "1": + return weight_tensors + + _validate_nvfp4_fake_qat_support(config, delay_wgrad_compute, weight_tensors) + + # Keep CuTe DSL optional for every Megatron process that does not enable this path. + from megatron.core.fusions.fused_nvfp4_qdq import ( + current_nvfp4_qdq_config, + fake_nvfp4_quantization_ste, + ) + + qdq_config = current_nvfp4_qdq_config() + return [fake_nvfp4_quantization_ste(weight, qdq_config) for weight in weight_tensors] + + +__all__ = ["NVFP4_FAKE_QAT_FLAG", "maybe_fake_quantize_nvfp4_weight_tensors"] diff --git a/megatron/core/fusions/fused_nvfp4_qdq.py b/megatron/core/fusions/fused_nvfp4_qdq.py new file mode 100644 index 00000000000..0f95c65306d --- /dev/null +++ b/megatron/core/fusions/fused_nvfp4_qdq.py @@ -0,0 +1,1293 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Fused CuTe DSL NVFP4 quantize-dequantize for fake QAT. + +The kernel keeps the E4M3 block scale and packed E2M1 values in registers and +writes only the dequantized BF16/FP16 result. Its arithmetic order mirrors +Transformer Engine's 1D, 1x16, per-tensor NVFP4 implementation. The vectorized +load, FP4 conversion, and Four Over Six structure are adapted from FlashInfer's +CuTe DSL NVFP4 quantizer. + +Supported contract: + +* contiguous rank-2 BF16 or FP16 input on SM10x; +* 1x16 block scaling and a caller-provided FP32 per-tensor amax; +* round-to-nearest quantization with ordinary quant fast math disabled; +* standard NVFP4, plus the full Four Over Six MAE/MSE, E4M3-max 256/448, + and exact/FP16-error matrix; +* no stochastic rounding, RHT, 2D quantization, transpose, or row scaling. +""" + +from __future__ import annotations + +import functools +import os +from dataclasses import dataclass +from enum import IntEnum +from typing import Any, Optional + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T, dsl_user_op + +_FP32_MAX = 3.4028234663852886e38 +_FP4_BLOCK_SIZE = 16 +_STANDARD_THREADS = 256 +# Preserve the 4-CTA compile launch bound while using a deeper runtime grid to +# reduce each thread's grid-stride work on model-sized tensors. +_STANDARD_MIN_BLOCKS_PER_SM = 4 +_STANDARD_GRID_BLOCKS_PER_SM = 24 +_4OVER6_THREADS = 128 +# The largest specialization uses 56 registers/thread, so 8x128 threads stays +# below the SM10x 64K-register budget without spills while doubling active CTAs. +_4OVER6_BLOCKS_PER_SM = 8 +_INT32_MAX = 2**31 - 1 + + +class NVFP4QDQErrorMode(IntEnum): + """Four Over Six candidate error metric.""" + + MAE = 0 + MSE = 1 + + +@dataclass(frozen=True) +class NVFP4QDQConfig: + """Compile-time numerical configuration for fused NVFP4 QDQ.""" + + use_4over6: bool = False + e4m3_max: int = 448 + error_mode: NVFP4QDQErrorMode = NVFP4QDQErrorMode.MAE + error_use_fp16: bool = False + + def __post_init__(self) -> None: + if self.e4m3_max not in (256, 448): + raise ValueError(f"NVFP4 E4M3 max must be 256 or 448, got {self.e4m3_max}.") + if not self.use_4over6 and self.e4m3_max != 448: + raise ValueError("E4M3 max 256 is only supported by Four Over Six.") + if not self.use_4over6 and self.error_use_fp16: + raise ValueError("The FP16 error contract only applies to Four Over Six.") + + +def _env_flag(name: str, default: str = "0") -> bool: + value = os.getenv(name, default).strip().lower() + if value in ("1", "true", "yes", "on"): + return True + if value in ("0", "false", "no", "off", ""): + return False + raise ValueError(f"{name} must be a boolean value, got {value!r}.") + + +def _env_applies_to_weights(name: str, default: str) -> bool: + value = os.getenv(name, default).strip().lower() + if value not in ("none", "weights", "activations", "all"): + raise ValueError( + f"{name} must be one of none, weights, activations, or all; got {value!r}." + ) + return value in ("weights", "all") + + +def current_nvfp4_qdq_config() -> NVFP4QDQConfig: + """Resolve the weight QDQ contract from Transformer Engine environment variables.""" + if _env_flag("NVTE_USE_FAST_MATH"): + raise ValueError( + "Fused NVFP4 QDQ requires NVTE_USE_FAST_MATH=0; ordinary quant fast math " + "is outside its numerical contract." + ) + + use_4over6 = _env_applies_to_weights("NVTE_NVFP4_4OVER6", "none") + use_e4m3_256 = _env_applies_to_weights("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + error_mode_name = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").strip().upper() + try: + error_mode = NVFP4QDQErrorMode[error_mode_name] + except KeyError as exc: + raise ValueError( + f"NVTE_NVFP4_4OVER6_ERR_MODE must be MAE or MSE, got {error_mode_name!r}." + ) from exc + + return NVFP4QDQConfig( + use_4over6=use_4over6, + e4m3_max=256 if use_4over6 and use_e4m3_256 else 448, + error_mode=error_mode, + # Despite the legacy variable name, this selects TE's FP16 candidate-error contract. + error_use_fp16=(use_4over6 and _env_flag("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH")), + ) + + +@dsl_user_op +def _get_ptr(tensor: cute.Tensor, offset: Int32, *, loc=None, ip=None) -> Int64: + elem_ptr = tensor.iterator + offset + return Int64(llvm.ptrtoint(T.i64(), elem_ptr.llvm_ptr, loc=loc, ip=ip)) + + +@dsl_user_op +def _load_v4_u32(base_ptr: Int64, *, loc=None, ip=None) -> tuple[Uint32, Uint32, Uint32, Uint32]: + result = llvm.inline_asm( + llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]), + [Int64(base_ptr).ir_value(loc=loc, ip=ip)], + "ld.global.v4.u32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Uint32(llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)), + Uint32(llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)), + Uint32(llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)), + Uint32(llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _store_v4_u32( + base_ptr: Int64, v0: Uint32, v1: Uint32, v2: Uint32, v3: Uint32, *, loc=None, ip=None +) -> None: + llvm.inline_asm( + None, + [ + Int64(base_ptr).ir_value(loc=loc, ip=ip), + Uint32(v0).ir_value(loc=loc, ip=ip), + Uint32(v1).ir_value(loc=loc, ip=ip), + Uint32(v2).ir_value(loc=loc, ip=ip), + Uint32(v3).ir_value(loc=loc, ip=ip), + ], + "st.global.v4.u32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _fadd_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "add.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fsub_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "sub.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fmul_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _scale_f32x2_rn( + value0: Float32, value1: Float32, scale: Float32, *, loc=None, ip=None +) -> tuple[Float32, Float32]: + """Multiply two FP32 values by one common scale with packed RN arithmetic.""" + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [ + Float32(value0).ir_value(loc=loc, ip=ip), + Float32(value1).ir_value(loc=loc, ip=ip), + Float32(scale).ir_value(loc=loc, ip=ip), + ], + """ + { + .reg .b64 values, scales; + mov.b64 values, {$2, $3}; + mov.b64 scales, {$4, $4}; + mul.f32x2 values, values, scales; + mov.b64 {$0, $1}, values; + } + """, + "=f,=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _square_f32x2_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + """Square two FP32 values with the packed RN instruction used by TE.""" + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b64 values; + mov.b64 values, {$2, $3}; + mul.f32x2 values, values, values; + mov.b64 {$0, $1}, values; + } + """, + "=f,=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _fdiv_rn(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "div.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fmin(a: Float32, b: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + "min.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _fabs(a: Float32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(a).ir_value(loc=loc, ip=ip)], + "abs.f32 $0, $1;", + "=f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _half2_abs(x: Uint32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [Uint32(x).ir_value(loc=loc, ip=ip)], + "and.b32 $0, $1, 0x7FFF7FFF;", + "=r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _half2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)], + "max.f16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bfloat2_max(a: Uint32, b: Uint32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [Uint32(a).ir_value(loc=loc, ip=ip), Uint32(b).ir_value(loc=loc, ip=ip)], + "max.bf16x2 $0, $1, $2;", + "=r,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _half2_max_to_f32(x: Uint32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Uint32(x).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 h0, h1; + .reg .f32 f0, f1; + mov.b32 {h0, h1}, $1; + cvt.f32.f16 f0, h0; + cvt.f32.f16 f1, h1; + max.f32 $0, f0, f1; + } + """, + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _bfloat2_max_to_f32(x: Uint32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Uint32(x).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b32 lo, hi; + .reg .f32 f0, f1; + and.b32 lo, $1, 0xFFFF; + shr.b32 hi, $1, 16; + shl.b32 lo, lo, 16; + shl.b32 hi, hi, 16; + mov.b32 f0, lo; + mov.b32 f1, hi; + max.f32 $0, f0, f1; + } + """, + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _half2_to_f32x2(x: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(x).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 lo, hi; + mov.b32 {lo, hi}, $2; + cvt.f32.f16 $0, lo; + cvt.f32.f16 $1, hi; + } + """, + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _bfloat2_to_f32x2(x: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(x).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b32 lo, hi; + and.b32 lo, $2, 0xFFFF; + shr.b32 hi, $2, 16; + shl.b32 lo, lo, 16; + shl.b32 hi, hi, 16; + mov.b32 $0, lo; + mov.b32 $1, hi; + } + """, + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _cvt_f32_to_e4m3(a: Float32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [Float32(a).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 fp8_pair; + .reg .f32 zero; + mov.f32 zero, 0f00000000; + cvt.rn.satfinite.e4m3x2.f32 fp8_pair, zero, $1; + cvt.u32.u16 $0, fp8_pair; + } + """, + "=r,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _cvt_e4m3_to_f32(a: Uint32, *, loc=None, ip=None) -> Float32: + return Float32( + llvm.inline_asm( + T.f32(), + [Uint32(a).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 fp8_pair; + .reg .b32 h2; + .reg .b16 lo, hi; + cvt.u16.u32 fp8_pair, $1; + cvt.rn.f16x2.e4m3x2 h2, fp8_pair; + mov.b32 {lo, hi}, h2; + cvt.f32.f16 $0, lo; + } + """, + "=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _cvt_e2m1x8( + v0: Float32, + v1: Float32, + v2: Float32, + v3: Float32, + v4: Float32, + v5: Float32, + v6: Float32, + v7: Float32, + *, + loc=None, + ip=None, +) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [ + Float32(v0).ir_value(loc=loc, ip=ip), + Float32(v1).ir_value(loc=loc, ip=ip), + Float32(v2).ir_value(loc=loc, ip=ip), + Float32(v3).ir_value(loc=loc, ip=ip), + Float32(v4).ir_value(loc=loc, ip=ip), + Float32(v5).ir_value(loc=loc, ip=ip), + Float32(v6).ir_value(loc=loc, ip=ip), + Float32(v7).ir_value(loc=loc, ip=ip), + ], + """ + { + .reg .b8 b0, b1, b2, b3; + cvt.rn.satfinite.e2m1x2.f32 b0, $2, $1; + cvt.rn.satfinite.e2m1x2.f32 b1, $4, $3; + cvt.rn.satfinite.e2m1x2.f32 b2, $6, $5; + cvt.rn.satfinite.e2m1x2.f32 b3, $8, $7; + mov.b32 $0, {b0, b1, b2, b3}; + } + """, + "=r,f,f,f,f,f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _e2m1x2_to_f32x2(a: Uint32, *, loc=None, ip=None) -> tuple[Float32, Float32]: + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(a).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b8 byte0, byte1, byte2, byte3; + .reg .b32 h2; + .reg .b16 lo, hi; + .reg .b32 code_lo, code_hi, bits_lo, bits_hi; + .reg .f32 f_lo, f_hi; + .reg .pred negzero_lo, negzero_hi; + + mov.b32 {byte0, byte1, byte2, byte3}, $2; + cvt.rn.f16x2.e2m1x2 h2, byte0; + mov.b32 {lo, hi}, h2; + cvt.f32.f16 f_lo, lo; + cvt.f32.f16 f_hi, hi; + mov.b32 bits_lo, f_lo; + mov.b32 bits_hi, f_hi; + and.b32 code_lo, $2, 0xF; + shr.u32 code_hi, $2, 4; + and.b32 code_hi, code_hi, 0xF; + setp.eq.u32 negzero_lo, code_lo, 0x8; + setp.eq.u32 negzero_hi, code_hi, 0x8; + selp.u32 bits_lo, 0x80000000, bits_lo, negzero_lo; + selp.u32 bits_hi, 0x80000000, bits_hi, negzero_hi; + mov.b32 $0, bits_lo; + mov.b32 $1, bits_hi; + } + """, + "=f,=f,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _scaled_e2m1x2_e4m3_to_f32x2( + packed: Uint32, scale: Uint32, *, loc=None, ip=None +) -> tuple[Float32, Float32]: + result = llvm.inline_asm( + llvm.StructType.get_literal([T.f32(), T.f32()]), + [Uint32(packed).ir_value(loc=loc, ip=ip), Uint32(scale).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b8 b0, b1, b2, b3; + .reg .b16 fp8_pair, scale_h, unused_h, lo, hi; + .reg .b32 q_h2, scale_h2, product_h2; + mov.b32 {b0, b1, b2, b3}, $2; + cvt.rn.f16x2.e2m1x2 q_h2, b0; + cvt.u16.u32 fp8_pair, $3; + cvt.rn.f16x2.e4m3x2 scale_h2, fp8_pair; + mov.b32 {scale_h, unused_h}, scale_h2; + mov.b32 scale_h2, {scale_h, scale_h}; + mul.rn.f16x2 product_h2, q_h2, scale_h2; + mov.b32 {lo, hi}, product_h2; + cvt.f32.f16 $0, lo; + cvt.f32.f16 $1, hi; + } + """, + "=f,=f,r,r", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return ( + Float32(llvm.extractvalue(T.f32(), result, [0], loc=loc, ip=ip)), + Float32(llvm.extractvalue(T.f32(), result, [1], loc=loc, ip=ip)), + ) + + +@dsl_user_op +def _pack_f32x2_to_half2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 lo, hi; + cvt.rn.f16.f32 lo, $1; + cvt.rn.f16.f32 hi, $2; + mov.b32 $0, {lo, hi}; + } + """, + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _pack_f32x2_to_bfloat2(a: Float32, b: Float32, *, loc=None, ip=None) -> Uint32: + return Uint32( + llvm.inline_asm( + T.i32(), + [Float32(a).ir_value(loc=loc, ip=ip), Float32(b).ir_value(loc=loc, ip=ip)], + """ + { + .reg .b16 lo, hi; + cvt.rn.bf16.f32 lo, $1; + cvt.rn.bf16.f32 hi, $2; + mov.b32 $0, {lo, hi}; + } + """, + "=r,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dsl_user_op +def _normal_block_scale( + block_amax: Float32, global_encode_scale: Float32, *, loc=None, ip=None +) -> Float32: + """TE's intentionally associated ``amax * (S_enc * (1/6))`` expression.""" + return Float32( + llvm.inline_asm( + T.f32(), + [ + Float32(block_amax).ir_value(loc=loc, ip=ip), + Float32(global_encode_scale).ir_value(loc=loc, ip=ip), + ], + """ + { + .reg .pred zero; + .reg .f32 scale_mul, result; + setp.eq.f32 zero, $1, 0f00000000; + mul.rn.f32 scale_mul, $2, 0f3E2AAAAB; + mul.rn.f32 result, $1, scale_mul; + selp.f32 $0, 0f00000000, result, zero; + } + """, + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.jit +def _input_values(words: tuple, is_bfloat16: bool) -> tuple: + values = () + for i in cutlass.range_constexpr(8): + if cutlass.const_expr(is_bfloat16): + lo, hi = _bfloat2_to_f32x2(words[i]) + else: + lo, hi = _half2_to_f32x2(words[i]) + values = values + (lo, hi) + return values + + +@cute.jit +def _block_amax(words: tuple, is_bfloat16: bool) -> Float32: + maxima = () + for i in cutlass.range_constexpr(8): + value = _half2_abs(words[i]) + maxima = maxima + (value,) + + if cutlass.const_expr(is_bfloat16): + max01 = _bfloat2_max(maxima[0], maxima[1]) + max23 = _bfloat2_max(maxima[2], maxima[3]) + max45 = _bfloat2_max(maxima[4], maxima[5]) + max67 = _bfloat2_max(maxima[6], maxima[7]) + max_value = _bfloat2_max(_bfloat2_max(max01, max23), _bfloat2_max(max45, max67)) + return _bfloat2_max_to_f32(max_value) + + max01 = _half2_max(maxima[0], maxima[1]) + max23 = _half2_max(maxima[2], maxima[3]) + max45 = _half2_max(maxima[4], maxima[5]) + max67 = _half2_max(maxima[6], maxima[7]) + max_value = _half2_max(_half2_max(max01, max23), _half2_max(max45, max67)) + return _half2_max_to_f32(max_value) + + +@cute.jit +def _scale_pack_e2m1(values: tuple, scale: Float32) -> tuple[Uint32, Uint32]: + """Scale and immediately pack each eight-value half to limit live FP32 state.""" + scaled_lo = () + for pair_idx in cutlass.range_constexpr(4): + value0, value1 = _scale_f32x2_rn(values[2 * pair_idx], values[2 * pair_idx + 1], scale) + scaled_lo = scaled_lo + (value0, value1) + lo = _cvt_e2m1x8( + scaled_lo[0], + scaled_lo[1], + scaled_lo[2], + scaled_lo[3], + scaled_lo[4], + scaled_lo[5], + scaled_lo[6], + scaled_lo[7], + ) + + scaled_hi = () + for pair_idx in cutlass.range_constexpr(4, 8): + value0, value1 = _scale_f32x2_rn(values[2 * pair_idx], values[2 * pair_idx + 1], scale) + scaled_hi = scaled_hi + (value0, value1) + hi = _cvt_e2m1x8( + scaled_hi[0], + scaled_hi[1], + scaled_hi[2], + scaled_hi[3], + scaled_hi[4], + scaled_hi[5], + scaled_hi[6], + scaled_hi[7], + ) + return lo, hi + + +@cute.jit +def _scale_pack_input_words( + words: tuple, scale: Float32, is_bfloat16: bool +) -> tuple[Uint32, Uint32]: + """Standard-path input conversion with at most eight scaled FP32 values live.""" + scaled_lo = () + for i in cutlass.range_constexpr(4): + if cutlass.const_expr(is_bfloat16): + value0, value1 = _bfloat2_to_f32x2(words[i]) + else: + value0, value1 = _half2_to_f32x2(words[i]) + value0, value1 = _scale_f32x2_rn(value0, value1, scale) + scaled_lo = scaled_lo + (value0, value1) + lo = _cvt_e2m1x8( + scaled_lo[0], + scaled_lo[1], + scaled_lo[2], + scaled_lo[3], + scaled_lo[4], + scaled_lo[5], + scaled_lo[6], + scaled_lo[7], + ) + + scaled_hi = () + for i in cutlass.range_constexpr(4, 8): + if cutlass.const_expr(is_bfloat16): + value0, value1 = _bfloat2_to_f32x2(words[i]) + else: + value0, value1 = _half2_to_f32x2(words[i]) + value0, value1 = _scale_f32x2_rn(value0, value1, scale) + scaled_hi = scaled_hi + (value0, value1) + hi = _cvt_e2m1x8( + scaled_hi[0], + scaled_hi[1], + scaled_hi[2], + scaled_hi[3], + scaled_hi[4], + scaled_hi[5], + scaled_hi[6], + scaled_hi[7], + ) + return lo, hi + + +@cute.jit +def _global_encode_scale(amax: Float32, e4m3_max: int) -> Float32: + scale = _fdiv_rn(Float32(float(e4m3_max * 6)), amax) + scale = _fmin(scale, Float32(_FP32_MAX)) + if amax == Float32(0.0): + scale = Float32(1.0) + if scale == Float32(0.0): + scale = Float32(1.0) + return scale + + +@cute.jit +def _candidate_inverse_scale(scale_f32: Float32, global_decode_scale: Float32) -> Float32: + product = _fmul_rn(scale_f32, global_decode_scale) + return _fmin(_fdiv_rn(Float32(1.0), product), Float32(_FP32_MAX)) + + +@cute.jit +def _standard_quantize( + words: tuple, + block_amax: Float32, + global_encode_scale: Float32, + global_decode_scale: Float32, + is_bfloat16: bool, +) -> tuple[Uint32, Uint32, Uint32]: + scale_high_precision = _normal_block_scale(block_amax, global_encode_scale) + scale = _cvt_f32_to_e4m3(scale_high_precision) + scale_f32 = _cvt_e4m3_to_f32(scale) + inverse = _candidate_inverse_scale(scale_f32, global_decode_scale) + lo, hi = _scale_pack_input_words(words, inverse, is_bfloat16) + return scale, lo, hi + + +@cute.jit +def _candidate_error( + original: tuple, + lo: Uint32, + hi: Uint32, + scale: Uint32, + global_amax: Float32, + global_encode_scale: Float32, + config: NVFP4QDQConfig, +) -> Float32: + error = Float32(0.0) + if cutlass.const_expr(config.error_use_fp16): + for pair_idx in cutlass.range_constexpr(8): + if cutlass.const_expr(pair_idx < 4): + packed_pair = lo >> Uint32(8 * pair_idx) + else: + packed_pair = hi >> Uint32(8 * (pair_idx - 4)) + candidate0, candidate1 = _scaled_e2m1x2_e4m3_to_f32x2(packed_pair, scale) + original0, original1 = _scale_f32x2_rn( + original[2 * pair_idx], original[2 * pair_idx + 1], global_encode_scale + ) + diff0 = _fsub_rn(candidate0, original0) + diff1 = _fsub_rn(candidate1, original1) + if cutlass.const_expr(config.error_mode == NVFP4QDQErrorMode.MSE): + term0, term1 = _square_f32x2_rn(diff0, diff1) + else: + term0 = _fabs(diff0) + term1 = _fabs(diff1) + error = _fadd_rn(error, term0) + error = _fadd_rn(error, term1) + return error + + denominator = Float32(float(6 * config.e4m3_max)) + scale_f32 = _cvt_e4m3_to_f32(scale) + for pair_idx in cutlass.range_constexpr(8): + if cutlass.const_expr(pair_idx < 4): + packed_pair = lo >> Uint32(8 * pair_idx) + else: + packed_pair = hi >> Uint32(8 * (pair_idx - 4)) + candidate0, candidate1 = _e2m1x2_to_f32x2(packed_pair) + + dequant0 = _fmul_rn(candidate0, scale_f32) + dequant0 = _fmul_rn(dequant0, global_amax) + dequant0 = _fdiv_rn(dequant0, denominator) + diff0 = _fsub_rn(dequant0, original[2 * pair_idx]) + if cutlass.const_expr(config.error_mode == NVFP4QDQErrorMode.MSE): + term0 = _fmul_rn(diff0, diff0) + else: + term0 = _fabs(diff0) + error = _fadd_rn(error, term0) + + dequant1 = _fmul_rn(candidate1, scale_f32) + dequant1 = _fmul_rn(dequant1, global_amax) + dequant1 = _fdiv_rn(dequant1, denominator) + diff1 = _fsub_rn(dequant1, original[2 * pair_idx + 1]) + if cutlass.const_expr(config.error_mode == NVFP4QDQErrorMode.MSE): + term1 = _fmul_rn(diff1, diff1) + else: + term1 = _fabs(diff1) + error = _fadd_rn(error, term1) + return error + + +@cute.jit +def _four_over_six_quantize( + values: tuple, + block_amax: Float32, + global_amax: Float32, + global_encode_scale: Float32, + global_decode_scale: Float32, + config: NVFP4QDQConfig, +) -> tuple[Uint32, Uint32, Uint32]: + # TE intentionally associates this differently from standard NVFP4. Keep + # the zero-amax path in the same expression stream as TE as well: packing + # the candidates is what preserves negative-zero E2M1 sign bits. + scale6_hp = _fmul_rn(_fdiv_rn(block_amax, Float32(6.0)), global_encode_scale) + scale4_hp = _fmul_rn(scale6_hp, Float32(1.5)) + scale4 = _cvt_f32_to_e4m3(_fmin(scale4_hp, Float32(448.0))) + scale6 = _cvt_f32_to_e4m3(_fmin(scale6_hp, Float32(448.0))) + scale4_f32 = _cvt_e4m3_to_f32(scale4) + scale6_f32 = _cvt_e4m3_to_f32(scale6) + inv4 = _candidate_inverse_scale(scale4_f32, global_decode_scale) + inv6 = _candidate_inverse_scale(scale6_f32, global_decode_scale) + lo4, hi4 = _scale_pack_e2m1(values, inv4) + lo6, hi6 = _scale_pack_e2m1(values, inv6) + error4 = _candidate_error(values, lo4, hi4, scale4, global_amax, global_encode_scale, config) + error6 = _candidate_error(values, lo6, hi6, scale6, global_amax, global_encode_scale, config) + + # Strict comparison is part of the contract: ties select map-to-6. + selected_scale = scale6 + selected_lo = lo6 + selected_hi = hi6 + if error4 < error6: + selected_scale = scale4 + selected_lo = lo4 + selected_hi = hi4 + + return selected_scale, selected_lo, selected_hi + + +@cute.jit +def _dequantize_pack_pair(packed_pair: Uint32, final_scale: Float32, is_bfloat16: bool) -> Uint32: + q0, q1 = _e2m1x2_to_f32x2(packed_pair) + out0, out1 = _scale_f32x2_rn(q0, q1, final_scale) + if cutlass.const_expr(is_bfloat16): + return _pack_f32x2_to_bfloat2(out0, out1) + return _pack_f32x2_to_half2(out0, out1) + + +@cute.jit +def _dequantize_store( + output: cute.Tensor, + offset: Int32, + lo: Uint32, + hi: Uint32, + scale: Uint32, + global_amax: Float32, + e4m3_max: int, + is_bfloat16: bool, +) -> None: + """Dequantize, cast, and store one half-block at a time.""" + scale_f32 = _cvt_e4m3_to_f32(scale) + final_scale = _fmul_rn(scale_f32, global_amax) + final_scale = _fmul_rn(final_scale, Float32(1.0 / float(6 * e4m3_max))) + + ptr0 = _get_ptr(output, offset) + out0 = _dequantize_pack_pair(lo, final_scale, is_bfloat16) + out1 = _dequantize_pack_pair(lo >> Uint32(8), final_scale, is_bfloat16) + out2 = _dequantize_pack_pair(lo >> Uint32(16), final_scale, is_bfloat16) + out3 = _dequantize_pack_pair(lo >> Uint32(24), final_scale, is_bfloat16) + _store_v4_u32(ptr0, out0, out1, out2, out3) + + ptr1 = _get_ptr(output, offset + Int32(8)) + out4 = _dequantize_pack_pair(hi, final_scale, is_bfloat16) + out5 = _dequantize_pack_pair(hi >> Uint32(8), final_scale, is_bfloat16) + out6 = _dequantize_pack_pair(hi >> Uint32(16), final_scale, is_bfloat16) + out7 = _dequantize_pack_pair(hi >> Uint32(24), final_scale, is_bfloat16) + _store_v4_u32(ptr1, out4, out5, out6, out7) + + +class _NVFP4QDQKernel: + """One thread processes one contiguous 1x16 quantization block.""" + + def __init__(self, is_bfloat16: bool, config: NVFP4QDQConfig) -> None: + self.is_bfloat16 = is_bfloat16 + self.config = config + if config.use_4over6: + self.threads = _4OVER6_THREADS + self.min_blocks_per_sm = _4OVER6_BLOCKS_PER_SM + self.grid_blocks_per_sm = _4OVER6_BLOCKS_PER_SM + else: + self.threads = _STANDARD_THREADS + self.min_blocks_per_sm = _STANDARD_MIN_BLOCKS_PER_SM + self.grid_blocks_per_sm = _STANDARD_GRID_BLOCKS_PER_SM + + @cute.jit + def __call__( + self, + input_tensor: cute.Tensor, + output_tensor: cute.Tensor, + global_amax: cute.Tensor, + total_blocks: Int32, + num_ctas: Int32, + stream, + ) -> None: + self.kernel(input_tensor, output_tensor, global_amax, total_blocks).launch( + grid=[num_ctas, 1, 1], + block=[self.threads, 1, 1], + max_number_threads=[self.threads, 1, 1], + min_blocks_per_mp=self.min_blocks_per_sm, + smem=0, + stream=stream, + ) + + @cute.kernel + def kernel( + self, + input_tensor: cute.Tensor, + output_tensor: cute.Tensor, + global_amax: cute.Tensor, + total_blocks: Int32, + ) -> None: + """Quantize and immediately dequantize grid-stride 1x16 blocks.""" + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + grid_dim, _, _ = cute.arch.grid_dim() + + amax = Float32(global_amax[Int32(0)]) + global_encode_scale = _global_encode_scale(amax, self.config.e4m3_max) + global_decode_scale = _fdiv_rn(Float32(1.0), global_encode_scale) + block = block_idx * Int32(self.threads) + thread_idx + stride = grid_dim * Int32(self.threads) + while block < total_blocks: + offset = block * Int32(_FP4_BLOCK_SIZE) + ptr0 = _get_ptr(input_tensor, offset) + ptr1 = _get_ptr(input_tensor, offset + Int32(8)) + w0, w1, w2, w3 = _load_v4_u32(ptr0) + w4, w5, w6, w7 = _load_v4_u32(ptr1) + words = (w0, w1, w2, w3, w4, w5, w6, w7) + block_amax = _block_amax(words, self.is_bfloat16) + + if cutlass.const_expr(self.config.use_4over6): + values = _input_values(words, self.is_bfloat16) + scale, lo, hi = _four_over_six_quantize( + values, block_amax, amax, global_encode_scale, global_decode_scale, self.config + ) + else: + scale, lo, hi = _standard_quantize( + words, block_amax, global_encode_scale, global_decode_scale, self.is_bfloat16 + ) + + _dequantize_store( + output_tensor, offset, lo, hi, scale, amax, self.config.e4m3_max, self.is_bfloat16 + ) + block = block + stride + + +@dataclass(frozen=True) +class _NVFP4QDQSpecialization: + """Compiled callable and its statically selected launch geometry.""" + + launch: Any + threads: int + grid_blocks_per_sm: int + + +_KERNEL_CACHE: dict[tuple[Any, ...], _NVFP4QDQSpecialization] = {} + + +@functools.cache +def _device_info(device_index: int) -> tuple[tuple[int, int], int]: + """Cache immutable capability and SM-count metadata outside the QAT hot path.""" + with torch.cuda.device(device_index): + capability = torch.cuda.get_device_capability(device_index) + multiprocessors = torch.cuda.get_device_properties(device_index).multi_processor_count + return capability, multiprocessors + + +def _validate_input(x: torch.Tensor, amax: torch.Tensor) -> tuple[int, tuple[int, int], int, int]: + if not x.is_cuda: + raise ValueError("Fused NVFP4 QDQ requires a CUDA tensor.") + if x.dtype not in (torch.bfloat16, torch.float16): + raise TypeError(f"Fused NVFP4 QDQ supports BF16 and FP16, got {x.dtype}.") + if x.ndim != 2: + raise ValueError(f"Fused NVFP4 QDQ requires a rank-2 tensor, got shape {tuple(x.shape)}.") + if not x.is_contiguous(): + raise ValueError("Fused NVFP4 QDQ requires a contiguous tensor.") + if x.data_ptr() % 16 != 0: + raise ValueError("Fused NVFP4 QDQ requires a 16-byte-aligned input tensor.") + if x.shape[1] % _FP4_BLOCK_SIZE != 0: + raise ValueError( + f"Fused NVFP4 QDQ requires K divisible by {_FP4_BLOCK_SIZE}, got {x.shape[1]}." + ) + num_elements = x.numel() + if num_elements == 0: + raise ValueError("Fused NVFP4 QDQ does not support empty tensors.") + if num_elements > _INT32_MAX: + raise ValueError( + f"Fused NVFP4 QDQ supports at most {_INT32_MAX} elements, got {num_elements}." + ) + if not amax.is_cuda or amax.device != x.device: + raise ValueError("The FP32 per-tensor amax must be on the input tensor's CUDA device.") + if amax.dtype != torch.float32 or amax.numel() != 1: + raise TypeError("The per-tensor amax must contain exactly one FP32 value.") + device_index = x.device.index + if device_index is None: + raise RuntimeError("CUDA tensor does not have a concrete device index.") + capability, multiprocessors = _device_info(device_index) + if capability[0] != 10: + raise ValueError(f"Fused NVFP4 QDQ requires SM10x, got compute capability {capability}.") + return device_index, capability, multiprocessors, num_elements + + +def _compile_specialization(dtype: torch.dtype, config: NVFP4QDQConfig) -> _NVFP4QDQSpecialization: + """Compile one dtype/config specialization outside the steady-state path.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Warm up fused NVFP4 QDQ before CUDA graph capture.") + kernel = _NVFP4QDQKernel(dtype == torch.bfloat16, config) + element_type = cutlass.BFloat16 if dtype == torch.bfloat16 else cutlass.Float16 + dynamic_elements = cute.sym_int() + input_fake = cute.runtime.make_fake_compact_tensor( + element_type, (dynamic_elements,), assumed_align=16 + ) + output_fake = cute.runtime.make_fake_compact_tensor( + element_type, (dynamic_elements,), assumed_align=16 + ) + amax_fake = cute.runtime.make_fake_compact_tensor(cutlass.Float32, (1,), assumed_align=4) + stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + compiled = cute.compile( + kernel, + input_fake, + output_fake, + amax_fake, + Int32(1), + Int32(1), + stream_fake, + options="--enable-tvm-ffi", + ) + return _NVFP4QDQSpecialization( + launch=compiled, threads=kernel.threads, grid_blocks_per_sm=kernel.grid_blocks_per_sm + ) + + +def _launch_fused_nvfp4_qdq( + x: torch.Tensor, + amax: torch.Tensor, + config: NVFP4QDQConfig, + capability: tuple[int, int], + multiprocessors: int, + num_elements: int, +) -> torch.Tensor: + """Launch on the current CUDA device with cached static dispatch.""" + key = (capability, x.dtype, config) + specialization = _KERNEL_CACHE.get(key) + if specialization is None: + specialization = _compile_specialization(x.dtype, config) + _KERNEL_CACHE[key] = specialization + + output = torch.empty_like(x) + input_flat = x.detach().view(-1) + output_flat = output.view(-1) + amax_flat = amax.detach().reshape(1) + total_blocks = num_elements // _FP4_BLOCK_SIZE + num_ctas = min( + (total_blocks + specialization.threads - 1) // specialization.threads, + multiprocessors * specialization.grid_blocks_per_sm, + ) + specialization.launch(input_flat, output_flat, amax_flat, total_blocks, num_ctas) + return output + + +def compute_nvfp4_amax(x: torch.Tensor) -> torch.Tensor: + """Compute the TE-compatible FP32 per-tensor amax with PyTorch.""" + if x.numel() == 0: + raise ValueError("Cannot compute NVFP4 amax for an empty tensor.") + return torch.linalg.vector_norm(x.detach(), ord=float("inf"), dtype=torch.float32) + + +def fused_nvfp4_qdq( + x: torch.Tensor, amax: torch.Tensor, config: Optional[NVFP4QDQConfig] = None +) -> torch.Tensor: + """Run register-resident NVFP4 QDQ and return a detached high-precision tensor.""" + if config is None: + config = current_nvfp4_qdq_config() + device_index, capability, multiprocessors, num_elements = _validate_input(x, amax) + + if torch.cuda.current_device() == device_index: + return _launch_fused_nvfp4_qdq(x, amax, config, capability, multiprocessors, num_elements) + with torch.cuda.device(device_index): + return _launch_fused_nvfp4_qdq(x, amax, config, capability, multiprocessors, num_elements) + + +class _FusedNVFP4QDQSTE(torch.autograd.Function): + """Identity backward around the non-differentiable fused QDQ kernel.""" + + @staticmethod + def forward( + ctx: Any, x: torch.Tensor, amax: torch.Tensor, config: NVFP4QDQConfig + ) -> torch.Tensor: + """Apply fused QDQ in the STE forward pass.""" + del ctx + return fused_nvfp4_qdq(x, amax, config) + + @staticmethod + def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]: + """Propagate the weight gradient through fake quantization unchanged.""" + del ctx + return grad_output, None, None + + +def fake_nvfp4_quantization_ste( + x: torch.Tensor, config: Optional[NVFP4QDQConfig] = None +) -> torch.Tensor: + """Apply fused NVFP4 QDQ in forward and the straight-through estimator in backward.""" + if config is None: + config = current_nvfp4_qdq_config() + amax = compute_nvfp4_amax(x) + output = _FusedNVFP4QDQSTE.apply(x, amax, config) + if hasattr(x, "main_grad"): + output.main_grad = x.main_grad + return output + + +__all__ = [ + "NVFP4QDQConfig", + "NVFP4QDQErrorMode", + "compute_nvfp4_amax", + "current_nvfp4_qdq_config", + "fake_nvfp4_quantization_ste", + "fused_nvfp4_qdq", +] diff --git a/pyproject.toml b/pyproject.toml index ff718cfabb2..f66de13b112 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -263,6 +263,7 @@ markers = [ "internal: mark a test as a test to private/internal functions.", "flaky: mark flaky tests for LTS environment", "flaky_in_dev: mark flaky tests for DEV environment", + "launch_on_gb200: mark a test file for the marker-driven GB200 unit-test lane", ] [tool.coverage.run] diff --git a/tests/unit_tests/extension/test_transformer_engine_nvfp4_qat.py b/tests/unit_tests/extension/test_transformer_engine_nvfp4_qat.py new file mode 100644 index 00000000000..546769d6cc1 --- /dev/null +++ b/tests/unit_tests/extension/test_transformer_engine_nvfp4_qat.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +import megatron.core.extensions.transformer_engine_nvfp4_fake_qat as nvfp4_qat +from megatron.core.extensions.transformer_engine_int4_fake_qat import INT4_FAKE_QAT_FLAG + + +def _config(gradient_accumulation_fusion: bool = False, moe_single_grouped_weight: bool = False): + return SimpleNamespace( + gradient_accumulation_fusion=gradient_accumulation_fusion, + moe_single_grouped_weight=moe_single_grouped_weight, + ) + + +@pytest.fixture(autouse=True) +def _supported_te(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(nvfp4_qat, "is_te_min_version", lambda _version: True) + monkeypatch.setenv(INT4_FAKE_QAT_FLAG, "0") + + +class TestTransformerEngineNVFP4FakeQAT: + def test_supported_with_discrete_rank2_weights(self): + weights = [torch.nn.Parameter(torch.empty(4, 16)) for _ in range(3)] + + nvfp4_qat._validate_nvfp4_fake_qat_support(_config(), False, weights) + + def test_rejects_transformer_engine_before_2_17(self, monkeypatch: pytest.MonkeyPatch): + weight = torch.nn.Parameter(torch.empty(4, 16)) + monkeypatch.setattr(nvfp4_qat, "is_te_min_version", lambda _version: False) + + with pytest.raises(RuntimeError, match="Transformer Engine >= 2.17.0"): + nvfp4_qat._validate_nvfp4_fake_qat_support(_config(), False, [weight]) + + def test_rejects_simultaneous_int4_and_nvfp4(self, monkeypatch: pytest.MonkeyPatch): + weight = torch.nn.Parameter(torch.empty(4, 16)) + monkeypatch.setenv(INT4_FAKE_QAT_FLAG, "1") + + with pytest.raises(RuntimeError, match="mutually exclusive"): + nvfp4_qat._validate_nvfp4_fake_qat_support(_config(), False, [weight]) + + def test_rejects_gradient_accumulation_fusion(self): + weight = torch.nn.Parameter(torch.empty(4, 16)) + + with pytest.raises(RuntimeError, match="gradient_accumulation_fusion"): + nvfp4_qat._validate_nvfp4_fake_qat_support( + _config(gradient_accumulation_fusion=True), False, [weight] + ) + + def test_rejects_delayed_wgrad_compute(self): + weight = torch.nn.Parameter(torch.empty(4, 16)) + + with pytest.raises(RuntimeError, match="delayed wgrad"): + nvfp4_qat._validate_nvfp4_fake_qat_support(_config(), True, [weight]) + + def test_rejects_fsdp_weight_attributes(self): + weight = torch.nn.Parameter(torch.empty(4, 16)) + weight.__fsdp_param__ = True + weight.get_main_grad = lambda: torch.empty_like(weight) + + with pytest.raises(RuntimeError, match="Megatron FSDP"): + nvfp4_qat._validate_nvfp4_fake_qat_support(_config(), False, [weight]) + + def test_rejects_single_grouped_weight(self): + weight = torch.nn.Parameter(torch.empty(3, 4, 16)) + + with pytest.raises(RuntimeError, match="moe_single_grouped_weight"): + nvfp4_qat._validate_nvfp4_fake_qat_support( + _config(moe_single_grouped_weight=True), False, [weight] + ) + + def test_rejects_non_rank2_weight(self): + weight = torch.nn.Parameter(torch.empty(3, 4, 16)) + + with pytest.raises(RuntimeError, match="rank-2 tensor"): + nvfp4_qat._validate_nvfp4_fake_qat_support(_config(), False, [weight]) + + def test_disabled_path_returns_original_list(self, monkeypatch: pytest.MonkeyPatch): + weights = [torch.nn.Parameter(torch.empty(4, 16))] + monkeypatch.setenv(nvfp4_qat.NVFP4_FAKE_QAT_FLAG, "0") + monkeypatch.setattr(nvfp4_qat, "is_te_min_version", lambda _version: False) + + actual = nvfp4_qat.maybe_fake_quantize_nvfp4_weight_tensors(_config(), False, weights) + + assert actual is weights + + def test_enabled_path_resolves_config_once_and_maps_arbitrary_weight_count( + self, monkeypatch: pytest.MonkeyPatch + ): + weights = [torch.nn.Parameter(torch.empty(4, 16)) for _ in range(3)] + expected = [torch.empty_like(weight) for weight in weights] + qdq_config = object() + config_calls = 0 + calls = [] + + fake_module = ModuleType("megatron.core.fusions.fused_nvfp4_qdq") + + def current_config(): + nonlocal config_calls + config_calls += 1 + return qdq_config + + def fake_qdq(weight, config): + calls.append((weight, config)) + return expected[len(calls) - 1] + + setattr(fake_module, "current_nvfp4_qdq_config", current_config) + setattr(fake_module, "fake_nvfp4_quantization_ste", fake_qdq) + monkeypatch.setitem(sys.modules, "megatron.core.fusions.fused_nvfp4_qdq", fake_module) + monkeypatch.setenv(nvfp4_qat.NVFP4_FAKE_QAT_FLAG, "1") + + actual = nvfp4_qat.maybe_fake_quantize_nvfp4_weight_tensors(_config(), False, weights) + + assert config_calls == 1 + assert len(actual) == len(weights) + assert all(value is expected_value for value, expected_value in zip(actual, expected)) + assert all(weight is call[0] for weight, call in zip(weights, calls)) + assert all(call[1] is qdq_config for call in calls) diff --git a/tests/unit_tests/fusions/benchmark_fused_nvfp4_qdq.py b/tests/unit_tests/fusions/benchmark_fused_nvfp4_qdq.py new file mode 100644 index 00000000000..9af60fd91dc --- /dev/null +++ b/tests/unit_tests/fusions/benchmark_fused_nvfp4_qdq.py @@ -0,0 +1,303 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Benchmark fused NVFP4 QDQ against the former fake-QAT TE round trip. + +The ``torch.utils.benchmark`` methodology follows Transformer Engine commit +83e230873f00676d4966ca151de22c8bfc68a77f. The benchmark shape is written as +``block-axis x rows`` and stored as a contiguous ``[rows, block-axis]`` tensor, +so every 1x16 NVFP4 block lies along the first reported dimension. It mirrors +``TEGroupedLinear._get_weight_tensors`` by timing the complete +``for w in weight_tensors`` loop over independently stored, equal-shaped +Parameters. The number of weights is configurable and defaults to eight. +The primary speedup compares complete user-visible loops: + +* naive: pad + TE quantize + TE dequantize + crop + contiguous; +* fused: the production env/config validator and per-weight STE adapter, including + PyTorch FP32 per-tensor amax + register-resident CuTe DSL QDQ. + +The separately labeled precomputed-amax API number still includes output +allocation, TVM-FFI argument marshalling, and host launch overhead. It is a +diagnostic and is not used for the primary speedup claim. +""" + +from __future__ import annotations + +import argparse +import os +from collections.abc import Callable +from statistics import geometric_mean +from types import SimpleNamespace + +import cutlass +import torch +import torch.utils.benchmark as benchmark +import transformer_engine +import transformer_engine.pytorch as te + +from megatron.core.extensions.transformer_engine_nvfp4_fake_qat import ( + NVFP4_FAKE_QAT_FLAG, + maybe_fake_quantize_nvfp4_weight_tensors, +) +from megatron.core.fusions.fused_nvfp4_qdq import ( + NVFP4QDQConfig, + NVFP4QDQErrorMode, + compute_nvfp4_amax, + fused_nvfp4_qdq, +) + +DEFAULT_LOGICAL_SHAPES = [(6144, 4096)] +DEFAULT_NUM_WEIGHTS = 8 + + +def _configs() -> list[tuple[str, NVFP4QDQConfig]]: + # Exact-error remains in the zero-tolerance correctness matrix, but the + # performance target is standard NVFP4 plus TE's current FP16-error path. + configs = [("nvfp4", NVFP4QDQConfig())] + for error_mode in (NVFP4QDQErrorMode.MAE, NVFP4QDQErrorMode.MSE): + for e4m3_max in (448, 256): + configs.append( + ( + f"4over6-{error_mode.name.lower()}-e4m3-{e4m3_max}-fp16-error", + NVFP4QDQConfig( + use_4over6=True, + e4m3_max=e4m3_max, + error_mode=error_mode, + error_use_fp16=True, + ), + ) + ) + return configs + + +def _make_te_quantizer(config: NVFP4QDQConfig): + return te.NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_amax_reduction=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=False, + nvfp4_use_4over6=config.use_4over6, + nvfp4_e4m3_max=config.e4m3_max, + nvfp4_4over6_err_mode=config.error_mode.name, + with_random_sign_mask=False, + ) + + +def _naive_te_qdq(x: torch.Tensor, quantizer) -> torch.Tensor: + m, n = x.shape + padded_m = ((m + 15) // 16) * 16 + if padded_m == m: + x_padded = x.contiguous() + else: + padding = torch.zeros((padded_m - m, n), dtype=x.dtype, device=x.device) + x_padded = torch.cat((x.contiguous(), padding), dim=0) + return quantizer.quantize(x_padded).dequantize(dtype=x.dtype)[:m, :n].contiguous() + + +def _median_us(function: Callable[[], object], min_run_time: float) -> float: + timing = benchmark.Timer( + stmt="function()", globals={"function": function}, num_threads=1 + ).blocked_autorange(min_run_time=min_run_time) + return timing.median * 1e6 + + +def _benchmark_case( + logical_shape: tuple[int, int], + dtype: torch.dtype, + config: NVFP4QDQConfig, + num_weights: int, + min_run_time: float, + repeats: int, +) -> tuple[list[float], list[float], list[float]]: + os.environ["NVTE_USE_FAST_MATH"] = "0" + os.environ["NVTE_NVFP4_4OVER6"] = "weights" if config.use_4over6 else "none" + os.environ["NVTE_NVFP4_4OVER6_E4M3_USE_256"] = "weights" if config.e4m3_max == 256 else "none" + os.environ["NVTE_NVFP4_4OVER6_ERR_MODE"] = config.error_mode.name + os.environ["NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"] = "1" if config.error_use_fp16 else "0" + os.environ["OPEN_TRAINING_INT4_FAKE_QAT_FLAG"] = "0" + os.environ[NVFP4_FAKE_QAT_FLAG] = "1" + torch.manual_seed(42) + torch.cuda.manual_seed(42) + block_axis, rows = logical_shape + weight_tensors = [ + torch.nn.Parameter(torch.randn((rows, block_axis), dtype=dtype, device="cuda")) + for _ in range(num_weights) + ] + amax_tensors = [compute_nvfp4_amax(w) for w in weight_tensors] + quantizer = _make_te_quantizer(config) + qat_config = SimpleNamespace( + gradient_accumulation_fusion=False, moe_single_grouped_weight=False + ) + + def naive() -> list[torch.Tensor]: + return [_naive_te_qdq(w, quantizer) for w in weight_tensors] + + def fused_end_to_end() -> list[torch.Tensor]: + return maybe_fake_quantize_nvfp4_weight_tensors(qat_config, False, weight_tensors) + + def fused_precomputed_amax() -> list[torch.Tensor]: + return [fused_nvfp4_qdq(w, amax, config) for w, amax in zip(weight_tensors, amax_tensors)] + + # Warm native TE and every compiled CuTe specialization before timing. + # Both primary closures retain their production QAT autograd wrappers. + expected = naive() + actual = fused_end_to_end() + for weight_idx, (expected_weight, actual_weight) in enumerate(zip(expected, actual)): + if not torch.equal( + expected_weight.detach().view(torch.uint16), actual_weight.detach().view(torch.uint16) + ): + mismatch_count = torch.count_nonzero( + expected_weight.detach().view(torch.uint16) + != actual_weight.detach().view(torch.uint16) + ).item() + raise AssertionError( + f"fused QDQ weight {weight_idx} differs from TE in " f"{mismatch_count} elements" + ) + del expected, actual + fused_precomputed_amax() + torch.cuda.synchronize() + + naive_us = [] + fused_e2e_us = [] + precomputed_amax_us = [] + for _ in range(repeats): + # Interleaved A/B/A guards the primary comparison against clock drift + # and neighboring-workload interference on a shared GPU node. + naive_us.append(_median_us(naive, min_run_time)) + fused_e2e_us.append(_median_us(fused_end_to_end, min_run_time)) + naive_us.append(_median_us(naive, min_run_time)) + precomputed_amax_us.append(_median_us(fused_precomputed_amax, min_run_time)) + return naive_us, fused_e2e_us, precomputed_amax_us + + +def _format_samples(samples: list[float]) -> str: + raw = ", ".join(f"{sample:.3f}" for sample in samples) + return f"{_sample_median(samples):.3f} [{raw}]" + + +def _sample_median(samples: list[float]) -> float: + ordered = sorted(samples) + middle = len(ordered) // 2 + if len(ordered) % 2: + return ordered[middle] + return 0.5 * (ordered[middle - 1] + ordered[middle]) + + +def _parse_shape(value: str) -> tuple[int, int]: + try: + block_axis, rows = value.lower().split("x", maxsplit=1) + shape = int(block_axis), int(rows) + except (TypeError, ValueError) as exc: + raise argparse.ArgumentTypeError( + "shapes must use BLOCK_AXISxROWS, for example 6144x4096" + ) from exc + if shape[0] <= 0 or shape[1] <= 0 or shape[0] % 16 != 0: + raise argparse.ArgumentTypeError( + "BLOCK_AXIS and ROWS must be positive and BLOCK_AXIS must be divisible by 16" + ) + return shape + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--shape", action="append", type=_parse_shape, dest="shapes") + parser.add_argument("--num-weights", type=int, default=DEFAULT_NUM_WEIGHTS) + parser.add_argument("--dtype", choices=("bf16", "fp16", "both"), default="both") + parser.add_argument("--min-run-time", type=float, default=1.0) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument( + "--image", + default=os.getenv("MILES_IMAGE", "unknown"), + help="Explicit Miles image tag or digest for the output metadata.", + ) + parser.add_argument( + "--commit", + default=os.getenv("MEGATRON_COMMIT", "unknown"), + help="Tested Megatron commit for output metadata.", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if args.min_run_time <= 0.0: + raise ValueError("--min-run-time must be positive") + if args.repeats <= 0: + raise ValueError("--repeats must be positive") + if args.num_weights <= 0: + raise ValueError("--num-weights must be positive") + available, reason = te.is_nvfp4_available(return_reason=True) + if not available: + raise RuntimeError(reason) + + shapes = args.shapes or DEFAULT_LOGICAL_SHAPES + dtypes = { + "bf16": [torch.bfloat16], + "fp16": [torch.float16], + "both": [torch.bfloat16, torch.float16], + }[args.dtype] + + print(f"image={args.image}") + print(f"megatron_commit={args.commit}") + print(f"gpu={torch.cuda.get_device_name()}") + print(f"compute_capability={torch.cuda.get_device_capability()}") + print(f"torch={torch.__version__} cuda={torch.version.cuda}") + print(f"transformer_engine={transformer_engine.__version__}") + print(f"cutlass_dsl={getattr(cutlass, '__version__', 'unknown')}") + print(f"min_run_time_s={args.min_run_time}") + print(f"repeats={args.repeats}") + print(f"num_weights={args.num_weights}") + print(f"num_gemms={args.num_weights}") + print("weight_storage=discrete_parameters") + print("gradient_accumulation_fusion=false") + print("moe_single_grouped_weight=false") + print("fused_path=maybe_fake_quantize_nvfp4_weight_tensors") + print("shape_contract=logical_block_axis_x_rows") + print("tensor_layout=contiguous_[rows,block_axis]") + for block_axis, rows in shapes: + print(f"in_features={block_axis}") + print(f"out_features={rows}") + print(f"stored_weight_shape=[{rows},{block_axis}]") + print("primary_order=naive/fused/naive per repeat") + print("NVTE_USE_FAST_MATH=0") + print() + print( + "| dtype | logical shape (block-axis x rows) | mode | " + f"naive TE {args.num_weights}-weight loop median [A/B/A raw] (us) | " + f"fused QAT {args.num_weights}-weight loop median [raw] (us) | " + f"precomputed-amax {args.num_weights}-weight loop median [raw] (us) | " + "end-to-end speedup |" + ) + print("|---|---:|---|---:|---:|---:|---:|") + all_speedups = [] + speedups_by_dtype: dict[torch.dtype, list[float]] = {dtype: [] for dtype in dtypes} + for dtype in dtypes: + for shape in shapes: + for label, config in _configs(): + naive_us, fused_e2e_us, precomputed_amax_us = _benchmark_case( + shape, dtype, config, args.num_weights, args.min_run_time, args.repeats + ) + speedup = _sample_median(naive_us) / _sample_median(fused_e2e_us) + all_speedups.append(speedup) + speedups_by_dtype[dtype].append(speedup) + print( + f"| {str(dtype).removeprefix('torch.')} | {shape[0]}x{shape[1]} | " + f"{label} | {_format_samples(naive_us)} | " + f"{_format_samples(fused_e2e_us)} | " + f"{_format_samples(precomputed_amax_us)} | " + f"{speedup:.3f}x |", + flush=True, + ) + print() + print(f"geomean_speedup={geometric_mean(all_speedups):.3f}x") + for dtype, speedups in speedups_by_dtype.items(): + print( + f"geomean_speedup_{str(dtype).removeprefix('torch.')}=" + f"{geometric_mean(speedups):.3f}x" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unit_tests/fusions/test_fused_nvfp4_qdq.py b/tests/unit_tests/fusions/test_fused_nvfp4_qdq.py new file mode 100644 index 00000000000..2da785a1800 --- /dev/null +++ b/tests/unit_tests/fusions/test_fused_nvfp4_qdq.py @@ -0,0 +1,366 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Bit-exact tests for the fused CuTe DSL NVFP4 QDQ kernel. + +The data patterns and Four Over Six Cartesian matrix mirror FlashInfer's +``tests/utils/test_fp4_quantize.py::test_nvfp4_quantize_te_reference``. The +oracle follows Transformer Engine's strict +``tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py`` test and calls TE's native +quantize-then-dequantize path because FlashInfer's per-tensor path has a +different numerical contract. +""" + +from __future__ import annotations + +import os + +import pytest +import torch + +pytest.importorskip("cutlass") +te = pytest.importorskip("transformer_engine.pytorch") + +from megatron.core.fusions.fused_nvfp4_qdq import ( # noqa: E402 + NVFP4QDQConfig, + NVFP4QDQErrorMode, + compute_nvfp4_amax, + current_nvfp4_qdq_config, + fake_nvfp4_quantization_ste, + fused_nvfp4_qdq, +) +from megatron.core.utils import is_te_min_version # noqa: E402 + +_recipe_available, _recipe_unavailable_reason = te.is_nvfp4_available(return_reason=True) +pytestmark = [ + pytest.mark.launch_on_gb200, + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required"), + pytest.mark.skipif( + not is_te_min_version("2.17.0"), + reason="Fused NVFP4 fake QAT requires Transformer Engine >= 2.17.0", + ), + pytest.mark.skipif(not _recipe_available, reason=_recipe_unavailable_reason), +] + + +@pytest.fixture(scope="module", autouse=True) +def _select_local_cuda_device() -> None: + """Keep torchrun workers on their assigned GPUs without initializing collectives.""" + if torch.cuda.is_available() and "LOCAL_RANK" in os.environ: + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + + +SHAPES = [ + # Minimum-K and odd-row cases absent from FlashInfer's swizzled-layout matrix. + (1, 16), + (1, 32), + (3, 48), + # FlashInfer strict-test shapes and both TE BF16 dispatch routes after M padding. + (1, 64), + (3, 128), + (16, 64), + (31, 128), + (32, 128), + (128, 64), + (128, 1024), + (256, 256), + (1024, 2048), +] + + +CONFIGS = [pytest.param(NVFP4QDQConfig(), id="nvfp4")] +for _error_mode in (NVFP4QDQErrorMode.MAE, NVFP4QDQErrorMode.MSE): + for _e4m3_max in (448, 256): + for _error_use_fp16 in (False, True): + CONFIGS.append( + pytest.param( + NVFP4QDQConfig( + use_4over6=True, + e4m3_max=_e4m3_max, + error_mode=_error_mode, + error_use_fp16=_error_use_fp16, + ), + id=( + f"4over6-{_error_mode.name.lower()}-e4m3-{_e4m3_max}-" + f"{'fp16-error' if _error_use_fp16 else 'exact-error'}" + ), + ) + ) + + +def _make_input(shape: tuple[int, int], dtype: torch.dtype, init_data: str) -> torch.Tensor: + torch.manual_seed(42) + torch.cuda.manual_seed(42) + m, n = shape + if init_data == "random": + x = torch.randn(shape, dtype=dtype, device="cuda") + if m > 1: + x[0].zero_() + return x + if init_data == "boundary": + base = torch.linspace(-12.0, 12.0, steps=n // 2, dtype=torch.float32, device="cuda") + eps = torch.full_like(base, 1e-3) + eps = torch.maximum(eps, torch.full_like(base, 1e-4)) + row = torch.empty(n, dtype=torch.float32, device="cuda") + row[0::2] = base - eps + row[1::2] = base + eps + return row.unsqueeze(0).repeat(m, 1).to(dtype=dtype) + if init_data == "zeros": + # Alternate signed zeros so the integer-view equality below exercises + # TE's E2M1 sign-bit contract for zero-amax blocks. + return ( + torch.tensor([-0.0, 0.0], dtype=torch.float32, device="cuda") + .repeat(m, n // 2) + .to(dtype=dtype) + ) + if init_data == "maxes": + return torch.full(shape, torch.finfo(dtype).max, dtype=dtype, device="cuda") + raise ValueError(f"Unknown init_data: {init_data}") + + +def _make_te_quantizer(config: NVFP4QDQConfig): + return te.NVFP4Quantizer( + rowwise=True, + columnwise=False, + with_amax_reduction=False, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + row_scaled_nvfp4=False, + nvfp4_use_4over6=config.use_4over6, + nvfp4_e4m3_max=config.e4m3_max, + nvfp4_4over6_err_mode=config.error_mode.name, + with_random_sign_mask=False, + ) + + +def _te_reference(x: torch.Tensor, config: NVFP4QDQConfig) -> tuple[torch.Tensor, torch.Tensor]: + m, n = x.shape + padded_m = ((m + 15) // 16) * 16 + if padded_m == m: + x_padded = x.contiguous() + else: + padding = torch.zeros((padded_m - m, n), dtype=x.dtype, device=x.device) + x_padded = torch.cat((x.contiguous(), padding), dim=0) + + quantized = _make_te_quantizer(config).quantize(x_padded) + reference = quantized.dequantize(dtype=x.dtype)[:m, :n].contiguous() + assert quantized._amax_rowwise is not None + return reference, quantized._amax_rowwise.reshape(1) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) +@pytest.mark.parametrize("shape", SHAPES, ids=lambda shape: f"{shape[0]}x{shape[1]}") +@pytest.mark.parametrize("init_data", ["random", "boundary", "zeros", "maxes"]) +@pytest.mark.parametrize("config", CONFIGS) +@torch.inference_mode() +def test_fused_nvfp4_qdq_is_bit_exact_with_te( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + shape: tuple[int, int], + init_data: str, + config: NVFP4QDQConfig, +) -> None: + """Cover BF16/FP16 x shapes x data patterns x the full supported feature matrix.""" + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "1" if config.error_use_fp16 else "0") + x = _make_input(shape, dtype, init_data) + amax = compute_nvfp4_amax(x) + expected, te_amax = _te_reference(x, config) + actual = fused_nvfp4_qdq(x, amax, config) + + assert torch.equal(amax.reshape(1).view(torch.int32), te_amax.view(torch.int32)) + # Integer views distinguish signed zero; tolerance-zero floating comparison does not. + actual_bits = actual.view(torch.uint16) + expected_bits = expected.view(torch.uint16) + assert torch.equal( + actual_bits, expected_bits + ), f"bit mismatch count: {torch.count_nonzero(actual_bits != expected_bits).item()}" + torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0) + + +def test_fused_nvfp4_qdq_uses_straight_through_gradient_and_preserves_main_grad() -> None: + x = torch.randn((3, 32), dtype=torch.bfloat16, device="cuda", requires_grad=True) + main_grad = torch.empty_like(x) + x.main_grad = main_grad + output = fake_nvfp4_quantization_ste(x, NVFP4QDQConfig()) + output.backward(torch.ones_like(output)) + + torch.testing.assert_close(x.grad, torch.ones_like(x), rtol=0.0, atol=0.0) + assert output.main_grad is main_grad + + +@pytest.mark.parametrize( + ("four_over_six_scope", "e4m3_256_scope", "expected_enabled", "expected_max"), + [ + ( + four_over_six_scope, + e4m3_256_scope, + four_over_six_scope in ("weights", "all"), + expected_max, + ) + for four_over_six_scope in ("none", "activations", "weights", "all") + for e4m3_256_scope in ("none", "activations", "weights", "all") + for expected_max in [ + ( + 256 + if four_over_six_scope in ("weights", "all") + and e4m3_256_scope in ("weights", "all") + else 448 + ) + ] + ], +) +@pytest.mark.parametrize( + ("error_mode", "error_use_fp16"), [("MAE", False), ("MAE", True), ("MSE", False), ("MSE", True)] +) +def test_current_nvfp4_qdq_config_maps_full_latest_te_env_contract( + monkeypatch: pytest.MonkeyPatch, + four_over_six_scope: str, + e4m3_256_scope: str, + expected_enabled: bool, + expected_max: int, + error_mode: str, + error_use_fp16: bool, +) -> None: + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + monkeypatch.setenv("NVTE_NVFP4_4OVER6", four_over_six_scope) + monkeypatch.setenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", e4m3_256_scope) + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_MODE", error_mode) + monkeypatch.setenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH", "1" if error_use_fp16 else "0") + config = current_nvfp4_qdq_config() + assert config.use_4over6 is expected_enabled + assert config.e4m3_max == expected_max + assert config.error_mode is NVFP4QDQErrorMode[error_mode] + # The latest TE meaning is FP16-rounded candidate error, not a general + # instruction-level fast-math toggle. + assert config.error_use_fp16 is (expected_enabled and error_use_fp16) + + +@pytest.mark.parametrize("legacy_scope", ["inputs", "gradients"]) +def test_current_nvfp4_qdq_config_rejects_stale_te_scopes( + monkeypatch: pytest.MonkeyPatch, legacy_scope: str +) -> None: + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + monkeypatch.setenv("NVTE_NVFP4_4OVER6", legacy_scope) + with pytest.raises(ValueError, match="activations"): + current_nvfp4_qdq_config() + + +def test_current_nvfp4_qdq_config_rejects_quant_fast_math(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("NVTE_USE_FAST_MATH", "1") + with pytest.raises(ValueError, match="NVTE_USE_FAST_MATH=0"): + current_nvfp4_qdq_config() + + +def test_te_grouped_linear_real_discrete_weight_qdq_and_backward( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from megatron.core.extensions import transformer_engine as te_extension + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.transformer.transformer_config import TransformerConfig + + group_count, rows, columns = 3, 2, 32 + monkeypatch.setenv("NVTE_USE_FAST_MATH", "0") + monkeypatch.setenv("NVTE_NVFP4_4OVER6", "none") + monkeypatch.delenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", raising=False) + monkeypatch.setenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") + monkeypatch.setenv("OPEN_TRAINING_NVFP4_FAKE_QAT_FLAG", "1") + config = TransformerConfig( + num_layers=1, + hidden_size=columns, + num_attention_heads=1, + params_dtype=torch.bfloat16, + gradient_accumulation_fusion=False, + moe_single_grouped_weight=False, + ) + pg_collection = ProcessGroupCollection() + pg_collection.expt_tp = None + layer = te_extension.TEGroupedLinear( + num_gemms=group_count, + input_size=columns, + output_size=rows, + parallel_mode=None, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + is_expert=True, + pg_collection=pg_collection, + ) + + assert layer.fuse_wgrad_accumulation is False + assert not getattr(layer, "single_grouped_weight", False) + assert getattr(layer, "weight", None) is None + weights = [getattr(layer, f"weight{group_idx}") for group_idx in range(group_count)] + assert set(dict(layer.named_parameters())) == { + f"weight{group_idx}" for group_idx in range(group_count) + } + + actual_weights = te_extension.TEGroupedLinear._get_weight_tensors(layer) + qdq_config = current_nvfp4_qdq_config() + expected_weights = [_te_reference(weight, qdq_config)[0] for weight in weights] + + assert len(actual_weights) == group_count + assert all(weight.requires_grad for weight in actual_weights) + assert all( + torch.equal(actual.view(torch.uint16), expected.view(torch.uint16)) + for actual, expected in zip(actual_weights, expected_weights) + ) + + m_splits = [2, 1, 3] + inp = torch.ones( + (sum(m_splits), columns), dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + output, bias = layer(inp, m_splits) + output.backward(torch.ones_like(output)) + + assert bias is None + assert tuple(output.shape) == (sum(m_splits), rows) + assert inp.grad is not None and torch.isfinite(inp.grad).all() + for weight, m_split in zip(weights, m_splits): + assert weight.grad is not None + torch.testing.assert_close( + weight.grad, torch.full_like(weight.grad, m_split), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_fused_nvfp4_qdq_rejects_unsupported_input_dtype(dtype: torch.dtype) -> None: + x = torch.randn((2, 16), dtype=dtype, device="cuda") + with pytest.raises(TypeError, match="supports BF16 and FP16"): + fused_nvfp4_qdq(x, x.abs().amax().float(), NVFP4QDQConfig()) + + +def test_fused_nvfp4_qdq_rejects_non_block_aligned_k() -> None: + x = torch.randn((2, 17), dtype=torch.bfloat16, device="cuda") + with pytest.raises(ValueError, match="K divisible by 16"): + fused_nvfp4_qdq(x, compute_nvfp4_amax(x), NVFP4QDQConfig()) + + +def test_fused_nvfp4_qdq_rejects_misaligned_contiguous_storage() -> None: + storage = torch.randn(33, dtype=torch.bfloat16, device="cuda") + x = storage[1:].view(2, 16) + assert x.is_contiguous() + assert x.data_ptr() % 16 != 0 + with pytest.raises(ValueError, match="16-byte-aligned"): + fused_nvfp4_qdq(x, compute_nvfp4_amax(x), NVFP4QDQConfig()) + + +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two CUDA devices") +def test_fused_nvfp4_qdq_uses_and_restores_non_current_device() -> None: + if int(os.getenv("WORLD_SIZE", "1")) > 1: + pytest.skip("Run the cross-device state test in a dedicated single process") + primary_device = torch.cuda.current_device() + secondary_device = (primary_device + 1) % torch.cuda.device_count() + with torch.cuda.device(primary_device): + with torch.cuda.device(secondary_device): + x = _make_input((3, 32), torch.bfloat16, "boundary") + amax = compute_nvfp4_amax(x) + expected, _ = _te_reference(x, NVFP4QDQConfig()) + + assert torch.cuda.current_device() == primary_device + actual = fused_nvfp4_qdq(x, amax, NVFP4QDQConfig()) + assert torch.cuda.current_device() == primary_device + + assert torch.equal(actual.view(torch.uint16), expected.view(torch.uint16))