kernels: fp4: support blockwise outer scale

alto/kernels/dispatch/config.py

@@ -38,15 +38,19 @@ 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
     * tensorwise: apply a global scale factor to the entire tensor (for NVFP4 only)
     * 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])

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

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",
     )

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",
     )
 

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",
 )

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,
         )

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():