sm12x: hoist the E8M0 block-scale upcast out of the FP8 GEMM hot path

vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py

@@ -12,6 +12,7 @@
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     process_fp8_weight_block_strategy,
 )
+from vllm.logger import init_logger
 from vllm.model_executor.utils import replace_parameter
 
 from ..base import (
@@ -20,6 +21,8 @@
 )
 from .ScaledMMLinearKernel import FP8ScaledMMLinearLayerConfig
 
+logger = init_logger(__name__)
+
 
 @dataclass
 class FP8BlockParams(FP8Params):
@@ -46,6 +49,17 @@ class Fp8BlockScaledMMLinearKernel(
     # and therefore do not need the input quantization step in apply_weights.
     apply_input_quant: ClassVar[bool] = True
 
+    # Set to True in subclasses whose kernel cannot bind an E8M0 block-scale
+    # tensor and upcasts it to fp32 on every call. Such subclasses populate
+    # ``layer.<FP32_BLOCK_SCALE_ATTR>`` once in process_weights_after_loading;
+    # apply_weights then hands the kernel the cached fp32 copy instead of
+    # repeating the upcast per forward. The original E8M0 parameter is left in
+    # place because other consumers read it directly (e.g. DeepSeek-V4's fused
+    # o_proj and DSpark kernels read layer.weight_scale{,_inv}).
+    prefers_fp32_block_scale: ClassVar[bool] = False
+
+    FP32_BLOCK_SCALE_ATTR: ClassVar[str] = "_fp32_block_scale"
+
     def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
         super().__init__(config)
         act_scale_descriptor = config.activation_quant_key.scale
@@ -94,6 +108,42 @@ def process_weights_after_loading(self, layer: torch.nn.Module):
         replace_parameter(layer, params.WEIGHT, new_weight.data)
         replace_parameter(layer, scale_attr_name, new_weight_scale.data)
 
+        self._maybe_cache_fp32_block_scale(layer)
+
+    def _maybe_cache_fp32_block_scale(self, layer: torch.nn.Module) -> None:
+        """Stash an fp32 copy of an E8M0 block scale alongside the original.
+
+        E8M0 is exponent-only, so the upcast is an exact bit shift; the cached
+        copy is numerically identical to what the kernel would compute per call.
+        """
+        if not self.prefers_fp32_block_scale:
+            return
+        e8m0 = getattr(torch, "float8_e8m0fnu", None)
+        if e8m0 is None:
+            return
+        params = self._get_layer_params(layer)
+        scale = (
+            params.weight_scale
+            if params.weight_scale_inv is None
+            else params.weight_scale_inv
+        )
+        if scale is None or scale.dtype != e8m0:
+            return
+        from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+            _upcast_e8m0_to_fp32,
+        )
+
+        layer.register_buffer(
+            self.FP32_BLOCK_SCALE_ATTR,
+            _upcast_e8m0_to_fp32(scale.data).contiguous(),
+            persistent=False,
+        )
+        logger.info_once(
+            "Caching fp32 copies of E8M0 block scales for %s; the per-call "
+            "upcast is skipped.",
+            type(self).__name__,
+        )
+
     def apply_weights(
         self,
         layer: torch.nn.Module,
@@ -109,6 +159,10 @@ def apply_weights(
             if params.weight_scale_inv is None
             else params.weight_scale_inv
         )
+        if self.prefers_fp32_block_scale:
+            cached = getattr(layer, self.FP32_BLOCK_SCALE_ATTR, None)
+            if cached is not None:
+                weight_scale = cached
         input_scale = params.input_scale
         scale_up = params.input_scale_ub
 

vllm/model_executor/kernels/linear/scaled_mm/triton.py

@@ -157,6 +157,11 @@ def apply_weights(
 
 
 class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    # Triton cannot bind an E8M0 tensor, so w8a8_triton_block_scaled_mm upcasts
+    # the block scale to fp32 on every call. The weight scale is static, so let
+    # the base class cache the upcast at load time instead.
+    prefers_fp32_block_scale = True
+
     @classmethod
     def is_supported(cls, compute_capability=None):
         if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):