diff --git a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py index e6ec7926c4..92d6458bb3 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/xpu.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/xpu.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import math from collections.abc import Sequence import torch @@ -204,7 +205,34 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): ) scale = getattr(layer, scale_attr) - # Checkpoint scale is [n_blocks, k_blocks] (one value per 128x128 tile). + # Ragged N (N % block_n != 0): oneDNN needs n_blocks to divide N. + # Weight untouched; only repeat scale rows to a finer N-group gn that + # divides both N and block_n (gn = gcd(N, block_n)): + # scale [ceil(N/block_n), K/block_k] --> [N/gn, K/block_k] + # oneDNN only accepts gn that is a multiple of 16, and gcd(N, block_n) + # is a power of two (block_n=128), so gn must be >= 16. No-op when + # N % block_n == 0. + block_n, block_k = self.weight_group_shape + N, K = layer.weight.shape + if N % block_n != 0: + gn = math.gcd(N, block_n) + assert gn % 16 == 0, ( + f"XPU block-scaled FP8: N ({N}) yields group width {gn}, but " + f"oneDNN only supports multiples of 16; this weight shape is " + f"unsupported." + ) + col_start = torch.arange(N // gn, device=scale.device) * gn + src_idx = torch.div(col_start, block_n, rounding_mode="floor") + scale = scale.index_select(0, src_idx).contiguous() + + # Ragged K needs the runtime activation scale expanded too, which we + # don't handle; DeepSeek/GLM keep K block-aligned, so fail loudly. + assert K % block_k == 0, ( + f"XPU block-scaled FP8 requires K ({K}) to be a multiple of the " + f"weight block size ({block_k}); ragged-K weights are unsupported." + ) + + # Checkpoint scale is [n_blocks, k_blocks] (one value per block tile). # oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout. # We store the transposed contiguous buffer as a .t() view so that: # - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape