[PyTorch] Enable FP8 block scaling for fusible ops

tests/pytorch/test_fusible_ops.py

@@ -36,6 +36,7 @@
 )
 from transformer_engine.pytorch import (
     QuantizedTensor,
+    Float8BlockQuantizer,
     Float8CurrentScalingQuantizer,
     Float8Quantizer,
     MXFP8Quantizer,
@@ -58,6 +59,9 @@
 fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
 mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True)
 nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True)
+fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available(
+    return_reason=True
+)
 
 # Supported data types
 _dtypes: list[torch.dtype] = [torch.float32, torch.float16]
@@ -76,6 +80,8 @@
 if nvfp4_available:
     _quantization_list.append("nvfp4")
     _quantization_list.append("nvfp4_4over6")
+if fp8_block_scaling_available:
+    _quantization_list.append("fp8_block_scaling")
 
 
 @pytest.fixture(autouse=True, scope="function")
@@ -110,6 +116,8 @@ def maybe_skip_quantization(
         and not nvfp4_available
     ):
         pytest.skip(reason_for_no_nvfp4)
+    if quantization == "fp8_block_scaling" and not fp8_block_scaling_available:
+        pytest.skip(reason_for_no_fp8_block_scaling)
 
     # Check dims
     if dims is not None:
@@ -121,6 +129,9 @@ def maybe_skip_quantization(
         elif quantization == "mxfp8":
             if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0:
                 pytest.skip("MXFP8 GEMMs require dims that are divisible by 32")
+        elif quantization == "fp8_block_scaling":
+            if math.prod(dims[:-1]) % 128 != 0 or dims[-1] % 128 != 0:
+                pytest.skip("FP8 block scaling requires dims that are divisible by 128")
         elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"):
             if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0:
                 pytest.skip("NVFP4 GEMMs require dims that are divisible by 16")
@@ -186,6 +197,17 @@ def make_reference_and_test_tensors(
         test = quantizer(test)
     elif quantization == "mxfp8":
         test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test)
+    elif quantization == "fp8_block_scaling":
+        tensor_type = "input"
+        if quantizer_role is not None:
+            tensor_type = quantizer_role.tensor_type
+        # Weights use 2D (128x128) blocks; everything else 1D (1x128).
+        test = Float8BlockQuantizer(
+            fp8_dtype=te.DType.kFloat8E4M3,
+            rowwise=True,
+            columnwise=True,
+            block_scaling_dim=2 if tensor_type == "weight" else 1,
+        )(test)
     elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"):
         tensor_type = "input"
         if quantizer_role is not None:
@@ -1017,8 +1039,9 @@ def _test_basic_linear(
             )
         torch.testing.assert_close(dw_test, w_ref.grad, **tols)
 
-    @pytest.mark.parametrize("weight_shape", ((64, 32), (3, 5)))
-    @pytest.mark.parametrize("in_shape", ((-1,), (5, 1, -1), (4, 2, 4, -1)))
+    # (128, 128) + a 128-token in_shape keep FP8 block scaling (128-divisible dims) unskipped.
+    @pytest.mark.parametrize("weight_shape", ((64, 32), (3, 5), (128, 128)))
+    @pytest.mark.parametrize("in_shape", ((-1,), (5, 1, -1), (4, 2, 4, -1), (128, -1)))
     @pytest.mark.parametrize("dtype", _dtypes)
     @pytest.mark.parametrize("quantization", _quantization_list)
     @pytest.mark.parametrize("accumulate_into_main_grad", (False, True))
@@ -3360,7 +3383,7 @@ def test_layernorm_mlp(
         quantization: Optional[str],
         device: torch.device = "cuda",
         hidden_size: int = 256,
-        sequence_length: int = 48,
+        sequence_length: int = 64,
         batch_size: int = 4,
         ffn_hidden_size: int = 384,
         layernorm_epsilon: float = 1e-5,

tests/pytorch/utils.py

@@ -114,6 +114,7 @@ def quantization_tols(name: str) -> dict[str, float]:
         "fp8_delayed_scaling",
         "fp8_current_scaling",
         "fp8_blockwise",
+        "fp8_block_scaling",
         "mxfp8",
         "mxfp8_block_scaling",
     ):

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -27,6 +27,7 @@
 from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe
 from ...quantized_tensor import QuantizedTensorStorage
 from ...tensor import (
+    Float8BlockQuantizer,
     Float8CurrentScalingQuantizer,
     MXFP8Quantizer,
     MXFP8Tensor,
@@ -791,9 +792,10 @@ def _is_graph_safe_path_supported(
         * FP8 per-tensor current scaling is backed by grouped current-scaling quantization
           (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling,
           which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0).
-          Every other quantization recipe (fp8 delayed scaling, fp8 block scaling, ...)
-          falls back to the legacy flow because the corresponding grouped quantization kernels are
-          missing.
+        * FP8 block scaling uses the grouped-tensor path on Hopper (CC 9.0) with cuBLAS 13.4+;
+          it is Hopper-only (no MXFP8-broadcast emulation), so elsewhere it falls back.
+        * Every other quantization recipe (fp8 delayed scaling, ...) falls back to the legacy flow
+          because the corresponding grouped quantization kernels are missing.
         * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0)
           -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it.
         * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it
@@ -812,6 +814,12 @@ def _is_graph_safe_path_supported(
                 ):
                     return False
                 return True
+            if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers):
+                # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere
+                # fall back to the split-quantize (MXFP8-emulated) flow.
+                if get_device_compute_capability() >= (10, 0):
+                    return False
+                return tex.get_cublasLt_version() >= 130400
             # MXFP8 and NVFP4 grouped quantization kernels require Blackwell.
             if not (10, 0) <= get_device_compute_capability() <= (11, 0):
                 return False
@@ -1666,11 +1674,12 @@ def _fuser_backward_grouped_tensor(
             )
             grad_output_quantizer.optimize_for_gemm = True
 
-            if (
-                has_bias
-                and not self._scale_bias
-                and isinstance(grad_output_quantizer, MXFP8Quantizer)
-            ):
+            # FP8 block scaling computes dbias in the rowwise (dgrad) pass, so only fuse
+            # when dgrad is required.
+            fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or (
+                isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad
+            )
+            if has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
                     dy_2d, grad_output_quantizer, num_groups, split_sizes
                 )

