[XPU] Support MXFP8 linear weights for INC DeepSeek V4 model (#48476)
Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Signed-off-by: Xiaochang Wu <xiaochang.wu@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xiaochang Wu committed
Aug 6, 2026 at 13:17 UTC
d779835a1196d45225c8429ee23bff293517f77e
1 file changed
+36
-5
vllm/model_executor/kernels/linear/mxfp8/xpu.py
+36
-5
@@ -28,10 +28,38 @@ class XPUMxFp8LinearKernel(Mxfp8LinearKernel):
28
return True, None
29
30
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
31
+ # Checkpoint scale is [N, K//32] (one E8M0 scale per 32 elements).
32
+ # oneDNN fp8_gemm requires contiguous [K//32, N] layout. Store the
33
+ # transposed contiguous buffer as a .t() view so that:
34
+ # - dequantization consumers still see the checkpoint shape
35
+ # - apply_weights recovers the oneDNN layout via .t() at zero cost
36
weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
32
- weight_scale = weight_scale.t().contiguous()
33
- replace_parameter(layer, "weight", layer.weight.t())
34
- replace_parameter(layer, "weight_scale", weight_scale.data)
37
+ scale_kn = weight_scale.data.t().contiguous()
38
+ replace_parameter(layer, "weight_scale", scale_kn.t())
39
+
40
+ if getattr(layer, "is_bmm", False):
41
+ self._prepare_bmm_params(layer, scale_kn)
42
+
43
+ def _prepare_bmm_params(
44
+ self, layer: torch.nn.Module, scale_kn: torch.Tensor
45
+ ) -> None:
46
+ """Precompute batched weight and scale for grouped fp8_bmm (e.g. wo_a).
47
+
48
+ Splits scale [K//32, N_total] into [G, K//32, N_per_group] and weight
49
+ [N_total, K] into contiguous [G, K, N_per_group] for batch GEMM.
50
+ """
51
+ batch = layer.bmm_batch_size
52
+ k_blocks, n_blocks = scale_kn.shape
53
+ layer.bmm_scale = (
54
+ scale_kn.reshape(k_blocks, batch, n_blocks // batch)
55
+ .permute(1, 0, 2)
56
+ .contiguous()
57
+ )
58
+ w = layer.weight
59
+ n_total, k = w.shape
60
+ layer.bmm_weight = (
61
+ w.reshape(batch, n_total // batch, k).permute(0, 2, 1).contiguous()
62
+ )
63
64
def apply_weights(
65
self,
@@ -41,11 +69,14 @@ class XPUMxFp8LinearKernel(Mxfp8LinearKernel):
69
) -> torch.Tensor:
70
out_dtype = x.dtype
71
x_fp8, x_scale = quant_mxfp8(x)
72
+ # Weight is [N, K]; .t() gives a [K, N] view without copying.
73
+ # Scale is stored as [N, K//32]; .t() recovers the contiguous
74
+ # [K//32, N] buffer that oneDNN expects.
75
return torch.ops._xpu_C.fp8_gemm(
76
x_fp8,
46
- layer.weight,
77
+ layer.weight.t(),
78
out_dtype,
79
x_scale,
49
- layer.weight_scale,
80
+ layer.weight_scale.t(),
81
bias,
82
)