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 14, 05:30 UTC
Pruning kernel Core Tests (#26727)
Signed-off-by: Fardin Hoque <kfhfar@amazon.com>
Fardin Hoque committed
Oct 13, 2025 at 16:08 UTC
fa96fb9c707919b5a2a2d0b0e9feb610a3a1f75a
4 files changed
+5
-18
tests/kernels/core/test_fused_quant_layernorm.py
-1
@@ -15,7 +15,6 @@ VEC_HIDDEN_SIZES = range(1024, 1030)
15
# Avoid combinatorial explosion with full Cartesian product
16
NUM_TOKENS_HIDDEN_SIZES = [
17
*[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
18
- *[(83, i) for i in [1, 1033, 2048, 5120]],
18
*[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
19
*[(4096, i) for i in [1, 64, 5137]],
20
]
tests/kernels/core/test_layernorm.py
+2
-14
@@ -11,19 +11,7 @@ from vllm.platforms import current_platform
11
12
DTYPES = [torch.half, torch.bfloat16, torch.float]
13
NUM_TOKENS = [7, 83, 4096] # Arbitrary values for testing
14
-HIDDEN_SIZES = [
15
- 8,
16
- 768,
17
- 769,
18
- 770,
19
- 771,
20
- 5120,
21
- 5124,
22
- 5125,
23
- 5126,
24
- 8192,
25
- 8199,
26
-] # Arbitrary values for testing
14
+HIDDEN_SIZES = [8, 768, 769, 5120, 5125, 8192] # Arbitrary values for testing
15
ADD_RESIDUAL = [False, True]
16
SEEDS = [0]
17
CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
@@ -118,7 +106,7 @@ def test_poly_norm(
106
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
107
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
108
@pytest.mark.parametrize("dtype", DTYPES)
121
-@pytest.mark.parametrize("quant_scale", [1.0, 0.01, 10.0])
109
+@pytest.mark.parametrize("quant_scale", [0.01, 1.0, 10.0])
110
@pytest.mark.parametrize("seed", SEEDS)
111
@pytest.mark.parametrize("device", CUDA_DEVICES)
112
@pytest.mark.parametrize("strided_input", [False, True])
tests/kernels/core/test_permute_cols.py
+1
-1
@@ -9,7 +9,7 @@ from vllm._custom_ops import permute_cols
9
10
11
@pytest.mark.parametrize("shape", [(1, 512), (544, 4096), (67, 8192)])
12
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
12
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
13
def test_permute_cols(shape, dtype):
14
x = torch.randn(shape, dtype=dtype).cuda()
15
perm = torch.randperm(x.shape[1]).to(torch.int).cuda()
tests/kernels/core/test_pos_encoding.py
+2
-2
@@ -12,8 +12,8 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
12
from vllm.platforms import current_platform
13
14
IS_NEOX_STYLE = [True, False]
15
-DTYPES = [torch.half, torch.bfloat16, torch.float]
16
-HEAD_SIZES = [64, 80, 112, 120, 256]
15
+DTYPES = [torch.bfloat16, torch.float]
16
+HEAD_SIZES = [64, 80, 120, 256]
17
ROTARY_DIMS = [None, 32] # None means rotary dim == head size
18
NUM_HEADS = [17] # Arbitrary values for testing
19
BATCH_SIZES = [5] # Arbitrary values for testing