A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 02:30 UTC
[CUDA][CUTLASS] Enable cutlass scaled mm for non-compatible sizes (#41868)
Signed-off-by: John Calderon <jcalderon@nvidia.com>
John Calderon committed
May 8, 2026 at 18:58 UTC
8a2fc80b848af8f1e902d4b2b14afe378d9d4d89
2 files changed
+111
-1
tests/kernels/quantization/test_cutlass_scaled_mm.py
+68
@@ -40,6 +40,18 @@ MNK_FACTORS = [
40
(512, 24576, 128),
41
]
42
43
+# Shapes with N or K not divisible by 16. These exercise the padding path
44
+# inside CutlassFP8ScaledMMLinearKernel.apply_scaled_mm (e.g. Qwen2.5-VL
45
+# vision MLP dims).
46
+UNALIGNED_MNK_FACTORS = [
47
+ (32, 3420, 1280),
48
+ (32, 1280, 6840),
49
+ (1, 3420, 1280),
50
+ (64, 6840, 1280),
51
+ (16, 100, 200),
52
+ (33, 255, 513),
53
+]
54
+
55
CUDA_DEVICES = [
56
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
57
]
@@ -152,6 +164,62 @@ def test_cutlass_fp8_gemm(
164
cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape, use_bias)
165
166
167
+@pytest.mark.parametrize("m,n,k", UNALIGNED_MNK_FACTORS)
168
+@pytest.mark.parametrize(
169
+ "a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
170
+)
171
+@pytest.mark.parametrize(
172
+ "b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
173
+)
174
+@pytest.mark.parametrize("use_bias", [True, False])
175
+@pytest.mark.skipif(
176
+ not current_platform.has_device_capability(89),
177
+ reason="FP8 is not supported on this GPU type.",
178
+)
179
+def test_cutlass_fp8_gemm_padded(
180
+ m: int, n: int, k: int, a_scale_group_shape, b_scale_group_shape, use_bias: bool
181
+):
182
+ """Test CUTLASS FP8 GEMM with padding for non-16-aligned N/K dims.
183
+
184
+ Exercises CutlassFP8ScaledMMLinearKernel.apply_scaled_mm which pads
185
+ inputs to satisfy CUTLASS alignment requirements — the path taken by
186
+ models like Qwen2.5-VL whose vision MLP has non-16-aligned dims.
187
+ """
188
+ from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
189
+ CutlassFP8ScaledMMLinearKernel,
190
+ )
191
+
192
+ a = to_fp8(torch.randn((m, k), device="cuda"))
193
+ b = to_fp8(torch.randn((n, k), device="cuda").t())
194
+
195
+ a_scales_shape = scale_shape(a.shape, a_scale_group_shape)
196
+ b_scales_shape = scale_shape(b.shape, b_scale_group_shape)
197
+
198
+ scale_a = torch.randn(a_scales_shape, device="cuda", dtype=torch.float32)
199
+ scale_b = torch.randn(b_scales_shape, device="cuda", dtype=torch.float32)
200
+
201
+ scale_a = scale_a.t().contiguous().t()
202
+ scale_b = scale_b.t().contiguous().t()
203
+
204
+ out_dtype = torch.bfloat16
205
+ bias = torch.rand((n,), device="cuda", dtype=out_dtype) * 10 if use_bias else None
206
+
207
+ baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
208
+
209
+ kernel = object.__new__(CutlassFP8ScaledMMLinearKernel)
210
+ out = kernel.apply_scaled_mm(
211
+ A=a,
212
+ B=b,
213
+ out_dtype=out_dtype,
214
+ As=scale_a,
215
+ Bs=scale_b,
216
+ bias=bias,
217
+ output_shape=[m, n],
218
+ )
219
+
220
+ torch.testing.assert_close(out, baseline, rtol=5e-1, atol=1.5e-1)
221
+
222
+
223
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
224
@pytest.mark.parametrize(
225
"a_scale_group_shape,b_scale_group_shape", [((1, 128), (128, 128))]
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+43
-1
@@ -162,6 +162,20 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
162
def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
163
return True, None
164
165
+ @staticmethod
166
+ def _pad_to_alignment(
167
+ x: torch.Tensor, dim: int, alignment: int, value: float = 0.0
168
+ ) -> torch.Tensor:
169
+ """Pad tensor ``x`` along ``dim`` to the next multiple of
170
+ ``alignment``."""
171
+ remainder = x.shape[dim] % alignment
172
+ if remainder == 0:
173
+ return x
174
+ pad_size = alignment - remainder
175
+ pad_spec = [0] * (2 * x.dim())
176
+ pad_spec[-(2 * dim + 1)] = pad_size
177
+ return torch.nn.functional.pad(x, pad_spec, value=value)
178
+
179
def apply_scaled_mm(
180
self,
181
*,
@@ -173,10 +187,38 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
187
bias: torch.Tensor | None,
188
output_shape: list,
189
) -> torch.Tensor:
176
- # Fused GEMM_DQ
190
+ # Per-tensor/Per-channel padding to use Cutlass instead of Triton.
191
+ K, N = B.shape
192
+ pad_k = (16 - K % 16) % 16
193
+ pad_n = (16 - N % 16) % 16
194
+
195
+ if pad_k > 0 or pad_n > 0:
196
+ # B is column-major [K, N]. Transpose to row-major [N, K],
197
+ # pad both dims in one call, then transpose back so the
198
+ # result keeps column-major layout with stride (1, K_padded).
199
+ B = torch.nn.functional.pad(B.t().contiguous(), (0, pad_k, 0, pad_n)).t()
200
+
201
+ if pad_k > 0:
202
+ A = self._pad_to_alignment(A, dim=1, alignment=16)
203
+ if pad_n > 0:
204
+ if bias is not None:
205
+ bias = self._pad_to_alignment(bias, dim=0, alignment=16)
206
+ # Bs is per-tensor (numel==1) or per-channel (numel==N)
207
+ # in this kernel class — never 2D block-wise.
208
+ if Bs.numel() > 1:
209
+ Bs = self._pad_to_alignment(
210
+ Bs.view(-1), dim=0, alignment=16, value=1.0
211
+ )
212
+ if Bs.dim() == 1 and B.shape[1] > 1:
213
+ Bs = Bs.view(-1, 1)
214
+
215
output = ops.cutlass_scaled_mm(
216
A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
217
)
218
+
219
+ if pad_n > 0:
220
+ output = output[..., :N].contiguous()
221
+
222
return output.view(*output_shape)
223
224