[ROCm] Cache fp32 upcast of static e8m0 weight scale in AITER scaled_mm (#47773)
Signed-off-by: jiacao-amd <jiahui.cao@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
jiacao-amd committed
Jul 28, 2026 at 09:34 UTC
4f56321d7ec25dd041c7bf0aa47b4eadd075a7e5
1 file changed
+16
-8
vllm/model_executor/kernels/linear/scaled_mm/aiter.py
+16
-8
@@ -9,6 +9,9 @@ from vllm._aiter_ops import (
9
rocm_aiter_ops,
10
)
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13
+ _upcast_e8m0_to_fp32,
14
+)
15
from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
GroupShape,
17
)
@@ -16,6 +19,7 @@ from vllm.model_executor.utils import replace_parameter
19
from vllm.platforms import current_platform
20
21
from .BlockScaledMMLinearKernel import (
22
+ FP8BlockParams,
23
Fp8BlockScaledMMLinearKernel,
24
)
25
from .cutlass import CutlassInt8ScaledMMLinearKernel
@@ -375,6 +379,17 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
379
and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
380
)
381
382
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
383
+ super().process_weights_after_loading(layer)
384
+
385
+ params = FP8BlockParams.from_layer(layer)
386
+ if params.weight_scale_inv is not None:
387
+ ws, attr = params.weight_scale_inv, params.WEIGHT_SCALE_INV
388
+ else:
389
+ ws, attr = params.weight_scale, params.WEIGHT_SCALE
390
+ if ws is not None and ws.dtype == torch.float8_e8m0fnu:
391
+ replace_parameter(layer, attr, _upcast_e8m0_to_fp32(ws).contiguous())
392
+
393
@classmethod
394
def is_supported(cls, compute_capability=None):
395
return (
@@ -406,19 +421,12 @@ class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
421
Bs: torch.Tensor,
422
) -> torch.Tensor:
423
if As.dtype != Bs.dtype:
409
- from vllm.model_executor.layers.quantization.utils.fp8_utils import (
410
- _upcast_e8m0_to_fp32,
411
- )
412
-
424
if As.dtype == torch.float8_e8m0fnu:
425
As = _upcast_e8m0_to_fp32(As).contiguous()
426
else:
427
As = As.to(torch.float32)
428
418
- if Bs.dtype == torch.float8_e8m0fnu:
419
- Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
420
- else:
421
- Bs = Bs.to(torch.float32)
429
+ Bs = Bs.to(torch.float32)
430
431
out_dtype = self.config.out_dtype
432
if self.use_triton: