Skip to content

Commit 2d80bda

Browse files
zufangzhuwyettzeng
authored andcommitted
[XPU] [Bugfix] process ragged weights in xpu linear backend (vllm-project#52118)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Signed-off-by: Wyett <wyettzeng@gmail.com>
1 parent fea8d31 commit 2d80bda

1 file changed

Lines changed: 29 additions & 1 deletion

File tree

  • vllm/model_executor/kernels/linear/scaled_mm

vllm/model_executor/kernels/linear/scaled_mm/xpu.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# SPDX-License-Identifier: Apache-2.0
22
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
33

4+
import math
45
from collections.abc import Sequence
56

67
import torch
@@ -204,7 +205,34 @@ def process_weights_after_loading(self, layer: torch.nn.Module):
204205
)
205206
scale = getattr(layer, scale_attr)
206207

207-
# Checkpoint scale is [n_blocks, k_blocks] (one value per 128x128 tile).
208+
# Ragged N (N % block_n != 0): oneDNN needs n_blocks to divide N.
209+
# Weight untouched; only repeat scale rows to a finer N-group gn that
210+
# divides both N and block_n (gn = gcd(N, block_n)):
211+
# scale [ceil(N/block_n), K/block_k] --> [N/gn, K/block_k]
212+
# oneDNN only accepts gn that is a multiple of 16, and gcd(N, block_n)
213+
# is a power of two (block_n=128), so gn must be >= 16. No-op when
214+
# N % block_n == 0.
215+
block_n, block_k = self.weight_group_shape
216+
N, K = layer.weight.shape
217+
if N % block_n != 0:
218+
gn = math.gcd(N, block_n)
219+
assert gn % 16 == 0, (
220+
f"XPU block-scaled FP8: N ({N}) yields group width {gn}, but "
221+
f"oneDNN only supports multiples of 16; this weight shape is "
222+
f"unsupported."
223+
)
224+
col_start = torch.arange(N // gn, device=scale.device) * gn
225+
src_idx = torch.div(col_start, block_n, rounding_mode="floor")
226+
scale = scale.index_select(0, src_idx).contiguous()
227+
228+
# Ragged K needs the runtime activation scale expanded too, which we
229+
# don't handle; DeepSeek/GLM keep K block-aligned, so fail loudly.
230+
assert K % block_k == 0, (
231+
f"XPU block-scaled FP8 requires K ({K}) to be a multiple of the "
232+
f"weight block size ({block_k}); ragged-K weights are unsupported."
233+
)
234+
235+
# Checkpoint scale is [n_blocks, k_blocks] (one value per block tile).
208236
# oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout.
209237
# We store the transposed contiguous buffer as a .t() view so that:
210238
# - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape

0 commit comments

Comments
 (0)