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 12, 22:15 UTC
[Performance] Reduce DeepGEMM N dim restriction from 128 to 64 multiplier (#28687)
Signed-off-by: Alexander Matveev <amatveev@redhat.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Alexander Matveev committed
Nov 19, 2025 at 18:47 UTC
3aaa94ac99f4b295ba95f14b4968620b2127044f
3 files changed
+36
-6
.buildkite/test-pipeline.yaml
+20
@@ -550,6 +550,26 @@ steps:
550
commands:
551
- pytest -v -s kernels/mamba
552
553
+- label: Kernels DeepGEMM Test (H100)
554
+ timeout_in_minutes: 45
555
+ gpu: h100
556
+ num_gpus: 1
557
+ optional: true
558
+ source_file_dependencies:
559
+ - tools/install_deepgemm.sh
560
+ - vllm/utils/deep_gemm.py
561
+ - vllm/model_executor/layers/fused_moe
562
+ - vllm/model_executor/layers/quantization
563
+ - tests/kernels/quantization/test_block_fp8.py
564
+ - tests/kernels/moe/test_deepgemm.py
565
+ - tests/kernels/moe/test_batched_deepgemm.py
566
+ - tests/kernels/attention/test_deepgemm_attention.py
567
+ commands:
568
+ - pytest -v -s tests/kernels/quantization/test_block_fp8.py -k deep_gemm
569
+ - pytest -v -s tests/kernels/moe/test_deepgemm.py
570
+ - pytest -v -s tests/kernels/moe/test_batched_deepgemm.py
571
+ - pytest -v -s tests/kernels/attention/test_deepgemm_attention.py
572
+
573
- label: Model Executor Test # 23min
574
timeout_in_minutes: 35
575
torch_nightly: true
tests/kernels/quantization/test_block_fp8.py
+7
-4
@@ -22,6 +22,7 @@ from vllm.utils.deep_gemm import (
22
fp8_gemm_nt,
23
get_col_major_tma_aligned_tensor,
24
per_block_cast_to_fp8,
25
+ should_use_deepgemm_for_fp8_linear,
26
)
27
from vllm.utils.import_utils import has_deep_gemm
28
@@ -157,10 +158,6 @@ def test_w8a8_block_fp8_cutlass_matmul():
158
@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGemm kernels not available.")
159
@torch.inference_mode()
160
def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
160
- # only aligned sizes
161
- if M % 4 != 0 or K % 128 != 0 or N % 64 != 0:
162
- pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
163
-
161
torch.manual_seed(seed)
162
fp8_info = torch.finfo(torch.float8_e4m3fn)
163
fp8_max = fp8_info.max
@@ -168,6 +165,12 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
165
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
166
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
167
168
+ # only aligned sizes are supported by deepgemm
169
+ if not should_use_deepgemm_for_fp8_linear(
170
+ output_dtype=out_dtype, weight=B_fp32, supports_deep_gemm=True
171
+ ):
172
+ pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
173
+
174
A_fp8, As_fp8 = per_token_group_quant_fp8(A_fp32, block_size[1])
175
B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)
176
vllm/utils/deep_gemm.py
+9
-2
@@ -365,11 +365,18 @@ def should_use_deepgemm_for_fp8_linear(
365
):
366
if supports_deep_gemm is None:
367
supports_deep_gemm = is_deep_gemm_supported()
368
+
369
+ # Verify DeepGEMM N/K dims requirements
370
+ # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul
371
+ # test inside kernels/quatization/test_block_fp8.py
372
+ N_MULTIPLE = 64
373
+ K_MULTIPLE = 128
374
+
375
return (
376
supports_deep_gemm
377
and output_dtype == torch.bfloat16
371
- and weight.shape[0] % 128 == 0
372
- and weight.shape[1] % 128 == 0
378
+ and weight.shape[0] % N_MULTIPLE == 0
379
+ and weight.shape[1] % K_MULTIPLE == 0
380
)
381
382