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 10, 12:45 UTC
Decouple page_size_bytes calculation in AttentionSpec for TPU/RPA Compatibility. (#31635)
Signed-off-by: Lihao Ran <imlihao.ran@gmail.com> Signed-off-by: Lumosis <30372757+Lumosis@users.noreply.github.com>
Lumosis committed
Jan 8, 2026 at 01:00 UTC
b634e619bbcfed0abe4e01d0e2d97fb1fdfdbbd5
6 files changed
+75
-20
tests/v1/core/test_kv_sharing.py
+3
-1
@@ -11,7 +11,9 @@ pytestmark = pytest.mark.cpu_test
11
12
13
def new_kv_cache_spec():
14
- return FullAttentionSpec(16, 1, 1, torch.float32, False)
14
+ return FullAttentionSpec(
15
+ block_size=16, num_kv_heads=1, head_size=1, dtype=torch.float32
16
+ )
17
18
19
def test_initialize_kv_cache_for_kv_sharing_different_attn_groups():
tests/v1/core/test_prefix_caching.py
+32
-9
@@ -94,7 +94,12 @@ def make_kv_cache_config(block_size: int, num_blocks: int) -> KVCacheConfig:
94
kv_cache_groups=[
95
KVCacheGroupSpec(
96
["layer"],
97
- FullAttentionSpec(block_size, 1, 1, torch.float32),
97
+ FullAttentionSpec(
98
+ block_size=block_size,
99
+ num_kv_heads=1,
100
+ head_size=1,
101
+ dtype=torch.float32,
102
+ ),
103
)
104
],
105
)
@@ -109,18 +114,31 @@ def make_kv_cache_config_hybrid_model(
114
kv_cache_groups=[
115
KVCacheGroupSpec(
116
["layer1"],
112
- FullAttentionSpec(block_size, 1, 1, torch.float32),
117
+ FullAttentionSpec(
118
+ block_size=block_size,
119
+ num_kv_heads=1,
120
+ head_size=1,
121
+ dtype=torch.float32,
122
+ ),
123
),
124
KVCacheGroupSpec(
125
["layer2"],
126
SlidingWindowSpec(
117
- block_size, 1, 1, torch.float32, sliding_window=2 * block_size
127
+ block_size=block_size,
128
+ num_kv_heads=1,
129
+ head_size=1,
130
+ dtype=torch.float32,
131
+ sliding_window=2 * block_size,
132
),
133
),
134
KVCacheGroupSpec(
135
["layer3"],
136
SlidingWindowSpec(
123
- block_size, 1, 1, torch.float32, sliding_window=2 * block_size
137
+ block_size=block_size,
138
+ num_kv_heads=1,
139
+ head_size=1,
140
+ dtype=torch.float32,
141
+ sliding_window=2 * block_size,
142
),
143
),
144
],
@@ -1616,15 +1634,20 @@ def test_different_block_size():
1634
kv_cache_groups=[
1635
KVCacheGroupSpec(
1636
["layer1"],
1619
- FullAttentionSpec(block_size * 2, 1, 1, torch.float16),
1637
+ FullAttentionSpec(
1638
+ block_size=block_size * 2,
1639
+ num_kv_heads=1,
1640
+ head_size=1,
1641
+ dtype=torch.float16,
1642
+ ),
1643
),
1644
KVCacheGroupSpec(
1645
["layer2"],
1646
SlidingWindowSpec(
1624
- block_size,
1625
- 1,
1626
- 1,
1627
- torch.float32,
1647
+ block_size=block_size,
1648
+ num_kv_heads=1,
1649
+ head_size=1,
1650
+ dtype=torch.float32,
1651
sliding_window=2 * block_size,
1652
),
1653
),
tests/v1/core/test_scheduler.py
+7
-1
@@ -1573,7 +1573,13 @@ def create_scheduler_with_priority(
1573
kv_cache_tensors=[],
1574
kv_cache_groups=[
1575
KVCacheGroupSpec(
1576
- ["layer"], FullAttentionSpec(block_size, 1, 1, torch.float32, False)
1576
+ ["layer"],
1577
+ FullAttentionSpec(
1578
+ block_size=block_size,
1579
+ num_kv_heads=1,
1580
+ head_size=1,
1581
+ dtype=torch.float32,
1582
+ ),
1583
)
1584
],
1585
)
tests/v1/core/utils.py
+7
-1
@@ -142,7 +142,13 @@ def create_scheduler(
142
kv_cache_tensors=[],
143
kv_cache_groups=[
144
KVCacheGroupSpec(
145
- ["layer"], FullAttentionSpec(block_size, 1, 1, torch.float32, False)
145
+ ["layer"],
146
+ FullAttentionSpec(
147
+ block_size=block_size,
148
+ num_kv_heads=1,
149
+ head_size=1,
150
+ dtype=torch.float32,
151
+ ),
152
)
153
],
154
)
tests/v1/kv_connector/unit/utils.py
+7
-1
@@ -148,7 +148,13 @@ def create_scheduler(
148
kv_cache_tensors=[],
149
kv_cache_groups=[
150
KVCacheGroupSpec(
151
- ["layer"], FullAttentionSpec(block_size, 1, 1, torch.float32, False)
151
+ ["layer"],
152
+ FullAttentionSpec(
153
+ block_size=block_size,
154
+ num_kv_heads=1,
155
+ head_size=1,
156
+ dtype=torch.float32,
157
+ ),
158
)
159
],
160
)
vllm/v1/kv_cache_interface.py
+19
-7
@@ -61,14 +61,23 @@ class KVCacheSpec:
61
return copy.deepcopy(specs[0])
62
63
64
-@dataclass(frozen=True)
64
+@dataclass(frozen=True, kw_only=True)
65
class AttentionSpec(KVCacheSpec):
66
num_kv_heads: int
67
head_size: int
68
dtype: torch.dtype
69
+ page_size_padded: int | None = None
70
71
@property
72
def page_size_bytes(self) -> int:
73
+ real_page_size = self.real_page_size_bytes
74
+ if self.page_size_padded is not None:
75
+ assert self.page_size_padded >= real_page_size
76
+ return self.page_size_padded
77
+ return real_page_size
78
+
79
+ @property
80
+ def real_page_size_bytes(self) -> int:
81
return (
82
2
83
* self.block_size
@@ -78,7 +87,7 @@ class AttentionSpec(KVCacheSpec):
87
)
88
89
81
-@dataclass(frozen=True)
90
+@dataclass(frozen=True, kw_only=True)
91
class FullAttentionSpec(AttentionSpec):
92
"""
93
When hybrid allocator is disabled and the model contains both full
@@ -150,6 +159,7 @@ class FullAttentionSpec(AttentionSpec):
159
head_size=specs[0].head_size,
160
head_size_v=specs[0].head_size_v,
161
dtype=specs[0].dtype,
162
+ page_size_padded=specs[0].page_size_padded,
163
sliding_window=cls.merge_window_sizes(sliding_window),
164
attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
165
)
@@ -168,7 +178,7 @@ class FullAttentionSpec(AttentionSpec):
178
return merged_spec
179
180
@property
171
- def page_size_bytes(self) -> int:
181
+ def real_page_size_bytes(self) -> int:
182
return (
183
self.block_size
184
* self.num_kv_heads
@@ -177,13 +187,13 @@ class FullAttentionSpec(AttentionSpec):
187
)
188
189
180
-@dataclass(frozen=True)
190
+@dataclass(frozen=True, kw_only=True)
191
class MLAAttentionSpec(FullAttentionSpec):
192
# TODO(Lucas/Chen): less hacky way to do this
193
cache_dtype_str: str | None = None
194
195
@property
186
- def page_size_bytes(self) -> int:
196
+ def real_page_size_bytes(self) -> int:
197
if self.cache_dtype_str == "fp8_ds_mla":
198
# See `vllm/v1/attention/backends/mla/flashmla_sparse.py`
199
# for details.
@@ -210,11 +220,12 @@ class MLAAttentionSpec(FullAttentionSpec):
220
num_kv_heads=specs[0].num_kv_heads,
221
head_size=specs[0].head_size,
222
dtype=specs[0].dtype,
223
+ page_size_padded=specs[0].page_size_padded,
224
cache_dtype_str=cache_dtype_str_set.pop(),
225
)
226
227
217
-@dataclass(frozen=True)
228
+@dataclass(frozen=True, kw_only=True)
229
class ChunkedLocalAttentionSpec(AttentionSpec):
230
attention_chunk_size: int
231
@@ -233,7 +244,7 @@ class ChunkedLocalAttentionSpec(AttentionSpec):
244
return cdiv(num_tokens, self.block_size) * self.page_size_bytes
245
246
236
-@dataclass(frozen=True)
247
+@dataclass(frozen=True, kw_only=True)
248
class SlidingWindowSpec(AttentionSpec):
249
sliding_window: int
250
@@ -335,6 +346,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
346
head_size_v=specs[0].head_size_v,
347
sink_len=specs[0].sink_len,
348
dtype=specs[0].dtype,
349
+ page_size_padded=specs[0].page_size_padded,
350
sliding_window=cls.merge_window_sizes(sliding_window),
351
attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
352
)