[Bugfix] Pad Mamba page size instead of scaling block_size in unify_kv_cache_spec_page_size (#45207)
Signed-off-by: Sahil170595 <147995121+Sahil170595@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Sahil Kadadekar committed
Jul 7, 2026 at 18:01 UTC
55da232db6963613d34229dfd257236e6f3c8097
2 files changed
+101
-5
tests/v1/core/test_kv_cache_utils.py
+85
@@ -2484,6 +2484,91 @@ def test_unify_hybrid_kv_cache_specs():
2484
kv_cache_utils.unify_hybrid_kv_cache_specs(kv_cache_spec)
2485
2486
2487
+def test_unify_kv_cache_spec_page_size_mamba():
2488
+ """Regression test for https://github.com/vllm-project/vllm/issues/43626.
2489
+
2490
+ MambaSpec's page_size_bytes is determined by its state shapes and does not
2491
+ change with block_size, so unify_kv_cache_spec_page_size must pad the Mamba
2492
+ page instead of scaling its block_size. This situation arises when a layer
2493
+ with a page larger than the (already platform-aligned) Mamba page joins the
2494
+ specs, e.g. a dense draft model with more KV heads than the hybrid main
2495
+ model.
2496
+ """
2497
+ # 1. Hybrid main model (Mamba + full attention, pages already aligned at
2498
+ # 16KB) plus a dense draft model layer with a 2x larger page (32KB).
2499
+ # Reproduces the bare AssertionError from #43626: 32768 % 16384 == 0, so
2500
+ # the old code scaled the Mamba block_size, which left page_size_bytes
2501
+ # unchanged at 16384.
2502
+ mamba_spec = new_mamba_spec() # page_size_bytes = 16384
2503
+ main_attn_spec = new_kv_cache_spec() # page_size_bytes = 16384
2504
+ draft_attn_spec = new_kv_cache_spec(num_kv_heads=4) # page_size_bytes = 32768
2505
+ assert mamba_spec.page_size_bytes == main_attn_spec.page_size_bytes == 16384
2506
+ assert draft_attn_spec.page_size_bytes == 32768
2507
+
2508
+ unified = kv_cache_utils.unify_kv_cache_spec_page_size(
2509
+ {
2510
+ "mamba_layer": mamba_spec,
2511
+ "main_attn_layer": main_attn_spec,
2512
+ "draft_attn_layer": draft_attn_spec,
2513
+ }
2514
+ )
2515
+ # Mamba page is padded; block_size (caching granularity) is unchanged.
2516
+ assert unified["mamba_layer"].page_size_bytes == 32768
2517
+ assert unified["mamba_layer"].page_size_padded == 32768
2518
+ assert unified["mamba_layer"].block_size == mamba_spec.block_size
2519
+ # Attention layer with smaller page still unifies by scaling block_size.
2520
+ assert unified["main_attn_layer"].page_size_bytes == 32768
2521
+ assert unified["main_attn_layer"].block_size == 2 * main_attn_spec.block_size
2522
+ assert unified["main_attn_layer"].page_size_padded is None
2523
+ # Layer already at max page size is unchanged.
2524
+ assert unified["draft_attn_layer"] == draft_attn_spec
2525
+
2526
+ # 2. Mamba page already padded by the platform (state smaller than the
2527
+ # padded page); the padding is re-applied at the new maximum.
2528
+ padded_mamba_spec = new_mamba_spec(
2529
+ shapes=((2, 256), (3, 32, 32)), page_size_padded=16384
2530
+ )
2531
+ assert padded_mamba_spec.page_size_bytes == 16384
2532
+ unified = kv_cache_utils.unify_kv_cache_spec_page_size(
2533
+ {
2534
+ "mamba_layer": padded_mamba_spec,
2535
+ "draft_attn_layer": draft_attn_spec,
2536
+ }
2537
+ )
2538
+ assert unified["mamba_layer"].page_size_bytes == 32768
2539
+ assert unified["mamba_layer"].page_size_padded == 32768
2540
+
2541
+ # 3. Mamba page that does not evenly divide the maximum page size is
2542
+ # padded as well (the divisibility constraint only applies to block_size
2543
+ # scaling).
2544
+ odd_mamba_spec = new_mamba_spec(shapes=((6144,),))
2545
+ assert odd_mamba_spec.page_size_bytes == 24576
2546
+ assert 32768 % odd_mamba_spec.page_size_bytes != 0
2547
+ unified = kv_cache_utils.unify_kv_cache_spec_page_size(
2548
+ {
2549
+ "mamba_layer": odd_mamba_spec,
2550
+ "draft_attn_layer": draft_attn_spec,
2551
+ }
2552
+ )
2553
+ assert unified["mamba_layer"].page_size_bytes == 32768
2554
+
2555
+ # 4. Attention layers with non-divisible page sizes still raise.
2556
+ with pytest.raises(NotImplementedError):
2557
+ kv_cache_utils.unify_kv_cache_spec_page_size(
2558
+ {
2559
+ "attn_layer": new_kv_cache_spec(block_size=24), # 24576
2560
+ "draft_attn_layer": draft_attn_spec, # 32768
2561
+ }
2562
+ )
2563
+
2564
+ # 5. Uniform page sizes are returned unchanged.
2565
+ specs = {
2566
+ "mamba_layer": new_mamba_spec(),
2567
+ "attn_layer": new_kv_cache_spec(),
2568
+ }
2569
+ assert kv_cache_utils.unify_kv_cache_spec_page_size(specs) == specs
2570
+
2571
+
2572
def test_hma_not_disabled_when_kv_events_enabled():
2573
"""
2574
Test enabling KV events must not force disable_hybrid_kv_cache_manager to True.
vllm/v1/core/kv_cache_utils.py
+16
-5
@@ -1051,11 +1051,12 @@ def unify_kv_cache_spec_page_size(
1051
) -> dict[str, KVCacheSpec]:
1052
"""
1053
Unify the page size of the given KVCacheSpec. If the page size of all layers
1054
- are the same, return the original KVCacheSpec. If not same, first try to
1055
- unify page size by increasing the block size of layers with smaller page
1056
- size. If a smaller attention page does not evenly divide the maximum page
1057
- size, keep its logical block size and pad its physical page instead --- but
1058
- only for attention layers whose backend opts in via
1054
+ are the same, return the original KVCacheSpec. If not same, unify the page
1055
+ size by increasing the block size of layers with smaller page size. Two
1056
+ cases cannot be unified by block size alone and pad their physical page to
1057
+ the maximum instead: Mamba layers, whose page size comes from state shapes
1058
+ and is independent of block size; and attention layers whose page does not
1059
+ evenly divide the maximum and whose backend opts in via
1060
``AttentionSpec.indexes_kv_by_block_stride`` (the padded page is read through
1061
a strided view, which not every backend handles). Raise NotImplementedError
1062
if failed to unify the page size.
@@ -1076,6 +1077,16 @@ def unify_kv_cache_spec_page_size(
1077
for layer_name, layer_spec in kv_cache_spec.items():
1078
if layer_spec.page_size_bytes == max_page_size:
1079
new_kv_cache_spec[layer_name] = layer_spec
1080
+ elif isinstance(layer_spec, MambaSpec):
1081
+ # MambaSpec's page size is determined by its state shapes and does
1082
+ # not scale with block_size, so pad the page instead. This is the
1083
+ # same padding mechanism the platform uses to align Mamba pages
1084
+ # with the main model's attention page size; it is needed here
1085
+ # when another layer (e.g. from a draft model) has a larger page
1086
+ # than the already-aligned Mamba page.
1087
+ new_spec: KVCacheSpec = replace(layer_spec, page_size_padded=max_page_size)
1088
+ assert new_spec.page_size_bytes == max_page_size
1089
+ new_kv_cache_spec[layer_name] = new_spec
1090
else:
1091
layer_page_size = layer_spec.page_size_bytes
1092
if max_page_size % layer_page_size == 0: