[XPU] Fix FP8 block scale layout for MLA compatibility (#50349)
Signed-off-by: Ma Jian <jian1.ma@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Ma Jian committed
Jul 31, 2026 at 11:05 UTC
ef0d084b4b11141f009bf2210d5eac9b00b90ed0
1 file changed
+34
-25
vllm/model_executor/kernels/linear/scaled_mm/xpu.py
+34
-25
@@ -203,30 +203,38 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
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)
206
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]
207
+ # Checkpoint scale is [n_blocks, k_blocks] (one value per 128x128 tile).
208
+ # oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout.
209
+ # We store the transposed contiguous buffer as a .t() view so that:
210
+ # - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape
211
+ # - apply_block_scaled_mm recovers the contiguous buffer via .t()
212
+ scale_kn = scale.data.t().contiguous() # [k_blocks, n_blocks]
213
+ replace_parameter(layer, scale_attr, scale_kn.t()) # view: [n_blocks, k_blocks]
214
+
215
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]
216
+ self._prepare_bmm_params(layer, scale_kn)
217
+
218
+ def _prepare_bmm_params(
219
+ self, layer: torch.nn.Module, scale_kn: torch.Tensor
220
+ ) -> None:
221
+ """Precompute batched weight and scale for grouped fp8_bmm (e.g. wo_a).
222
+
223
+ Splits scale [k_blocks, n_blocks] into [G, k_blocks, n_blocks_per_group]
224
+ and weight [N_total, K] into [G, K, N_per_group] for batch GEMM.
225
+ """
226
+ batch = layer.bmm_batch_size
227
+ k_blocks, n_blocks = scale_kn.shape
228
+ layer.bmm_scale = (
229
+ scale_kn.reshape(k_blocks, batch, n_blocks // batch)
230
+ .permute(1, 0, 2)
231
+ .contiguous()
232
+ )
233
+ w = layer.weight
234
+ N_total, K = w.shape
235
+ layer.bmm_weight = w.reshape(batch, N_total // batch, K).permute(
236
+ 0, 2, 1
237
+ ) # [G, K, N_per_group]
238
239
def apply_block_scaled_mm(
240
self,
@@ -235,13 +243,14 @@ class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
243
As: torch.Tensor,
244
Bs: torch.Tensor,
245
) -> torch.Tensor:
238
- # Weight is [N, K]. Use .t() to create a [K, N] view without copying.
239
- # Bs is already [K/128, N/128] from process_weights_after_loading.
246
+ # B is [N, K]; .t() gives [K, N] view (no copy).
247
+ # Bs is stored as [n_blocks, k_blocks] view; .t() recovers the
248
+ # contiguous [k_blocks, n_blocks] buffer that oneDNN expects.
249
return torch.ops._xpu_C.fp8_gemm(
250
A,
251
B.t(),
252
self.config.out_dtype,
253
As,
245
- Bs,
254
+ Bs.t(),
255
torch.Tensor(),
256
)