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 15, 08:30 UTC
Speed up the `kernels/quantization/` tests (#18669)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
May 25, 2025 at 01:02 UTC
63934543a0f05edfc6a5f2afa235b5e026b27b71
3 files changed
+17
-25
tests/kernels/quantization/test_block_fp8.py
+7
-7
@@ -36,16 +36,16 @@ vllm_config.scheduler_config.max_model_len = 8192
36
37
# Test configurations
38
DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
39
-NUM_TOKENS = [7, 83, 2048]
39
+NUM_TOKENS = [7, 2050]
40
D = [512, 4096, 5120, 13824]
41
-GROUP_SIZE = [64, 128, 256, 512]
42
-M = [1, 7, 8, 83, 84, 512, 2048, 4096]
43
-N = [128, 512, 1024, 4096, 7168, 7748, 13824]
44
-K = [256, 4096, 5120, 3884, 13824, 16384]
41
+GROUP_SIZE = [64, 128, 512]
42
+M = [1, 7, 8, 83, 84, 4096]
43
+N = [128, 512, 7168, 7748, 13824]
44
+K = [256, 3884, 4096, 13824, 16384]
45
# Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8
46
# and its hidden size is 7168.
47
-M_moe = [1, 2, 7, 83, 128, 512, 2048]
48
-M_moe_dg = [128, 192, 512, 1335, 2048]
47
+M_moe = [1, 2, 7, 83, 128, 2048]
48
+M_moe_dg = [128, 192, 1335, 2048]
49
N_moe = [128, 256, 1024, 4608] # [13824]
50
K_moe = [256, 512, 7168] # [13824]
51
BLOCK_SIZE = [[128, 128]]
tests/kernels/quantization/test_gguf.py
+2
-2
@@ -35,11 +35,11 @@ def get_gguf_MoE_tensors(
35
return GGUFReader(sample_file).tensors
36
37
38
-DTYPES = [torch.half, torch.bfloat16, torch.float32]
38
+DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
39
# Hidden_size for testing, must match the sample file in HF repo,
40
# we have `hidden_size = 256, 1024` for test in HF repo currently.
41
HIDDEN_SIZES = [256, 1024]
42
-NUM_TOKENS = [7, 83, 128, 2048] # Arbitrary values for testing
42
+NUM_TOKENS = [7, 2050] # Arbitrary values for testing
43
SEEDS = [0]
44
QUANT_TYPES = [
45
# i-matrix
tests/kernels/quantization/test_triton_scaled_mm.py
+8
-16
@@ -13,8 +13,13 @@ from vllm.platforms import current_platform
13
14
device = "cuda"
15
16
+triton_scaled_mm_module = importlib.import_module(
17
+ "vllm.model_executor.layers.quantization.compressed_tensors."
18
+ "triton_scaled_mm")
19
+triton_scaled_mm = triton_scaled_mm_module.triton_scaled_mm
20
17
-def scaled_mm_torch(a: torch.Tensor,
21
+
22
+def torch_scaled_mm(a: torch.Tensor,
23
b: torch.Tensor,
24
scale_a: torch.Tensor,
25
scale_b: torch.Tensor,
@@ -101,21 +106,8 @@ def test_scaled_mm(M, N, K, in_dtype, out_dtype, use_scalar_scale_a,
106
if use_bias:
107
bias = torch.rand((N, ), device=device, dtype=out_dtype)
108
104
- triton_scaled_mm_module = importlib.import_module(
105
- "vllm.model_executor.layers.quantization.compressed_tensors."
106
- "triton_scaled_mm")
107
- triton_scaled_mm = triton_scaled_mm_module.triton_scaled_mm
108
-
109
c_check = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
110
111
- a_cpu = a.cpu()
112
- b_cpu = b.cpu()
113
- scale_a_cpu = scale_a.cpu()
114
- scale_b_cpu = scale_b.cpu()
115
- bias_cpu = None if bias is None else bias.cpu()
116
-
117
- c_actual = scaled_mm_torch(a_cpu, b_cpu, scale_a_cpu, scale_b_cpu,
118
- out_dtype, bias_cpu)
111
+ c_actual = torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
112
120
- c_check_cpu = c_check.cpu()
121
- torch.testing.assert_close(c_check_cpu, c_actual, rtol=1e-1, atol=1e-1)
113
+ torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1)