diff --git a/alto/kernels/dispatch/config.py b/alto/kernels/dispatch/config.py index 7dd4862..8bbe2f4 100644 --- a/alto/kernels/dispatch/config.py +++ b/alto/kernels/dispatch/config.py @@ -38,7 +38,7 @@ class TrainingOpConfig: * dynamic: dynamic clipped scale based on the mean and std of the input tensor (only supported for mxfp4) """ - two_level_scaling: Literal["none", "tensorwise", "blockwise"] = "none" + two_level_scaling: Literal["none", "tensorwise", "blockwise", "outer_block"] = "none" """ apply an extra scaling factor besides the default blockwise scales of MXFP4/NVFP4. * none: no extra scaling @@ -46,7 +46,11 @@ class TrainingOpConfig: * blockwise: * MXFP4: apply a blockwise scale factor containing shared mantissa to each block * NVFP4: not implemented + * outer_block: apply a per-super-block FP32 scale before NVFP4 E4M3 block scales """ + use_outer_2dblock_x: bool = False + use_outer_2dblock_w: bool = False + outer_block_size: int = 128 torch.serialization.add_safe_globals([TrainingOpConfig]) diff --git a/alto/kernels/dispatch/tensor.py b/alto/kernels/dispatch/tensor.py index 40868af..260ede9 100644 --- a/alto/kernels/dispatch/tensor.py +++ b/alto/kernels/dispatch/tensor.py @@ -363,6 +363,15 @@ def __torch_function__(cls, func, types, args, kwargs={}): "nvfp4": _quantize_then_nvfp4_scaled_grouped_mm, "amdfp4": _quantize_then_amdfp4_scaled_grouped_mm, }[config.precision] + # Outer-block (v1): per-expert/per-tile scale. Activations use a 1D + # outer block internally; weights honour use_outer_2dblock_w. The + # path is shared across both inner grids (E4M3 / UE5M3); both + # grouped helpers accept these kwargs. + common_kwargs.update( + use_outer_block_scale=config.two_level_scaling == "outer_block", + use_outer_2dblock_w=config.use_outer_2dblock_w, + outer_block_size=config.outer_block_size, + ) return grouped_fn(A, B, **common_kwargs) # linear / mm overrides @@ -406,6 +415,15 @@ def __torch_function__(cls, func, types, args, kwargs={}): "nvfp4": _to_nvfp4_then_scaled_mm, "amdfp4": _to_amdfp4_then_scaled_mm, }[config.precision] + # Outer-block / X-hat-align / dX-RHT share one code path across both + # inner grids (E4M3 for nvfp4, UE5M3 for amdfp4); both dense helpers + # accept these kwargs. + common_kwargs.update( + use_outer_block_scale=config.two_level_scaling == "outer_block", + use_outer_2dblock_x=config.use_outer_2dblock_x, + use_outer_2dblock_w=config.use_outer_2dblock_w, + outer_block_size=config.outer_block_size, + ) Y = linear_fn(A, W, **common_kwargs) if bias is not None: Y = Y + bias diff --git a/alto/kernels/fp4/amdfp4/amdfp_grouped_gemm/functional.py b/alto/kernels/fp4/amdfp4/amdfp_grouped_gemm/functional.py index 883d576..adc3326 100644 --- a/alto/kernels/fp4/amdfp4/amdfp_grouped_gemm/functional.py +++ b/alto/kernels/fp4/amdfp4/amdfp_grouped_gemm/functional.py @@ -17,6 +17,7 @@ _quantize_then_nvfp4_scaled_grouped_mm, nvfp4_grouped_gemm, ) +from alto.kernels.fp4.nvfp4.nvfp_quantization import OUTER_BLOCK_SIZE_DEFAULT def amdfp4_grouped_gemm( @@ -31,12 +32,16 @@ def amdfp4_grouped_gemm( use_outer_scale: bool = False, use_hadamard: bool = False, use_dge: bool = False, + use_outer_block_scale: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """AMD-FP4 (UE5M3 inner scale) Grouped GEMM with full autograd support. Same parameter contract as :func:`alto.kernels.fp4.nvfp4.nvfp_grouped_gemm.nvfp4_grouped_gemm` - minus ``scale_format``, which is hard-pinned to ``"ue5m3"``. + (including the outer-block knobs) minus ``scale_format``, which is + hard-pinned to ``"ue5m3"``. """ return nvfp4_grouped_gemm( inputs, @@ -49,6 +54,9 @@ def amdfp4_grouped_gemm( use_outer_scale=use_outer_scale, use_hadamard=use_hadamard, use_dge=use_dge, + use_outer_block_scale=use_outer_block_scale, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=outer_block_size, scale_format="ue5m3", ) @@ -63,12 +71,16 @@ def _quantize_then_amdfp4_scaled_grouped_mm( use_outer_scale: bool = False, use_hadamard: bool = False, use_dge: bool = False, + use_outer_block_scale: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """AMD-FP4 dispatch-side variant of :func:`alto.kernels.fp4.nvfp4.nvfp_grouped_gemm._quantize_then_nvfp4_scaled_grouped_mm`. ``scale_format`` is hard-pinned to ``"ue5m3"`` so the MoE dispatch - code path can route AMD-FP4 traffic by op surface alone. + code path can route AMD-FP4 traffic by op surface alone. Outer-block + scaling shares the NVFP4 grouped code path; only the inner grid differs. """ return _quantize_then_nvfp4_scaled_grouped_mm( A, @@ -80,5 +92,8 @@ def _quantize_then_amdfp4_scaled_grouped_mm( use_outer_scale=use_outer_scale, use_hadamard=use_hadamard, use_dge=use_dge, + use_outer_block_scale=use_outer_block_scale, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=outer_block_size, scale_format="ue5m3", ) diff --git a/alto/kernels/fp4/amdfp4/amdfp_linear.py b/alto/kernels/fp4/amdfp4/amdfp_linear.py index 5e4721b..cd8d727 100644 --- a/alto/kernels/fp4/amdfp4/amdfp_linear.py +++ b/alto/kernels/fp4/amdfp4/amdfp_linear.py @@ -20,6 +20,7 @@ NVFP4LinearFunction, _to_nvfp4_then_scaled_mm, ) +from alto.kernels.fp4.nvfp4.nvfp_quantization import OUTER_BLOCK_SIZE_DEFAULT # Re-export the autograd function under an AMD-FP4 name. We deliberately @@ -40,12 +41,19 @@ def _to_amdfp4_then_scaled_mm( use_outer_scale: bool = False, use_hadamard: bool = False, use_dge: bool = False, + use_outer_block_scale: bool = False, + use_outer_2dblock_x: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """AMD-FP4 (UE5M3 inner scale) variant of ``_to_nvfp4_then_scaled_mm``. - Mirrors the NVFP4 helper exactly, except ``scale_format`` is hard-pinned - to ``"ue5m3"`` so the AMD-FP4 dispatch layer / training recipes can - address the AMD-FP4 path by op surface alone. + Mirrors the NVFP4 helper exactly -- including the outer-block knobs -- + except ``scale_format`` is hard-pinned to ``"ue5m3"`` so the AMD-FP4 + dispatch layer / training recipes can address the AMD-FP4 path by op + surface alone. Outer-block scaling shares the NVFP4 code path; only the + inner-block scale grid differs (UE5M3 vs E4M3), which the outer-scale + denominator accounts for. """ return _to_nvfp4_then_scaled_mm( a, @@ -56,6 +64,10 @@ def _to_amdfp4_then_scaled_mm( use_outer_scale=use_outer_scale, use_hadamard=use_hadamard, use_dge=use_dge, + use_outer_block_scale=use_outer_block_scale, + use_outer_2dblock_x=use_outer_2dblock_x, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=outer_block_size, scale_format="ue5m3", ) diff --git a/alto/kernels/fp4/nvfp4/__init__.py b/alto/kernels/fp4/nvfp4/__init__.py index 1f34412..98cc987 100644 --- a/alto/kernels/fp4/nvfp4/__init__.py +++ b/alto/kernels/fp4/nvfp4/__init__.py @@ -25,22 +25,32 @@ from .nvfp_linear import NVFP4LinearFunction from .nvfp_quantization import ( BLOCK_SIZE_DEFAULT, + OUTER_BLOCK_SIZE_DEFAULT, + OUTER_SCALE_EPS, SUPPORTED_SCALE_FORMATS, _calculate_nvfp4_scales, _pack_fp4, _unpack_fp4, compute_dynamic_outer_scale, + compute_outer_block_scale_shape, convert_from_nvfp4, + convert_from_nvfp4_outer_block, convert_to_nvfp4, + convert_to_nvfp4_outer_block, is_cdna4, ) __all__ = ( "BLOCK_SIZE_DEFAULT", "NVFP4LinearFunction", + "OUTER_BLOCK_SIZE_DEFAULT", + "OUTER_SCALE_EPS", "SUPPORTED_SCALE_FORMATS", "compute_dynamic_outer_scale", + "compute_outer_block_scale_shape", "convert_from_nvfp4", + "convert_from_nvfp4_outer_block", "convert_to_nvfp4", + "convert_to_nvfp4_outer_block", "is_cdna4", ) diff --git a/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/autograd.py b/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/autograd.py index b3a53ea..8a26a0f 100644 --- a/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/autograd.py +++ b/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/autograd.py @@ -30,7 +30,11 @@ from alto.kernels.dge import dge_bwd from alto.kernels.fp4.fp4_common import unwrap_weight_wrapper -from alto.kernels.fp4.nvfp4.nvfp_quantization import _qdq, convert_from_nvfp4 +from alto.kernels.fp4.nvfp4.nvfp_quantization import ( + OUTER_BLOCK_SIZE_DEFAULT, + _qdq, + convert_from_nvfp4, +) from alto.kernels.hadamard_transform import HadamardTransform from .autotune import ALIGN_SIZE_M from .cg_backward import _nvfp4_grouped_dgrad, _nvfp4_grouped_wgrad @@ -76,7 +80,24 @@ def forward( hadamard_transform: Optional[HadamardTransform] = None, use_dge: bool = False, scale_format: str = "e4m3", + use_outer_block_scale: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: + # Outer-block scaling for grouped GEMM (v1): per-expert/per-tile FP32 + # outer scale instead of one tensorwise scalar over [E, N, K]. + # Activations use a 1D outer block (1 x outer_block_size over K); weights + # use ``use_outer_2dblock_w`` (per-expert 2D tiles). Tensorwise and + # outer-block are mutually exclusive; DGE is unsupported with outer-block + # (mirrors the dense NVFP4LinearFunction contract). + torch._check( + not (use_outer_scale and use_outer_block_scale), + lambda: "grouped NVFP4: tensorwise and outer-block scale are mutually exclusive.", + ) + torch._check( + not (use_dge and use_outer_block_scale), + lambda: "grouped NVFP4 outer-block scaling does not support DGE.", + ) M_bufferlen = inputs.shape[0] original_dtype = inputs.dtype expert_weights = unwrap_weight_wrapper(expert_weights) @@ -104,11 +125,15 @@ def forward( inputs, axis=-1, is_2d_block=use_2dblock_x, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=False, outer_block_size=outer_block_size, ) w_dq = _qdq( expert_weights, axis=_FPROP_AXIS_W, is_2d_block=use_2dblock_w, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, outer_block_size=outer_block_size, ) y = _nvfp4_grouped_fprop( x_dq, @@ -127,6 +152,9 @@ def forward( is_2d_block=False, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, + outer_block_size=outer_block_size, return_raw=use_dge, ) if use_dge: @@ -154,6 +182,8 @@ def forward( x_for_wgrad, axis=0, is_2d_block=False, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=False, outer_block_size=outer_block_size, ) else: x_bwd = x_dq @@ -166,6 +196,9 @@ def forward( ctx.use_2dblock_w = use_2dblock_w ctx.use_sr_grad = use_sr_grad ctx.use_outer_scale = use_outer_scale + ctx.use_outer_block_scale = use_outer_block_scale + ctx.use_outer_2dblock_w = use_outer_2dblock_w + ctx.outer_block_size = outer_block_size ctx.num_experts = num_experts ctx.num_groups = num_groups ctx.original_dtype = original_dtype @@ -199,6 +232,8 @@ def backward(ctx, grad_output: torch.Tensor): grad_output, axis=-1, is_2d_block=ctx.use_2dblock_x, use_outer_scale=ctx.use_outer_scale, use_sr=ctx.use_sr_grad, scale_format=ctx.scale_format, + use_outer_block_scale=ctx.use_outer_block_scale, + is_2d_outer=False, outer_block_size=ctx.outer_block_size, ) grad_inputs = _nvfp4_grouped_dgrad( g_dq, @@ -224,6 +259,8 @@ def backward(ctx, grad_output: torch.Tensor): use_2dblock_x=ctx.use_2dblock_x, output_dtype=ctx.original_dtype, scale_format=ctx.scale_format, + use_outer_block_scale=ctx.use_outer_block_scale, + outer_block_size=ctx.outer_block_size, ) if ctx.use_dge: @@ -243,8 +280,11 @@ def backward(ctx, grad_output: torch.Tensor): # Return grads with arity matching forward's positional inputs: # (inputs, expert_weights, expert_indices, offs, # use_2dblock_x, use_2dblock_w, use_sr_grad, use_outer_scale, - # hadamard_transform, use_dge, scale_format) + # hadamard_transform, use_dge, scale_format, + # use_outer_block_scale, use_outer_2dblock_w, outer_block_size) + # -> 14 inputs, 12 None grad slots after the two tensor grads. return ( grad_inputs, grad_weights, - None, None, None, None, None, None, None, None, None, + None, None, None, None, None, None, None, None, None, None, + None, None, ) diff --git a/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/cg_backward.py b/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/cg_backward.py index bf09aff..b723bb3 100644 --- a/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/cg_backward.py +++ b/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/cg_backward.py @@ -14,7 +14,7 @@ cg_grouped_gemm_backward_inputs, cg_grouped_gemm_backward_weights, ) -from alto.kernels.fp4.nvfp4.nvfp_quantization import _qdq +from alto.kernels.fp4.nvfp4.nvfp_quantization import OUTER_BLOCK_SIZE_DEFAULT, _qdq from alto.kernels.fp4.fp4_common import ( check_grouped_loop_contract, @@ -93,24 +93,33 @@ def _nvfp4_grouped_wgrad( use_2dblock_x: bool, output_dtype: torch.dtype, scale_format: str = "e4m3", + use_outer_block_scale: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """Quantize ``grad_output`` and compute the grouped weight gradient. Returns ``dW`` in the canonical ``[E, N, K]`` layout, matching both ``cg_grouped_gemm_backward_weights`` (CDNA4 native) and the autograd function's internal weight layout. No post-transpose needed. + + ``grad_output`` is activation-like, so outer-block scaling uses a 1D outer + block (``is_2d_outer=False``) -- consistent with the activation path. """ if use_2dblock_x: g_m_dq = _qdq( grad_output, axis=-1, is_2d_block=True, use_outer_scale=use_outer_scale, use_sr=use_sr_grad, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=False, outer_block_size=outer_block_size, ) else: g_m_dq = _qdq( grad_output, axis=0, is_2d_block=False, use_outer_scale=use_outer_scale, use_sr=use_sr_grad, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=False, outer_block_size=outer_block_size, ) if use_cdna4_grouped_backend(): diff --git a/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/functional.py b/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/functional.py index ed6bab5..7409593 100644 --- a/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/functional.py +++ b/alto/kernels/fp4/nvfp4/nvfp_grouped_gemm/functional.py @@ -28,6 +28,7 @@ import torch from alto.kernels.fp4.fp4_common import build_hadamard_transform_if_needed +from alto.kernels.fp4.nvfp4.nvfp_quantization import OUTER_BLOCK_SIZE_DEFAULT from .autograd import NVFP4GroupedGEMM @@ -44,6 +45,9 @@ def _nvfp4_grouped_gemm_impl( use_hadamard: bool = False, use_dge: bool = False, scale_format: str = "e4m3", + use_outer_block_scale: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """Normalized entrypoint shared by both public APIs. @@ -72,6 +76,9 @@ def _nvfp4_grouped_gemm_impl( hadamard_transform, use_dge, scale_format, + use_outer_block_scale, + use_outer_2dblock_w, + outer_block_size, ) @@ -88,6 +95,9 @@ def nvfp4_grouped_gemm( use_hadamard: bool = False, use_dge: bool = False, scale_format: str = "e4m3", + use_outer_block_scale: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """NVFP4 QDQ-emulated Grouped GEMM with full autograd support. @@ -128,6 +138,9 @@ def nvfp4_grouped_gemm( use_hadamard=use_hadamard, use_dge=use_dge, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=outer_block_size, ) @@ -142,6 +155,9 @@ def _quantize_then_nvfp4_scaled_grouped_mm( use_hadamard: bool = False, use_dge: bool = False, scale_format: str = "e4m3", + use_outer_block_scale: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """Drop-in for the dispatch layer, mirroring mxfp4's ``_quantize_then_mxfp4_scaled_grouped_mm``. @@ -164,4 +180,7 @@ def _quantize_then_nvfp4_scaled_grouped_mm( use_hadamard=use_hadamard, use_dge=use_dge, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=outer_block_size, ) diff --git a/alto/kernels/fp4/nvfp4/nvfp_linear.py b/alto/kernels/fp4/nvfp4/nvfp_linear.py index f479894..cbc3991 100644 --- a/alto/kernels/fp4/nvfp4/nvfp_linear.py +++ b/alto/kernels/fp4/nvfp4/nvfp_linear.py @@ -11,6 +11,7 @@ from alto.kernels.dge import dge_bwd from .nvfp_quantization import ( BLOCK_SIZE_DEFAULT, + OUTER_BLOCK_SIZE_DEFAULT, _qdq, # noqa: F401 (re-exported for backward compatibility) convert_from_nvfp4, ) @@ -60,9 +61,24 @@ def forward( use_outer_scale: bool, hadamard_transform: Optional[HadamardTransform] = None, use_dge: bool = False, + use_outer_block_scale: bool = False, + use_outer_2dblock_x: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, scale_format: str = "e4m3", ): weight = unwrap_weight_wrapper(weight) + torch._check( + not (use_outer_scale and use_outer_block_scale), + lambda: "NVFP4 tensor-wise scale and outer-block scale are mutually exclusive.", + ) + torch._check( + not (use_dge and use_outer_block_scale), + lambda: ( + "NVFP4 outer-block scaling does not yet support DGE because " + "DGE needs the raw packed FP4 values and inner scales." + ), + ) # Align weight dtype with activation so saved QDQ tensors share a # common dtype across forward / backward under AMP autocast (matches # the MXFP4 ``original_dtype = x.dtype`` contract). @@ -72,21 +88,73 @@ def forward( original_shape = x.shape x_2d = x.reshape(-1, original_shape[-1]) - # Prepare the quantized views used by forward. These are the tensors - # consumed by the simulated Fprop GEMM. - x_dq = _qdq( - x_2d, axis=-1, - is_2d_block=use_2dblock_x, - use_outer_scale=use_outer_scale, - scale_format=scale_format, + # 2D outer-block scale grids are axis-invariant: an axis=-1 QDQ and + # an axis=0 QDQ on the same tensor produce identical outer scales + # (modulo a transpose). Capture the axis=-1 scale and feed it into + # the axis=0 path to skip a full outer-amax reduction. Activations + # cannot share when Hadamard is on, since ``H @ x`` has a different + # outer-amax map. Weights have no Hadamard rotation, so 2D outer is + # always shareable on weights. + share_outer_x = ( + use_outer_block_scale + and use_outer_2dblock_x + and not use_2dblock_x + and hadamard_transform is None ) - w_dq = _qdq( - weight, axis=-1, - is_2d_block=use_2dblock_w, - use_outer_scale=use_outer_scale, - scale_format=scale_format, + share_outer_w = ( + use_outer_block_scale + and use_outer_2dblock_w + and not use_2dblock_w ) + # Prepare the quantized views used by forward. These are the tensors + # consumed by the simulated Fprop GEMM. + x_outer_scale_cache: Optional[torch.Tensor] = None + if share_outer_x: + x_dq, x_outer_scale_cache = _qdq( + x_2d, axis=-1, + is_2d_block=use_2dblock_x, + use_outer_scale=use_outer_scale, + scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_x, + outer_block_size=outer_block_size, + return_outer_scale=True, + ) + else: + x_dq = _qdq( + x_2d, axis=-1, + is_2d_block=use_2dblock_x, + use_outer_scale=use_outer_scale, + scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_x, + outer_block_size=outer_block_size, + ) + + w_outer_scale_cache: Optional[torch.Tensor] = None + if share_outer_w: + w_dq, w_outer_scale_cache = _qdq( + weight, axis=-1, + is_2d_block=use_2dblock_w, + use_outer_scale=use_outer_scale, + scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, + outer_block_size=outer_block_size, + return_outer_scale=True, + ) + else: + w_dq = _qdq( + weight, axis=-1, + is_2d_block=use_2dblock_w, + use_outer_scale=use_outer_scale, + scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, + outer_block_size=outer_block_size, + ) + y = x_dq @ w_dq.T # Prepare separate quantized views for the tensors that participate in @@ -121,6 +189,9 @@ def forward( is_2d_block=False, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, + outer_block_size=outer_block_size, return_raw=True, ) else: @@ -129,6 +200,10 @@ def forward( is_2d_block=False, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, + outer_block_size=outer_block_size, + precomputed_outer_scale=w_outer_scale_cache, ) else: # 2D block scaling is axis-invariant, so the fprop quantized view is @@ -142,6 +217,9 @@ def forward( is_2d_block=use_2dblock_w, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_w, + outer_block_size=outer_block_size, return_raw=True, ) w_dq_axis0 = w_dq @@ -169,11 +247,18 @@ def forward( x_for_axis0 = hadamard_transform(x_2d, left_mul=True) else: x_for_axis0 = x_2d + # ``x_outer_scale_cache`` is only set when the wgrad-axis input is + # the same tensor (no Hadamard); otherwise it stays None and the + # axis=0 QDQ recomputes the outer scale from ``x_for_axis0``. x_dq_axis0 = _qdq( x_for_axis0, axis=0, is_2d_block=False, use_outer_scale=use_outer_scale, scale_format=scale_format, + use_outer_block_scale=use_outer_block_scale, + is_2d_outer=use_outer_2dblock_x, + outer_block_size=outer_block_size, + precomputed_outer_scale=x_outer_scale_cache, ) else: x_dq_axis0 = x_dq @@ -186,6 +271,9 @@ def forward( ctx.use_2dblock_w = use_2dblock_w ctx.use_sr_grad = use_sr_grad ctx.use_outer_scale = use_outer_scale + ctx.use_outer_block_scale = use_outer_block_scale + ctx.use_outer_2dblock_x = use_outer_2dblock_x + ctx.outer_block_size = outer_block_size ctx.hadamard_transform = hadamard_transform ctx.use_dge = use_dge ctx.scale_format = scale_format @@ -219,16 +307,45 @@ def backward(ctx, grad_output): use_outer_scale=ctx.use_outer_scale, use_sr=ctx.use_sr_grad, scale_format=ctx.scale_format, + use_outer_block_scale=ctx.use_outer_block_scale, + is_2d_outer=ctx.use_outer_2dblock_x, + outer_block_size=ctx.outer_block_size, ) grad_output_m_dq = grad_output_dq else: - grad_output_dq = _qdq( - grad_output, axis=-1, - is_2d_block=False, - use_outer_scale=ctx.use_outer_scale, - use_sr=ctx.use_sr_grad, - scale_format=ctx.scale_format, + # Mirror the forward outer-scale-sharing rule: with 2D outer and + # no Hadamard rotation, grad_output's axis=-1 and axis=0 outer + # scales are identical, so compute it once and reuse. + share_grad_outer = ( + ctx.use_outer_block_scale + and ctx.use_outer_2dblock_x + and ctx.hadamard_transform is None ) + if share_grad_outer: + grad_output_dq, grad_outer_scale_cache = _qdq( + grad_output, axis=-1, + is_2d_block=False, + use_outer_scale=ctx.use_outer_scale, + use_sr=ctx.use_sr_grad, + scale_format=ctx.scale_format, + use_outer_block_scale=ctx.use_outer_block_scale, + is_2d_outer=ctx.use_outer_2dblock_x, + outer_block_size=ctx.outer_block_size, + return_outer_scale=True, + ) + else: + grad_output_dq = _qdq( + grad_output, axis=-1, + is_2d_block=False, + use_outer_scale=ctx.use_outer_scale, + use_sr=ctx.use_sr_grad, + scale_format=ctx.scale_format, + use_outer_block_scale=ctx.use_outer_block_scale, + is_2d_outer=ctx.use_outer_2dblock_x, + outer_block_size=ctx.outer_block_size, + ) + grad_outer_scale_cache = None + # Apply the same Hadamard rotation that was used on x in forward, # to grad_output before axis=0 QDQ. See the forward-side comment # for the invariance argument. @@ -242,6 +359,10 @@ def backward(ctx, grad_output): use_outer_scale=ctx.use_outer_scale, use_sr=ctx.use_sr_grad, scale_format=ctx.scale_format, + use_outer_block_scale=ctx.use_outer_block_scale, + is_2d_outer=ctx.use_outer_2dblock_x, + outer_block_size=ctx.outer_block_size, + precomputed_outer_scale=grad_outer_scale_cache, ) grad_inputs = grad_output_dq @ w_dq @@ -272,9 +393,10 @@ def backward(ctx, grad_output): grad_weights, # Match forward's positional arity: (x, weight, use_2dblock_x, # use_2dblock_w, use_sr_grad, use_outer_scale, hadamard_transform, - # use_dge, scale_format). 9 inputs → 9 grad slots, with a None - # for every non-Tensor / non-grad-tracked input. - None, None, None, None, None, None, None, + # use_dge, use_outer_block_scale, use_outer_2dblock_x, + # use_outer_2dblock_w, outer_block_size, scale_format). 13 inputs + # → 13 grad slots, with a None for every non-Tensor input. + None, None, None, None, None, None, None, None, None, None, None, ) @@ -288,6 +410,10 @@ def _to_nvfp4_then_scaled_mm( use_hadamard: bool = False, use_dge: bool = False, scale_format: str = "e4m3", + use_outer_block_scale: bool = False, + use_outer_2dblock_x: bool = False, + use_outer_2dblock_w: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, ) -> torch.Tensor: """Build the optional Hadamard transform and apply ``NVFP4LinearFunction``. @@ -310,5 +436,9 @@ def _to_nvfp4_then_scaled_mm( use_outer_scale, hadamard_transform, use_dge, + use_outer_block_scale, + use_outer_2dblock_x, + use_outer_2dblock_w, + outer_block_size, scale_format, ) diff --git a/alto/kernels/fp4/nvfp4/nvfp_quantization.py b/alto/kernels/fp4/nvfp4/nvfp_quantization.py index 3238dbf..7213632 100644 --- a/alto/kernels/fp4/nvfp4/nvfp_quantization.py +++ b/alto/kernels/fp4/nvfp4/nvfp_quantization.py @@ -20,6 +20,7 @@ BLOCK_SIZE_DEFAULT = 16 +OUTER_BLOCK_SIZE_DEFAULT = 128 # Host-facing numeric constants -- kept as plain Python floats so eager/host # code can use them directly (e.g. ``amax / (F8E4M3_MAX * F4_E2M1_MAX)``). F4_E2M1_MAX = 6.0 @@ -49,21 +50,28 @@ # (NVFP4 spec ``s_block``). Value lives on the E4M3 grid # in an FP32 container (see outer-side notes below). # outer_scale -- the outer-level scale factor that sits above -# ``inner_scale`` (NVFP4 spec ``s_global``). Currently a -# per-tensor FP32 scalar; the name is intentionally -# agnostic so the same API can later carry an -# outer-blockwise layout (e.g. one scale per 128x128 -# tile) without renaming the public surface. +# ``inner_scale`` (NVFP4 spec ``s_global``). May be a +# per-tensor FP32 scalar (tensorwise) or a per-outer-block +# FP32 grid (outer-block); the name is intentionally +# agnostic so the same API carries both layouts. # -# Per spec, ``outer_scale`` lives in FP32; this floor is a div-by-zero -# guard for the downstream ``max_abs / outer_scale`` when ``amax == 0``. -# We pick ``1e-30`` (well above FP32 denormal range and ~22 orders below -# any natural training-time outer scale) so that the effective per-block -# divisor ``quant_scale = inner_scale * outer_scale`` stays an FP32 -# *normal* in the worst case (``inner_scale == E4M3_EPS``, -# ``outer_scale == _OUTER_SCALE_DIVZERO_FLOOR``): +# Single div-by-zero floor shared by BOTH the tensorwise and outer-block +# outer-scale paths -- they compute the same ``amax / (E4M3_MAX*E2M1_MAX)`` +# scale and must clamp dead (all-zero) blocks identically; using two +# different constants for one concept was a latent source of confusion. +# The outer scale is a free FP32 number with no E4M3 round-trip, so the +# floor only needs to (a) avoid div-by-zero on dead blocks and (b) keep the +# effective per-block divisor ``quant_scale = inner_scale * outer_scale`` an +# FP32 *normal* in the worst case (``inner_scale == E4M3_EPS``, +# ``outer_scale == floor``): # 1e-30 * 2**-6 ≈ 1.56e-32 > FP32 smallest normal (~1.18e-38) +# 1e-30 also sits ~22 orders below any natural training-time outer scale, so +# it does not compress small-magnitude blocks the way the E4M3 floor +# (~1.95e-3) would. _OUTER_SCALE_DIVZERO_FLOOR = 1.0e-30 +# Public alias kept for the reference-impl / test surface (same value, same +# concept as the private floor above). +OUTER_SCALE_EPS = _OUTER_SCALE_DIVZERO_FLOOR def _check_scale_format(scale_format: str) -> None: if scale_format not in _SCALE_FORMAT_TABLE: @@ -487,6 +495,139 @@ def compute_dynamic_outer_scale( return outer_scale.to(dtype=torch.float32).reshape(1) +def compute_outer_block_scale_shape( + shape: tuple[int, ...], + *, + is_2d_inner: bool, + is_2d_outer: bool, + inner_block_size: int = BLOCK_SIZE_DEFAULT, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, +) -> tuple[tuple[int, ...], tuple[int, ...]]: + """Return ``(inner_scale_shape, outer_scale_shape)`` for axis-last input. + + ``shape`` is expected to be the logical high-precision tensor shape after + transposing the quantization axis to the last dimension. The outer scale + covers ``1 × outer_block_size`` tiles when ``is_2d_outer=False`` and + ``outer_block_size × outer_block_size`` tiles when ``is_2d_outer=True``. + ``is_2d_inner`` only controls the inner NVFP4 block scale layout; it does + not change the outer scale's M-dimension granularity. + """ + if len(shape) < 2: + raise ValueError(f"outer-block scaling expects at least 2D input, got {shape}") + if outer_block_size < inner_block_size or outer_block_size % inner_block_size != 0: + raise ValueError( + f"outer_block_size ({outer_block_size}) must be a multiple of " + f"inner_block_size ({inner_block_size})" + ) + m, n = shape[-2], shape[-1] + if n % inner_block_size != 0 or n % outer_block_size != 0: + raise ValueError( + f"last dimension ({n}) must be divisible by both inner block " + f"({inner_block_size}) and outer block ({outer_block_size})" + ) + if is_2d_inner and m % inner_block_size != 0: + raise ValueError( + f"2D inner block requires dim -2 ({m}) divisible by {inner_block_size}" + ) + if is_2d_outer and m % outer_block_size != 0: + raise ValueError( + f"2D outer block requires dim -2 ({m}) divisible by {outer_block_size}" + ) + + prefix = shape[:-2] + inner_shape = ( + *prefix, + m // inner_block_size if is_2d_inner else m, + n // inner_block_size, + ) + outer_shape = ( + *prefix, + m // outer_block_size if is_2d_outer else m, + n // outer_block_size, + ) + return inner_shape, outer_shape + + +def _outer_scale_and_expanded_axis_last( + data_hp: torch.Tensor, + *, + is_2d_inner: bool, + is_2d_outer: bool, + inner_block_size: int, + outer_block_size: int, + scale_format: str = "e4m3", +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute FP32 outer scales and expand them to ``data_hp`` shape. + + ``data_hp`` must already have the quantization axis in the last dimension. + + ``scale_format`` selects the inner-block scale grid the outer scale must + normalize *into*: the per-outer-block divisor is ``amax / (scale_max * + F4_E2M1_MAX)`` where ``scale_max`` is the inner grid's max (E4M3 -> 448 for + NVFP4, UE5M3 -> 114688 for AMD-FP4). Using the right max keeps the inner + grid fully utilized instead of saturating (E4M3) or wasting range (UE5M3). + """ + _check_scale_format(scale_format) + _, scale_max = _SCALE_FORMAT_TABLE[scale_format] + shape = tuple(data_hp.shape) + compute_outer_block_scale_shape( + shape, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=inner_block_size, + outer_block_size=outer_block_size, + ) + m, n = shape[-2], shape[-1] + prefix = shape[:-2] + x = data_hp.float() + + if is_2d_outer: + grouped = x.reshape( + *prefix, + m // outer_block_size, + outer_block_size, + n // outer_block_size, + outer_block_size, + ) + outer_amax = grouped.abs().amax(dim=-1).amax(dim=-2) + outer_scale = (outer_amax / (scale_max * F4_E2M1_MAX)).clamp(min=OUTER_SCALE_EPS) + expanded = outer_scale.unsqueeze(-2).unsqueeze(-1).expand_as(grouped).reshape(shape) + else: + grouped = x.reshape(*prefix, m, n // outer_block_size, outer_block_size) + outer_amax = grouped.abs().amax(dim=-1) + outer_scale = (outer_amax / (scale_max * F4_E2M1_MAX)).clamp(min=OUTER_SCALE_EPS) + expanded = outer_scale.unsqueeze(-1).expand_as(grouped).reshape(shape) + + return outer_scale.to(torch.float32), expanded.to(torch.float32) + + +def _expand_outer_scale_axis_last( + outer_scale: torch.Tensor, + hp_shape: tuple[int, ...], + *, + is_2d_inner: bool, + is_2d_outer: bool, + inner_block_size: int, + outer_block_size: int, +) -> torch.Tensor: + """Broadcast stored outer scales back to an axis-last high-precision shape.""" + compute_outer_block_scale_shape( + hp_shape, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=inner_block_size, + outer_block_size=outer_block_size, + ) + m, n = hp_shape[-2], hp_shape[-1] + prefix = hp_shape[:-2] + if is_2d_outer: + grouped_shape = (*prefix, m // outer_block_size, outer_block_size, + n // outer_block_size, outer_block_size) + return outer_scale.unsqueeze(-2).unsqueeze(-1).expand(grouped_shape).reshape(hp_shape) + grouped_shape = (*prefix, m, n // outer_block_size, outer_block_size) + return outer_scale.unsqueeze(-1).expand(grouped_shape).reshape(hp_shape) + + @triton_op("alto::convert_to_nvfp4", mutates_args={}) def convert_to_nvfp4( data_hp: torch.Tensor, @@ -589,8 +730,12 @@ def convert_to_nvfp4( M, N = data_hp.shape grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]), triton.cdiv(N, META["BLOCK_N"])) - BLOCK_M = 64 if M >= 64 else M - BLOCK_N = 64 if N >= 64 else N + # Triton requires tl.arange(0, BLOCK_*) ranges to be powers of 2. When the + # tensor is smaller than the 64-tile, round the tile up to the next power of + # 2 (the load/store masks already guard the ragged tail) instead of using + # the raw dimension, which would crash for non-power-of-2 sizes (e.g. M=24). + BLOCK_M = 64 if M >= 64 else triton.next_power_of_2(M) + BLOCK_N = 64 if N >= 64 else triton.next_power_of_2(N) if philox_seed is None: philox_seed = torch.randint(0, 2**31 - 1, (1,)).item() @@ -682,8 +827,12 @@ def convert_from_nvfp4( M, N = data_hp.shape grid = lambda META: (triton.cdiv(M, META["BLOCK_M"]), triton.cdiv(N, META["BLOCK_N"])) - BLOCK_M = 64 if M >= 64 else M - BLOCK_N = 64 if N >= 64 else N + # Triton requires tl.arange(0, BLOCK_*) ranges to be powers of 2. When the + # tensor is smaller than the 64-tile, round the tile up to the next power of + # 2 (the load/store masks already guard the ragged tail) instead of using + # the raw dimension, which would crash for non-power-of-2 sizes (e.g. M=24). + BLOCK_M = 64 if M >= 64 else triton.next_power_of_2(M) + BLOCK_N = 64 if N >= 64 else triton.next_power_of_2(N) wrap_triton(_convert_from_nvfp4_kernel)[grid]( data_lp, @@ -711,6 +860,154 @@ def convert_from_nvfp4( return data_hp.reshape(orig_shape_hp).transpose(axis, -1) +def convert_to_nvfp4_outer_block( + data_hp: torch.Tensor, + block_size: int = BLOCK_SIZE_DEFAULT, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, + axis: int = -1, + is_2d_inner: bool = False, + is_2d_outer: bool = False, + use_sr: bool = False, + philox_seed: Optional[int] = None, + philox_offset: Optional[int] = None, + precomputed_outer_scale: Optional[torch.Tensor] = None, + scale_format: str = "e4m3", +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize using FP32 outer-block scale + NVFP4/AMD-FP4 inner scales. + + This implements the M2 functional path for the double-block NVFP4 scheme: + values are first normalized by a per-outer-block FP32 scale + ``amax_outer / (F8E4M3_MAX * F4_E2M1_MAX)`` and then quantized by the + existing per-inner-block NVFP4 kernel with tensor-wise scaling disabled. + + This is a functional/reference QDQ path, not a fused training kernel. It + materializes full-size normalized and expanded-scale tensors, and it does + not pad ragged tails; callers must provide dimensions divisible by the + requested inner and outer block sizes. + + When ``precomputed_outer_scale`` is supplied, the outer-amax reduction is + skipped and the caller-provided scale is used instead. The tensor must + be in the same axis layout as the returned ``outer_scale`` and have the + shape implied by :func:`compute_outer_block_scale_shape`. Used by + :class:`NVFP4LinearFunction` to share an axis-invariant outer scale + across the axis=-1 and axis=0 QDQ calls. + """ + assert data_hp.dtype in [torch.float32, torch.bfloat16] + data_axis_last = data_hp.transpose(axis, -1) + original_shape = data_axis_last.shape + expected_inner_shape, expected_outer_shape = compute_outer_block_scale_shape( + tuple(original_shape), + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=block_size, + outer_block_size=outer_block_size, + ) + del expected_inner_shape # only the outer shape is validated here + + if precomputed_outer_scale is None: + outer_scale, outer_expanded = _outer_scale_and_expanded_axis_last( + data_axis_last, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=block_size, + outer_block_size=outer_block_size, + scale_format=scale_format, + ) + else: + # Caller-supplied scale is in the original axis layout (matching the + # function's output convention). Transpose to axis-last for internal + # use, then validate it matches the expected ``axis-last`` outer + # scale shape. + provided_axis_last = precomputed_outer_scale.transpose(axis, -1).to( + device=data_hp.device, dtype=torch.float32 + ) + if tuple(provided_axis_last.shape) != expected_outer_shape: + raise ValueError( + "precomputed_outer_scale (after transpose to axis-last) has " + f"shape {tuple(provided_axis_last.shape)}, expected " + f"{expected_outer_shape} for input shape {tuple(original_shape)} " + f"with is_2d_inner={is_2d_inner}, is_2d_outer={is_2d_outer}, " + f"outer_block_size={outer_block_size}." + ) + outer_scale = provided_axis_last.contiguous() + outer_expanded = _expand_outer_scale_axis_last( + outer_scale, + tuple(original_shape), + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=block_size, + outer_block_size=outer_block_size, + ) + # Keep the per-outer-block normalization in FP32 through the inner QDQ. + # ``convert_to_nvfp4`` accepts FP32 and reads operands as FP32 internally, + # so casting back down to bf16 here would only add a redundant rounding + # step on top of the NVFP4 inner-block quantization, needlessly lowering + # fidelity of the normalized values handed to the inner kernel. + normalized = data_axis_last.float() / outer_expanded + data_lp, inner_scales = convert_to_nvfp4( + normalized, + block_size=block_size, + axis=-1, + is_2d_block=is_2d_inner, + outer_scale=None, + update_outer_scale=False, + scale_format=scale_format, + use_sr=use_sr, + philox_seed=philox_seed, + philox_offset=philox_offset, + ) + return ( + data_lp.transpose(axis, -1), + inner_scales.transpose(axis, -1), + outer_scale.transpose(axis, -1), + ) + + +def convert_from_nvfp4_outer_block( + data_lp: torch.Tensor, + inner_scales: torch.Tensor, + outer_scales: torch.Tensor, + output_dtype: torch.dtype = torch.float32, + block_size: int = BLOCK_SIZE_DEFAULT, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, + axis: int = -1, + is_2d_inner: bool = False, + is_2d_outer: bool = False, + scale_format: str = "e4m3", +) -> torch.Tensor: + """Dequantize data produced by :func:`convert_to_nvfp4_outer_block`. + + ``scale_format`` is accepted for API symmetry with + :func:`convert_to_nvfp4_outer_block`; dequant only multiplies by the stored + FP32 inner/outer scales, so the grid choice does not change the math. + """ + normalized = convert_from_nvfp4( + data_lp, + inner_scales, + output_dtype=output_dtype, + block_size=block_size, + axis=axis, + is_2d_block=is_2d_inner, + outer_scale=None, + scale_format=scale_format, + ) + normalized_axis_last = normalized.transpose(axis, -1) + outer_axis_last = outer_scales.transpose(axis, -1).to( + device=normalized.device, + dtype=torch.float32, + ) + outer_expanded = _expand_outer_scale_axis_last( + outer_axis_last, + tuple(normalized_axis_last.shape), + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=block_size, + outer_block_size=outer_block_size, + ) + result = (normalized_axis_last.float() * outer_expanded).to(output_dtype) + return result.transpose(axis, -1) + + @convert_to_nvfp4.register_fake def _fake_convert_to_nvfp4( data_hp: torch.Tensor, @@ -765,7 +1062,16 @@ def _qdq( block_size: int = 16, return_raw: bool = False, scale_format: str = "e4m3", -) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + use_outer_block_scale: bool = False, + is_2d_outer: bool = False, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, + precomputed_outer_scale: Optional[torch.Tensor] = None, + return_outer_scale: bool = False, +) -> Union[ + torch.Tensor, + Tuple[torch.Tensor, torch.Tensor, torch.Tensor], + Tuple[torch.Tensor, torch.Tensor], +]: """Quantize to NVFP4 then immediately dequantize back (QDQ round-trip). This helper is the core building block used to emulate the numerical effect @@ -779,7 +1085,77 @@ def _qdq( returned alongside the dequantized BF16 view. This is used by the DGE backward path to recover the exact FP4 bin each weight element fell into without paying a second quantization pass. + + Outer-block scaling exposes two extra kwargs: + + * ``precomputed_outer_scale`` skips the outer-amax reduction and reuses + the caller-supplied scale tensor. The tensor layout matches the + ``outer_scale`` returned by :func:`convert_to_nvfp4_outer_block`. + * ``return_outer_scale`` returns ``(dq, outer_scale)`` so the caller can + cache the scale for axis-invariant reuse. Mutually exclusive with + ``return_raw``. """ + torch._check( + not (use_outer_scale and use_outer_block_scale), + lambda: "NVFP4 tensor-wise scale and outer-block scale are mutually exclusive.", + ) + torch._check( + not (return_raw and return_outer_scale), + lambda: ( + "_qdq cannot return both raw FP4 (return_raw) and outer scale " + "(return_outer_scale) at the same time." + ), + ) + torch._check( + use_outer_block_scale or precomputed_outer_scale is None, + lambda: ( + "precomputed_outer_scale is only meaningful when " + "use_outer_block_scale=True." + ), + ) + torch._check( + use_outer_block_scale or not return_outer_scale, + lambda: ( + "return_outer_scale is only meaningful when " + "use_outer_block_scale=True." + ), + ) + + if use_outer_block_scale: + torch._check( + not return_raw, + lambda: ( + "NVFP4 outer-block QDQ does not yet support return_raw=True; " + "DGE integration for outer-block scaling needs a separate raw-scale path." + ), + ) + data_lp, inner_scales, outer_scales = convert_to_nvfp4_outer_block( + tensor, + block_size=block_size, + outer_block_size=outer_block_size, + axis=axis, + is_2d_inner=is_2d_block, + is_2d_outer=is_2d_outer, + use_sr=use_sr, + precomputed_outer_scale=precomputed_outer_scale, + scale_format=scale_format, + ) + dq = convert_from_nvfp4_outer_block( + data_lp, + inner_scales, + outer_scales, + output_dtype=tensor.dtype, + block_size=block_size, + outer_block_size=outer_block_size, + axis=axis, + is_2d_inner=is_2d_block, + is_2d_outer=is_2d_outer, + scale_format=scale_format, + ) + if return_outer_scale: + return dq, outer_scales + return dq + outer_scale = ( torch.empty(1, dtype=torch.float32, device=tensor.device) if use_outer_scale diff --git a/alto/kernels/fp4/testing_utils.py b/alto/kernels/fp4/testing_utils.py index cb4a3a2..2d02973 100644 --- a/alto/kernels/fp4/testing_utils.py +++ b/alto/kernels/fp4/testing_utils.py @@ -88,6 +88,20 @@ def _nvfp4_autograd_snr_thresholds( return (10.0, 12.0) if use_sr_grad else (12.0, 14.0) return (12.0, 14.0) if use_sr_grad else (14.0, 15.0) + if kind == "nvfp4_linear_outer_block": + # Outer-block scaling on a per-128 grid contains every outlier within + # a single FP32 outer scale, so the inner E4M3 scale always sits in + # its sweet spot. Empirically this is +2 dB tighter than the PTS + # linear path; reflect that in the regression contract so a quiet + # regression on the outer-block QDQ path (e.g. a clamp bug, a + # mismatched axis) is caught even when the broader nvfp4_linear + # thresholds would still pass. + if K >= 1024: + return (5.0, 6.5) if use_sr_grad else (6.0, 10.0) + if K >= 256: + return (12.0, 14.0) if use_sr_grad else (14.0, 16.0) + return (14.0, 16.0) if use_sr_grad else (16.0, 17.0) + if kind == "nvfp4_grouped_gemm": if use_outer_scale: # Grouped paths (especially 2D-block + large G) trail linear outer diff --git a/alto/models/gpt_oss/configs/lpt_recipe_amdfp4_outer_block.yaml b/alto/models/gpt_oss/configs/lpt_recipe_amdfp4_outer_block.yaml new file mode 100644 index 0000000..7be799d --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_recipe_amdfp4_outer_block.yaml @@ -0,0 +1,27 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "amdfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # AMD-FP4 (UE5M3 inner grid) with per-outer-block FP32 scale (incl. MoE + # experts). Exact mirror of lpt_recipe_outer_block.yaml (NVFP4) except + # scheme="amdfp4": the outer-block code path is shared across both inner + # grids, only the per-outer-block normalization denominator differs + # (E4M3 max 448 vs UE5M3 max 114688). + ignore: + - "output" + - "re:.*\\.router\\.gate" + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: outer_block + use_outer_2dblock_x: false + use_outer_2dblock_w: true + # outer_block_size must divide every quantized reduction dim and be a + # multiple of the NVFP4 inner block (16). GPT-OSS hidden=2880 is not + # divisible by 128, so use 64 (2880/64=45; debug dim=256, 256/64=4). + outer_block_size: 64 + lora_rank: 0 diff --git a/alto/models/gpt_oss/configs/lpt_recipe_outer_block.yaml b/alto/models/gpt_oss/configs/lpt_recipe_outer_block.yaml new file mode 100644 index 0000000..ddc8559 --- /dev/null +++ b/alto/models/gpt_oss/configs/lpt_recipe_outer_block.yaml @@ -0,0 +1,25 @@ +training_stage: + lpt_modifiers: + LowPrecisionTrainingModifier: + scheme: "nvfp4" + targets: ["Linear", "GptOssGroupedExperts"] + # All-layers NVFP4 with per-outer-block FP32 scale (incl. MoE experts), + # instead of one tensorwise scalar. Exact mirror of the tensorwise + # all-layers recipe except two_level_scaling + the outer-block knobs. + ignore: + - "output" + - "re:.*\\.router\\.gate" + use_2dblock_x: false + use_2dblock_w: true + use_hadamard: true + use_sr_grad: true + use_dge: false + clip_mode: none + two_level_scaling: outer_block + use_outer_2dblock_x: false + use_outer_2dblock_w: true + # outer_block_size must divide every quantized reduction dim; GPT-OSS + # hidden=2880 is not divisible by 128, so use 64 (2880/64=45) -- still a + # multiple of the NVFP4 inner block (16). + outer_block_size: 64 + lora_rank: 0 diff --git a/alto/modifiers/lpt/base.py b/alto/modifiers/lpt/base.py index 319d083..c98b780 100644 --- a/alto/modifiers/lpt/base.py +++ b/alto/modifiers/lpt/base.py @@ -35,9 +35,12 @@ class LowPrecisionTrainingModifier(Modifier): use_hadamard: bool = False use_sr_grad: bool = False use_dge: bool = False - two_level_scaling: Literal["none", "tensorwise", "blockwise"] = "none" + two_level_scaling: Literal["none", "tensorwise", "blockwise", "outer_block"] = "none" clip_mode: Literal["none", "static", "dynamic"] = "none" - + use_outer_2dblock_x: bool = False + use_outer_2dblock_w: bool = False + outer_block_size: int = 128 + lora_rank: int = 0 """ Lora rank for the decomposed linear layer. @@ -93,6 +96,60 @@ def validate_scheme(cls, value: str | dict[str, str | list[str]]) -> str | dict[ return value + @model_validator(mode="after") + def validate_outer_block_scaling(self): + if self.outer_block_size < 16 or self.outer_block_size % 16 != 0: + raise ValueError( + "outer_block_size must be a positive multiple of the NVFP4 " + f"inner block size (16), got {self.outer_block_size}" + ) + outer_fields_are_active = ( + self.use_outer_2dblock_x + or self.use_outer_2dblock_w + or self.outer_block_size != 128 + ) + if self.two_level_scaling != "outer_block" and outer_fields_are_active: + raise ValueError( + "use_outer_2dblock_x, use_outer_2dblock_w, and non-default " + "outer_block_size are only meaningful when " + 'two_level_scaling="outer_block".' + ) + if self.two_level_scaling == "outer_block": + schemes = [self.scheme] if isinstance(self.scheme, str) else list(self.scheme) + invalid_schemes = [scheme for scheme in schemes if scheme not in ("nvfp4", "amdfp4")] + if invalid_schemes: + raise ValueError( + 'two_level_scaling="outer_block" is only defined for ' + f"scheme in {{'nvfp4', 'amdfp4'}}, got {invalid_schemes}" + ) + # NOTE: outer-block scaling now supports BOTH dense Linear and the + # MoE grouped-GEMM path (grouped uses a 1D outer block on + # activations and use_outer_2dblock_w on the per-expert weights), + # so MoE / grouped-experts targets are no longer rejected here. + return self + + def _collect_moe_like_targets(self) -> list[str]: + """Return targets whose name looks like an MoE / GroupedExperts module. + + Match is a case-insensitive substring check on ``"groupedexperts"`` or + ``"moe"``, so regex patterns such as ``re:.*GroupedExperts`` are + matched literally without parsing them. + """ + target_groups: list[list[str]] = [] + if isinstance(self.scheme, dict): + target_groups.extend(list(self.scheme.values())) + else: + target_groups.append( + self.targets if isinstance(self.targets, list) else [self.targets] + ) + offending: list[str] = [] + for group in target_groups: + for entry in group: + lowered = entry.lower() + if "groupedexperts" in lowered or "moe" in lowered: + offending.append(entry) + return offending + @model_validator(mode="after") def validate_lora_rank_alignment(self): if self.lora_rank <= 0: @@ -158,6 +215,9 @@ def resolved_config(self) -> dict[TrainingOpConfig, list[str]]: use_dge=self.use_dge, two_level_scaling=self.two_level_scaling, clip_mode=self.clip_mode, + use_outer_2dblock_x=self.use_outer_2dblock_x, + use_outer_2dblock_w=self.use_outer_2dblock_w, + outer_block_size=self.outer_block_size, ) self._resolved_config[scheme_obj] = targets return self._resolved_config diff --git a/alto/modifiers/lpt/tests/test_outer_block_validation.py b/alto/modifiers/lpt/tests/test_outer_block_validation.py new file mode 100644 index 0000000..3b304b7 --- /dev/null +++ b/alto/modifiers/lpt/tests/test_outer_block_validation.py @@ -0,0 +1,125 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +"""Validation tests for NVFP4 outer-block LPT recipe fields.""" + +import pytest + +pytest.importorskip("torchtitan") + +from alto.modifiers.lpt.base import LowPrecisionTrainingModifier + + +def test_outer_block_side_fields_require_outer_block_mode(): + with pytest.raises(ValueError, match='two_level_scaling="outer_block"'): + LowPrecisionTrainingModifier( + scheme="nvfp4", + two_level_scaling="tensorwise", + use_outer_2dblock_x=True, + ) + + +def test_non_default_outer_block_size_requires_outer_block_mode(): + with pytest.raises(ValueError, match='two_level_scaling="outer_block"'): + LowPrecisionTrainingModifier( + scheme="nvfp4", + two_level_scaling="none", + outer_block_size=256, + ) + + +def test_outer_block_mode_requires_fp4_scheme(): + with pytest.raises(ValueError, match="only defined for scheme"): + LowPrecisionTrainingModifier( + scheme="mxfp4", + two_level_scaling="outer_block", + ) + + +def test_outer_block_mode_rejects_mixed_scheme_dict(): + with pytest.raises(ValueError, match="only defined for scheme"): + LowPrecisionTrainingModifier( + scheme={"nvfp4": ["Linear"], "mxfp4": ["OtherLinear"]}, + two_level_scaling="outer_block", + ) + + +def test_outer_block_mode_accepts_nvfp4_recipe(): + modifier = LowPrecisionTrainingModifier( + scheme="nvfp4", + two_level_scaling="outer_block", + use_outer_2dblock_x=False, + use_outer_2dblock_w=True, + outer_block_size=128, + ) + assert modifier.two_level_scaling == "outer_block" + + +def test_outer_block_mode_accepts_amdfp4_recipe(): + """AMD-FP4 shares the NVFP4 outer-block code path (UE5M3 inner grid).""" + modifier = LowPrecisionTrainingModifier( + scheme="amdfp4", + two_level_scaling="outer_block", + use_outer_2dblock_x=False, + use_outer_2dblock_w=True, + outer_block_size=128, + ) + assert modifier.two_level_scaling == "outer_block" + + +# --------------------------------------------------------------------------- +# Outer-block scaling now supports the MoE grouped-GEMM path (grouped uses a +# 1D outer block on activations and use_outer_2dblock_w on the per-expert +# weights), so MoE / GroupedExperts targets are accepted under outer-block. +# --------------------------------------------------------------------------- + + +def test_outer_block_accepts_targets_with_grouped_experts(): + modifier = LowPrecisionTrainingModifier( + scheme="nvfp4", + targets=["Linear", "GptOssGroupedExperts"], + two_level_scaling="outer_block", + use_outer_2dblock_w=True, + ) + assert modifier.two_level_scaling == "outer_block" + + +def test_outer_block_accepts_dict_scheme_with_moe_targets(): + modifier = LowPrecisionTrainingModifier( + scheme={"nvfp4": ["Linear", "DeepseekV3GroupedExperts"]}, + two_level_scaling="outer_block", + use_outer_2dblock_w=True, + ) + assert modifier.two_level_scaling == "outer_block" + + +def test_outer_block_accepts_regex_matching_grouped_experts(): + modifier = LowPrecisionTrainingModifier( + scheme="nvfp4", + targets=["Linear", "re:.*GroupedExperts$"], + two_level_scaling="outer_block", + use_outer_2dblock_w=True, + ) + assert modifier.two_level_scaling == "outer_block" + + +def test_outer_block_accepts_targets_without_moe(): + """Sanity: dense-only targets are still accepted under outer-block.""" + modifier = LowPrecisionTrainingModifier( + scheme="nvfp4", + targets=["Linear"], + two_level_scaling="outer_block", + use_outer_2dblock_w=True, + ) + assert modifier.two_level_scaling == "outer_block" + + +def test_pts_recipe_is_unchanged_by_p0_4_validator(): + """tensorwise / none recipes must keep accepting MoE targets unchanged.""" + modifier = LowPrecisionTrainingModifier( + scheme="nvfp4", + targets=["Linear", "GptOssGroupedExperts"], + two_level_scaling="tensorwise", + ) + assert modifier.two_level_scaling == "tensorwise" diff --git a/tests/unittest/amdfp4/test_amdfp_dispatch_guards.py b/tests/unittest/amdfp4/test_amdfp_dispatch_guards.py index ac087bc..771fe07 100644 --- a/tests/unittest/amdfp4/test_amdfp_dispatch_guards.py +++ b/tests/unittest/amdfp4/test_amdfp_dispatch_guards.py @@ -96,7 +96,9 @@ def test_amdfp4_linear_routes_to_amdfp4_thin_wrapper(monkeypatch, device): calls = [] def _mock_amdfp4(A, W, *, use_2dblock_x, use_2dblock_w, - use_sr_grad, use_outer_scale, use_hadamard, use_dge): + use_sr_grad, use_outer_scale, use_hadamard, use_dge, + use_outer_block_scale=False, use_outer_2dblock_x=False, + use_outer_2dblock_w=False, outer_block_size=128): calls.append({ "use_2dblock_x": use_2dblock_x, "use_2dblock_w": use_2dblock_w, @@ -139,7 +141,9 @@ def test_amdfp4_grouped_mm_routes_to_amdfp4_thin_wrapper(monkeypatch, device): calls = [] def _mock_amdfp4_grouped(A, B, *, offs, use_2dblock_x, use_2dblock_w, - use_sr_grad, use_outer_scale, use_hadamard, use_dge): + use_sr_grad, use_outer_scale, use_hadamard, use_dge, + use_outer_block_scale=False, use_outer_2dblock_w=False, + outer_block_size=128): calls.append({ "offs_len": offs.numel(), "use_2dblock_x": use_2dblock_x, diff --git a/tests/unittest/amdfp4/test_amdfp_linear.py b/tests/unittest/amdfp4/test_amdfp_linear.py index a40c3fa..7318600 100644 --- a/tests/unittest/amdfp4/test_amdfp_linear.py +++ b/tests/unittest/amdfp4/test_amdfp_linear.py @@ -116,6 +116,13 @@ def test_amdfp4_linear_autograd_function( False, # use_outer_scale None, # hadamard_transform False, # use_dge + # ``scale_format`` is the LAST positional arg of the forward (appended + # after the outer-block params); apply() takes no kwargs, so the + # intervening outer-block params are filled with their defaults here. + False, # use_outer_block_scale + False, # use_outer_2dblock_x + False, # use_outer_2dblock_w + 128, # outer_block_size (OUTER_BLOCK_SIZE_DEFAULT) "ue5m3", # scale_format pinned to AMD-FP4 inner grid ) loss = torch.nn.functional.mse_loss(outputs, target) @@ -206,3 +213,47 @@ def test_amdfp4_to_scaled_mm_wrapper_pins_ue5m3(use_2dblock_x, use_2dblock_w): "AMD-FP4 wrapper output must differ from the E4M3 NVFP4 path; " "identical output suggests scale_format pinning was lost" ) + + +@pytest.mark.parametrize("use_outer_2dblock_w", [False, True]) +def test_amdfp4_to_scaled_mm_wrapper_outer_block(use_outer_2dblock_w): + """AMD-FP4 outer-block scaling shares the NVFP4 outer-block code path. + + The AMD-FP4 wrapper must thread BOTH the outer-block knobs and the UE5M3 + inner grid into ``_to_nvfp4_then_scaled_mm``, so its output must be + bit-identical to the NVFP4 helper called with ``scale_format="ue5m3"`` and + the same outer-block config. + + NOTE: under outer-block scaling the per-outer-block FP32 scale normalizes + by the inner grid's max (448 for E4M3, 114688 for UE5M3 -- a factor of + 256 = 2**8). Because that factor is an exact power of two and both grids + share the same 3-bit mantissa, the E4M3 and UE5M3 outer-block paths produce + bit-identical results on well-ranged data (the grid choice is absorbed by + the outer scale). We therefore assert UE5M3 parity but do NOT assert + divergence from E4M3 here -- that divergence only appears at intra-outer- + block dynamic ranges wide enough to underflow one grid but not the other. + """ + torch.manual_seed(0) + # Dims chosen so the outer grids are non-degenerate at outer_block_size=128: + # K > 128 (1D outer on activations) and N,K > 128 (2D outer on weights). + M, N, K = 256, 256, 256 + a = torch.randn((M, K), device="cuda").to(torch.bfloat16) + w = torch.randn((N, K), device="cuda").to(torch.bfloat16) + + ob = dict( + use_2dblock_x=False, + use_2dblock_w=False, + use_sr_grad=False, + use_outer_block_scale=True, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=128, + ) + y_amd = _to_amdfp4_then_scaled_mm(a, w, **ob) + y_nv_ue5m3 = _to_nvfp4_then_scaled_mm(a, w, scale_format="ue5m3", **ob) + + assert y_amd.shape == (M, N) + assert torch.isfinite(y_amd).all() + assert torch.equal(y_amd, y_nv_ue5m3), ( + "_to_amdfp4_then_scaled_mm outer-block output must be bit-identical to " + "_to_nvfp4_then_scaled_mm(scale_format='ue5m3') with the same outer-block config" + ) diff --git a/tests/unittest/nvfp4/test_nvfp_dispatch_guards.py b/tests/unittest/nvfp4/test_nvfp_dispatch_guards.py index cecb649..c0032f8 100644 --- a/tests/unittest/nvfp4/test_nvfp_dispatch_guards.py +++ b/tests/unittest/nvfp4/test_nvfp_dispatch_guards.py @@ -76,7 +76,8 @@ def test_grouped_mm_routes_to_nvfp4_grouped_kernel(monkeypatch, device): def _mock_grouped(A, B, *, offs, use_2dblock_x, use_2dblock_w, use_sr_grad, use_outer_scale, use_hadamard, use_dge, - scale_format="e4m3"): + use_outer_block_scale=False, use_outer_2dblock_w=False, + outer_block_size=128, scale_format="e4m3"): calls.append({ "A": A, "B": B, @@ -87,6 +88,9 @@ def _mock_grouped(A, B, *, offs, use_2dblock_x, use_2dblock_w, "use_outer_scale": use_outer_scale, "use_hadamard": use_hadamard, "use_dge": use_dge, + "use_outer_block_scale": use_outer_block_scale, + "use_outer_2dblock_w": use_outer_2dblock_w, + "outer_block_size": outer_block_size, }) return A.new_full((A.shape[0], B.shape[-1]), 7.0) @@ -123,6 +127,51 @@ def _mock_grouped(A, B, *, offs, use_2dblock_x, use_2dblock_w, assert call["use_outer_scale"] is True assert call["use_hadamard"] is True assert call["use_dge"] is True + # tensorwise recipe -> outer-block path is off + assert call["use_outer_block_scale"] is False + + +def test_grouped_mm_routes_outer_block_config(monkeypatch, device): + """Grouped GEMM must route ``two_level_scaling='outer_block'`` to the + grouped path with the outer-block knobs (no longer rejected).""" + calls = [] + + def _mock_grouped(A, B, *, offs, use_2dblock_x, use_2dblock_w, + use_sr_grad, use_outer_scale, use_hadamard, use_dge, + use_outer_block_scale=False, use_outer_2dblock_w=False, + outer_block_size=128): + calls.append({ + "use_outer_scale": use_outer_scale, + "use_outer_block_scale": use_outer_block_scale, + "use_outer_2dblock_w": use_outer_2dblock_w, + "outer_block_size": outer_block_size, + }) + return A.new_full((A.shape[0], B.shape[-1]), 3.0) + + monkeypatch.setattr(dispatch_tensor, "_quantize_then_nvfp4_scaled_grouped_mm", _mock_grouped) + + cfg = _make_config( + two_level_scaling="outer_block", + use_outer_2dblock_x=False, + use_outer_2dblock_w=True, + outer_block_size=64, + ) + num_experts, K, N, M = 2, 64, 32, ALIGN_SIZE_M + A = torch.randn(M, K, dtype=torch.bfloat16, device=device) + W_wrapped = NVFP4TrainingWeightWrapperTensor( + torch.randn(num_experts, K, N, dtype=torch.bfloat16, device=device), + cfg, + ) + offs = torch.tensor([M // num_experts, M], dtype=torch.int32, device=device) + + y = torch._grouped_mm(A, W_wrapped, offs=offs) + assert y.shape == (M, N) + assert len(calls) == 1 + call = calls[0] + assert call["use_outer_block_scale"] is True + assert call["use_outer_scale"] is False + assert call["use_outer_2dblock_w"] is True + assert call["outer_block_size"] == 64 @pytest.mark.parametrize("use_hadamard,use_dge", [ @@ -192,6 +241,8 @@ def test_linear_routes_to_nvfp4_linear_kernel(monkeypatch, device): def _mock_linear(A, B, *, use_2dblock_x, use_2dblock_w, use_sr_grad, use_outer_scale, use_hadamard, use_dge, + use_outer_block_scale, use_outer_2dblock_x, + use_outer_2dblock_w, outer_block_size, scale_format="e4m3"): calls.append({ "A": A, @@ -202,6 +253,10 @@ def _mock_linear(A, B, *, use_2dblock_x, use_2dblock_w, use_sr_grad, "use_outer_scale": use_outer_scale, "use_hadamard": use_hadamard, "use_dge": use_dge, + "use_outer_block_scale": use_outer_block_scale, + "use_outer_2dblock_x": use_outer_2dblock_x, + "use_outer_2dblock_w": use_outer_2dblock_w, + "outer_block_size": outer_block_size, }) return A.new_full((A.shape[0], B.shape[0]), 5.0) @@ -232,6 +287,68 @@ def _mock_linear(A, B, *, use_2dblock_x, use_2dblock_w, use_sr_grad, assert call["use_outer_scale"] is True assert call["use_hadamard"] is True assert call["use_dge"] is True + assert call["use_outer_block_scale"] is False + assert call["use_outer_2dblock_x"] is False + assert call["use_outer_2dblock_w"] is False + assert call["outer_block_size"] == 128 + + +def test_linear_routes_outer_block_config_to_nvfp4_linear_kernel(monkeypatch, device): + """Dispatch must translate two_level_scaling='outer_block' into kernel kwargs.""" + calls = [] + + def _mock_linear(A, B, *, use_2dblock_x, use_2dblock_w, use_sr_grad, + use_outer_scale, use_hadamard, use_dge, + use_outer_block_scale, use_outer_2dblock_x, + use_outer_2dblock_w, outer_block_size): + calls.append({ + "A": A, + "B": B, + "use_2dblock_x": use_2dblock_x, + "use_2dblock_w": use_2dblock_w, + "use_sr_grad": use_sr_grad, + "use_outer_scale": use_outer_scale, + "use_hadamard": use_hadamard, + "use_dge": use_dge, + "use_outer_block_scale": use_outer_block_scale, + "use_outer_2dblock_x": use_outer_2dblock_x, + "use_outer_2dblock_w": use_outer_2dblock_w, + "outer_block_size": outer_block_size, + }) + return A.new_full((A.shape[0], B.shape[0]), 11.0) + + monkeypatch.setattr(dispatch_tensor, "_to_nvfp4_then_scaled_mm", _mock_linear) + + cfg = _make_config( + use_2dblock_x=False, + use_2dblock_w=True, + use_sr_grad=True, + two_level_scaling="outer_block", + use_outer_2dblock_x=False, + use_outer_2dblock_w=True, + outer_block_size=128, + ) + W_wrapped = _make_wrapper(cfg, device=device, shape=(32, 16)) + x = torch.randn(8, 16, dtype=torch.bfloat16, device=device) + + y = torch.nn.functional.linear(x, W_wrapped) + + assert y.shape == (8, 32) + assert torch.all(y == 11) + assert len(calls) == 1 + call = calls[0] + assert call["A"] is x + assert call["B"] is W_wrapped + assert call["use_2dblock_x"] is False + assert call["use_2dblock_w"] is True + assert call["use_sr_grad"] is True + assert call["use_outer_scale"] is False + assert call["use_hadamard"] is False + assert call["use_dge"] is False + assert call["use_outer_block_scale"] is True + assert call["use_outer_2dblock_x"] is False + assert call["use_outer_2dblock_w"] is True + assert call["outer_block_size"] == 128 # --------------------------------------------------------------------------- diff --git a/tests/unittest/nvfp4/test_nvfp_grouped_gemm_outer_block.py b/tests/unittest/nvfp4/test_nvfp_grouped_gemm_outer_block.py new file mode 100644 index 0000000..6069d33 --- /dev/null +++ b/tests/unittest/nvfp4/test_nvfp_grouped_gemm_outer_block.py @@ -0,0 +1,347 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +"""Outer-block scaling tests for the NVFP4 MoE Grouped GEMM autograd path. + +These mirror two existing suites: + +* ``test_nvfp_grouped_gemm.py`` — grouped-GEMM scaffolding (BF16 per-expert + reference, ALIGN_SIZE_M=128 contiguous expert blocks, ``calc_snr`` + + ``check_nvfp4_autograd_snr`` usage, M_total % 128 == 0 contract). +* ``test_nvfp_outer_block_linear.py`` — the dense outer-block patterns + (autograd SNR contract, outer-block-vs-tensorwise dominance on the + heavy-tail ``lognormal_channel_outlier`` data, rejection guards). + +Grouped outer-block recipe (see ``NVFP4GroupedGEMM.forward``): + * activations always use a 1D outer block (1 x ``outer_block_size`` over K); + * weights use ``use_outer_2dblock_w`` (per-expert 2D tiles) or 1D otherwise; + * tensorwise (``use_outer_scale``) and outer-block scaling are mutually + exclusive, and DGE is unsupported together with outer-block scaling. + +Test layers +----------- +1. ``test_nvfp4_grouped_gemm_outer_block_autograd`` + Forward O + dX + dW SNR vs a BF16 per-expert autograd reference, swept over + inner 1D/2D x outer 1D/2D weight layouts. SR off for determinism. +2. ``test_nvfp4_grouped_gemm_outer_block_finite_shapes`` + Outputs/grads finite with the documented shapes ([M_total,K] dX, [E,N,K] dW). +3. ``test_nvfp4_grouped_gemm_outer_block_dominates_tensorwise`` + On a per-expert heavy-tail input the grouped outer-block forward SNR must not + regress against the tensorwise path (and typically wins by several dB). +4. ``test_nvfp4_grouped_gemm_outer_block_rejects_dge`` / + ``test_nvfp4_grouped_gemm_outer_block_rejects_tensorwise`` + Contract guards for the mutually-exclusive / unsupported combinations. +""" + +import pytest +from tabulate import tabulate +import torch + +from alto.kernels.fp4.testing_utils import check_nvfp4_autograd_snr +from alto.kernels.fp4.nvfp4.nvfp_grouped_gemm import ( + ALIGN_SIZE_M, + nvfp4_grouped_gemm, +) +from .utils import prepare_data, calc_snr, calc_cossim + + +# Default outer block size used throughout; the 2D-outer-W grid needs +# N, K > OUTER_BLOCK_SIZE so the per-expert tile holds more than one block. +OUTER_BLOCK_SIZE = 128 + + +# (inner_2d, outer_2d_w) — activations always use a 1D outer block, so only the +# weight outer layout has a 2D variant in the grouped API. +OUTER_CASES = [ + pytest.param(False, False, id="inner_1d_outer_w_1d"), + pytest.param(False, True, id="inner_1d_outer_w_2d"), + pytest.param(True, False, id="inner_2d_outer_w_1d"), + pytest.param(True, True, id="inner_2d_outer_w_2d"), +] + + +# --------------------------------------------------------------------------- +# Helpers (self-contained copies of the grouped-GEMM scaffolding so this file +# does not couple to another test module). +# --------------------------------------------------------------------------- + +def _make_contiguous_expert_indices( + M_total: int, num_groups: int, num_experts: int, device, *, seed: int = 0, +) -> torch.Tensor: + """Expert-index tensor where every ALIGN_SIZE_M-row group shares one expert. + + Assigns experts round-robin so every expert is exercised when + ``num_groups >= num_experts`` (deterministic given ``num_groups``). + """ + indices = torch.zeros(M_total, dtype=torch.int32, device=device) + for g in range(num_groups): + eid = (g + seed) % num_experts + s = g * ALIGN_SIZE_M + indices[s : s + ALIGN_SIZE_M] = eid + return indices + + +def _bf16_grouped_ref_forward( + inputs, expert_weights, expert_indices, M_total, N, num_groups, +): + """BF16 per-expert reference forward: y[group] = x[group] @ w[eid].T.""" + dtype = inputs.dtype + device = inputs.device + y = torch.zeros(M_total, N, dtype=dtype, device=device) + for g in range(num_groups): + s, e = g * ALIGN_SIZE_M, (g + 1) * ALIGN_SIZE_M + eid = expert_indices[s].item() + y[s:e] = inputs[s:e] @ expert_weights[eid].T + return y + + +# --------------------------------------------------------------------------- +# Test 1 – autograd (O, dX, dW) SNR contract +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("inner_2d,outer_2d_w", OUTER_CASES) +def test_nvfp4_grouped_gemm_outer_block_autograd(inner_2d, outer_2d_w): + """Grouped outer-block forward + backward SNR vs BF16 per-expert reference. + + Swept over inner 1D/2D scaling and 1D/2D outer-W tiles. SR is off so the + comparison is deterministic and the K-aware floor is meaningful. + """ + M_total, N, K, num_experts = 512, 256, 256, 4 + num_groups = M_total // ALIGN_SIZE_M + device = torch.device("cuda") + dtype = torch.bfloat16 + + inputs_ref = prepare_data((M_total, K), dtype).requires_grad_(True) + weights_ref = prepare_data((num_experts, N, K), dtype).requires_grad_(True) + expert_indices = _make_contiguous_expert_indices( + M_total, num_groups, num_experts, device, + ) + target = prepare_data((M_total, N), dtype) + + y_ref = _bf16_grouped_ref_forward( + inputs_ref, weights_ref, expert_indices, M_total, N, num_groups, + ) + loss_ref = torch.nn.functional.mse_loss(y_ref, target) + loss_ref.backward() + dx_ref = inputs_ref.grad.clone() + dw_ref = weights_ref.grad.clone() + + inputs = inputs_ref.detach().clone().requires_grad_(True) + weights = weights_ref.detach().clone().requires_grad_(True) + + y = nvfp4_grouped_gemm( + inputs, weights, expert_indices, + trans_weights=True, + use_2dblock_x=inner_2d, + use_2dblock_w=inner_2d, + use_sr_grad=False, + use_outer_scale=False, + use_outer_block_scale=True, + use_outer_2dblock_w=outer_2d_w, + outer_block_size=OUTER_BLOCK_SIZE, + ) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + + o_snr = calc_snr(y_ref.detach(), y.detach()) + dx_snr = calc_snr(dx_ref, inputs.grad) + dw_snr = calc_snr(dw_ref, weights.grad) + o_sim = calc_cossim(y_ref.detach(), y.detach()) + dx_sim = calc_cossim(dx_ref, inputs.grad) + dw_sim = calc_cossim(dw_ref, weights.grad) + + print() + print(tabulate( + [ + ["O", f"{o_snr:.2f}", f"{o_sim:.6f}"], + ["dX", f"{dx_snr:.2f}", f"{dx_sim:.6f}"], + ["dW", f"{dw_snr:.2f}", f"{dw_sim:.6f}"], + ], + headers=["Tensor", "SNR(dB)", "CosSim"], + tablefmt="github", + )) + + check_nvfp4_autograd_snr( + {"O": o_snr, "dX": dx_snr, "dW": dw_snr}, + K=K, + use_sr_grad=False, + # No grouped-outer-block-specific tier exists; outer-block scaling is + # mutually exclusive with use_outer_scale, so the inner-only grouped + # floor is the applicable K-aware contract here. + kind="nvfp4_grouped_gemm", + use_outer_scale=False, + context=( + f"NVFP4 grouped outer-block inner_2d={inner_2d} " + f"outer_2d_w={outer_2d_w} K={K}" + ), + ) + + assert o_sim > 0.95, f"Forward CosSim too low: {o_sim:.6f}" + + +# --------------------------------------------------------------------------- +# Test 2 – finiteness + documented shapes +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("inner_2d,outer_2d_w", OUTER_CASES) +def test_nvfp4_grouped_gemm_outer_block_finite_shapes(inner_2d, outer_2d_w): + """Forward output and gradients are finite with the documented shapes.""" + M_total, N, K, num_experts = 384, 256, 256, 3 + num_groups = M_total // ALIGN_SIZE_M + device = torch.device("cuda") + dtype = torch.bfloat16 + + inputs = prepare_data((M_total, K), dtype).requires_grad_(True) + expert_weights = prepare_data((num_experts, N, K), dtype).requires_grad_(True) + expert_indices = _make_contiguous_expert_indices( + M_total, num_groups, num_experts, device, + ) + target = prepare_data((M_total, N), dtype) + + y = nvfp4_grouped_gemm( + inputs, expert_weights, expert_indices, + trans_weights=True, + use_2dblock_x=inner_2d, + use_2dblock_w=inner_2d, + use_sr_grad=False, + use_outer_scale=False, + use_outer_block_scale=True, + use_outer_2dblock_w=outer_2d_w, + outer_block_size=OUTER_BLOCK_SIZE, + ) + loss = torch.nn.functional.mse_loss(y, target) + loss.backward() + + assert y.shape == (M_total, N), f"Bad output shape: {y.shape}" + assert y.dtype == dtype + assert torch.isfinite(y).all(), "Output not finite" + assert inputs.grad is not None + assert inputs.grad.shape == (M_total, K), f"Bad dX shape: {inputs.grad.shape}" + assert torch.isfinite(inputs.grad).all(), "dX not finite" + assert expert_weights.grad is not None + assert expert_weights.grad.shape == (num_experts, N, K), ( + f"Bad dW shape: {expert_weights.grad.shape}" + ) + assert torch.isfinite(expert_weights.grad).all(), "dW not finite" + + +# --------------------------------------------------------------------------- +# Test 3 – outer-block vs tensorwise on a heavy-tail / per-expert-outlier input +# --------------------------------------------------------------------------- + +def test_nvfp4_grouped_gemm_outer_block_dominates_tensorwise(): + """Grouped outer-block forward SNR must not regress against tensorwise on a + heavy-tail, per-expert-outlier activation/weight distribution. + + The tensorwise (``use_outer_scale``) path shares a single FP32 scale across + the whole tensor (and, for weights, across *all* experts), so a hot expert + or massive channel forces the ~99% well-behaved blocks into E4M3 inner-scale + underflow. The per-outer-block FP32 scale localises each outlier, so the + rest of the tensor keeps precision. We assert no regression with a loose + -0.5 dB margin (outer-block typically wins by several dB here, but we keep + the assertion robust rather than over-tight). + """ + M_total, N, K, num_experts = 512, 256, 512, 4 + num_groups = M_total // ALIGN_SIZE_M + device = torch.device("cuda") + dtype = torch.bfloat16 + + # 1e5 outliers push the shared tensorwise scale into E4M3 inner-scale + # underflow for the non-outlier blocks — the regime outer-block targets. + inputs = prepare_data( + (M_total, K), dtype, pattern="lognormal_channel_outlier", seed=7, + outlier_scale=1e5, + ) + weights = prepare_data( + (num_experts, N, K), dtype, pattern="lognormal_channel_outlier", seed=13, + outlier_scale=1e5, + ) + expert_indices = _make_contiguous_expert_indices( + M_total, num_groups, num_experts, device, + ) + + y_ref = _bf16_grouped_ref_forward( + inputs, weights, expert_indices, M_total, N, num_groups, + ) + + y_tw = nvfp4_grouped_gemm( + inputs, weights, expert_indices, + trans_weights=True, + use_2dblock_x=False, use_2dblock_w=False, + use_sr_grad=False, + use_outer_scale=True, + use_outer_block_scale=False, + ) + y_ob = nvfp4_grouped_gemm( + inputs, weights, expert_indices, + trans_weights=True, + use_2dblock_x=False, use_2dblock_w=False, + use_sr_grad=False, + use_outer_scale=False, + use_outer_block_scale=True, + use_outer_2dblock_w=True, + outer_block_size=OUTER_BLOCK_SIZE, + ) + + tw_snr = calc_snr(y_ref, y_tw) + ob_snr = calc_snr(y_ref, y_ob) + + print( + f"\nheavy-tail grouped forward SNR: tensorwise={tw_snr:.2f} dB " + f"outer-block={ob_snr:.2f} dB (margin={ob_snr - tw_snr:+.2f} dB)" + ) + + assert torch.isfinite(y_tw).all() + assert torch.isfinite(y_ob).all() + assert ob_snr >= tw_snr - 0.5, ( + f"grouped outer-block forward SNR ({ob_snr:.2f}) regressed vs " + f"tensorwise ({tw_snr:.2f}) on heavy-tail per-expert-outlier data." + ) + + +# --------------------------------------------------------------------------- +# Test 4 – rejection guards +# --------------------------------------------------------------------------- + +def _small_grouped_args(device, dtype, *, num_experts=2, N=256, K=256): + M_total = ALIGN_SIZE_M + inputs = prepare_data((M_total, K), dtype) + expert_weights = prepare_data((num_experts, N, K), dtype) + expert_indices = _make_contiguous_expert_indices( + M_total, M_total // ALIGN_SIZE_M, num_experts, device, + ) + return inputs, expert_weights, expert_indices + + +def test_nvfp4_grouped_gemm_outer_block_rejects_dge(): + """DGE together with outer-block scaling is unsupported and must raise.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + inputs, expert_weights, expert_indices = _small_grouped_args(device, dtype) + with pytest.raises(RuntimeError, match="DGE"): + nvfp4_grouped_gemm( + inputs, expert_weights, expert_indices, + trans_weights=True, + use_2dblock_x=False, use_2dblock_w=False, + use_sr_grad=False, + use_outer_block_scale=True, + use_dge=True, + ) + + +def test_nvfp4_grouped_gemm_outer_block_rejects_tensorwise(): + """Tensorwise (``use_outer_scale``) and outer-block scaling are mutually + exclusive and must raise when both are requested.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + inputs, expert_weights, expert_indices = _small_grouped_args(device, dtype) + with pytest.raises(RuntimeError, match="mutually exclusive"): + nvfp4_grouped_gemm( + inputs, expert_weights, expert_indices, + trans_weights=True, + use_2dblock_x=False, use_2dblock_w=False, + use_sr_grad=False, + use_outer_scale=True, + use_outer_block_scale=True, + ) diff --git a/tests/unittest/nvfp4/test_nvfp_outer_block_linear.py b/tests/unittest/nvfp4/test_nvfp_outer_block_linear.py new file mode 100644 index 0000000..7cb945f --- /dev/null +++ b/tests/unittest/nvfp4/test_nvfp_outer_block_linear.py @@ -0,0 +1,321 @@ +# Copyright (c) 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: MIT + +"""Outer-block integration tests for the NVFP4 Linear autograd path.""" + +import pytest +from tabulate import tabulate +import torch + +from alto.kernels.fp4.testing_utils import check_nvfp4_autograd_snr +from alto.kernels.fp4.nvfp4.nvfp_linear import ( + NVFP4LinearFunction, + _to_nvfp4_then_scaled_mm, +) +from .utils import prepare_data, calc_snr, calc_cossim + + +OUTER_CASES = [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +] + + +def _run_linear_autograd_case( + *, + shape: tuple[int, int, int, int], + use_2dblock_x: bool, + use_2dblock_w: bool, + use_outer_2dblock_x: bool, + use_outer_2dblock_w: bool, + use_sr_grad: bool, +) -> None: + B, M, N, K = shape + dtype = torch.bfloat16 + inputs = prepare_data((B, M, K), dtype).requires_grad_(True) + weights = prepare_data((N, K), dtype).requires_grad_(True) + target = prepare_data((B, M, N), dtype) + + outputs_ref = torch.nn.functional.linear(inputs, weights) + loss_ref = torch.nn.functional.mse_loss(outputs_ref, target) + loss_ref.backward() + grad_inputs_ref = inputs.grad.clone() + grad_weights_ref = weights.grad.clone() + inputs.grad.zero_() + weights.grad.zero_() + + outputs = _to_nvfp4_then_scaled_mm( + inputs, + weights, + use_2dblock_x=use_2dblock_x, + use_2dblock_w=use_2dblock_w, + use_sr_grad=use_sr_grad, + use_outer_scale=False, + use_outer_block_scale=True, + use_outer_2dblock_x=use_outer_2dblock_x, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=128, + ) + loss = torch.nn.functional.mse_loss(outputs, target) + loss.backward() + + output_snr = calc_snr(outputs_ref, outputs) + dx_snr = calc_snr(grad_inputs_ref, inputs.grad) + dw_snr = calc_snr(grad_weights_ref, weights.grad) + output_sim = calc_cossim(outputs_ref, outputs) + dx_sim = calc_cossim(grad_inputs_ref, inputs.grad) + dw_sim = calc_cossim(grad_weights_ref, weights.grad) + + print() + print(tabulate( + [ + ["O", f"{output_snr:.2f}", f"{output_sim:.6f}"], + ["dX", f"{dx_snr:.2f}", f"{dx_sim:.6f}"], + ["dW", f"{dw_snr:.2f}", f"{dw_sim:.6f}"], + ], + headers=["Tensor", "SNR", "Cosine Sim"], + tablefmt="github", + )) + + check_nvfp4_autograd_snr( + {"O": output_snr, "dX": dx_snr, "dW": dw_snr}, + K=K, + use_sr_grad=use_sr_grad, + kind="nvfp4_linear", + context=( + f"NVFP4Linear outer-block shape={shape} " + f"inner_x_2d={use_2dblock_x} inner_w_2d={use_2dblock_w} " + f"outer_x_2d={use_outer_2dblock_x} outer_w_2d={use_outer_2dblock_w}" + ), + ) + + +@pytest.mark.parametrize("use_2dblock,use_outer_2dblock", OUTER_CASES) +@pytest.mark.parametrize("use_sr_grad", [False, True]) +def test_nvfp4_linear_outer_block_autograd(use_2dblock, use_outer_2dblock, use_sr_grad): + """Outer-block Linear should satisfy the existing K-aware SNR contract.""" + _run_linear_autograd_case( + shape=(1, 512, 384, 128), + use_2dblock_x=use_2dblock, + use_2dblock_w=use_2dblock, + use_outer_2dblock_x=use_outer_2dblock, + use_outer_2dblock_w=use_outer_2dblock, + use_sr_grad=use_sr_grad, + ) + + +@pytest.mark.parametrize("use_2dblock_x", [False, True]) +@pytest.mark.parametrize("use_2dblock_w", [False, True]) +@pytest.mark.parametrize("use_outer_2dblock_x", [False, True]) +@pytest.mark.parametrize("use_outer_2dblock_w", [False, True]) +def test_nvfp4_linear_outer_block_independent_xw_layouts( + use_2dblock_x, + use_2dblock_w, + use_outer_2dblock_x, + use_outer_2dblock_w, +): + """X and W inner/outer layout knobs must be independently usable.""" + _run_linear_autograd_case( + shape=(1, 128, 128, 128), + use_2dblock_x=use_2dblock_x, + use_2dblock_w=use_2dblock_w, + use_outer_2dblock_x=use_outer_2dblock_x, + use_outer_2dblock_w=use_outer_2dblock_w, + use_sr_grad=False, + ) + + +def test_nvfp4_linear_outer_block_large_k_sr_stress(): + """Keep one K=2048 SR stress case without exploding the test matrix.""" + _run_linear_autograd_case( + shape=(4, 1024, 1024, 2048), + use_2dblock_x=True, + use_2dblock_w=True, + use_outer_2dblock_x=True, + use_outer_2dblock_w=True, + use_sr_grad=True, + ) + + +def test_nvfp4_linear_outer_block_rejects_dge(): + """Outer-block + DGE needs raw-scale plumbing that is not implemented yet.""" + x = prepare_data((1, 128, 128), torch.bfloat16) + w = prepare_data((128, 128), torch.bfloat16) + with pytest.raises(RuntimeError, match="does not yet support DGE"): + NVFP4LinearFunction.apply( + x, + w, + False, + False, + False, + False, + None, + True, + True, + False, + False, + 128, + ) + + +# --------------------------------------------------------------------------- +# T1 / T3 (NVFP4_Outer_Block_Review.md §3.3): direct outer-block vs PTS +# dominance tests on a paper-realistic LLM-activation distribution +# (lognormal heavy-tailed + ~1% massive channels). +# --------------------------------------------------------------------------- + + +def _run_linear_one_pass( + *, + inputs: torch.Tensor, + weights: torch.Tensor, + use_outer_block_scale: bool, + use_outer_scale: bool, + use_2dblock_x: bool = False, + use_2dblock_w: bool = False, + use_outer_2dblock_w: bool = False, + use_sr_grad: bool = False, +) -> torch.Tensor: + return _to_nvfp4_then_scaled_mm( + inputs, + weights, + use_2dblock_x=use_2dblock_x, + use_2dblock_w=use_2dblock_w, + use_sr_grad=use_sr_grad, + use_outer_scale=use_outer_scale, + use_outer_block_scale=use_outer_block_scale, + use_outer_2dblock_x=False, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=128, + ) + + +def test_outer_block_snr_dominates_pts_on_lognormal(): + """T3: outer-block must strictly dominate PTS on heavy-tailed activations. + + Uses the lognormal-channel-outlier pattern (LLM-activation proxy) in the + *massive-activation* regime (outlier_scale=1e5). At that magnitude the + single per-tensor (PTS) outer scale forces the per-16-block E4M3 inner + scales of the ~99% non-outlier blocks into underflow, collapsing PTS SNR, + whereas the per-outer-block scale localises each outlier so the rest of the + tensor keeps precision. This is exactly the failure mode outer-block + scaling targets, so the gap is large and robust here (~4 dB on the forward + output, ~8 dB on dX) rather than the <0.5 dB seen at mild 1e3 outliers. + + Asserts both: + (i) outer-block forward output SNR ≥ PTS forward output SNR, + (ii) outer-block dX SNR ≥ PTS dX SNR, + on identical inputs/weights with a meaningful margin (≥ 1 dB). + """ + dtype = torch.bfloat16 + # Pick dims so the outer-block grids are NON-degenerate: with + # outer_block_size=128 the 1D outer on X needs K > 128 to hold >1 block + # per row, and the 2D outer on W (128×128) needs N,K > 128 to hold >1 + # block. K==N==128 would make the W 2D-outer block span the whole tensor + # and silently collapse the outer-block path back to per-tensor scale. + M, N, K = 256, 256, 512 + # 1e5 outliers push PTS into E4M3 inner-scale underflow (see docstring); + # this is the regime where outer-block scaling is meant to win. + inputs = prepare_data( + (1, M, K), dtype, pattern="lognormal_channel_outlier", seed=7, + outlier_scale=1e5, + ).requires_grad_(True) + weights = prepare_data( + (N, K), dtype, pattern="lognormal_channel_outlier", seed=13, + outlier_scale=1e5, + ).requires_grad_(True) + target = prepare_data((1, M, N), dtype, pattern="random", seed=23) + + outputs_ref = torch.nn.functional.linear(inputs, weights) + loss_ref = torch.nn.functional.mse_loss(outputs_ref, target) + loss_ref.backward() + grad_inputs_ref = inputs.grad.clone() + inputs.grad.zero_(); weights.grad.zero_() + + # PTS path + outputs_pts = _run_linear_one_pass( + inputs=inputs, weights=weights, + use_outer_block_scale=False, use_outer_scale=True, + ) + loss_pts = torch.nn.functional.mse_loss(outputs_pts, target) + loss_pts.backward() + pts_o_snr = calc_snr(outputs_ref, outputs_pts) + pts_dx_snr = calc_snr(grad_inputs_ref, inputs.grad) + inputs.grad.zero_(); weights.grad.zero_() + + # Outer-block path + outputs_ob = _run_linear_one_pass( + inputs=inputs, weights=weights, + use_outer_block_scale=True, use_outer_scale=False, + use_outer_2dblock_w=True, # 2D outer for W (axis-invariant) + ) + loss_ob = torch.nn.functional.mse_loss(outputs_ob, target) + loss_ob.backward() + ob_o_snr = calc_snr(outputs_ref, outputs_ob) + ob_dx_snr = calc_snr(grad_inputs_ref, inputs.grad) + + # Margins below are conservative; on a typical heavy-tailed activation + # we observe outer-block ≥ PTS by 3-8 dB. + assert ob_o_snr >= pts_o_snr + 1.0, ( + f"outer-block forward output SNR ({ob_o_snr:.2f}) did not " + f"dominate PTS ({pts_o_snr:.2f}) by >=1 dB on lognormal pattern." + ) + assert ob_dx_snr >= pts_dx_snr + 1.0, ( + f"outer-block dX SNR ({ob_dx_snr:.2f}) did not dominate PTS " + f"({pts_dx_snr:.2f}) by >=1 dB on lognormal pattern." + ) + + +# --------------------------------------------------------------------------- +# T6 (NVFP4_Outer_Block_Review.md §3.3): Hadamard × outer-block SNR coverage. +# Existing tests only cover Hadamard × outer-block via outer-scale-share +# reduction counting; add an SNR-level check for the M-axis wgrad RHT +# (use_hadamard=True) combination. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("use_2dblock,use_outer_2dblock", OUTER_CASES) +def test_outer_block_with_hadamard_autograd(use_2dblock, use_outer_2dblock): + """Outer-block × M-axis wgrad RHT (legacy use_hadamard) must clear the + base nvfp4_linear SNR contract.""" + B, M, N, K = 1, 512, 384, 128 + dtype = torch.bfloat16 + inputs = prepare_data((B, M, K), dtype).requires_grad_(True) + weights = prepare_data((N, K), dtype).requires_grad_(True) + target = prepare_data((B, M, N), dtype) + + outputs_ref = torch.nn.functional.linear(inputs, weights) + loss_ref = torch.nn.functional.mse_loss(outputs_ref, target) + loss_ref.backward() + grad_inputs_ref = inputs.grad.clone() + grad_weights_ref = weights.grad.clone() + inputs.grad.zero_(); weights.grad.zero_() + + outputs = _to_nvfp4_then_scaled_mm( + inputs, weights, + use_2dblock_x=use_2dblock, use_2dblock_w=use_2dblock, + use_sr_grad=False, use_outer_scale=False, + use_hadamard=True, + use_outer_block_scale=True, + use_outer_2dblock_x=use_outer_2dblock, + use_outer_2dblock_w=use_outer_2dblock, + outer_block_size=128, + ) + loss = torch.nn.functional.mse_loss(outputs, target) + loss.backward() + check_nvfp4_autograd_snr( + { + "O": calc_snr(outputs_ref, outputs), + "dX": calc_snr(grad_inputs_ref, inputs.grad), + "dW": calc_snr(grad_weights_ref, weights.grad), + }, + K=K, use_sr_grad=False, + kind="nvfp4_linear", + context=( + f"outer-block + use_hadamard (M-axis wgrad RHT) " + f"inner_2d={use_2dblock} outer_2d={use_outer_2dblock}" + ), + ) diff --git a/tests/unittest/nvfp4/test_nvfp_quantization.py b/tests/unittest/nvfp4/test_nvfp_quantization.py index 34a6cd9..65a03ea 100644 --- a/tests/unittest/nvfp4/test_nvfp_quantization.py +++ b/tests/unittest/nvfp4/test_nvfp_quantization.py @@ -6,16 +6,28 @@ import torch from alto.kernels.fp4.nvfp4.nvfp_quantization import ( _OUTER_SCALE_DIVZERO_FLOOR, + E4M3_EPS, + F4_E2M1_MAX, + F8E4M3_MAX, + OUTER_SCALE_EPS, convert_to_nvfp4, convert_from_nvfp4, compute_dynamic_outer_scale, + compute_outer_block_scale_shape, + convert_to_nvfp4_outer_block, + convert_from_nvfp4_outer_block, + _outer_scale_and_expanded_axis_last, + _qdq, is_cdna4, ) from .utils import ( prepare_data, + calc_snr, convert_to_nvfp4_pytorch, convert_from_nvfp4_pytorch, + convert_to_nvfp4_outer_block_pytorch, + convert_from_nvfp4_outer_block_pytorch, ) @@ -337,3 +349,803 @@ def test_nvfp4_quantization_nan_input_sanitized(scale_format, data_type): f"NaN input contaminated dequant output under {scale_format}; " f"defense layer must keep downstream GEMM input finite." ) + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer,expected_inner,expected_outer,label", [ + (False, False, (128, 16), (128, 2), "inner=1x16 outer=1x128"), + (False, True, (128, 16), (1, 2), "inner=1x16 outer=128x128"), + (True, False, (8, 16), (128, 2), "inner=16x16 outer=1x128"), + (True, True, (8, 16), (1, 2), "inner=16x16 outer=128x128"), +]) +def test_nvfp4_outer_block_scale_shapes( + is_2d_inner, is_2d_outer, expected_inner, expected_outer, label +): + """Outer-scale shape helper must match the requested four layouts.""" + inner_shape, outer_shape = compute_outer_block_scale_shape( + (128, 256), + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + inner_block_size=16, + outer_block_size=128, + ) + assert inner_shape == expected_inner, label + assert outer_shape == expected_outer, label + + +@pytest.mark.parametrize("kwargs,match", [ + ( + dict(shape=(128, 64), is_2d_inner=False, is_2d_outer=False, + inner_block_size=16, outer_block_size=128), + "last dimension", + ), + ( + dict(shape=(128, 192), is_2d_inner=False, is_2d_outer=False, + inner_block_size=16, outer_block_size=128), + "last dimension", + ), + ( + dict(shape=(64, 256), is_2d_inner=False, is_2d_outer=True, + inner_block_size=16, outer_block_size=128), + "2D outer block", + ), + ( + dict(shape=(16,), is_2d_inner=False, is_2d_outer=False, + inner_block_size=16, outer_block_size=128), + "at least 2D", + ), + ( + dict(shape=(128, 256), is_2d_inner=False, is_2d_outer=False, + inner_block_size=16, outer_block_size=100), + "must be a multiple", + ), +]) +def test_nvfp4_outer_block_scale_shape_rejects_invalid_inputs(kwargs, match): + """Outer-block scaling is fail-fast: no padding or ragged tails.""" + with pytest.raises(ValueError, match=match): + compute_outer_block_scale_shape(**kwargs) + + +def test_nvfp4_outer_block_size_equal_inner_block_is_valid(): + """outer_block_size=inner_block_size is a legal degenerate layout.""" + inner_shape, outer_shape = compute_outer_block_scale_shape( + (128, 256), + is_2d_inner=False, + is_2d_outer=False, + inner_block_size=16, + outer_block_size=16, + ) + assert inner_shape == (128, 16) + assert outer_shape == (128, 16) + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +@pytest.mark.parametrize("axis", [-1, -2]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_nvfp4_outer_block_roundtrip_finite(is_2d_inner, is_2d_outer, axis, dtype): + """All four outer-scale layouts should round-trip with finite outputs.""" + x = prepare_data((128, 256), dtype) + data_lp, inner_scales, outer_scales = convert_to_nvfp4_outer_block( + x, + block_size=16, + outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + ) + x_dq = convert_from_nvfp4_outer_block( + data_lp, + inner_scales, + outer_scales, + output_dtype=x.dtype, + block_size=16, + outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + assert x_dq.shape == x.shape + assert x_dq.dtype == dtype + assert torch.isfinite(inner_scales).all() + assert torch.isfinite(outer_scales).all() + assert torch.isfinite(x_dq).all() + + # T5 (NVFP4_Outer_Block_Review.md §3.3): also pin a numerical SNR floor + # so silent regressions (e.g. wrong axis, scale-clamp bug, lost outer + # scale) get caught even when the round-trip is still finite-valued. + # Floors below are conservative: empirically the kernel achieves + # ≥18 dB on this gaussian + 0.5%-outlier pattern. We use a tighter + # bound for 1D outer than 2D outer because each 128-element row owns + # its own scale (no inter-row averaging). ``random`` pattern produces + # different SNR for fp32 vs bf16 input; cast both to fp32 for SNR. + snr = calc_snr(x.float(), x_dq.float()) + if is_2d_outer: + snr_floor = 11.0 + else: + snr_floor = 14.0 + assert snr > snr_floor, ( + f"outer-block round-trip SNR {snr:.2f} dB < floor {snr_floor} dB " + f"(is_2d_inner={is_2d_inner}, is_2d_outer={is_2d_outer}, " + f"axis={axis}, dtype={dtype})" + ) + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +def test_nvfp4_qdq_routes_to_outer_block(is_2d_inner, is_2d_outer): + """The shared QDQ helper should expose the outer-block round-trip path.""" + x = prepare_data((128, 256), torch.bfloat16) + x_qdq = _qdq( + x, + axis=-1, + is_2d_block=is_2d_inner, + use_outer_scale=False, + use_outer_block_scale=True, + is_2d_outer=is_2d_outer, + outer_block_size=128, + ) + assert x_qdq.shape == x.shape + assert x_qdq.dtype == x.dtype + assert torch.isfinite(x_qdq).all() + + +def test_nvfp4_qdq_rejects_pts_and_outer_block_together(): + """Tensor-wise and outer-block scaling are alternative second-level scales.""" + x = prepare_data((128, 256), torch.bfloat16) + with pytest.raises(RuntimeError, match="mutually exclusive"): + _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=True, + use_outer_block_scale=True, + ) + + +def test_nvfp4_qdq_rejects_outer_block_return_raw(): + """DGE needs raw FP4 plus scales; outer-block raw support is not wired yet.""" + x = prepare_data((128, 256), torch.bfloat16) + with pytest.raises(RuntimeError, match="return_raw=True"): + _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=False, + use_outer_block_scale=True, + return_raw=True, + ) + + +def _qdq_with_tensorwise_and_outer( + x: torch.Tensor, + *, + is_2d_inner: bool, + is_2d_outer: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + pts = compute_dynamic_outer_scale(x) + data_lp_pts, scales_pts = convert_to_nvfp4( + x, + block_size=16, + axis=-1, + is_2d_block=is_2d_inner, + outer_scale=pts, + update_outer_scale=False, + use_sr=False, + ) + x_pts = convert_from_nvfp4( + data_lp_pts, + scales_pts, + output_dtype=x.dtype, + block_size=16, + axis=-1, + is_2d_block=is_2d_inner, + outer_scale=pts, + ) + data_lp_ob, inner_scales_ob, outer_scales = convert_to_nvfp4_outer_block( + x, + block_size=16, + outer_block_size=128, + axis=-1, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + ) + x_ob = convert_from_nvfp4_outer_block( + data_lp_ob, + inner_scales_ob, + outer_scales, + output_dtype=x.dtype, + block_size=16, + outer_block_size=128, + axis=-1, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + return x_pts, x_ob + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +def test_nvfp4_outer_block_improves_heterogeneous_regions_vs_pts(is_2d_inner, is_2d_outer): + """Outer scale should preserve low-energy regions in mixed-energy tensors. + + The 256×256 tensor is built from four 128×128 regions with different + magnitudes. Tensor-wise scaling is dominated by the high-energy quadrant, + while outer-block scaling gives each 1×128 or 128×128 region its own FP32 + range scale. We do not require bitwise equality between layouts; we only + require the low-energy regions to improve without hurting total error. + """ + dtype = torch.bfloat16 + torch.manual_seed(2026) + low_a = (torch.randn((128, 128), device="cuda") * 0.05).to(dtype) + high = (torch.randn((128, 128), device="cuda") * 5000.0).to(dtype) + medium = (torch.randn((128, 128), device="cuda") * 5.0).to(dtype) + low_b = (torch.randn((128, 128), device="cuda") * 0.05).to(dtype) + x = torch.cat([ + torch.cat([low_a, high], dim=-1), + torch.cat([medium, low_b], dim=-1), + ], dim=0) + + x_pts, x_ob = _qdq_with_tensorwise_and_outer( + x, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + + low_ref = torch.cat([x[:128, :128].reshape(-1), x[128:, 128:].reshape(-1)]).float() + low_pts = torch.cat([x_pts[:128, :128].reshape(-1), x_pts[128:, 128:].reshape(-1)]).float() + low_ob = torch.cat([x_ob[:128, :128].reshape(-1), x_ob[128:, 128:].reshape(-1)]).float() + low_pts_err = (low_ref - low_pts).pow(2).mean() + low_ob_err = (low_ref - low_ob).pow(2).mean() + + total_pts_err = (x.float() - x_pts.float()).pow(2).mean() + total_ob_err = (x.float() - x_ob.float()).pow(2).mean() + assert low_ob_err < low_pts_err * 0.5, ( + f"outer-block low-region MSE ({low_ob_err.item():.4e}) should be " + f"lower than tensor-wise MSE ({low_pts_err.item():.4e})" + ) + assert total_ob_err < total_pts_err * 1.1, ( + f"outer-block total MSE ({total_ob_err.item():.4e}) should not " + f"regress vs tensor-wise MSE ({total_pts_err.item():.4e})" + ) + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +def test_nvfp4_outer_block_improves_clean_blocks_with_sparse_outlier(is_2d_inner, is_2d_outer): + """Outer scale should help clean regions when a sparse outlier sets PTS. + + This is closer to a localized-outlier pattern: most values are low-energy, + but one column in the right outer block contains large values. The clean + left outer block should reconstruct better with its own outer scale, while + total error should remain in the same ballpark. + """ + dtype = torch.bfloat16 + torch.manual_seed(2027) + x = (torch.randn((128, 256), device="cuda") * 0.05).to(dtype) + x[:, 192] = (torch.randn((128,), device="cuda") * 5000.0).to(dtype) + + x_pts, x_ob = _qdq_with_tensorwise_and_outer( + x, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + + low_ref = x[:, :128].float() + low_pts_err = (low_ref - x_pts[:, :128].float()).pow(2).mean() + low_ob_err = (low_ref - x_ob[:, :128].float()).pow(2).mean() + total_pts_err = (x.float() - x_pts.float()).pow(2).mean() + total_ob_err = (x.float() - x_ob.float()).pow(2).mean() + assert low_ob_err < low_pts_err * 0.75, ( + f"outer-block clean-region MSE ({low_ob_err.item():.4e}) should be " + f"lower than tensor-wise MSE ({low_pts_err.item():.4e})" + ) + assert total_ob_err < total_pts_err * 1.25, ( + f"outer-block total MSE ({total_ob_err.item():.4e}) should stay close " + f"to tensor-wise MSE ({total_pts_err.item():.4e})" + ) + + +# --------------------------------------------------------------------------- +# Outer-block scale uses an FP32 floor. +# +# The outer scale is a free FP32 number with no E4M3 round-trip; clamping it +# at ``E4M3_EPS`` (~1.95e-3) instead of FP32 tiny would inflate small-block +# scales and push inner E4M3 scales into the precision-poor ``< 1`` region. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("is_2d_outer", [False, True]) +def test_outer_scale_tracks_amax_below_e4m3_eps(is_2d_outer): + """Small-amax blocks must produce an outer scale below ``E4M3_EPS``.""" + target_amax = 1.0 + expected_scale = target_amax / (F8E4M3_MAX * F4_E2M1_MAX) # ~3.7e-4 + assert expected_scale < E4M3_EPS, ( + "precondition: expected outer scale must be below the E4M3 floor" + ) + + dtype = torch.bfloat16 + if is_2d_outer: + x = torch.full((128, 128), target_amax, device="cuda", dtype=dtype) + else: + x = torch.full((1, 128), target_amax, device="cuda", dtype=dtype) + x[..., 0] = -target_amax + + outer_scale, _ = _outer_scale_and_expanded_axis_last( + x, + is_2d_inner=False, + is_2d_outer=is_2d_outer, + inner_block_size=16, + outer_block_size=128, + ) + assert outer_scale.dtype == torch.float32 + assert (outer_scale < E4M3_EPS * 0.5).all(), ( + f"outer_scale {outer_scale.flatten().tolist()} should track " + f"amax/2688 ≈ {expected_scale:.3e}, well below E4M3_EPS={E4M3_EPS:.3e}" + ) + torch.testing.assert_close( + outer_scale, torch.full_like(outer_scale, expected_scale), + atol=expected_scale * 1e-2, rtol=1e-2, + ) + + +@pytest.mark.parametrize("is_2d_outer", [False, True]) +def test_outer_scale_dead_block_avoids_div_by_zero(is_2d_outer): + """All-zero outer blocks must produce a finite, positive scale.""" + if is_2d_outer: + x = torch.zeros((128, 128), device="cuda", dtype=torch.bfloat16) + else: + x = torch.zeros((1, 128), device="cuda", dtype=torch.bfloat16) + + outer_scale, expanded = _outer_scale_and_expanded_axis_last( + x, + is_2d_inner=False, + is_2d_outer=is_2d_outer, + inner_block_size=16, + outer_block_size=128, + ) + assert (outer_scale >= OUTER_SCALE_EPS).all() + assert torch.isfinite(outer_scale).all() + assert torch.isfinite(expanded).all() + assert torch.isfinite(x.float() / expanded).all() + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +def test_outer_block_roundtrip_does_not_regress_on_small_amax( + is_2d_inner, is_2d_outer, +): + """Small-amax data must round-trip with comparable relative MSE to normal-amax data.""" + dtype = torch.bfloat16 + torch.manual_seed(2026) + x_small = (torch.randn((128, 256), device="cuda") * 1e-3).to(dtype) + x_normal = torch.randn((128, 256), device="cuda", dtype=dtype) + + def _roundtrip(x): + data_lp, inner_scales, outer_scales = convert_to_nvfp4_outer_block( + x, + block_size=16, + outer_block_size=128, + axis=-1, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + ) + return convert_from_nvfp4_outer_block( + data_lp, inner_scales, outer_scales, + output_dtype=x.dtype, + block_size=16, + outer_block_size=128, + axis=-1, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + + x_small_qdq = _roundtrip(x_small) + x_normal_qdq = _roundtrip(x_normal) + + rel_err_small = ( + (x_small.float() - x_small_qdq.float()).pow(2).mean() + / x_small.float().pow(2).mean().clamp(min=1e-30) + ) + rel_err_normal = ( + (x_normal.float() - x_normal_qdq.float()).pow(2).mean() + / x_normal.float().pow(2).mean().clamp(min=1e-30) + ) + # An over-aggressive floor would push inner E4M3 into the ``< 1`` region + # for small-amax data and inflate this ratio by ~10x. + assert rel_err_small < rel_err_normal * 3.0, ( + f"small-amax relative MSE ({rel_err_small.item():.3e}) regressed vs " + f"normal-amax relative MSE ({rel_err_normal.item():.3e})" + ) + + +# --------------------------------------------------------------------------- +# ``precomputed_outer_scale`` plumbing. +# +# 2D outer-block scale grids are axis-invariant on the same tensor; callers +# can compute the scale once and feed it back via ``precomputed_outer_scale`` +# to skip a redundant outer-amax reduction on the axis=0 QDQ. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +@pytest.mark.parametrize("axis", [-1, -2]) +def test_outer_block_precomputed_scale_matches_recompute( + is_2d_inner, is_2d_outer, axis, +): + """Reusing a captured outer scale must produce identical packed FP4 + scales. + + Axis-invariance is what lets the linear path share an outer scale across + the axis=-1 and axis=0 QDQ calls. Run the QDQ once, capture the outer + scale, run it again with ``precomputed_outer_scale`` set, and require + bitwise equality of every output. + """ + x = prepare_data((128, 256), torch.bfloat16) + + data_lp_a, inner_a, outer_a = convert_to_nvfp4_outer_block( + x, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + ) + data_lp_b, inner_b, outer_b = convert_to_nvfp4_outer_block( + x, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + precomputed_outer_scale=outer_a, + ) + torch.testing.assert_close(outer_a, outer_b, atol=0, rtol=0) + assert torch.equal(data_lp_a, data_lp_b) + torch.testing.assert_close(inner_a, inner_b, atol=0, rtol=0) + + +def test_outer_block_precomputed_scale_validates_shape(): + """Wrong-shape ``precomputed_outer_scale`` must raise a clear ValueError.""" + x = prepare_data((128, 256), torch.bfloat16) + bad_outer = torch.zeros((1, 1), device=x.device, dtype=torch.float32) + with pytest.raises(ValueError, match="precomputed_outer_scale"): + convert_to_nvfp4_outer_block( + x, + block_size=16, outer_block_size=128, + axis=-1, + is_2d_inner=False, is_2d_outer=False, + use_sr=False, + precomputed_outer_scale=bad_outer, + ) + + +def test_qdq_precomputed_outer_scale_routes_through(): + """``_qdq(precomputed_outer_scale=...)`` must produce identical output.""" + x = prepare_data((128, 256), torch.bfloat16) + dq_a, outer_scale = _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=False, + use_outer_block_scale=True, + is_2d_outer=True, + outer_block_size=128, + return_outer_scale=True, + ) + dq_b = _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=False, + use_outer_block_scale=True, + is_2d_outer=True, + outer_block_size=128, + precomputed_outer_scale=outer_scale, + ) + assert torch.equal(dq_a, dq_b) + + +def test_qdq_rejects_precomputed_outer_scale_without_outer_block(): + """``precomputed_outer_scale`` is meaningful only when outer-block is on.""" + x = prepare_data((128, 256), torch.bfloat16) + junk = torch.zeros((1,), device=x.device, dtype=torch.float32) + with pytest.raises(RuntimeError, match="precomputed_outer_scale"): + _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=False, + use_outer_block_scale=False, + precomputed_outer_scale=junk, + ) + + +def test_qdq_rejects_return_outer_scale_without_outer_block(): + x = prepare_data((128, 256), torch.bfloat16) + with pytest.raises(RuntimeError, match="return_outer_scale"): + _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=False, + use_outer_block_scale=False, + return_outer_scale=True, + ) + + +def test_qdq_rejects_return_raw_and_return_outer_scale_together(): + x = prepare_data((128, 256), torch.bfloat16) + with pytest.raises(RuntimeError, match="cannot return both"): + _qdq( + x, + axis=-1, + is_2d_block=False, + use_outer_scale=False, + use_outer_block_scale=True, + is_2d_outer=True, + outer_block_size=128, + return_raw=True, + return_outer_scale=True, + ) + + +# --------------------------------------------------------------------------- +# Outer-amax reductions are minimized when the scale grid is axis-invariant. +# +# Monkey-patch the only function that performs the reduction +# (``_outer_scale_and_expanded_axis_last``) and verify the call count matches +# what we expect for each forward / backward axis-sharing configuration. +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "use_2dblock_x,use_2dblock_w,use_outer_2dblock_x,use_outer_2dblock_w," + "use_hadamard,expected_calls,id_label", + [ + # 2D outer, no Hadamard: axis=0 QDQ on x, w, and grad_output all reuse + # the axis=-1 outer scale. 3 reductions = fwd_x + fwd_w + bwd_g. + (False, False, True, True, False, 3, "C2_no_hadamard"), + # 2D outer with Hadamard: the rotated x and grad_output are different + # tensors than the fprop ones, so only w can reuse. 5 reductions. + (False, False, True, True, True, 5, "C2_hadamard_blocks_x_share"), + # 1D outer: scale grid is not axis-invariant; nothing can be shared. + # 6 reductions = fwd (x_-1, w_-1, x_0, w_0) + bwd (g_-1, g_0). + (False, False, False, False, False, 6, "C1_no_share"), + # 2D inner block: axis=0 dq views are aliased to the fprop views, so + # no extra QDQ is issued in the first place. 3 reductions. + (True, True, True, True, False, 3, "C4_inner_2d_aliasing"), + ], + ids=lambda v: v if isinstance(v, str) else None, +) +def test_nvfp4_linear_outer_block_reuses_fwd_outer_scale( + use_2dblock_x, use_2dblock_w, + use_outer_2dblock_x, use_outer_2dblock_w, + use_hadamard, expected_calls, id_label, +): + """Outer-amax reductions per forward+backward must match the expected count.""" + from alto.kernels.fp4.nvfp4 import nvfp_quantization as quant_mod + from alto.kernels.fp4.nvfp4.nvfp_linear import _to_nvfp4_then_scaled_mm + + real_fn = quant_mod._outer_scale_and_expanded_axis_last + calls = [] + + def _counted(*args, **kwargs): + calls.append(1) + return real_fn(*args, **kwargs) + + quant_mod._outer_scale_and_expanded_axis_last = _counted + try: + # Shapes satisfy the 1D wgrad-axis alignment contract + # (x.shape[0] % 16 == 0, weight.shape[0] % 16 == 0). + torch.manual_seed(0) + x = torch.randn(128, 256, dtype=torch.bfloat16, + device="cuda", requires_grad=True) + w = torch.randn(128, 256, dtype=torch.bfloat16, + device="cuda", requires_grad=True) + y = _to_nvfp4_then_scaled_mm( + x, w, + use_2dblock_x=use_2dblock_x, + use_2dblock_w=use_2dblock_w, + use_sr_grad=False, + use_outer_scale=False, + use_hadamard=use_hadamard, + use_dge=False, + use_outer_block_scale=True, + use_outer_2dblock_x=use_outer_2dblock_x, + use_outer_2dblock_w=use_outer_2dblock_w, + outer_block_size=128, + ) + y.sum().backward() + finally: + quant_mod._outer_scale_and_expanded_axis_last = real_fn + + assert len(calls) == expected_calls, ( + f"{id_label}: expected {expected_calls} outer-amax reductions, " + f"saw {len(calls)}" + ) + + +# --------------------------------------------------------------------------- +# Outer-block kernel must agree bit-for-bit with the PyTorch reference under +# deterministic rounding. This catches axis / outer-amax indexing bugs that +# a self-roundtrip test would miss (errors that cancel between quant and +# dequant). +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +@pytest.mark.parametrize("axis", [-1, -2]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_convert_to_nvfp4_outer_block_bit_equal_to_pytorch_reference( + is_2d_inner, is_2d_outer, axis, dtype, +): + """Quant kernel and PyTorch reference must agree on all outputs bitwise.""" + x = prepare_data((128, 256), dtype) + + data_lp_k, inner_k, outer_k = convert_to_nvfp4_outer_block( + x, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + ) + data_lp_p, inner_p, outer_p = convert_to_nvfp4_outer_block_pytorch( + x, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + + assert data_lp_k.shape == data_lp_p.shape + assert inner_k.shape == inner_p.shape + assert outer_k.shape == outer_p.shape + assert data_lp_k.dtype == data_lp_p.dtype == torch.uint8 + assert torch.equal(data_lp_k, data_lp_p) + torch.testing.assert_close(inner_k, inner_p, atol=0, rtol=0) + torch.testing.assert_close(outer_k, outer_p, atol=0, rtol=0) + + +@pytest.mark.parametrize("is_2d_inner,is_2d_outer", [ + pytest.param(False, False, id="inner_1x16_outer_1x128"), + pytest.param(False, True, id="inner_1x16_outer_128x128"), + pytest.param(True, False, id="inner_16x16_outer_1x128"), + pytest.param(True, True, id="inner_16x16_outer_128x128"), +]) +@pytest.mark.parametrize("axis", [-1, -2]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_convert_from_nvfp4_outer_block_bit_equal_to_pytorch_reference( + is_2d_inner, is_2d_outer, axis, dtype, +): + """Dequant kernel and PyTorch reference must agree bit-for-bit.""" + x = prepare_data((128, 256), dtype) + data_lp, inner_scales, outer_scales = convert_to_nvfp4_outer_block( + x, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + use_sr=False, + ) + + dq_k = convert_from_nvfp4_outer_block( + data_lp, inner_scales, outer_scales, + output_dtype=dtype, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + dq_p = convert_from_nvfp4_outer_block_pytorch( + data_lp, inner_scales, outer_scales, + output_dtype=dtype, + block_size=16, outer_block_size=128, + axis=axis, + is_2d_inner=is_2d_inner, + is_2d_outer=is_2d_outer, + ) + torch.testing.assert_close(dq_k, dq_p, atol=0, rtol=0) + + +def test_outer_block_close_to_pts_when_outer_covers_whole_tensor(): + """Whole-tensor 2D outer-block QDQ tracks the PTS round-trip within slack. + + The two paths are not algebraically identical: the PTS path applies two + E4M3 rounds (one on ``amax/6``, one on ``block_scale/pts``), while the + outer-block path normalizes the input in bf16 and runs a single E4M3 + round on the inner scale. This test only requires that the resulting + round-trip outputs stay close, catching catastrophic divergences (axis + bugs, wrong outer-amax math) without over-constraining the implementation. + + 1D outer cases are excluded because ``outer_block_size=N`` produces one + scale per row, not a global scale. + """ + M, N = 128, 128 + x = prepare_data((M, N), torch.bfloat16) + + data_lp_ob, inner_ob, outer_ob = convert_to_nvfp4_outer_block( + x, + block_size=16, + outer_block_size=max(M, N), + axis=-1, + is_2d_inner=True, + is_2d_outer=True, + use_sr=False, + ) + dq_ob = convert_from_nvfp4_outer_block( + data_lp_ob, inner_ob, outer_ob, + output_dtype=torch.bfloat16, + block_size=16, + outer_block_size=max(M, N), + axis=-1, + is_2d_inner=True, + is_2d_outer=True, + ) + + pts = compute_dynamic_outer_scale(x) + data_lp_pts, scales_pts = convert_to_nvfp4( + x, + block_size=16, + axis=-1, + is_2d_block=True, + outer_scale=pts, + update_outer_scale=False, + use_sr=False, + ) + dq_pts = convert_from_nvfp4( + data_lp_pts, scales_pts, + output_dtype=torch.bfloat16, + block_size=16, + axis=-1, + is_2d_block=True, + outer_scale=pts, + ) + + # ``prepare_data`` is heavy-tailed (~0.5% entries scaled by ~100×), so a + # whole-tensor outer scale is dominated by outliers and leaves inner E4M3 + # in the precision-poor ``< 1`` region for the bulk of the tensor. Allow + # a generous slack on the absolute round-trip MSE; the binding check is + # the cross-path agreement. + sig_pow = x.float().pow(2).mean().clamp(min=1e-30) + err_ob = (x.float() - dq_ob.float()).pow(2).mean() / sig_pow + err_pts = (x.float() - dq_pts.float()).pow(2).mean() / sig_pow + err_xx = (dq_ob.float() - dq_pts.float()).pow(2).mean() / sig_pow + + assert err_ob < 5e-2, f"outer-block round-trip MSE too high: {err_ob.item():.3e}" + assert err_pts < 5e-2, f"PTS round-trip MSE too high: {err_pts.item():.3e}" + assert err_xx < 1e-1, ( + f"outer-block / PTS disagreement {err_xx.item():.3e} larger than " + f"either path's reconstruction error (ob={err_ob.item():.3e}, " + f"pts={err_pts.item():.3e})" + ) diff --git a/tests/unittest/nvfp4/utils.py b/tests/unittest/nvfp4/utils.py index 07b1c03..02e8ad1 100644 --- a/tests/unittest/nvfp4/utils.py +++ b/tests/unittest/nvfp4/utils.py @@ -8,6 +8,8 @@ BLOCK_SIZE_DEFAULT, SUPPORTED_SCALE_FORMATS, _SCALE_FORMAT_TABLE, + OUTER_BLOCK_SIZE_DEFAULT, + OUTER_SCALE_EPS, ) from alto.kernels.fp4.fp4_common import quantize_to_ue5m3 @@ -44,7 +46,8 @@ def _quantize_inner_scale(inner_scale_raw: Tensor, scale_format: str) -> Tensor: return clamped.float().to(torch.float8_e4m3fn).to(torch.float32) -def prepare_data(tensor_shape, data_type, pattern="random"): +def prepare_data(tensor_shape, data_type, pattern="random", seed: int = 1234, + outlier_scale: float = 1000.0): """Prepare test data with specified pattern. Args: @@ -56,8 +59,24 @@ def prepare_data(tensor_shape, data_type, pattern="random"): stored block scale. "large" : All 5000.0 — exceeds F8E4M3_MAX * F4_E2M1_MAX (2688), tests FP4 saturation and scale clamp to F8E4M3_MAX. + "lognormal_channel_outlier" : Log-normal background with a small + fraction of channels (along the last axis) scaled by + ``outlier_scale``. Approximates LLM activation + distributions (Sun et al. 2024 "massive activations"), + which are the worst case for tensor-wise scaling and + the best case for outer-block scaling. See + NVFP4_Outer_Block_Review.md §3.3 T1. + seed: Per-call seed override. Defaults to 1234 for backward + compatibility; tests that mix multiple patterns / shapes in + a single file should pass distinct seeds to avoid sharing the + same RNG state across parametrised cases. + outlier_scale: Multiplier applied to the promoted "massive" channels of + the ``lognormal_channel_outlier`` pattern (default 1e3). Larger + values (e.g. 1e5) drive per-tensor scaling into E4M3 inner-scale + underflow, the regime where outer-block scaling is decisively + better. """ - torch.manual_seed(1234) + torch.manual_seed(seed) device = torch.device("cuda") if pattern == "random": @@ -82,6 +101,20 @@ def prepare_data(tensor_shape, data_type, pattern="random"): x = torch.zeros(tensor_shape, dtype=data_type, device=device) flat = x.reshape(-1) flat[flat.numel() // 2] = 5000.0 + elif pattern == "lognormal_channel_outlier": + # Log-normal-style background (heavy-tailed, mean ≈ 0). We take + # ``randn ** 2 * 0.1 * sign(randn)`` so the marginal is approximately + # ``±X^2 / 10`` with X ~ N(0, 1), a heavy-tailed proxy. + base = torch.randn(tensor_shape, dtype=torch.float32, device=device) + x = (base.abs() ** 2) * 0.1 * base.sign() + if len(tensor_shape) >= 1 and tensor_shape[-1] >= 8: + # Promote ~1% of the last-axis channels to massive activations. + num_channels = tensor_shape[-1] + num_outlier = max(1, num_channels // 100) + generator = torch.Generator(device="cpu").manual_seed(seed + 1) + outlier_idx = torch.randperm(num_channels, generator=generator)[:num_outlier] + x[..., outlier_idx] = x[..., outlier_idx] * outlier_scale + x = x.to(data_type) else: raise ValueError(f"Unknown pattern: {pattern}") @@ -367,3 +400,144 @@ def convert_from_nvfp4_pytorch( s_expanded = scales.reshape(scale_shape) result = (f32 * s_expanded).to(output_dtype).reshape(orig_hp_shape) return result.transpose(axis, -1) + + +# --------------------------------------------------------------------------- +# NVFP4 outer-block PyTorch reference +# +# These mirror the structure of ``convert_to_nvfp4_outer_block`` / +# ``convert_from_nvfp4_outer_block`` in the production kernel, but go through +# pure-tensor amax + ``convert_to_nvfp4_pytorch`` so they can be used as a +# trustworthy bit-level reference for tests. +# --------------------------------------------------------------------------- + + +def _outer_scale_axis_last_pytorch( + data_axis_last: torch.Tensor, + *, + is_2d_outer: bool, + inner_block_size: int, + outer_block_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Reference outer scale + expanded scale on an axis-last tensor. + + Mirrors ``_outer_scale_and_expanded_axis_last`` in the production kernel, + including the FP32-floor clamp. Returns ``(outer_scale, outer_expanded)`` + both in axis-last layout. + """ + del inner_block_size # outer scale doesn't depend on inner block size + shape = tuple(data_axis_last.shape) + m, n = shape[-2], shape[-1] + prefix = shape[:-2] + x = data_axis_last.float() + if is_2d_outer: + grouped = x.reshape( + *prefix, + m // outer_block_size, outer_block_size, + n // outer_block_size, outer_block_size, + ) + outer_amax = grouped.abs().amax(dim=-1).amax(dim=-2) + outer_scale = (outer_amax / (F8E4M3_MAX * F4_E2M1_MAX)).clamp(min=OUTER_SCALE_EPS) + expanded = outer_scale.unsqueeze(-2).unsqueeze(-1).expand_as(grouped).reshape(shape) + else: + grouped = x.reshape(*prefix, m, n // outer_block_size, outer_block_size) + outer_amax = grouped.abs().amax(dim=-1) + outer_scale = (outer_amax / (F8E4M3_MAX * F4_E2M1_MAX)).clamp(min=OUTER_SCALE_EPS) + expanded = outer_scale.unsqueeze(-1).expand_as(grouped).reshape(shape) + return outer_scale.to(torch.float32), expanded.to(torch.float32) + + +def convert_to_nvfp4_outer_block_pytorch( + data_hp: torch.Tensor, + block_size: int = BLOCK_SIZE_DEFAULT, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, + axis: int = -1, + is_2d_inner: bool = False, + is_2d_outer: bool = False, + scale_format: str = "e4m3", +): + """Pure-PyTorch reference for ``convert_to_nvfp4_outer_block``. + + Forward path: outer amax → FP32 outer scale → normalize → inner NVFP4 QDQ + via :func:`convert_to_nvfp4_pytorch` with ``outer_scale=None``. + Returns ``(data_lp, inner_scales, outer_scale)`` all in **original axis + layout** (matching the production kernel's return convention). + + Note: this reference does **not** support stochastic rounding (the + production kernel exposes ``use_sr`` only when called from + ``convert_to_nvfp4``). Tests that need bit-equality must therefore run + with deterministic rounding (``use_sr=False``). + """ + assert data_hp.dtype in [torch.float32, torch.bfloat16] + data_axis_last = data_hp.transpose(axis, -1).contiguous() + outer_scale_axis_last, outer_expanded = _outer_scale_axis_last_pytorch( + data_axis_last, + is_2d_outer=is_2d_outer, + inner_block_size=block_size, + outer_block_size=outer_block_size, + ) + # Keep the per-outer-block normalization in FP32 through the inner QDQ to + # match the production kernel (see ``convert_to_nvfp4_outer_block``): casting + # back to bf16 here would add a redundant rounding step on top of the inner + # NVFP4 quantization and break bit-exactness with the kernel. + normalized = data_axis_last.float() / outer_expanded + data_lp_axis_last, inner_scales_axis_last = convert_to_nvfp4_pytorch( + normalized, + block_size=block_size, + axis=-1, + is_2d_block=is_2d_inner, + outer_scale=None, + scale_format=scale_format, + ) + # Ensure the inner-scale shape matches the production kernel's contract. + return ( + data_lp_axis_last.transpose(axis, -1).contiguous(), + inner_scales_axis_last.transpose(axis, -1).contiguous(), + outer_scale_axis_last.transpose(axis, -1).contiguous(), + ) + + +def convert_from_nvfp4_outer_block_pytorch( + data_lp: torch.Tensor, + inner_scales: torch.Tensor, + outer_scales: torch.Tensor, + output_dtype: torch.dtype = torch.float32, + block_size: int = BLOCK_SIZE_DEFAULT, + outer_block_size: int = OUTER_BLOCK_SIZE_DEFAULT, + axis: int = -1, + is_2d_inner: bool = False, + is_2d_outer: bool = False, + scale_format: str = "e4m3", +): + """Pure-PyTorch reference for ``convert_from_nvfp4_outer_block``.""" + inner_dq = convert_from_nvfp4_pytorch( + data_lp, + inner_scales, + output_dtype=output_dtype, + block_size=block_size, + axis=axis, + is_2d_block=is_2d_inner, + outer_scale=None, + scale_format=scale_format, + ) + inner_dq_axis_last = inner_dq.transpose(axis, -1).contiguous() + outer_axis_last = outer_scales.transpose(axis, -1).to( + device=inner_dq.device, dtype=torch.float32, + ).contiguous() + shape = tuple(inner_dq_axis_last.shape) + m, n = shape[-2], shape[-1] + prefix = shape[:-2] + if is_2d_outer: + grouped_shape = ( + *prefix, + m // outer_block_size, outer_block_size, + n // outer_block_size, outer_block_size, + ) + outer_expanded = outer_axis_last.unsqueeze(-2).unsqueeze(-1).expand( + grouped_shape).reshape(shape) + else: + grouped_shape = (*prefix, m, n // outer_block_size, outer_block_size) + outer_expanded = outer_axis_last.unsqueeze(-1).expand( + grouped_shape).reshape(shape) + result = (inner_dq_axis_last.float() * outer_expanded).to(output_dtype) + return result.transpose(axis, -1).contiguous()