[Bugfix] Scale KV page size for hidden states extraction with TP (#50894)
Signed-off-by: Claude <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis-z@users.noreply.github.com> Signed-off-by: Orestis Zambounis <23146389+orestis-z@users.noreply.github.com> Signed-off-by: Speculators Bot <bot@speculators.ai> Signed-off-by: Speculators Bot <speculatorsbot@gmail.com> Co-authored-by: Orestis Zambounis <orestis-z@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Speculators Bot <speculatorsbot@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Orestis Zambounis committed
Sep 14, 2026 at 16:07 UTC
ff5f6d41b1ccf1e305e425db41bc7f52dbce8621
2 files changed
+73
tests/v1/core/test_kv_cache_utils.py
+39
@@ -3208,6 +3208,45 @@ def test_mla_with_incompatible_swa_uses_one_full_allocation_group(caplog_vllm):
3208
assert "attention compute is unchanged" in caplog_vllm.text
3209
3210
3211
+def test_hidden_states_with_tp_scales_page_size():
3212
+ """When TP shrinks KV pages below the hidden-state per-token cost,
3213
+ get_kv_cache_groups must scale up target block sizes so that the
3214
+ common page accommodates the unsharded hidden states."""
3215
+ # Simulate TP=4 sharding a model with 8 KV heads → 2 per rank.
3216
+ # KV page = block_size(16) * num_kv_heads(2) * head_size(64) * dtype(2)
3217
+ # = 16 * 2 * 64 * 2 = 4096 bytes.
3218
+ kv_spec = new_kv_cache_spec(
3219
+ block_size=16,
3220
+ num_kv_heads=2,
3221
+ head_size=64,
3222
+ dtype=torch.bfloat16,
3223
+ )
3224
+ # Hidden-state per-token cost = num_hidden_states(6) * hidden_size(512)
3225
+ # * dtype(2) = 6144 bytes, which exceeds the 4096-byte KV page.
3226
+ hs_spec = HiddenStateCacheSpec(
3227
+ block_size=16,
3228
+ num_kv_heads=6,
3229
+ head_size=512,
3230
+ dtype=torch.bfloat16,
3231
+ )
3232
+ specs = {
3233
+ "target.0.attn": kv_spec,
3234
+ "target.1.attn": kv_spec,
3235
+ "cache_only_layers.48": hs_spec,
3236
+ }
3237
+
3238
+ groups = get_kv_cache_groups(_grouping_config(), specs)
3239
+
3240
+ # The hidden-state layer should be present and no assertion should fire.
3241
+ all_layers = {name for g in groups for name in g.layer_names}
3242
+ assert "cache_only_layers.48" in all_layers
3243
+
3244
+ # The target group block sizes must have been scaled up.
3245
+ for g in groups:
3246
+ if "cache_only_layers.48" not in g.layer_names:
3247
+ assert g.kv_cache_spec.block_size > kv_spec.block_size
3248
+
3249
+
3250
def test_get_kv_cache_spec_kind_prefers_specific_attention_subclasses():
3251
assert get_kv_cache_spec_kind(new_mla_spec()) == KVCacheSpecKind.MLA_ATTENTION
3252
vllm/v1/core/kv_cache_utils.py
+34
@@ -2217,6 +2217,38 @@ def _largest_divisor_at_most(value: int, limit: int) -> int:
2217
return 1
2218
2219
2220
+def _ensure_min_page_size(
2221
+ groups: list[KVCacheGroupSpec],
2222
+ common_page: int,
2223
+ hidden_specs: dict[str, HiddenStateCacheSpec],
2224
+) -> tuple[list[KVCacheGroupSpec], int]:
2225
+ """Scale up group block sizes so the common page is at least as large as
2226
+ the biggest hidden-state per-token cost.
2227
+
2228
+ This protects the hidden state extraction feature, where we store hidden states in
2229
+ the KV cache, and return them to clients using a custom hidden state connector.
2230
+
2231
+ Returns the (possibly scaled) groups and updated common page size.
2232
+ """
2233
+ max_per_token = max(
2234
+ spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)
2235
+ for spec in hidden_specs.values()
2236
+ )
2237
+ if max_per_token <= common_page:
2238
+ return groups, common_page
2239
+
2240
+ scale = cdiv(max_per_token, common_page)
2241
+ common_page *= scale
2242
+ scaled: list[KVCacheGroupSpec] = []
2243
+ for g in groups:
2244
+ s = g.kv_cache_spec
2245
+ kw: dict[str, int] = {"block_size": s.block_size * scale}
2246
+ if isinstance(s, (AttentionSpec, MambaSpec)) and s.page_size_padded is not None:
2247
+ kw["page_size_padded"] = s.page_size_padded * scale
2248
+ scaled.append(KVCacheGroupSpec(g.layer_names, replace(s, **kw)))
2249
+ return scaled, common_page
2250
+
2251
+
2252
def get_kv_cache_groups(
2253
vllm_config: VllmConfig,
2254
kv_cache_spec: dict[str, KVCacheSpec],
@@ -2288,6 +2320,8 @@ def get_kv_cache_groups(
2320
# Add hidden-state layers back with page aligned to the common page.
2321
if hidden_specs:
2322
common_page = get_uniform_page_size([g.kv_cache_spec for g in groups])
2323
+ # TP may shrink the common page below hidden-state per-token cost.
2324
+ groups, common_page = _ensure_min_page_size(groups, common_page, hidden_specs)
2325
group_block_size = math.gcd(*(g.kv_cache_spec.block_size for g in groups))
2326
for name, spec in hidden_specs.items():
2327
per_token = spec.num_kv_heads * spec.head_size * get_dtype_size(spec.dtype)