[Linear] [Kernel] add block-wise scaled_mm (#49932)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
zofia committed
Aug 6, 2026 at 14:40 UTC
8060260e1d1718d831419b5e9fe7a1757d6cc27b
3 files changed
+171
tests/kernels/quantization/test_block_fp8.py
+60
@@ -3,6 +3,7 @@
3
4
# Adapted from https://github.com/sgl-project/sglang/pull/2575
5
import itertools
6
+import types
7
8
import pytest
9
import torch
@@ -200,6 +201,65 @@ def test_w8a8_block_fp8_cutlass_matmul():
201
assert rel_diff < 0.001
202
203
204
+@pytest.mark.skipif(
205
+ not (current_platform.is_cuda() and current_platform.has_device_capability(90)),
206
+ reason="torch._scaled_mm DeepSeek-style block scaling only supports SM90.",
207
+)
208
+def test_w8a8_block_fp8_torch_scaled_mm_matmul():
209
+ # BlockWiseTorchFP8ScaledMMLinearKernel: 1x128 activation + 128x128 weight
210
+ # block scaling routed through torch._scaled_mm. M=83 is not a multiple of
211
+ # 4 so this also exercises the M-padding path.
212
+ from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
213
+ BlockWiseTorchFP8ScaledMMLinearKernel,
214
+ )
215
+
216
+ M = 83
217
+ N = 576
218
+ K = 7168
219
+ block_size = [128, 128]
220
+ out_dtype = torch.bfloat16
221
+ seed = 0
222
+
223
+ torch.manual_seed(seed)
224
+ factor_for_scale = 1e-2
225
+ fp8_info = torch.finfo(torch.float8_e4m3fn)
226
+ fp8_max, fp8_min = fp8_info.max, fp8_info.min
227
+
228
+ A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
229
+ B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
230
+ B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
231
+
232
+ block_n, block_k = block_size[0], block_size[1]
233
+ n_tiles = (N + block_n - 1) // block_n
234
+ k_tiles = (K + block_k - 1) // block_k
235
+ Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
236
+
237
+ # Reference uses row-major activation scales.
238
+ A_fp8, As = per_token_group_quant_fp8(
239
+ A_fp32, block_size[1], column_major_scales=False
240
+ )
241
+ ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
242
+
243
+ # The kernel expects column-major activation scales on CUDA.
244
+ A_fp8_cuda, As_cuda = per_token_group_quant_fp8(
245
+ A_fp32, block_size[1], column_major_scales=True
246
+ )
247
+
248
+ stub = BlockWiseTorchFP8ScaledMMLinearKernel.__new__(
249
+ BlockWiseTorchFP8ScaledMMLinearKernel
250
+ )
251
+ stub.config = types.SimpleNamespace(out_dtype=out_dtype)
252
+ out = stub.apply_block_scaled_mm(
253
+ A_fp8_cuda.cuda(), B_fp8.cuda(), As_cuda.cuda(), Bs.cuda()
254
+ )
255
+
256
+ ref_out = ref_out.cuda()
257
+ rel_diff = torch.mean(
258
+ torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
259
+ ) / torch.mean(torch.abs(ref_out.to(torch.float32)))
260
+ assert rel_diff < 0.001
261
+
262
+
263
@pytest.mark.skipif(
264
current_platform.is_fp8_fnuz(),
265
reason="This platform supports e4m3fnuz, not e4m3fn.",
vllm/model_executor/kernels/linear/__init__.py
+3
@@ -187,6 +187,7 @@ from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
187
MarlinFP8ScaledMMLinearKernel,
188
)
189
from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
190
+ BlockWiseTorchFP8ScaledMMLinearKernel,
191
ChannelWiseTorchFP8ScaledMMLinearKernel,
192
PerTensorTorchFP8ScaledMMLinearKernel,
193
RowWiseTorchFP8ScaledMMLinearKernel,
@@ -281,6 +282,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
282
PerTensorTorchFP8ScaledMMLinearKernel,
283
ChannelWiseTorchFP8ScaledMMLinearKernel,
284
RowWiseTorchFP8ScaledMMLinearKernel,
285
+ BlockWiseTorchFP8ScaledMMLinearKernel,
286
},
287
"aiter": {
288
AiterInt8ScaledMMLinearKernel,
@@ -380,6 +382,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
382
MarlinFP8ScaledMMLinearKernel,
383
TritonFp8BlockScaledMMKernel,
384
HummingFP8ScaledMMLinearKernel,
385
+ BlockWiseTorchFP8ScaledMMLinearKernel,
386
],
387
PlatformEnum.ROCM: [
388
AiterFp8BlockScaledMMKernel,
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py
+108
@@ -6,8 +6,11 @@ import math
6
import torch
7
8
from vllm.config import CompilationMode, get_current_vllm_config
9
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
10
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
11
from vllm.platforms import current_platform
12
13
+from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
14
from .ScaledMMLinearKernel import (
15
FP8ScaledMMLinearKernel,
16
FP8ScaledMMLinearLayerConfig,
@@ -244,3 +247,108 @@ class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
247
if bias is not None:
248
output = output + bias
249
return output.to(out_dtype).view(*output_shape)
250
+
251
+
252
+class BlockWiseTorchFP8ScaledMMLinearKernel(Fp8BlockScaledMMLinearKernel):
253
+ """FP8 block-scaled linear kernel using ``torch._scaled_mm``.
254
+
255
+ Implements the block-scaled path of ``torch._scaled_mm``, which
256
+ dispatches on the shapes of the scale tensors. For ``A = [M, K]`` and
257
+ ``B = [K, N]`` (both fp8) the op's block path requires, with float32
258
+ scales:
259
+ * 1x128 activation: ``scale_a = [M, ceil(K / 128)]``
260
+ * 128x128 weight: ``scale_b = [ceil(K / 128), ceil(N / 128)]``
261
+ """
262
+
263
+ def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
264
+ super().__init__(config)
265
+ act_scale_descriptor = config.activation_quant_key.scale
266
+ self.quant_fp8 = QuantFP8(
267
+ static=act_scale_descriptor.static,
268
+ group_shape=act_scale_descriptor.group_shape,
269
+ num_token_padding=self.get_output_padding(),
270
+ use_ue8m0=False,
271
+ column_major_scales=current_platform.is_cuda_alike(),
272
+ )
273
+
274
+ @classmethod
275
+ def is_supported(
276
+ cls, compute_capability: int | None = None
277
+ ) -> tuple[bool, str | None]:
278
+ if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
279
+ return False, "requires CUDA, ROCm or XPU."
280
+ # torch._scaled_mm DeepSeek-style (1x128, 128x128) block scaling is
281
+ # only implemented for SM90 (Hopper) on NVIDIA CUDA. See PyTorch's
282
+ # _check_deepseek_support():
283
+ # https://github.com/pytorch/pytorch/blob/33812ece06e2b0d597f73fbe41de03a83f9109f9/aten/src/ATen/native/cuda/ScaledBlas.cpp#L804-L820 # noqa: E501
284
+ if current_platform.is_cuda():
285
+ is_sm90 = (
286
+ compute_capability == 90
287
+ if compute_capability is not None
288
+ else current_platform.is_device_capability(90)
289
+ )
290
+ if not is_sm90:
291
+ return (
292
+ False,
293
+ "DeepSeek-style (1x128, 128x128) block scaling on CUDA "
294
+ "requires compute capability 90 (Hopper).",
295
+ )
296
+ return True, None
297
+
298
+ @classmethod
299
+ def can_implement(
300
+ cls, config: FP8ScaledMMLinearLayerConfig
301
+ ) -> tuple[bool, str | None]:
302
+ can_implement_base, reason = super().can_implement(config)
303
+ if not can_implement_base:
304
+ return can_implement_base, reason
305
+
306
+ act_group_shape = config.activation_quant_key.scale.group_shape
307
+ if act_group_shape != GroupShape(1, 128):
308
+ return (
309
+ False,
310
+ "requires 1x128 activation quantization.",
311
+ )
312
+ weight_group_shape = config.weight_quant_key.scale.group_shape
313
+ if weight_group_shape != GroupShape(128, 128):
314
+ return (
315
+ False,
316
+ "requires 128x128 block weight quantization.",
317
+ )
318
+ return True, None
319
+
320
+ def apply_block_scaled_mm(
321
+ self,
322
+ A: torch.Tensor,
323
+ B: torch.Tensor,
324
+ As: torch.Tensor,
325
+ Bs: torch.Tensor,
326
+ ) -> torch.Tensor:
327
+ # torch._scaled_mm implemented by cuBLASLt requires the M dimension to be a
328
+ # multiple of 4; pad A (and its column-major scale) up and slice back.
329
+ # https://docs.nvidia.com/cuda/cublas/index.html?highlight=128%2520element%25201D%2520and%2520128x128%25202D#scaling-factors-layouts # noqa: E501
330
+ M = A.shape[0]
331
+ pad = (-M) % 4 if current_platform.is_cuda() else 0
332
+ if pad:
333
+ Ap = A.new_empty((M + pad, A.shape[1]))
334
+ Ap[:M] = A
335
+ A = Ap
336
+ # As is column-major [M, ceil(K/128)] (stride (1, M)); build a padded
337
+ # column-major buffer and copy only the valid rows. Padded rows are left
338
+ # uninitialized — they don't affect the sliced-back output.
339
+ Asp = As.new_empty((As.shape[1], M + pad)).t()
340
+ Asp[:M] = As
341
+ As = Asp
342
+
343
+ output = torch._scaled_mm(
344
+ A,
345
+ B.t(),
346
+ scale_a=As,
347
+ scale_b=Bs.t(),
348
+ out_dtype=self.config.out_dtype,
349
+ )
350
+ if type(output) is tuple and len(output) == 2:
351
+ output = output[0]
352
+ if pad:
353
+ output = output[:M, :]
354
+ return output