[XPU][Bugfix] Do not transpose weight_scale_inv at load time (#47116)
Signed-off-by: Ma Jian <jian1.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Ma Jian committed
Jul 6, 2026 at 17:15 UTC
26c754d8475e753928d57b71c2398515eabe0a0d
1 file changed
+2
-24
vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+2
-24
@@ -197,29 +197,6 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
197
return False, "XPUFp8BlockScaledMM only support on XPU"
198
return True, None
199
200
- def process_weights_after_loading(self, layer: torch.nn.Module):
201
- super().process_weights_after_loading(layer)
202
- scale_attr = (
203
- "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale"
204
- )
205
- scale = getattr(layer, scale_attr)
206
- # Models with scale_fmt=ue8m0 (e.g. DeepSeek-V4) store weight scales
207
- # as float8_e8m0fnu. The oneDNN fp8_gemm kernel dispatches to its
208
- # "block quant" path only when NEITHER scale is e8m0:
209
- #
210
- # is_block_quant = (m1_sc != e8m0) && (m2_sc != e8m0) && ...
211
- #
212
- # Since activation scales are always float32 (use_ue8m0=False on XPU,
213
- # DeepGEMM requires Hopper/Blackwell), an e8m0 weight scale causes
214
- # is_block_quant=false and falls into the wrong per-channel path,
215
- # producing NaN. Converting e8m0→float32 here at load time (one-time,
216
- # negligible overhead for small scale tensors) ensures the kernel sees
217
- # matching dtypes and correctly enters the block-quant path with the
218
- # actual group_size derived from scale tensor shapes.
219
- if scale.dtype == torch.float8_e8m0fnu:
220
- scale = scale.to(torch.float32)
221
- replace_parameter(layer, scale_attr, scale.data.t().contiguous())
222
-
200
def apply_block_scaled_mm(
201
self,
202
A: torch.Tensor,
@@ -228,11 +205,12 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
205
Bs: torch.Tensor,
206
) -> torch.Tensor:
207
# Weight is [N, K]. Use .t() to create a [K, N] view without copying.
208
+ # Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN.
209
return torch.ops._xpu_C.fp8_gemm(
210
A,
211
B.t(),
212
self.config.out_dtype,
213
As,
236
- Bs,
214
+ Bs.t().contiguous(),
215
torch.Tensor(),
216
)