[Bug] Fix Weight Loading for Block FP8 Cutlass SM90

[yewentao256]

Purpose

Fix Weight Loading for Cutlass SM90.

Test Plan

vllm serve deepseek-ai/DeepSeek-V3.2-Exp -tp 8 --max-num-seqs 128 --load-format dummy --enforce_eager

Original

16:37:34 [multiproc_executor.py:671]     output = w8a8_blockscale_func(q_input, weight, x_scale, weight_scale,
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]   File "/opt/vllm-source/vllm/model_executor/layers/quantization/utils/fp8_utils.py", line 46, in cutlass_scaled_mm
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]     return ops.cutlass_scaled_mm(
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]            ^^^^^^^^^^^^^^^^^^^^^^
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]   File "/opt/vllm-source/vllm/_custom_ops.py", line 667, in cutlass_scaled_mm
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]     torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]   File "/opt/vllm/lib64/python3.12/site-packages/torch/_ops.py", line 1243, in __call__
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]     return self._op(*args, **kwargs)
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671]            ^^^^^^^^^^^^^^^^^^^^^^^^^
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671] RuntimeError: b_scale_group_shape must be [128, 128].
(APIServer pid=1) (EngineCore_DP0 pid=293) (Worker_TP3 pid=305) ERROR 09-29 16:37:34 [multiproc_executor.py:671] 

Now

(APIServer pid=3511336) INFO 09-29 22:03:11 [launcher.py:42] Route: /invocations, Methods: POST
(APIServer pid=3511336) INFO 09-29 22:03:11 [launcher.py:42] Route: /metrics, Methods: GET
(APIServer pid=3511336) INFO:     Started server process [3511336]
(APIServer pid=3511336) INFO:     Waiting for application startup.
(APIServer pid=3511336) INFO:     Application startup complete.

[gemini-code-assist]

Code Review

This pull request correctly fixes a bug in the weight loading logic for FP8 block-quantized models on SM90 GPUs, which could cause a runtime error. The refactoring also improves code clarity. However, I've identified a related issue where torch.bfloat16 is hardcoded when checking if DeepGEMM should be used. This could lead to incorrect behavior for models using other data types like float16. I've suggested a fix to use the layer's original data type instead.

[mgoin]

Thanks for fixing this, makes sense why it failed

[youkaichao]

thanks for fixing it!