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
[CI Perf] Prune tests in `tests/kernels/attention/` (#22936)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 14, 2025 at 23:34 UTC
0fe85087a950f3ca94d60293c865c0e6c05e6eff
8 files changed
+39
-38
tests/kernels/attention/test_aiter_flash_attn.py
+3
-3
@@ -9,10 +9,10 @@ import torch
9
import vllm.v1.attention.backends.rocm_aiter_fa # noqa: F401
10
from vllm.platforms import current_platform
11
12
-NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
12
+NUM_HEADS = [(4, 4), (8, 2)]
13
HEAD_SIZES = [128, 256]
14
-BLOCK_SIZES = [16, 32]
15
-DTYPES = [torch.float16, torch.bfloat16]
14
+BLOCK_SIZES = [16]
15
+DTYPES = [torch.bfloat16]
16
QDTYPES = [None]
17
# one value large enough to test overflow in index calculation.
18
# one value small enough to test the schema op check
tests/kernels/attention/test_attention.py
+2
-5
@@ -29,17 +29,14 @@ MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
29
NUM_BLOCKS = 4321 # Arbitrary values for testing
30
PARTITION_SIZE = 512
31
PARTITION_SIZE_ROCM = 256
32
-# flshattF and tritonflashattF supported: {torch.float16, torch.bfloat16}
33
-DTYPES = [
34
- torch.half, torch.bfloat16, torch.float
35
-] if not current_platform.is_rocm() else [torch.half, torch.bfloat16]
32
+DTYPES = [torch.bfloat16]
33
NUM_GEN_SEQS = [7] # Arbitrary values for testing
34
NUM_PREFILL_SEQS = [3] # Arbitrary values for testing
35
NUM_HEADS = [(40, 40), (64, 8)] # Arbitrary values for testing
36
37
# This should be sync with get_supported_head_sizes() in
38
# vllm.attention.ops.paged_attn.PagedAttention
42
-HEAD_SIZES = [32, 64, 80, 96, 112, 120, 128, 192, 256]
39
+HEAD_SIZES = [32, 80, 128, 256]
40
41
BLOCK_SIZES = [16, 32]
42
USE_ALIBI = [False, True]
tests/kernels/attention/test_cache.py
+2
-2
@@ -11,11 +11,11 @@ from vllm import _custom_ops as ops
11
from vllm.platforms import current_platform
12
13
COPYING_DIRECTION = [('cuda', 'cpu'), ('cuda', 'cuda'), ('cpu', 'cuda')]
14
-DTYPES = [torch.half, torch.bfloat16, torch.float]
14
+DTYPES = [torch.bfloat16, torch.float]
15
NUM_TOKENS = [42] # Arbitrary values for testing
16
NUM_LAYERS = [1] # Arbitrary values for testing
17
NUM_HEADS = [8] # Arbitrary values for testing
18
-HEAD_SIZES = [64, 80, 120, 256]
18
+HEAD_SIZES = [64, 80, 256]
19
BLOCK_SIZES = [8, 16, 32]
20
CACHE_LAYOUTS = ["NHD", "HND"]
21
tests/kernels/attention/test_flash_attn.py
+9
-7
@@ -12,14 +12,16 @@ from vllm.vllm_flash_attn import (fa_version_unsupported_reason,
12
flash_attn_with_kvcache,
13
is_fa_version_supported)
14
15
-NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
15
+NUM_HEADS = [(4, 4), (8, 2)]
16
HEAD_SIZES = [128, 256]
17
-BLOCK_SIZES = [16, 32]
18
-DTYPES = [torch.float16, torch.bfloat16]
17
+BLOCK_SIZES = [16]
18
+DTYPES = [torch.bfloat16]
19
QDTYPES = [None, torch.float8_e4m3fn]
20
# one value large enough to test overflow in index calculation.
21
# one value small enough to test the schema op check
22
NUM_BLOCKS = [32768, 2048]
23
+SOFT_CAPS = [None, 50.0]
24
+SLIDING_WINDOWS = [None, 256]
25
26
27
def ref_paged_attn(
@@ -83,9 +85,9 @@ def ref_paged_attn(
85
@pytest.mark.parametrize("head_size", HEAD_SIZES)
86
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
87
@pytest.mark.parametrize("dtype", DTYPES)
86
-@pytest.mark.parametrize("soft_cap", [None, 10.0, 50.0])
88
+@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
89
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
88
-@pytest.mark.parametrize("sliding_window", [None, 256])
90
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
91
@pytest.mark.parametrize("fa_version", [2, 3])
92
@pytest.mark.parametrize("q_dtype", QDTYPES)
93
@torch.inference_mode()
@@ -198,9 +200,9 @@ def test_flash_attn_with_paged_kv(
200
@pytest.mark.parametrize("num_heads", NUM_HEADS)
201
@pytest.mark.parametrize("head_size", HEAD_SIZES)
202
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
201
-@pytest.mark.parametrize("sliding_window", [None, 256])
203
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
204
@pytest.mark.parametrize("dtype", DTYPES)
203
-@pytest.mark.parametrize("soft_cap", [None, 10.0, 50.0])
205
+@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
206
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
207
@pytest.mark.parametrize("fa_version", [2, 3])
208
@pytest.mark.parametrize("q_dtype", QDTYPES)
tests/kernels/attention/test_flashinfer.py
+13
-11
@@ -9,11 +9,13 @@ import torch
9
10
from vllm.platforms import current_platform
11
12
-NUM_HEADS = [(16, 16), (32, 8), (64, 8), (6, 1)]
12
+NUM_HEADS = [(32, 8), (6, 1)]
13
HEAD_SIZES = [128, 256]
14
BLOCK_SIZES = [16, 32]
15
-DTYPES = [torch.float16, torch.bfloat16]
15
+DTYPES = [torch.bfloat16]
16
NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
17
+SOFT_CAPS = [None, 30.0]
18
+SLIDING_WINDOWS = [None, 64]
19
20
21
def ref_paged_attn(
@@ -76,8 +78,8 @@ def ref_paged_attn(
78
@pytest.mark.parametrize("head_size", HEAD_SIZES)
79
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
80
@pytest.mark.parametrize("dtype", DTYPES)
79
-@pytest.mark.parametrize("soft_cap", [None, 30.0, 50.0])
80
-@pytest.mark.parametrize("sliding_window", [None, 64])
81
+@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
82
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
83
@torch.inference_mode
84
def test_flashinfer_decode_with_paged_kv(
85
kv_lens: list[int],
@@ -173,8 +175,8 @@ def test_flashinfer_decode_with_paged_kv(
175
@pytest.mark.parametrize("head_size", HEAD_SIZES)
176
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
177
@pytest.mark.parametrize("dtype", DTYPES)
176
-@pytest.mark.parametrize("soft_cap", [None, 30.0, 50.0])
177
-@pytest.mark.parametrize("sliding_window", [None, 64])
178
+@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
179
+@pytest.mark.parametrize("sliding_window", SLIDING_WINDOWS)
180
@torch.inference_mode
181
def test_flashinfer_prefill_with_paged_kv(
182
seq_lens: list[tuple[int, int]],
@@ -278,11 +280,11 @@ def test_flashinfer_prefill_with_paged_kv(
280
281
282
@pytest.mark.parametrize("seq_lens", [[(1, 132), (5, 18)]])
281
-@pytest.mark.parametrize("num_heads", [(32, 8), (6, 1)])
283
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
284
@pytest.mark.parametrize("head_size", HEAD_SIZES)
285
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
286
@pytest.mark.parametrize("dtype", DTYPES)
285
-@pytest.mark.parametrize("soft_cap", [None, 30.0, 50.0])
287
+@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
288
def test_flashinfer_prefill_with_paged_fp8_kv(
289
seq_lens: list[tuple[int, int]], num_heads: tuple[int, int],
290
head_size: int, dtype: torch.dtype, block_size: int,
@@ -385,11 +387,12 @@ def test_flashinfer_prefill_with_paged_fp8_kv(
387
388
389
@pytest.mark.parametrize("kv_lens", [[1328, 18, 463], [1, 54, 293, 70]])
388
-@pytest.mark.parametrize("num_heads", [(32, 8), (64, 8), (6, 1)])
390
+@pytest.mark.parametrize("num_heads", NUM_HEADS)
391
@pytest.mark.parametrize("head_size", HEAD_SIZES)
392
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
393
@pytest.mark.parametrize("dtype", DTYPES)
392
-@pytest.mark.parametrize("soft_cap", [None, 30.0, 50.0])
394
+@pytest.mark.parametrize("soft_cap", SOFT_CAPS)
395
+@pytest.mark.skip(reason="TODO: fix the accuracy issue")
396
@torch.inference_mode
397
def test_flashinfer_decode_with_paged_fp8_kv(
398
kv_lens: list[int],
@@ -399,7 +402,6 @@ def test_flashinfer_decode_with_paged_fp8_kv(
402
block_size: int,
403
soft_cap: Optional[float],
404
) -> None:
402
- pytest.skip("TODO: fix the accuracy issue")
405
# test doesn't work for num_heads = (16,16)
406
torch.set_default_device("cuda")
407
current_platform.seed_everything(0)
tests/kernels/attention/test_flashinfer_trtllm_attention.py
+3
-3
@@ -20,11 +20,11 @@ FLOAT32_BYTES = torch.finfo(torch.float).bits // 8
20
MAX_Q_LEN = 1024
21
MAX_KV_LEN = 4096
22
BATCH_SIZES = [4, 12]
23
-NUM_HEADS = [(64, 8), (16, 16), (40, 8), (32, 8)]
23
+NUM_HEADS = [(16, 16), (40, 8)]
24
HEAD_SIZES = [128]
25
-BLOCK_SIZES = [16, 32]
25
+BLOCK_SIZES = [16]
26
KV_LAYOUTS = ["HND"]
27
-DTYPES = [torch.float16, torch.bfloat16]
27
+DTYPES = [torch.bfloat16]
28
KV_CACHE_DTYPES = [None, current_platform.fp8_dtype()]
29
NUM_BLOCKS = 32768 # Large enough to test overflow in index calculation.
30
SOFT_CAPS = [None, 50.0]
tests/kernels/attention/test_prefix_prefill.py
+3
-3
@@ -19,13 +19,13 @@ from vllm.platforms import current_platform
19
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
20
21
NUM_HEADS = [64]
22
-NUM_QUERIES_PER_KV = [1, 8, 64]
23
-HEAD_SIZES = [128, 96, 24]
22
+NUM_QUERIES_PER_KV = [1, 64]
23
+HEAD_SIZES = [24, 128]
24
DTYPES = [torch.float16]
25
CUDA_DEVICES = [
26
f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)
27
]
28
-SLIDING_WINDOW = [0, 16, 64, 128, 256, 512, 2048]
28
+SLIDING_WINDOW = [0, 16, 2048]
29
KV_CACHE_DTYPES = ["auto", "fp8", "fp8_e5m2"]
30
31
OPS = [chunked_prefill_paged_decode, context_attention_fwd]
tests/kernels/attention/test_triton_unified_attention.py
+4
-4
@@ -9,11 +9,11 @@ import torch
9
from vllm.attention.ops.triton_unified_attention import unified_attention
10
from vllm.platforms import current_platform
11
12
-NUM_HEADS = [(4, 4), (8, 2), (16, 2)]
12
+NUM_HEADS = [(4, 4), (8, 2)]
13
HEAD_SIZES = [128, 256]
14
-BLOCK_SIZES = [16, 32]
14
+BLOCK_SIZES = [16]
15
16
-DTYPES = [torch.float16, torch.bfloat16]
16
+DTYPES = [torch.bfloat16]
17
QDTYPES = [None, torch.float8_e4m3fn] if not current_platform.is_rocm() else [
18
None, torch.float8_e4m3fnuz
19
]
@@ -85,7 +85,7 @@ def ref_paged_attn(
85
@pytest.mark.parametrize("block_size", BLOCK_SIZES)
86
@pytest.mark.parametrize("sliding_window", [None, 256])
87
@pytest.mark.parametrize("dtype", DTYPES)
88
-@pytest.mark.parametrize("soft_cap", [None, 10.0, 50.0])
88
+@pytest.mark.parametrize("soft_cap", [None, 50.0])
89
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
90
@pytest.mark.parametrize("q_dtype", QDTYPES)
91
@torch.inference_mode()