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, 22:35 UTC
[Misc] Add unit tests for chunked local attention (#21692)
Signed-off-by: Yong Hoon Shin <yhshin@meta.com>
Yong Hoon Shin committed
Jul 31, 2025 at 11:39 UTC
71470bc4afdab89eccc232b668a69571ffede1dc
2 files changed
+219
-13
tests/v1/attention/test_chunked_local_attention.py
new
+196
@@ -0,0 +1,196 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+import numpy as np
6
+import pytest
7
+import torch
8
+
9
+from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
10
+from vllm.v1.attention.backends.utils import (
11
+ make_local_attention_virtual_batches)
12
+
13
+
14
+@dataclass
15
+class LocalAttentionTestData:
16
+ # Input parameters
17
+ batch_spec: BatchSpec
18
+ attn_chunk_size: int
19
+ block_size: int
20
+ # Expected return values
21
+ expected_q_seqlens: list[int]
22
+ expected_k_seqlens: list[int]
23
+ expected_local_block_table: list[list[int]]
24
+
25
+
26
+test_data_list = [
27
+ # Same as example in docstring of make_local_attention_virtual_batches
28
+ # except block table has 9 columns instead of 10
29
+ LocalAttentionTestData(
30
+ batch_spec=BatchSpec(
31
+ query_lens=[4, 10, 5],
32
+ seq_lens=[6, 17, 9],
33
+ ),
34
+ attn_chunk_size=4,
35
+ block_size=2,
36
+ expected_q_seqlens=[2, 2, 1, 4, 4, 1, 4, 1],
37
+ expected_k_seqlens=[4, 2, 4, 4, 4, 1, 4, 1],
38
+ # 2 pages per local branch
39
+ # (chunk size 4 // block size 2)
40
+ expected_local_block_table=[
41
+ [0, 1], # local-batch 0, (batch 0, starting from k[0])
42
+ [2, 3], # local-batch 1, (batch 0, starting from k[4])
43
+ [11, 12], # local-batch 2, (batch 1, starting from k[4])
44
+ [13, 14], # local-batch 3, (batch 1, starting from k[8])
45
+ [15, 16], # local-batch 4, (batch 1, starting from k[12])
46
+ [17, 17], # local-batch 5, (batch 1, starting from k[16])
47
+ [20, 21], # local-batch 6, (batch 2, starting from k[4])
48
+ [22, 23], # local-batch 7, (batch 2, starting from k[8])
49
+ ]),
50
+ # Case where block indices are not clipped to block table ncols-1
51
+ # because tokens_in_last_block == attn_chunk_size
52
+ LocalAttentionTestData(batch_spec=BatchSpec(
53
+ query_lens=[8],
54
+ seq_lens=[12],
55
+ ),
56
+ attn_chunk_size=4,
57
+ block_size=2,
58
+ expected_q_seqlens=[4, 4],
59
+ expected_k_seqlens=[4, 4],
60
+ expected_local_block_table=[
61
+ [2, 3],
62
+ [4, 5],
63
+ ]),
64
+ # Case where all kv_seq positions are involved in attn
65
+ LocalAttentionTestData(
66
+ batch_spec=BatchSpec(
67
+ query_lens=[7],
68
+ # 10 - 7 = 3 previously computed tokens
69
+ seq_lens=[10],
70
+ ),
71
+ attn_chunk_size=4,
72
+ block_size=2,
73
+ expected_q_seqlens=[1, 4, 2],
74
+ expected_k_seqlens=[4, 4, 2],
75
+ expected_local_block_table=[
76
+ [0, 1],
77
+ [2, 3],
78
+ [4, 4],
79
+ ]),
80
+ # Case where attn_chunk_size > kv_seq_len
81
+ # so no extra mini virtual batches are created
82
+ LocalAttentionTestData(
83
+ batch_spec=BatchSpec(
84
+ query_lens=[4],
85
+ seq_lens=[6],
86
+ ),
87
+ # Larger than kv_seq_len
88
+ attn_chunk_size=10,
89
+ block_size=2,
90
+ # No change to q_seqlens and k_seqlens
91
+ expected_q_seqlens=[4],
92
+ expected_k_seqlens=[6],
93
+ # In this case, we only need a block-table like:
94
+ # block_table = [ [0, 1, 2] ] # 1 batch, 3 pages
95
+ # But we need to pad it to 5 pages per local batch
96
+ # because currently the pages_per_local_batch
97
+ # is calculated as (attn_chunk_size // block_size)
98
+ expected_local_block_table=[
99
+ [0, 1, 2, 2, 2],
100
+ ]),
101
+ # Block size equal to chunk size
102
+ # Expect single page per batch in local batch table
103
+ LocalAttentionTestData(
104
+ batch_spec=BatchSpec(
105
+ query_lens=[6, 6],
106
+ seq_lens=[8, 8],
107
+ ),
108
+ attn_chunk_size=4,
109
+ block_size=4,
110
+ expected_q_seqlens=[2, 4, 2, 4],
111
+ expected_k_seqlens=[4, 4, 4, 4],
112
+ # Initial block table = [
113
+ # [0, 1], < batch 0
114
+ # [2, 3], < batch 1
115
+ # ]
116
+ expected_local_block_table=[
117
+ [0], # local-batch 0, (batch 0, starting from k[0])
118
+ [1], # local-batch 1, (batch 0, starting from k[4])
119
+ [2], # local-batch 1, (batch 0, starting from k[0])
120
+ [3], # local-batch 1, (batch 0, starting from k[4])
121
+ ]),
122
+ # Case where query falls in the second attention chunk
123
+ # k_toks > 0 1 2 3 4
124
+ # q_toks v _____________
125
+ # 0 | 1
126
+ # 1 | 1 1
127
+ # 2 | 1 1 1
128
+ # 3 | 1 1 1 1
129
+ # 4 | 1
130
+ # where tokens 0,1,2,3 have been pre-computed
131
+ LocalAttentionTestData(batch_spec=BatchSpec(
132
+ query_lens=[1],
133
+ seq_lens=[5],
134
+ ),
135
+ attn_chunk_size=4,
136
+ block_size=2,
137
+ expected_q_seqlens=[1],
138
+ expected_k_seqlens=[1],
139
+ expected_local_block_table=[
140
+ [2, 2],
141
+ ]),
142
+]
143
+
144
+
145
+@pytest.mark.parametrize("test_data", test_data_list)
146
+def test_local_attention_virtual_batches(test_data: LocalAttentionTestData):
147
+ device = torch.device("cuda:0")
148
+ batch_spec = test_data.batch_spec
149
+ attn_chunk_size = test_data.attn_chunk_size
150
+ block_size = test_data.block_size
151
+ expected_q_seqlens = test_data.expected_q_seqlens
152
+ expected_k_seqlens = test_data.expected_k_seqlens
153
+ expected_local_block_table = test_data.expected_local_block_table
154
+
155
+ # Create common attention metadata
156
+ common_attn_metadata = create_common_attn_metadata(
157
+ batch_spec,
158
+ block_size,
159
+ device,
160
+ # Use torch.arange instead of torch.randint so we can assert on
161
+ # block table tensor values. The block table will have shape
162
+ # (num_batches, cdiv(max_seq_len, block_size)) and the values will be
163
+ # aranged from 0 to cdiv(max_seq_len, block_size)-1
164
+ arange_block_indices=True,
165
+ )
166
+
167
+ # Call the function
168
+ result = make_local_attention_virtual_batches(attn_chunk_size,
169
+ common_attn_metadata,
170
+ block_size)
171
+
172
+ # Convert to numpy for easier comparison
173
+ actual_q_seqlens = np.diff(result.query_start_loc_cpu.numpy())
174
+ actual_k_seqlens = result.seq_lens_cpu.numpy()
175
+
176
+ # Check that all query lengths are less than or equal to attn_chunk_size
177
+ assert all(q_len <= attn_chunk_size for q_len in actual_q_seqlens)
178
+ # Check that all key lengths are less than or equal to attn_chunk_size
179
+ assert all(k_len <= attn_chunk_size for k_len in actual_k_seqlens)
180
+ # Check that the total number of query tokens is preserved
181
+ assert sum(actual_q_seqlens) == sum(batch_spec.query_lens)
182
+
183
+ # Verify results
184
+ np.testing.assert_array_equal(actual_q_seqlens, expected_q_seqlens)
185
+ np.testing.assert_array_equal(actual_k_seqlens, expected_k_seqlens)
186
+
187
+ expected_block_table_tensor =\
188
+ torch.tensor(expected_local_block_table,
189
+ dtype=torch.int32,
190
+ device=device)
191
+
192
+ print(f"Expected block table:\n{expected_block_table_tensor}")
193
+ print(f"Actual block table:\n{result.block_table_tensor}")
194
+
195
+ torch.testing.assert_close(result.block_table_tensor,
196
+ expected_block_table_tensor)
tests/v1/attention/utils.py
+23
-13
@@ -40,7 +40,8 @@ def create_common_attn_metadata(
40
batch_spec: BatchSpec,
41
block_size: int,
42
device: torch.device,
43
- max_block_idx: int = 1000) -> CommonAttentionMetadata:
43
+ max_block_idx: int = 1000,
44
+ arange_block_indices: bool = False) -> CommonAttentionMetadata:
45
"""Create CommonAttentionMetadata from a BatchSpec and ModelParams."""
46
# Create query start locations
47
query_start_loc = torch.zeros(batch_spec.batch_size + 1,
@@ -65,19 +66,28 @@ def create_common_attn_metadata(
66
]
67
num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
68
68
- # Create block table (random for testing)
69
+ # Create block table and slot mapping
70
max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
70
- block_table_tensor = torch.randint(0,
71
- max_block_idx,
72
- (batch_spec.batch_size, max_blocks),
73
- dtype=torch.int32,
74
- device=device)
75
-
76
- # Create slot mapping
77
- slot_mapping = torch.randint(0,
78
- max_block_idx, (num_tokens, ),
79
- dtype=torch.int64,
80
- device=device)
71
+ if arange_block_indices:
72
+ num_blocks = batch_spec.batch_size * max_blocks
73
+ block_table_tensor = torch.arange(num_blocks,
74
+ dtype=torch.int32,
75
+ device=device).view(
76
+ batch_spec.batch_size,
77
+ max_blocks)
78
+ slot_mapping = torch.arange(num_tokens,
79
+ dtype=torch.int64,
80
+ device=device).view(num_tokens)
81
+ else:
82
+ block_table_tensor = torch.randint(0,
83
+ max_block_idx,
84
+ (batch_spec.batch_size, max_blocks),
85
+ dtype=torch.int32,
86
+ device=device)
87
+ slot_mapping = torch.randint(0,
88
+ max_block_idx, (num_tokens, ),
89
+ dtype=torch.int64,
90
+ device=device)
91
92
# Calculate max query length
93
max_query_len = max(batch_spec.query_lens)