@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 188 lines 6 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from tests.v1.attention.utils import (
7 BatchSpec,
8 create_common_attn_metadata,
9 create_vllm_config,
10 )
11 from vllm.config import CUDAGraphMode, SpeculativeConfig
12 from vllm.v1.attention.backend import AttentionCGSupport
13 from vllm.v1.attention.backends.linear_attn import (
14 BailingLinearAttentionMetadataBuilder,
15 LinearAttentionMetadataBuilder,
16 )
17 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
18 from vllm.v1.kv_cache_interface import MambaSpec
19
20 BLOCK_SIZE = 16
21 DEVICE = torch.device("cpu")
22
23
24 def _create_mamba_spec(num_speculative_blocks: int = 1) -> MambaSpec:
25 return MambaSpec(
26 block_size=BLOCK_SIZE,
27 shapes=((16, 64),),
28 dtypes=(torch.float16,),
29 num_speculative_blocks=num_speculative_blocks,
30 )
31
32
33 def test_bailing_linear_attention_reports_uniform_batch_cudagraph_support():
34 vllm_config = create_vllm_config(
35 hf_config_override={
36 "architectures": ["BailingMoeV2_5ForCausalLM"],
37 "model_type": "bailing_hybrid",
38 }
39 )
40
41 support = BailingLinearAttentionMetadataBuilder.get_cudagraph_support(
42 vllm_config, _create_mamba_spec()
43 )
44
45 assert support == AttentionCGSupport.UNIFORM_BATCH
46
47
48 def test_non_bailing_linear_attention_keeps_single_token_cudagraph_support():
49 vllm_config = create_vllm_config(
50 hf_config_override={
51 "architectures": ["MiniMaxText01ForCausalLM"],
52 "model_type": "minimax_text_01",
53 }
54 )
55
56 support = LinearAttentionMetadataBuilder.get_cudagraph_support(
57 vllm_config, _create_mamba_spec()
58 )
59
60 assert support == AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
61
62
63 def test_linear_attention_spec_decode_full_graph_metadata_pads_cache_slots():
64 vllm_config = create_vllm_config(
65 hf_config_override={
66 "architectures": ["BailingMoeV2_5ForCausalLM"],
67 "model_type": "bailing_hybrid",
68 }
69 )
70 vllm_config.speculative_config = SpeculativeConfig(
71 method="ngram",
72 num_speculative_tokens=1,
73 )
74 vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
75
76 builder = BailingLinearAttentionMetadataBuilder(
77 kv_cache_spec=_create_mamba_spec(),
78 layer_names=["model.layers.0.self_attn"],
79 vllm_config=vllm_config,
80 device=DEVICE,
81 )
82
83 common = create_common_attn_metadata(
84 BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]),
85 BLOCK_SIZE,
86 DEVICE,
87 )
88 common.block_table_tensor[2].fill_(-1)
89
90 metadata = builder.build(
91 common_prefix_len=0,
92 common_attn_metadata=common,
93 num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32),
94 )
95
96 assert metadata.num_decodes == 3
97 assert metadata.num_prefills == 0
98 assert metadata.num_decode_tokens == 4
99 assert metadata.state_indices_tensor_d is not None
100 assert metadata.state_indices_tensor_d.shape == (3, 2)
101 assert torch.equal(
102 metadata.state_indices_tensor_d[2],
103 torch.full((2,), PAD_SLOT_ID, dtype=torch.int32),
104 )
105 assert torch.equal(
106 metadata.state_indices_tensor[2],
107 torch.tensor(PAD_SLOT_ID, dtype=torch.int32),
108 )
109 assert metadata.query_start_loc_d is not None
110 assert metadata.query_start_loc_d.tolist() == [0, 2, 4, 4]
111 assert metadata.num_accepted_tokens is not None
112 assert metadata.num_accepted_tokens.tolist() == [1, 2, 1]
113
114
115 def test_linear_attention_full_graph_metadata_uses_stable_decode_buffers():
116 vllm_config = create_vllm_config(
117 hf_config_override={
118 "architectures": ["BailingMoeV2_5ForCausalLM"],
119 "model_type": "bailing_hybrid",
120 }
121 )
122 vllm_config.speculative_config = SpeculativeConfig(
123 method="ngram",
124 num_speculative_tokens=1,
125 )
126 vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY
127
128 builder = BailingLinearAttentionMetadataBuilder(
129 kv_cache_spec=_create_mamba_spec(),
130 layer_names=["model.layers.0.self_attn"],
131 vllm_config=vllm_config,
132 device=DEVICE,
133 )
134
135 common = create_common_attn_metadata(
136 BatchSpec(seq_lens=[20, 20, 0], query_lens=[2, 2, 0]),
137 BLOCK_SIZE,
138 DEVICE,
139 arange_block_indices=True,
140 )
141 common.block_table_tensor = torch.tensor(
142 [[10, 11], [12, 13], [-1, -1]],
143 dtype=torch.int32,
144 device=DEVICE,
145 )
146
147 first = builder.build(
148 common_prefix_len=0,
149 common_attn_metadata=common,
150 num_accepted_tokens=torch.tensor([1, 2, 1], dtype=torch.int32),
151 )
152 assert first.state_indices_tensor_d is not None
153 assert first.query_start_loc_d is not None
154 assert first.num_accepted_tokens is not None
155 state_ptr = first.state_indices_tensor_d.data_ptr()
156 query_ptr = first.query_start_loc_d.data_ptr()
157 accepted_ptr = first.num_accepted_tokens.data_ptr()
158
159 common2 = create_common_attn_metadata(
160 BatchSpec(seq_lens=[36, 0, 0], query_lens=[2, 0, 0]),
161 BLOCK_SIZE,
162 DEVICE,
163 arange_block_indices=True,
164 )
165 common2.block_table_tensor = torch.tensor(
166 [[20, 21], [-1, -1], [-1, -1]],
167 dtype=torch.int32,
168 device=DEVICE,
169 )
170 second = builder.build(
171 common_prefix_len=0,
172 common_attn_metadata=common2,
173 num_accepted_tokens=torch.tensor([2, 1, 1], dtype=torch.int32),
174 )
175
176 assert second.state_indices_tensor_d is not None
177 assert second.query_start_loc_d is not None
178 assert second.num_accepted_tokens is not None
179 assert second.state_indices_tensor_d.data_ptr() == state_ptr
180 assert second.query_start_loc_d.data_ptr() == query_ptr
181 assert second.num_accepted_tokens.data_ptr() == accepted_ptr
182 assert second.state_indices_tensor_d.tolist() == [
183 [20, 21],
184 [PAD_SLOT_ID, PAD_SLOT_ID],
185 [PAD_SLOT_ID, PAD_SLOT_ID],
186 ]
187 assert second.query_start_loc_d.tolist() == [0, 2, 2, 2]
188 assert second.num_accepted_tokens.tolist() == [2, 1, 1]