|
1 | 1 | # SPDX-License-Identifier: Apache-2.0 |
2 | 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
3 | 3 |
|
| 4 | +import math |
4 | 5 | from collections.abc import Sequence |
5 | 6 |
|
6 | 7 | import torch |
@@ -204,7 +205,34 @@ def process_weights_after_loading(self, layer: torch.nn.Module): |
204 | 205 | ) |
205 | 206 | scale = getattr(layer, scale_attr) |
206 | 207 |
|
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). |
208 | 236 | # oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout. |
209 | 237 | # We store the transposed contiguous buffer as a .t() view so that: |
210 | 238 | # - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape |
|
0 commit comments