[XPU] FP8 o_proj with fp8_bmm and load-time scale transpose (#48334)
Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Xiaochang Wu committed
Jul 20, 2026 at 16:32 UTC
823eaf667d33125259d2d280b0ed00747fa92440
5 files changed
+124
-15
vllm/_xpu_ops.py
+63
@@ -219,6 +219,63 @@ def _xpu_ops_deepseek_scaling_rope_fake(
219
return query, key
220
221
222
+def _xpu_fp8_bmm_impl(
223
+ a: torch.Tensor,
224
+ b: torch.Tensor,
225
+ out_dtype: torch.dtype,
226
+ a_scale: torch.Tensor,
227
+ b_scale: torch.Tensor,
228
+ bias: torch.Tensor | None,
229
+) -> torch.Tensor:
230
+ """XPU FP8 batched GEMM implementation for ``torch.ops.vllm.xpu_fp8_bmm``.
231
+
232
+ Computes batched matrix multiplication over the leading group dimension:
233
+ ``[G, M, K] @ [G, K, N] -> [G, M, N]``.
234
+
235
+ Args:
236
+ a: FP8 activation tensor with shape ``[G, M, K]``.
237
+ Does not need to be contiguous.
238
+ b: FP8 weight tensor with shape ``[G, K, N]``.
239
+ Does not need to be contiguous.
240
+ out_dtype: Output dtype accepted by the kernel (typically
241
+ ``torch.bfloat16`` for the DeepSeek-V4 O-proj path).
242
+ a_scale: Activation scale tensor for ``a``.
243
+ In current DeepSeek-V4 XPU usage it is block-scaled with shape
244
+ ``[G, M, K/bs]`` (``bs`` is the quant block size, e.g. 128).
245
+ Must be contiguous.
246
+ b_scale: Weight scale tensor for ``b``.
247
+ In current DeepSeek-V4 XPU usage it is block-scaled with shape
248
+ ``[G, K/bs, N/bs]`` (``bs`` is the quant block size, e.g. 128).
249
+ Must be contiguous.
250
+ bias: Optional bias tensor. Pass ``None`` when no bias is required.
251
+
252
+ Returns:
253
+ Output tensor with shape ``[G, M, N]`` and dtype ``out_dtype``.
254
+
255
+ Notes:
256
+ This implementation centralizes access to
257
+ ``torch.ops._xpu_C.fp8_bmm``. Both scales must be contiguous, while
258
+ ``a`` and ``b`` may be non-contiguous views.
259
+ """
260
+ return torch.ops._xpu_C.fp8_bmm(a, b, out_dtype, a_scale, b_scale, bias)
261
+
262
+
263
+def _xpu_fp8_bmm_fake(
264
+ a: torch.Tensor,
265
+ b: torch.Tensor,
266
+ out_dtype: torch.dtype,
267
+ a_scale: torch.Tensor,
268
+ b_scale: torch.Tensor,
269
+ bias: torch.Tensor | None,
270
+) -> torch.Tensor:
271
+ # [G, M, K] @ [G, K, N] => [G, M, N]
272
+ return torch.empty(
273
+ (a.shape[0], a.shape[1], b.shape[2]),
274
+ dtype=out_dtype,
275
+ device=a.device,
276
+ )
277
+
278
+
279
def _xpu_fp8_mqa_logits_impl(
280
q: torch.Tensor,
281
k_quant: torch.Tensor,
@@ -1053,6 +1110,12 @@ class xpu_ops:
1110
fake_impl=_xpu_mxfp4_quantize_fake,
1111
)
1112
1113
+ direct_register_custom_op(
1114
+ op_name="xpu_fp8_bmm",
1115
+ op_func=_xpu_fp8_bmm_impl,
1116
+ fake_impl=_xpu_fp8_bmm_fake,
1117
+ )
1118
+
1119
direct_register_custom_op(
1120
op_name="xpu_fp8_mqa_logits",
1121
op_func=_xpu_fp8_mqa_logits_impl,
vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+33
-2
@@ -197,6 +197,37 @@ 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
+ # Transpose scale from checkpoint layout [N/128, K/128] to
207
+ # oneDNN expected layout [K/128, N/128] at load time (one-time cost).
208
+ scale_t = scale.data.t().contiguous()
209
+ replace_parameter(layer, scale_attr, scale_t)
210
+
211
+ # For BMM layers (e.g. wo_a), precompute 3D scale and weight:
212
+ # [K/bs, N/bs] -> [batch, K/bs, N_per_batch/bs]
213
+ if getattr(layer, "is_bmm", False):
214
+ batch = layer.bmm_batch_size
215
+ k_blocks = scale_t.shape[0]
216
+ n_per_batch_blocks = scale_t.shape[1] // batch
217
+ layer.bmm_scale = (
218
+ scale_t.reshape(k_blocks, batch, n_per_batch_blocks)
219
+ .permute(1, 0, 2)
220
+ .contiguous()
221
+ )
222
+ # Precompute [G, K, N] weight for fp8_bmm.
223
+ # Original weight is [N_total, K] where N_total = G * N_per_group.
224
+ w = layer.weight.data
225
+ N_total, K = w.shape
226
+ N_per_group = N_total // batch
227
+ layer.bmm_weight = w.reshape(batch, N_per_group, K).permute(
228
+ 0, 2, 1
229
+ ) # [G, K, N]
230
+
231
def apply_block_scaled_mm(
232
self,
233
A: torch.Tensor,
@@ -205,12 +236,12 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
236
Bs: torch.Tensor,
237
) -> torch.Tensor:
238
# 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.
239
+ # Bs is already [K/128, N/128] from process_weights_after_loading.
240
return torch.ops._xpu_C.fp8_gemm(
241
A,
242
B.t(),
243
self.config.out_dtype,
244
As,
214
- Bs.t().contiguous(),
245
+ Bs,
246
torch.Tensor(),
247
)
vllm/models/deepseek_v4/common/ops/fused_inv_rope_fp8_quant.py
+1
-2
@@ -257,7 +257,6 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
257
)
258
grid = (tma_aligned_T, n_groups * heads_per_group)
259
use_gdc = current_platform.is_arch_support_pdl()
260
- pdl_kwargs = {"launch_pdl": True} if use_gdc else {}
260
_fused_inv_rope_fp8_quant_per_head[grid](
261
o,
262
positions,
@@ -281,8 +280,8 @@ def _fused_inv_rope_fp8_quant_kernel_impl(
280
HALF_ROPE=half_rope,
281
TMA_ALIGNED_SCALES=tma_aligned_scales,
282
USE_GDC=use_gdc,
283
+ launch_pdl=use_gdc,
284
num_stages=1,
285
- **pdl_kwargs,
285
num_warps=1,
286
)
287
return fp8_buf, scale_buf
vllm/models/deepseek_v4/xpu/model.py
-2
@@ -8,7 +8,6 @@ import regex as re
8
import torch
9
import torch.nn as nn
10
11
-from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed import (
13
get_ep_group,
@@ -978,7 +977,6 @@ class DeepseekV4DecoderLayer(nn.Module):
977
return x, residual, post_mix, res_mix
978
979
981
-@support_torch_compile
980
class DeepseekV4Model(nn.Module, EagleModelMixin):
981
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
982
super().__init__()
vllm/models/deepseek_v4/xpu/xpu_sparse.py
+27
-9
@@ -89,19 +89,37 @@ class DeepseekV4XPUAttention(DeepseekV4Attention):
89
return num_heads
90
91
def _o_proj(self, o: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
92
- # XPU uses BF16 reference wo_a path (same as ROCm).
93
- from vllm.models.deepseek_v4.amd.rocm import rocm_inv_rope_einsum
92
+ from vllm.models.deepseek_v4.common.ops.fused_inv_rope_fp8_quant import (
93
+ fused_inv_rope_fp8_quant,
94
+ )
95
95
- z = rocm_inv_rope_einsum(
96
- self.rotary_emb,
96
+ o_fp8, o_scale = fused_inv_rope_fp8_quant(
97
o,
98
positions,
99
- self.rope_head_dim,
100
- self.n_local_groups,
101
- self.o_lora_rank,
102
- self.wo_a,
99
+ self.rotary_emb.cos_sin_cache,
100
+ n_groups=self.n_local_groups,
101
+ heads_per_group=self.n_local_heads // self.n_local_groups,
102
+ nope_dim=self.nope_head_dim,
103
+ rope_dim=self.rope_head_dim,
104
+ tma_aligned_scales=False,
105
)
104
- return self.wo_b(z.flatten(1))
106
+
107
+ # Precomputed contiguous [G, K, N] weight and [G, K/bs, N/bs] scale.
108
+ wo_a_weight = self.wo_a.bmm_weight
109
+ wo_a_scale = self.wo_a.bmm_scale
110
+
111
+ # TODO: optimize fused_inv_rope_fp8_quant for xpu bmm to
112
+ # eliminate o_scale transpose + contiguous
113
+ z = torch.ops.vllm.xpu_fp8_bmm(
114
+ o_fp8.transpose(0, 1),
115
+ wo_a_weight,
116
+ torch.bfloat16,
117
+ o_scale.transpose(0, 1).contiguous(),
118
+ wo_a_scale,
119
+ None,
120
+ )
121
+
122
+ return self.wo_b(z.transpose(0, 1).flatten(1))
123
124
def forward_mqa(
125
self,