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 2, 07:15 UTC
[Perf] Optimize cutlass fp8 scaled mm bypassing padding, 20% kernel performance improvement (#43706)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
Jun 1, 2026 at 09:05 UTC
985c97a6a884f1a29e7584c05e11214f7fb96dbf
1 file changed
+57
-9
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+57
-9
@@ -312,7 +312,7 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
312
) -> torch.Tensor:
313
out_dtype = self.config.out_dtype
314
if self.is_hopper:
315
- return torch.ops.vllm.padded_cutlass(
315
+ return torch.ops.vllm.dynamic_padded_cutlass(
316
A,
317
B,
318
As,
@@ -320,14 +320,14 @@ class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
320
list(self.weight_group_shape),
321
out_dtype,
322
)
323
- else:
324
- return ops.cutlass_scaled_mm(
325
- A,
326
- B.T,
327
- out_dtype=out_dtype,
328
- scale_a=As,
329
- scale_b=Bs.T,
330
- )
323
+
324
+ return ops.cutlass_scaled_mm(
325
+ A,
326
+ B.T,
327
+ out_dtype=out_dtype,
328
+ scale_a=As,
329
+ scale_b=Bs.T,
330
+ )
331
332
333
def cutlass_scaled_mm(
@@ -397,8 +397,56 @@ def _padded_cutlass_fake(
397
)
398
399
400
+def _dynamic_padded_cutlass(
401
+ qx: torch.Tensor,
402
+ weight: torch.Tensor,
403
+ x_scale: torch.Tensor,
404
+ weight_scale: torch.Tensor,
405
+ block_size: list[int],
406
+ output_dtype: torch.dtype,
407
+) -> torch.Tensor:
408
+ def run_padded(
409
+ qx: torch.Tensor,
410
+ weight: torch.Tensor,
411
+ x_scale: torch.Tensor,
412
+ weight_scale: torch.Tensor,
413
+ ) -> torch.Tensor:
414
+ return _padded_cutlass(
415
+ qx, weight, x_scale, weight_scale, block_size, output_dtype
416
+ )
417
+
418
+ def run_direct(
419
+ qx: torch.Tensor,
420
+ weight: torch.Tensor,
421
+ x_scale: torch.Tensor,
422
+ weight_scale: torch.Tensor,
423
+ ) -> torch.Tensor:
424
+ return cutlass_scaled_mm(
425
+ qx, weight, x_scale, weight_scale, block_size, output_dtype
426
+ )
427
+
428
+ if torch.compiler.is_compiling():
429
+ return torch.cond(
430
+ qx.shape[0] % 4 != 0,
431
+ run_padded,
432
+ run_direct,
433
+ (qx, weight, x_scale, weight_scale),
434
+ )
435
+
436
+ if qx.shape[0] % 4 != 0:
437
+ return run_padded(qx, weight, x_scale, weight_scale)
438
+
439
+ return run_direct(qx, weight, x_scale, weight_scale)
440
+
441
+
442
direct_register_custom_op(
443
"padded_cutlass",
444
_padded_cutlass,
445
fake_impl=_padded_cutlass_fake,
446
)
447
+
448
+direct_register_custom_op(
449
+ "dynamic_padded_cutlass",
450
+ _dynamic_padded_cutlass,
451
+ fake_impl=_padded_cutlass_fake,
452
+)