transformer_engine/pytorch/ops/op.py

@@ -275,11 +275,6 @@ def reset_recipe_state(
                 if num_quantizers == 0:
                     continue
 
-                if recipe.float8_block_scaling():
-                    raise NotImplementedError(
-                        "Fusible operations do not support FP8 block scaling recipe"
-                    )
-
                 # Construct quantization recipe state
                 roles = self.get_quantizer_roles(mode)  # pylint: disable=assignment-from-none
                 if roles is not None:

transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py

@@ -5,6 +5,7 @@
 """Mixin class holding data specific for Float8BlockwiseQTensor"""
 
 from __future__ import annotations
+from collections.abc import Iterable
 import math
 from typing import Optional, Dict, Any, Tuple, Union
 import torch
@@ -323,6 +324,61 @@ def size(self, *args, **kwargs):
         reordered.append(dims[0])
         return torch.Size(reordered)
 
+    def view(self, shape):
+        """Reshape the leading (token) dims without dequantizing.
+
+        Mirrors ``MXFP8TensorStorage.view``. Only leading dims may change: block tiling fixes
+        the inner dim (1D) or inner two (2D), so scale-inv stays valid. Columnwise data is
+        stored transposed, viewed as ``[inner, *leading]``.
+        """
+        cur_shape = self.size()
+        if shape is None or shape == cur_shape:
+            return self
+        # Canonicalize shape
+        if not isinstance(shape, Iterable):
+            shape = [shape]
+        elif len(shape) == 1 and isinstance(shape[0], Iterable):
+            shape = shape[0]
+        shape = list(shape)
+        if -1 in shape:
+            d_inferred = -math.prod(cur_shape) // math.prod(shape)
+            for i, d in enumerate(shape):
+                if d == -1:
+                    shape[i] = d_inferred
+                    break
+        if shape == list(cur_shape):
+            return self
+
+        if self._is_2D_scaled:
+            if shape[-2:] != list(cur_shape)[-2:]:
+                raise RuntimeError(
+                    "Float8BlockwiseQTensorStorage (2D block scaling) cannot reshape the inner "
+                    f"two dimensions (attempted {tuple(cur_shape)} -> {tuple(shape)})"
+                )
+        elif shape[-1] != cur_shape[-1]:
+            raise RuntimeError(
+                "Float8BlockwiseQTensorStorage (1D block scaling) cannot reshape the inner "
+                f"dimension (attempted {tuple(cur_shape)} -> {tuple(shape)})"
+            )
+
+        new_rowwise_data = None
+        if self._rowwise_data is not None:
+            new_rowwise_data = self._rowwise_data.view(*shape)
+        new_columnwise_data = None
+        if self._columnwise_data is not None:
+            new_columnwise_data = self._columnwise_data.view([shape[-1], *shape[:-1]])
+
+        return Float8BlockwiseQTensorStorage(
+            new_rowwise_data,
+            self._rowwise_scale_inv,
+            new_columnwise_data,
+            self._columnwise_scale_inv,
+            self._fp8_dtype,
+            self._quantizer,
+            self._is_2D_scaled,
+            fake_dtype=self._dtype,
+        )
+
     @property
     def device(self):
         """Return the device of the tensor. Define this to avoid expensive PyObject lookups."""