@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 140 lines 4.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for attention configuration and per-KV-group backend selection."""
4
5 from types import SimpleNamespace
6
7 import pytest
8
9 from vllm.config.attention import AttentionConfig, HiSparseConfig
10 from vllm.model_executor.layers.attention.attention import (
11 _largest_kernel_block_within,
12 )
13 from vllm.v1.attention.backend import AttentionType, MultipleOf
14 from vllm.v1.attention.backends.registry import AttentionBackendEnum
15 from vllm.v1.attention.selector import get_attn_spec_kind
16 from vllm.v1.hisparse.runtime import ResolvedHiSparseConfig
17 from vllm.v1.kv_cache_interface import KVCacheSpecKind
18
19
20 @pytest.mark.parametrize(
21 "supported_sizes,expected",
22 [
23 ([MultipleOf(16)], 1536),
24 ([16, 32], 32),
25 ([2048], 2048),
26 ([MultipleOf(2048)], 2048),
27 ],
28 )
29 def test_largest_kernel_block_within(supported_sizes, expected):
30 class Backend:
31 @staticmethod
32 def get_supported_kernel_block_sizes(kv_cache_spec=None):
33 return supported_sizes
34
35 assert _largest_kernel_block_within(Backend, 1024, 1024 * 1536, 2048) == expected
36
37
38 @pytest.mark.parametrize(
39 "signals,expected",
40 [
41 (dict(use_mla=False, has_sliding_window=False), "full"),
42 (dict(use_mla=True, has_sliding_window=False), "mla"),
43 (dict(use_mla=True, has_sliding_window=True), "sw_mla"),
44 (dict(use_mla=False, has_sliding_window=True), "sw"),
45 ],
46 )
47 def test_get_attn_spec_kind_decoder(signals, expected):
48 kind_by_name = {
49 "full": KVCacheSpecKind.FULL_ATTENTION,
50 "mla": KVCacheSpecKind.MLA_ATTENTION,
51 "sw_mla": KVCacheSpecKind.SLIDING_WINDOW_MLA,
52 "sw": KVCacheSpecKind.SLIDING_WINDOW,
53 }
54 kind = get_attn_spec_kind(attn_type=AttentionType.DECODER, **signals)
55 assert kind is kind_by_name[expected]
56
57
58 @pytest.mark.parametrize(
59 "attn_type,expected",
60 [
61 (AttentionType.ENCODER_ONLY, KVCacheSpecKind.ENCODER_ONLY_ATTENTION),
62 (AttentionType.ENCODER_DECODER, KVCacheSpecKind.CROSS_ATTENTION),
63 ],
64 )
65 def test_get_attn_spec_kind_attn_type(attn_type, expected):
66 kind = get_attn_spec_kind(
67 use_mla=False,
68 has_sliding_window=False,
69 attn_type=attn_type,
70 )
71 assert kind is expected
72
73
74 def test_backend_per_kind_parses_strings():
75 cfg = AttentionConfig(
76 backend_per_kind={
77 "mla_attention": "FLASHINFER_MLA",
78 "sliding_window_mla": "triton_mla", # case-insensitive
79 }
80 )
81 assert cfg.backend_per_kind["mla_attention"] is AttentionBackendEnum.FLASHINFER_MLA
82 assert cfg.backend_per_kind["sliding_window_mla"] is AttentionBackendEnum.TRITON_MLA
83
84
85 def test_backend_per_kind_rejects_unknown_kind():
86 with pytest.raises(ValueError, match="Unknown KV cache group kind"):
87 AttentionConfig(backend_per_kind={"not_a_kind": "TRITON_MLA"})
88
89
90 def test_backend_per_kind_defaults_empty():
91 assert AttentionConfig().backend_per_kind == {}
92
93
94 def test_hisparse_device_buffer_size_boundaries():
95 vllm_config = SimpleNamespace(
96 attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
97 speculative_config=None,
98 )
99 resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
100 assert resolved is not None
101 assert resolved.device_buffer_size == 256
102
103 vllm_config.attention_config.hisparse_config = HiSparseConfig(
104 device_buffer_size=127
105 )
106 with pytest.raises(ValueError, match="expected at least 128"):
107 ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
108
109 vllm_config.attention_config.hisparse_config = HiSparseConfig(
110 device_buffer_size=32768
111 )
112 resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
113 assert resolved is not None
114 assert resolved.device_buffer_size == 32768
115
116 vllm_config.attention_config.hisparse_config = HiSparseConfig(
117 device_buffer_size=32769
118 )
119 with pytest.raises(ValueError, match="int16 slot-index limit"):
120 ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
121
122
123 def test_hisparse_device_buffer_covers_speculative_window():
124 vllm_config = SimpleNamespace(
125 attention_config=AttentionConfig(hisparse_config=HiSparseConfig()),
126 speculative_config=SimpleNamespace(
127 num_speculative_tokens=3,
128 parallel_drafting=False,
129 ),
130 )
131
132 resolved = ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)
133 assert resolved is not None
134 assert resolved.device_buffer_size == 5 * 128
135
136 vllm_config.attention_config.hisparse_config = HiSparseConfig(
137 device_buffer_size=4 * 128 - 1
138 )
139 with pytest.raises(ValueError, match="expected at least 512"):
140 ResolvedHiSparseConfig.from_vllm_config(vllm_config, model_top_k=128)