@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6b110badbb

[Bugfix][Mooncake] Save exact Mamba boundary states (#51358)

Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Dao Le <daole@inferact.ai>

Yifan Qiao committed Aug 28, 2026 at 19:40 UTC 6b110badbb22d3f66c7218b71138f13b7a6b3419

20 files changed +1494 -273

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+264 -29

@@ -69,6 +69,19 @@ def test_capable_connector_uses_divergent_partial_hit_lookup():

69 manager.get_computed_blocks.assert_not_called()

72 +def drain_boundary_state_offloads(manager):

73 + """Drain exact boundary-state block ids offered to a connector."""

74 + return manager.take_boundary_state_offloads()

75 +

76 +

77 +def _free_block_ids(manager):

78 + """Block ids the pool would hand out to the next allocation."""

79 + return {

80 + block.block_id

81 + for block in manager.block_pool.free_block_queue.get_all_free_blocks()

82 + }

83 +

84 +

85 def make_full_mamba_manager(

86 *,

87 dcp_world_size: int,

@@ -618,10 +631,9 @@ def test_external_mamba_hit_same_block_uses_running_cow_on_continue():

631 assert moved[0].block_id == cow_copy.dst_block_id

632

633

621 -def test_take_partial_tail_offloads_returns_cow_target():

622 - """The connector offload hand-off exposes the mamba CoW *target* block Y

623 - (the durable boundary state), not the overwritten source X, and only at

624 - the CoW step."""

634 +def test_boundary_state_offloads_returns_cow_target():

635 + """Boundary hand-offs expose aligned snapshots and the partial-tail CoW

636 + target, never the overwritten CoW source."""

637 hash_block_size = 2

638 block_size = 2 * hash_block_size

639 kv_cache_config = KVCacheConfig(

@@ -660,8 +672,18 @@ def test_take_partial_tail_offloads_returns_cow_target():

672 computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

673 assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None

674

663 - # Step A registered the partial tail but has not CoW'd yet: no offload.

664 - assert manager.take_partial_tail_offloads() == {}

675 + # Step A offers the materialized aligned checkpoint immediately.

676 + ((group_id, block_id, boundary_tokens),) = drain_boundary_state_offloads(manager)[

677 + "0"

678 + ]

679 + assert group_id == 1

680 + assert boundary_tokens == 4

681 + aligned_hash = req0.block_hashes[4 // hash_block_size - 1]

682 + aligned_block = manager.block_pool.get_cached_block(

683 + aligned_hash, kv_cache_group_ids=[1]

684 + )

685 + assert aligned_block is not None

686 + assert block_id == aligned_block[0].block_id

687

688 partial_mamba_hash = req0.block_hashes[6 // hash_block_size - 1]

689 source_block = manager.block_pool.get_cached_block(

@@ -675,34 +697,27 @@ def test_take_partial_tail_offloads_returns_cow_target():

697 req0.append_output_token_ids([3])

698 assert manager.allocate_slots(req0, 1) is not None

699

678 - offloads = manager.take_partial_tail_offloads()

700 + offloads = drain_boundary_state_offloads(manager)

701 assert list(offloads.keys()) == ["0"]

702 assert len(offloads["0"]) == 1

703 group_id, block_id, boundary_tokens = offloads["0"][0]

704 assert group_id == 1 # the mamba group

705 assert boundary_tokens == 6

684 - copies, _ = manager.take_kv_cache_block_copies()

706 + copies, retained = manager.take_kv_cache_block_copies()

707 cow_copy = next(c for c in copies if c.src_block_id == source_block_id)

708 # The offload points at the durable CoW target Y, not the overwritten X.

709 assert block_id == cow_copy.dst_block_id

710 assert block_id != source_block_id

711 # Draining clears it.

690 - assert manager.take_partial_tail_offloads() == {}

712 + assert manager.take_boundary_state_offloads() == {}

713

692 - # The hand-off pinned Y (its CoW retention is released after this step,

693 - # and Y is off the request block table); freeing the request unpins it.

694 - cow_block = manager.block_pool.blocks[block_id]

695 - pinned_ref = cow_block.ref_cnt

696 - assert pinned_ref >= 1

714 + manager.block_pool.free_blocks(retained)

715 manager.free(req0)

698 - assert cow_block.ref_cnt == pinned_ref - 1

716

717

701 -def test_partial_tail_pin_survives_released_cow_retention():

702 - """If the CoW retention is released before the hand-off is drained

703 - (immediate-free mode), the drain must rescue the cow block from the free

704 - queue: a raw ref increment would leave a ref>0 block allocatable, and the

705 - next allocation would pop it and assert."""

718 +def test_block_pool_touch_pins_released_cow_target():

719 + """The connector can rescue an offered CoW target after its step-scoped

720 + retention is released by using the bound BlockPool's touch method."""

721 hash_block_size = 2

722 block_size = 2 * hash_block_size

723 kv_cache_config = KVCacheConfig(

@@ -746,13 +761,18 @@ def test_partial_tail_pin_survives_released_cow_retention():

761 _copies, retained = manager.take_kv_cache_block_copies()

762 manager.block_pool.free_blocks(retained)

763

749 - offloads = manager.take_partial_tail_offloads()

764 + offloads = drain_boundary_state_offloads(manager)

765 ((_group_id, block_id, boundary_tokens),) = offloads["0"]

766 assert boundary_tokens == 6

767 cow_block = manager.block_pool.blocks[block_id]

768 + assert cow_block.ref_cnt == 0

769 + assert block_id in _free_block_ids(manager)

770 +

771 + manager.block_pool.touch([cow_block])

772 assert cow_block.ref_cnt == 1

773 + assert block_id not in _free_block_ids(manager)

774

755 - # The pinned block is out of the free queue: draining every free block

775 + # The connector-pinned block is out of the free queue: draining every free block

776 # neither trips the allocator's ref_cnt assert nor hands it out.

777 new_blocks = manager.block_pool.get_new_blocks(

778 manager.block_pool.get_num_free_blocks()

@@ -760,7 +780,7 @@ def test_partial_tail_pin_survives_released_cow_retention():

780 assert block_id not in {b.block_id for b in new_blocks}

781

782

763 -def test_partial_tail_offload_dropped_when_request_freed_before_drain():

783 +def test_boundary_state_offload_dropped_when_request_freed_before_drain():

784 """A hand-off recorded in the same scheduling pass as the request's death

785 must not be drained: its release hook has already run, so draining would

786 leak a pinned block."""

@@ -805,12 +825,12 @@ def test_partial_tail_offload_dropped_when_request_freed_before_drain():

825

826 # The request dies (preempt/abort) before the scheduler drains.

827 manager.block_pool.free_blocks(manager.pop_blocks_for_free(req0))

808 - assert manager.take_partial_tail_offloads() == {}

828 + assert manager.take_boundary_state_offloads() == {}

829

830

811 -def test_take_partial_tail_offloads_empty_without_partial_tail():

812 - """A prompt ending on a block boundary registers no partial tail, so there

813 - is nothing to offload."""

831 +def test_boundary_state_offloads_block_aligned_prompt():

832 + """A prompt ending on a block boundary registers no CoW partial tail; its

833 + boundary state block is handed off as a snapshot instead (once)."""

834 hash_block_size = 2

835 block_size = 2 * hash_block_size

836 kv_cache_config = KVCacheConfig(

@@ -848,12 +868,17 @@ def test_take_partial_tail_offloads_empty_without_partial_tail():

868 req0 = make_request("0", [0, 0, 1, 1], hash_block_size, sha256)

869 computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

870 assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None

851 - assert manager.take_partial_tail_offloads() == {}

871 + mamba_blocks = manager.coordinator.single_type_managers[1].req_to_blocks["0"]

872 + ((group_id, block_id, boundary),) = drain_boundary_state_offloads(manager)["0"]

873 + assert group_id == 1

874 + assert boundary == block_size

875 + assert block_id == mamba_blocks[0].block_id

876

877 req0.num_computed_tokens = 4

878 req0.append_output_token_ids([2])

879 assert manager.allocate_slots(req0, 1) is not None

856 - assert manager.take_partial_tail_offloads() == {}

880 + # The boundary was already handed off; decoding emits nothing new.

881 + assert manager.take_boundary_state_offloads() == {}

882

883

884 def test_truncate_computed_blocks_preserves_sparse_prefix_positions():

@@ -1786,3 +1811,213 @@ def test_dcp_partial_hit_with_eagle_rewinds_one_hash_unit():

1811 assert num_computed == 4

1812 assert [len(group) for group in computed_blocks.blocks] == [1, 1]

1813 assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None

1814 +

1815 +

1816 +def _snapshot_offload_kv_cache_config(

1817 + hash_block_size: int, block_size: int, num_blocks: int = 24

1818 +):

1819 + return KVCacheConfig(

1820 + num_blocks=num_blocks,

1821 + kv_cache_tensors=[],

1822 + kv_cache_groups=[

1823 + KVCacheGroupSpec(

1824 + ["full"],

1825 + FullAttentionSpec(

1826 + block_size=hash_block_size,

1827 + num_kv_heads=1,

1828 + head_size=1,

1829 + dtype=torch.float32,

1830 + ),

1831 + ),

1832 + KVCacheGroupSpec(

1833 + ["mamba"],

1834 + MambaSpec(

1835 + block_size=block_size,

1836 + shapes=(1, 1),

1837 + dtypes=(torch.float32,),

1838 + mamba_cache_mode="align",

1839 + ),

1840 + ),

1841 + ],

1842 + )

1843 +

1844 +

1845 +def test_retention_snapshots_handed_off_with_exact_block_ids():

1846 + """Under sparse retention, each retained mamba boundary state block is

1847 + handed to the connector with its exact block id. Connectors must not resolve

1848 + align-mode state blocks positionally because the block table is not

1849 + append-only."""

1850 + hash_block_size = 2

1851 + block_size = 4

1852 + manager = make_kv_cache_manager(

1853 + kv_cache_config=_snapshot_offload_kv_cache_config(hash_block_size, block_size),

1854 + max_model_len=8192,

1855 + enable_caching=True,

1856 + hash_block_size=hash_block_size,

1857 + )

1858 + manager.coordinator.retention_interval = 2 * block_size

1859 +

1860 + req0 = make_request("0", [0] * 16, hash_block_size, sha256)

1861 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

1862 + assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None

1863 +

1864 + mamba_blocks = manager.coordinator.single_type_managers[1].req_to_blocks["0"]

1865 + offloads = drain_boundary_state_offloads(manager)

1866 + ((group_id, block_id, boundary),) = offloads["0"]

1867 + assert group_id == 1

1868 + assert boundary == 2 * block_size

1869 + assert block_id == mamba_blocks[1].block_id

1870 +

1871 + req0.num_computed_tokens = 8

1872 + manager.new_step_starts()

1873 + assert manager.allocate_slots(req0, 8) is not None

1874 +

1875 + offloads = drain_boundary_state_offloads(manager)

1876 + ((group_id, block_id2, boundary2),) = offloads["0"]

1877 + assert group_id == 1

1878 + assert boundary2 == 4 * block_size

1879 + assert block_id2 == mamba_blocks[3].block_id

1880 + # Interior align-mode positions stay null and are never handed off.

1881 + assert mamba_blocks[0].is_null and mamba_blocks[2].is_null

1882 +

1883 + manager.free(req0)

1884 +

1885 +

1886 +def test_snapshot_handoff_dense_default_retention():

1887 + """With dense (default) retention, every materialized mamba boundary

1888 + state block is handed off — regular mamba-align + prefix-match-unit

1889 + deployments get store-able boundary snapshots without setting

1890 + VLLM_PREFIX_CACHE_RETENTION_INTERVAL."""

1891 + hash_block_size = 2

1892 + block_size = 4

1893 + manager = make_kv_cache_manager(

1894 + kv_cache_config=_snapshot_offload_kv_cache_config(hash_block_size, block_size),

1895 + max_model_len=8192,

1896 + enable_caching=True,

1897 + hash_block_size=hash_block_size,

1898 + )

1899 + assert manager.coordinator.retention_interval is None

1900 +

1901 + req0 = make_request("0", [0] * 16, hash_block_size, sha256)

1902 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

1903 + assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None

1904 +

1905 + mamba_blocks = manager.coordinator.single_type_managers[1].req_to_blocks["0"]

1906 + ((group_id, block_id, boundary),) = drain_boundary_state_offloads(manager)["0"]

1907 + assert group_id == 1

1908 + assert boundary == 2 * block_size

1909 + assert block_id == mamba_blocks[1].block_id

1910 +

1911 + req0.num_computed_tokens = 8

1912 + manager.new_step_starts()

1913 + assert manager.allocate_slots(req0, 8) is not None

1914 + ((group_id, block_id2, boundary2),) = drain_boundary_state_offloads(manager)["0"]

1915 + assert group_id == 1

1916 + assert boundary2 == 4 * block_size

1917 + assert block_id2 == mamba_blocks[3].block_id

1918 +

1919 +

1920 +def _run_chunked_prefill(manager, req, chunk_size, num_chunks):

1921 + """Prefill ``req`` one ``chunk_size`` chunk per scheduler step, yielding the

1922 + boundary-state hand-offs offered (and claimed) at each step."""

1923 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req)

1924 + for step in range(num_chunks):

1925 + manager.new_step_starts()

1926 + req.num_computed_tokens = step * chunk_size

1927 + if step == 0:

1928 + allocated = manager.allocate_slots(

1929 + req, chunk_size, num_computed, computed_blocks

1930 + )

1931 + else:

1932 + allocated = manager.allocate_slots(req, chunk_size)

1933 + assert allocated is not None

1934 + yield drain_boundary_state_offloads(manager).get(req.request_id, [])

1935 +

1936 +

1937 +def test_boundary_state_offer_includes_more_mamba_groups_than_two():

1938 + """One offer batch carries one exact block entry per mamba group."""

1939 + block_size = 8

1940 + num_mamba_groups = 3

1941 + manager = make_kv_cache_manager(

1942 + kv_cache_config=KVCacheConfig(

1943 + num_blocks=200,

1944 + kv_cache_tensors=[],

1945 + kv_cache_groups=[

1946 + KVCacheGroupSpec(

1947 + ["full"],

1948 + FullAttentionSpec(

1949 + block_size=block_size,

1950 + num_kv_heads=1,

1951 + head_size=1,

1952 + dtype=torch.float32,

1953 + ),

1954 + ),

1955 + *(

1956 + KVCacheGroupSpec(

1957 + [f"mamba{i}"],

1958 + MambaSpec(

1959 + block_size=block_size,

1960 + shapes=(1, 1),

1961 + dtypes=(torch.float32,),

1962 + mamba_cache_mode="align",

1963 + ),

1964 + )

1965 + for i in range(num_mamba_groups)

1966 + ),

1967 + ],

1968 + ),

1969 + max_model_len=8192,

1970 + enable_caching=True,

1971 + hash_block_size=block_size,

1972 + )

1973 +

1974 + req0 = make_request("0", [0] * (2 * block_size), block_size, sha256)

1975 + entries = next(_run_chunked_prefill(manager, req0, block_size, 1))

1976 +

1977 + assert [group_id for group_id, _, _ in entries] == [1, 2, 3]

1978 + assert {boundary for _, _, boundary in entries} == {block_size}

1979 +

1980 +

1981 +def test_boundary_states_offered_past_prompt_for_resumed_prefill():

1982 + """The core offers every committed boundary, including past

1983 + ``num_prompt_tokens``. A resumed request re-prefills its generated tokens

1984 + and every group re-saves them, so filtering on the original prompt length

1985 + would silently strip the mamba key for boundaries full attention still

1986 + stores; only the connector knows where its save window ends."""

1987 + block_size = 8

1988 + prompt_len = block_size

1989 + manager = make_kv_cache_manager(

1990 + kv_cache_config=_snapshot_offload_kv_cache_config(

1991 + hash_block_size=block_size, block_size=block_size, num_blocks=200

1992 + ),

1993 + max_model_len=8192,

1994 + enable_caching=True,

1995 + hash_block_size=block_size,

1996 + )

1997 +

1998 + req0 = make_request("0", [0] * prompt_len, block_size, sha256)

1999 + assert [

2000 + boundary

2001 + for entries in _run_chunked_prefill(manager, req0, block_size, 1)

2002 + for _, _, boundary in entries

2003 + ] == [prompt_len]

2004 +

2005 + # Generate past the next mamba boundary, then replay it as a resumed

2006 + # prefill: its committed boundary is offered even though it is past the

2007 + # prompt, matching what full attention saves for the same range.

2008 + for i in range(block_size):

2009 + manager.new_step_starts()

2010 + req0.num_computed_tokens = prompt_len + i

2011 + req0.append_output_token_ids([1])

2012 + assert manager.allocate_slots(req0, 1) is not None

2013 + drain_boundary_state_offloads(manager)

2014 + assert req0.num_tokens == 2 * block_size

2015 +

2016 + manager.free(req0)

2017 + manager.new_step_starts()

2018 + req0.num_computed_tokens = 0

2019 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)

2020 + assert manager.allocate_slots(req0, req0.num_tokens - num_computed, num_computed)

2021 + offered = [b for _, _, b in drain_boundary_state_offloads(manager).get("0", [])]

2022 + assert 2 * block_size in offered

2023 + assert req0.num_prompt_tokens < 2 * block_size

tests/v1/core/test_single_type_kv_cache_manager.py

+32

@@ -15,12 +15,14 @@ from vllm.v1.core.kv_cache_utils import (

15 from vllm.v1.core.single_type_kv_cache_manager import (

16 ChunkedLocalAttentionManager,

17 FullAttentionManager,

18 + MambaManager,

19 RSWAManager,

20 SlidingWindowManager,

21 )

22 from vllm.v1.kv_cache_interface import (

23 ChunkedLocalAttentionSpec,

24 FullAttentionSpec,

25 + MambaSpec,

26 RSWASpec,

27 SlidingWindowSpec,

28 )

@@ -69,6 +71,36 @@ def test_external_computed_blocks_do_not_corrupt_free_pool():

71 assert len(manager.req_to_blocks[request_id]) == 3

74 +def test_mamba_speculative_block_relocation_requires_exclusive_ownership():

75 + spec = MambaSpec(

76 + block_size=4,

77 + shapes=((1, 1),),

78 + dtypes=(torch.float32,),

79 + mamba_cache_mode="align",

80 + num_speculative_blocks=1,

81 + )

82 + block_pool = BlockPool(num_gpu_blocks=4, enable_caching=True, hash_block_size=4)

83 + manager = MambaManager(

84 + spec,

85 + block_pool=block_pool,

86 + enable_caching=True,

87 + kv_cache_group_id=0,

88 + scheduler_block_size=4,

89 + )

90 + block = block_pool.get_new_blocks(1)[0]

91 + blocks = [block]

92 +

93 + manager._relocate_speculative_block(blocks, 0)

94 +

95 + assert blocks == [block_pool.null_block, block]

96 + assert block.ref_cnt == 1

97 +

98 + pinned_block = block_pool.get_new_blocks(1)[0]

99 + block_pool.touch((pinned_block,))

100 + with pytest.raises(AssertionError, match="exclusively owned and unhashed"):

101 + manager._relocate_speculative_block([pinned_block], 0)

102 +

103 +

104 def get_sliding_window_manager(

105 sliding_window_spec,

106 block_pool,

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+90 -4

@@ -36,7 +36,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (

36 is_store_reachable_swa_chunk,

37 )

38 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

39 -from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock

39 +from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock, KVCacheBlockCopy

40 +from vllm.v1.core.sched.output import (

41 + KVConnectorBlockState,

42 + SchedulerOutput,

43 +)

44 from vllm.v1.kv_cache_interface import (

45 ChunkedLocalAttentionSpec,

46 FullAttentionSpec,

@@ -107,11 +111,16 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():

111 req_status = scheduler._req_status["req"]

112 req_status.group_states[0].block_ids[:] = [11, 12]

113 req_status.group_states[1].block_ids[:] = [0, 21]

110 - scheduler.manager.prepare_store.side_effect = (

111 - lambda keys, req_context: generate_store_output(keys)

114 + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (

115 + generate_store_output(keys)

116 )

117

114 - output = SimpleNamespace(partial_tail_offloads={"req": [(1, 99, 28)]})

118 + output = SimpleNamespace(

119 + kv_connector_block_state=KVConnectorBlockState(

120 + block_ids={},

121 + boundary_state_offloads={"req": [(1, 99, 28)]},

122 + )

123 + )

124 jobs = scheduler._build_partial_tail_store_jobs(output)

125

126 assert len(jobs) == 1

@@ -155,6 +164,83 @@ def test_partial_tail_store_uses_attention_and_recurrent_cow_sources():

164 assert recurrent_event.parent_block_hash is None

165

166

167 +def test_aligned_boundary_store_uses_exact_source_with_partial_tail():

168 + scheduler = _make_partial_tail_scheduler()

169 + _make_partial_tail_request(scheduler)

170 + req_status = scheduler._req_status["req"]

171 + req_status.group_states[0].block_ids[:] = [11, 12]

172 + req_status.group_states[1].block_ids[:] = [0, 21]

173 + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (

174 + generate_store_output(keys)

175 + )

176 +

177 + output = SimpleNamespace(

178 + kv_connector_block_state=KVConnectorBlockState(

179 + block_ids={},

180 + boundary_state_offloads={"req": [(1, 98, 16), (1, 99, 28)]},

181 + )

182 + )

183 + jobs = scheduler._build_partial_tail_store_jobs(output)

184 +

185 + assert len(jobs) == 2

186 + src_spec = next(

187 + job.src_spec for job in jobs.values() if len(job.src_spec.block_ids) == 1

188 + )

189 + assert isinstance(src_spec, GPULoadStoreSpec)

190 + assert src_spec.block_ids.tolist() == [98]

191 + assert src_spec.group_sizes == [0, 1]

192 + assert src_spec.block_indices == [0, 0]

193 +

194 +

195 +def test_aligned_boundary_store_flushes_before_cow_destination_reuse():

196 + scheduler = _make_partial_tail_scheduler()

197 + _make_partial_tail_request(scheduler)

198 + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (

199 + generate_store_output(keys)

200 + )

201 +

202 + output = SchedulerOutput.make_empty()

203 + output.kv_connector_block_state = KVConnectorBlockState(

204 + block_ids={},

205 + boundary_state_offloads={"req": [(1, 99, 16)]},

206 + )

207 + meta = scheduler.build_connector_meta(output)

208 + [job_id] = meta.store_jobs

209 +

210 + output = SchedulerOutput.make_empty()

211 + output.kv_cache_block_copies = [KVCacheBlockCopy(98, 99)]

212 + meta = scheduler.build_connector_meta(output)

213 +

214 + assert meta.jobs_to_flush == {job_id}

215 +

216 +

217 +def test_normal_store_excludes_align_mode_mamba_sources():

218 + scheduler = _make_partial_tail_scheduler()

219 + request = _make_partial_tail_request(scheduler)

220 + request.num_computed_tokens = 0

221 + request.status = RequestStatus.RUNNING

222 + req_status = scheduler._req_status["req"]

223 + req_status.group_states[0].block_ids[:] = [11]

224 + req_status.group_states[1].block_ids[:] = [99]

225 + req_status.update_offload_keys()

226 + scheduler.manager.prepare_store.side_effect = lambda keys, req_context: (

227 + generate_store_output(keys)

228 + )

229 +

230 + output = SimpleNamespace(

231 + num_scheduled_tokens={"req": 16},

232 + finished_req_ids=set(),

233 + )

234 + jobs = scheduler._build_store_jobs(output)

235 +

236 + assert len(jobs) == 1

237 + [job] = jobs.values()

238 + src_spec = job.src_spec

239 + assert isinstance(src_spec, GPULoadStoreSpec)

240 + assert src_spec.block_ids.tolist() == [11]

241 + assert src_spec.group_sizes == [1, 0]

242 +

243 +

244 def test_partial_lookup_returns_exact_boundary_and_group_load_keys():

245 scheduler = _make_partial_tail_scheduler()

246 request = _make_partial_tail_request(scheduler)

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

+37 -1

@@ -5,6 +5,9 @@ import threading

5 import time

6 from unittest.mock import MagicMock, patch

8 +import pytest

9 +import torch

10 +

11 from vllm.config import set_current_vllm_config

12 from vllm.distributed.kv_events import BlockStored

13 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

@@ -31,6 +34,7 @@ from vllm.v1.kv_cache_interface import (

34 KVCacheConfig,

35 KVCacheGroupSpec,

36 KVCacheTensor,

37 + MambaSpec,

38 )

39 from vllm.v1.outputs import KVConnectorOutput

@@ -62,6 +66,30 @@ def _make_kv_cache_config() -> KVCacheConfig:

66 )

69 +def test_scheduler_requires_align_mode_for_mamba():

70 + vllm_config = _make_vllm_config()

71 + mamba_spec = MambaSpec(

72 + block_size=16,

73 + shapes=((1, 1),),

74 + dtypes=(torch.float32,),

75 + mamba_cache_mode="none",

76 + )

77 + kv_cache_config = KVCacheConfig(

78 + num_blocks=4,

79 + kv_cache_tensors=[],

80 + kv_cache_groups=[KVCacheGroupSpec(["mamba"], mamba_spec)],

81 + )

82 +

83 + with (

84 + patch(

85 + "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."

86 + "scheduler.LookupKeyClient"

87 + ),

88 + pytest.raises(AssertionError, match="requires mamba_cache_mode='align'"),

89 + ):

90 + scheduler.MooncakeStoreScheduler(vllm_config, kv_cache_config)

91 +

92 +

93 def _make_block_stored() -> BlockStored:

94 return BlockStored(

95 block_hashes=[b"hash"],

@@ -97,6 +125,10 @@ def test_scheduler_role_initializes_store_scheduler_only():

125 mock_worker.assert_not_called()

126 assert connector.connector_scheduler is mock_scheduler.return_value

127 assert connector.connector_worker is None

128 + block_pool = MagicMock()

129 +

130 + connector.bind_gpu_block_pool(block_pool)

131 + mock_scheduler.return_value.bind_gpu_block_pool.assert_called_once_with(block_pool)

132

133

134 def test_worker_methods_delegate_to_store_worker():

@@ -241,9 +273,13 @@ def test_update_connector_output_and_take_events():

273

274 kv_events = mooncake_store_connector.MooncakeStoreKVEvents(num_workers=1)

275 kv_events.add_events([event])

244 - connector.update_connector_output(KVConnectorOutput(kv_cache_events=kv_events))

276 + output = KVConnectorOutput(kv_cache_events=kv_events)

277 + connector.update_connector_output(output)

278

279 assert connector._kv_cache_events is kv_events

280 + connector.connector_scheduler.update_connector_output.assert_called_once_with(

281 + output

282 + )

283 assert list(connector.take_events()) == [event]

284 assert connector._kv_cache_events is None

285

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

+21

@@ -451,6 +451,27 @@ def test_store_mask_retention_prefix_stable_as_aligned_length_grows():

451 assert longer[: len(shorter)] == shorter

452

453

454 +def test_store_mask_excludes_mamba_groups_lookup_unaffected():

455 + """Align-mode mamba block tables are not append-only (interior states are

456 + nulled/freed; speculative blocks relocate), so the positional normal save

457 + must never cover mamba chunks — regardless of retention. Lookups still

458 + probe mamba boundaries: their keys come from the pinned snapshot/CoW

459 + hand-off path instead."""

460 + groups = [

461 + KVCacheGroupSpec(["L0"], _full(16)),

462 + KVCacheGroupSpec(["L1"], _mamba_align(16)),

463 + ]

464 + coord = _make_coord(groups, hash_block_size=16)

465 + masks = coord.store_mask(64)

466 + assert masks[0] is None # full attn stays dense

467 + assert masks[1] == [False] * 4

468 +

469 + coord = _make_coord(groups, hash_block_size=16, retention_interval=32)

470 + assert coord.store_mask(64)[1] == [False] * 4

471 +

472 + assert coord.lookup_mask(64)[1] is None

473 +

474 +

475 # ----- Eagle / MTP interaction with load_mask -----

476

477

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

+6 -5

@@ -236,6 +236,7 @@ def test_e2e_swa_plus_full_save_then_lookup_hits():

236 )

237

238 hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]

239 + # 1-based: block 0 is the reserved null block and the save skips it.

240 save_req = ReqMeta(

241 req_id="r0",

242 token_len_chunk=64,

@@ -441,10 +442,10 @@ def test_sub_block_partial_tail_offload_reads_cow_block():

442 block_hashes=hs,

443 can_save=True,

444 num_prompt_tokens=20,

444 - partial_tail_offloads=[(1, mamba_cow_block, 12)],

445 + boundary_state_offloads=[(1, mamba_cow_block, 12)],

446 )

447

447 - send._maybe_offload_partial_tail(req)

448 + send._maybe_offload_boundary_states(req)

449

450 # boundary = 12 // 4 * 4 = 12 -> keyed by hs[12 // 4 - 1] = hs[2].

451 partial_hash = hs[2]

@@ -512,8 +513,8 @@ def test_offload_syncs_event_before_put():

513 block_hashes=hs,

514 can_save=True,

515 num_prompt_tokens=12,

515 - partial_tail_offloads=[(1, 7, 12)],

516 store_job_id=1,

517 + boundary_state_offloads=[(1, 7, 12)],

518 )

519 req.current_event = event

520

@@ -589,10 +590,10 @@ def test_sub_block_partial_tail_offload_covers_smaller_group_blocks():

590 block_hashes=hs,

591 can_save=True,

592 num_prompt_tokens=12,

592 - partial_tail_offloads=[(1, mamba_cow_block, 12)],

593 + boundary_state_offloads=[(1, mamba_cow_block, 12)],

594 )

595

595 - send._maybe_offload_partial_tail(req)

596 + send._maybe_offload_boundary_states(req)

597

598 # FA (block 4): full blocks ending at 4, 8 and 12, keyed by their normal

599 # block-end hashes; mamba (block 16): the partial boundary block under

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

+276 -8

@@ -16,6 +16,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor

16 MooncakeStoreScheduler,

17 )

18 from vllm.v1.core.block_pool import BlockPool

19 +from vllm.v1.core.sched.output import KVConnectorBlockState

22 def _make_bare_scheduler(

@@ -47,9 +48,20 @@ def _make_bare_scheduler(

48 scheduler._num_workers = 1

49 scheduler._next_store_job_id = 0

50 scheduler._pinned_saves = {}

51 + scheduler._boundary_state_group_ids = frozenset({1})

52 return scheduler

55 +def _make_connector_block_state(

56 + block_ids: tuple[list[int], ...] | None = None,

57 + offloads: list[tuple[int, int, int]] | None = None,

58 +) -> KVConnectorBlockState:

59 + return KVConnectorBlockState(

60 + block_ids={} if block_ids is None else {"req-0": block_ids},

61 + boundary_state_offloads=({} if offloads is None else {"req-0": offloads}),

62 + )

63 +

64 +

65 def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None):

66 return SimpleNamespace(

67 finished_req_ids=set(),

@@ -65,6 +77,7 @@ def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None):

77 scheduled_spec_decode_tokens=(

78 {"req-0": scheduled_spec_tokens} if scheduled_spec_tokens else {}

79 ),

80 + kv_connector_block_state=_make_connector_block_state(block_ids=([0, 1, 2],)),

81 )

@@ -84,6 +97,7 @@ def _make_decode_scheduler_output(

97 ),

98 num_scheduled_tokens={"req-0": num_scheduled_tokens},

99 scheduled_spec_decode_tokens={},

100 + kv_connector_block_state=_make_connector_block_state(block_ids=([0, 1, 2],)),

101 )

102

103

@@ -109,6 +123,7 @@ def _make_new_scheduler_output() -> SimpleNamespace:

123 ),

124 num_scheduled_tokens={"req-0": 32},

125 scheduled_spec_decode_tokens={},

126 + kv_connector_block_state=_make_connector_block_state(block_ids=([0, 1],)),

127 )

128

129

@@ -491,6 +506,7 @@ def _make_resumed_scheduler_output(*, num_scheduled_tokens: int) -> SimpleNamesp

506 ),

507 num_scheduled_tokens={"req-0": num_scheduled_tokens},

508 scheduled_spec_decode_tokens={},

509 + kv_connector_block_state=_make_connector_block_state(block_ids=([0, 1, 2],)),

510 )

511

512

@@ -915,7 +931,10 @@ def _add_pending_partial_tail_request(

931 ),

932 num_scheduled_tokens={},

933 scheduled_spec_decode_tokens={},

918 - partial_tail_offloads={"req-0": [(1, 7, 12)]},

934 + kv_connector_block_state=_make_connector_block_state(

935 + block_ids=block_ids,

936 + offloads=[(1, 7, 12)],

937 + ),

938 )

939

940

@@ -940,15 +959,262 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():

959 assert req_meta.req_id == "req-0"

960 assert req_meta.can_save is True

961 assert req_meta.token_len_chunk == 0

943 - assert req_meta.partial_tail_offloads == [(1, 7, 12)]

962 + assert req_meta.boundary_state_offloads == [(1, 7, 12)]

963 assert req_meta.num_prompt_tokens == 12

964 assert req_meta.block_ids == ([0],)

965 + store_job_id = req_meta.store_job_id

966 + assert scheduler._pinned_saves[store_job_id][0] == [7]

967 + assert scheduler._gpu_block_pool.blocks[0].ref_cnt == 0

968 + assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 1

969 tracker = scheduler._request_trackers["req-0"]

970 assert tracker.num_saved_tokens == 0

971 assert tracker.has_pending_offload is True

972

973

951 -def test_resumed_partial_tail_uses_handoff_boundary():

974 +def test_decode_boundary_state_offload_dropped_unclaimed():

975 + # A hand-off past the prefill end can never complete a joint hybrid hit

976 + # (every other group stops saving there), so it is neither transferred nor

977 + # claimed — leaving the core free to release the block immediately.

978 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

979 + request = SimpleNamespace(

980 + all_token_ids=list(range(24)),

981 + block_hashes=[b"h0", b"h1", b"h2", b"h3", b"h4", b"h5"],

982 + num_output_placeholders=0,

983 + num_prompt_tokens=12,

984 + )

985 + scheduler._unfinished_requests["req-0"] = (request, ([0],))

986 + scheduler._request_trackers["req-0"] = RequestTracker(

987 + req_id="req-0",

988 + token_len=12,

989 + allocated_block_ids=([0],),

990 + num_saved_tokens=12,

991 + token_ids=list(range(12)),

992 + prefill_end_tokens=12,

993 + )

994 +

995 + out = SimpleNamespace(

996 + finished_req_ids=set(),

997 + preempted_req_ids=set(),

998 + scheduled_new_reqs=[],

999 + scheduled_cached_reqs=SimpleNamespace(

1000 + req_ids=[],

1001 + new_block_ids=[],

1002 + num_computed_tokens=[],

1003 + resumed_req_ids=set(),

1004 + ),

1005 + num_scheduled_tokens={},

1006 + scheduled_spec_decode_tokens={},

1007 + # Boundary 16 is a decode boundary for a 12-token prefill.

1008 + kv_connector_block_state=_make_connector_block_state(offloads=[(1, 7, 16)]),

1009 + )

1010 +

1011 + meta = scheduler.build_connector_meta(out)

1012 +

1013 + assert meta.requests == []

1014 + assert scheduler._pinned_saves == {}

1015 + assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0

1016 + assert scheduler._request_trackers["req-0"].has_pending_offload is False

1017 +

1018 +

1019 +def _register_offload_request(scheduler, *, prefill_end_tokens, num_prompt_tokens):

1020 + request = SimpleNamespace(

1021 + all_token_ids=list(range(64)),

1022 + block_hashes=[bytes([i]) for i in range(16)],

1023 + num_output_placeholders=0,

1024 + num_prompt_tokens=num_prompt_tokens,

1025 + )

1026 + scheduler._unfinished_requests["req-0"] = (request, ([0],))

1027 + scheduler._request_trackers["req-0"] = RequestTracker(

1028 + req_id="req-0",

1029 + token_len=prefill_end_tokens,

1030 + allocated_block_ids=([0],),

1031 + num_saved_tokens=prefill_end_tokens,

1032 + token_ids=list(range(prefill_end_tokens)),

1033 + prefill_end_tokens=prefill_end_tokens,

1034 + )

1035 +

1036 +

1037 +def _make_offload_only_output(entries, block_ids=([0],)):

1038 + return SimpleNamespace(

1039 + finished_req_ids=set(),

1040 + preempted_req_ids=set(),

1041 + scheduled_new_reqs=[],

1042 + scheduled_cached_reqs=SimpleNamespace(

1043 + req_ids=[],

1044 + new_block_ids=[],

1045 + num_computed_tokens=[],

1046 + resumed_req_ids=set(),

1047 + ),

1048 + num_scheduled_tokens={},

1049 + scheduled_spec_decode_tokens={},

1050 + kv_connector_block_state=_make_connector_block_state(

1051 + block_ids=block_ids,

1052 + offloads=entries,

1053 + ),

1054 + )

1055 +

1056 +

1057 +def test_resumed_prefill_claims_boundaries_past_prompt_length():

1058 + # A resumed request re-prefills its previously generated tokens, so its

1059 + # save window (`prefill_end_tokens`) extends past `num_prompt_tokens`.

1060 + # Boundaries in that range must still be claimed, or the mamba key would be

1061 + # missing for boundaries full attention does store and no joint hit could

1062 + # complete there.

1063 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1064 + _register_offload_request(scheduler, prefill_end_tokens=20, num_prompt_tokens=12)

1065 + out = _make_offload_only_output([(1, 7, 16), (1, 9, 20), (1, 11, 24)])

1066 +

1067 + meta = scheduler.build_connector_meta(out)

1068 +

1069 + # 16 and 20 are inside the resumed prefill; 24 is past it.

1070 + assert meta.requests[0].boundary_state_offloads == [(1, 7, 16), (1, 9, 20)]

1071 + store_job_id = meta.requests[0].store_job_id

1072 + assert scheduler._pinned_saves[store_job_id][0] == [7, 9]

1073 +

1074 +

1075 +def test_boundary_state_job_pins_exact_blocks_once():

1076 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1077 + _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64)

1078 + scheduler._request_trackers["req-0"].allocated_block_ids = ([7],)

1079 + out = _make_offload_only_output(

1080 + [(1, 7, 16), (1, 8, 32), (1, 9, 48)], block_ids=([7],)

1081 + )

1082 +

1083 + meta = scheduler.build_connector_meta(out)

1084 +

1085 + req_meta = meta.requests[0]

1086 + assert req_meta.boundary_state_offloads == [

1087 + (1, 7, 16),

1088 + (1, 8, 32),

1089 + (1, 9, 48),

1090 + ]

1091 + store_job_id = req_meta.store_job_id

1092 + assert scheduler._pinned_saves[store_job_id][0] == [7, 8, 9]

1093 + assert [scheduler._gpu_block_pool.blocks[i].ref_cnt for i in (7, 8, 9)] == [

1094 + 1,

1095 + 1,

1096 + 1,

1097 + ]

1098 +

1099 + scheduler.update_connector_output(_make_worker_output({store_job_id: 1}))

1100 + assert [scheduler._gpu_block_pool.blocks[i].ref_cnt for i in (7, 8, 9)] == [

1101 + 0,

1102 + 0,

1103 + 0,

1104 + ]

1105 +

1106 +

1107 +def test_store_job_pins_current_non_null_non_mamba_blocks():

1108 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1109 + request = SimpleNamespace(

1110 + all_token_ids=list(range(48)),

1111 + block_hashes=[bytes([i]) for i in range(12)],

1112 + num_output_placeholders=0,

1113 + )

1114 + scheduler._unfinished_requests["req-0"] = (request, ([7, 2, 0], [21, 0, 22]))

1115 + scheduler._request_trackers["req-0"] = RequestTracker(

1116 + req_id="req-0",

1117 + token_len=44,

1118 + allocated_block_ids=([7, 2, 0], [21, 0, 22]),

1119 + num_saved_tokens=32,

1120 + token_ids=list(range(44)),

1121 + prefill_end_tokens=48,

1122 + )

1123 + stale_block_ids = ([7, 2, 0, 5], [21, 0, 22, 23])

1124 + current_block_ids = ([7, 2, 0, 6], [24, 0, 25, 26])

1125 + out = SimpleNamespace(

1126 + finished_req_ids=set(),

1127 + preempted_req_ids=set(),

1128 + scheduled_new_reqs=[],

1129 + scheduled_cached_reqs=SimpleNamespace(

1130 + req_ids=["req-0"],

1131 + new_block_ids=[([5], [23])],

1132 + num_computed_tokens=[44],

1133 + resumed_req_ids=set(),

1134 + ),

1135 + num_scheduled_tokens={"req-0": 4},

1136 + scheduled_spec_decode_tokens={},

1137 + kv_connector_block_state=_make_connector_block_state(

1138 + block_ids=current_block_ids,

1139 + offloads=[(1, 8, 16)],

1140 + ),

1141 + )

1142 + pool = scheduler._gpu_block_pool

1143 +

1144 + meta = scheduler.build_connector_meta(out)

1145 +

1146 + assert scheduler._request_trackers["req-0"].allocated_block_ids == stale_block_ids

1147 + req_meta = meta.requests[0]

1148 + assert req_meta.block_ids == current_block_ids

1149 + store_job_id = req_meta.store_job_id

1150 + assert scheduler._pinned_saves[store_job_id][0] == [8, 7, 2, 6]

1151 + assert pool.blocks[0].ref_cnt == 0

1152 + assert pool.blocks[5].ref_cnt == 0

1153 + assert [pool.blocks[i].ref_cnt for i in (21, 22, 23, 24, 25, 26)] == [

1154 + 0,

1155 + 0,

1156 + 0,

1157 + 0,

1158 + 0,

1159 + 0,

1160 + ]

1161 + assert [pool.blocks[i].ref_cnt for i in (8, 7, 2, 6)] == [1, 1, 1, 1]

1162 +

1163 + scheduler.update_connector_output(_make_worker_output({store_job_id: 1}))

1164 +

1165 + assert [pool.blocks[i].ref_cnt for i in (8, 7, 2, 6)] == [0, 0, 0, 0]

1166 +

1167 +

1168 +def test_boundary_state_release_is_per_store_job():

1169 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1170 + _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64)

1171 + first = scheduler.build_connector_meta(_make_offload_only_output([(1, 7, 16)]))

1172 + second = scheduler.build_connector_meta(_make_offload_only_output([(1, 8, 32)]))

1173 + first_job_id = first.requests[0].store_job_id

1174 + second_job_id = second.requests[0].store_job_id

1175 +

1176 + scheduler.update_connector_output(_make_worker_output({first_job_id: 1}))

1177 +

1178 + assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0

1179 + assert scheduler._gpu_block_pool.blocks[8].ref_cnt == 1

1180 + assert first_job_id not in scheduler._pinned_saves

1181 + assert second_job_id in scheduler._pinned_saves

1182 +

1183 +

1184 +def test_preemption_and_request_id_reuse_do_not_release_inflight_job():

1185 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1186 + _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64)

1187 + first = scheduler.build_connector_meta(_make_offload_only_output([(1, 7, 16)]))

1188 + first_job_id = first.requests[0].store_job_id

1189 +

1190 + scheduler.build_connector_meta(_make_preemption_scheduler_output())

1191 + assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 1

1192 + assert first_job_id in scheduler._pinned_saves

1193 +

1194 + _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64)

1195 + second = scheduler.build_connector_meta(_make_offload_only_output([(1, 8, 32)]))

1196 + second_job_id = second.requests[0].store_job_id

1197 + assert second.requests[0].boundary_state_offloads == [(1, 8, 32)]

1198 +

1199 + scheduler.update_connector_output(_make_worker_output({first_job_id: 1}))

1200 + assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0

1201 + assert scheduler._gpu_block_pool.blocks[8].ref_cnt == 1

1202 + assert second_job_id in scheduler._pinned_saves

1203 +

1204 +

1205 +def test_boundary_state_never_claimed_without_a_send_thread():

1206 + # A kv_consumer has no store job that can acknowledge the block reference.

1207 + scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1208 + scheduler.kv_role = "kv_consumer"

1209 + _register_offload_request(scheduler, prefill_end_tokens=64, num_prompt_tokens=64)

1210 + out = _make_offload_only_output([(1, 7, 16)])

1211 +

1212 + assert scheduler.build_connector_meta(out).requests == []

1213 + assert scheduler._pinned_saves == {}

1214 + assert scheduler._gpu_block_pool.blocks[7].ref_cnt == 0

1215 +

1216 +

1217 +def test_resumed_partial_tail_uses_exact_boundary():

1218 scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1219 # Resumption replays prompt + previously generated tokens.

1220 out = _add_pending_partial_tail_request(

@@ -961,7 +1227,7 @@ def test_resumed_partial_tail_uses_handoff_boundary():

1227 meta = scheduler.build_connector_meta(out)

1228

1229 assert len(meta.requests) == 1

964 - assert meta.requests[0].partial_tail_offloads == [(1, 7, 12)]

1230 + assert meta.requests[0].boundary_state_offloads == [(1, 7, 12)]

1231 # Ordinary metadata retains the full resumed prefill range.

1232 assert meta.requests[0].num_prompt_tokens == 20

1233 tracker = scheduler._request_trackers["req-0"]

@@ -969,8 +1235,9 @@ def test_resumed_partial_tail_uses_handoff_boundary():

1235 assert tracker.has_pending_offload is True

1236

1237

972 -def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary():

1238 +def test_resumed_partial_tail_attached_to_save_keeps_exact_boundary():

1239 scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)

1240 + scheduler._boundary_state_group_ids = frozenset({0})

1241 request = SimpleNamespace(

1242 all_token_ids=list(range(48)),

1243 block_hashes=[b"h0", b"h1", b"h2"],

@@ -987,13 +1254,13 @@ def test_resumed_partial_tail_attached_to_save_keeps_handoff_boundary():

1254 prefill_end_tokens=48,

1255 )

1256 out = _make_scheduler_output(scheduled_spec_tokens=None)

990 - out.partial_tail_offloads = {"req-0": [(0, 7, 36)]}

1257 + out.kv_connector_block_state.boundary_state_offloads = {"req-0": [(0, 7, 36)]}

1258

1259 meta = scheduler.build_connector_meta(out)

1260

1261 assert len(meta.requests) == 1

1262 assert meta.requests[0].can_save is True

996 - assert meta.requests[0].partial_tail_offloads == [(0, 7, 36)]

1263 + assert meta.requests[0].boundary_state_offloads == [(0, 7, 36)]

1264 assert meta.requests[0].num_prompt_tokens == 48

1265 # Ordinary saving still covers the full resumed prefill range.

1266 tracker = scheduler._request_trackers["req-0"]

@@ -1019,7 +1286,8 @@ def test_partial_tail_cow_block_is_referenced_for_the_job():

1286 store_job_id = meta.requests[0].store_job_id

1287 # It leads the list, as in `pop_blocks_for_free`, so that the reversed free

1288 # puts it last in eviction priority.

1022 - assert scheduler._pinned_saves[store_job_id][0] == [7, 0]

1289 + assert scheduler._pinned_saves[store_job_id][0] == [7]

1290 + assert pool.blocks[0].ref_cnt == 0

1291 assert pool.blocks[7].ref_cnt == 1

1292

1293 scheduler.update_connector_output(_make_worker_output({store_job_id: 1}))

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

+266 -23

@@ -181,6 +181,7 @@ def _run_store_req(thread, req_meta: ReqMeta) -> None:

181

182

183 def _make_store_req(req_id: str, block_hashes: list[bytes]) -> ReqMeta:

184 + # 1-based: block 0 is the reserved null block and the save skips it.

185 return ReqMeta(

186 req_id=req_id,

187 token_len_chunk=32,

@@ -651,6 +652,7 @@ def _make_partial_tail_send_thread(

652 enable_partial_hash_hits=True,

653 hash_block_size=4,

654 lcm_block_size=16,

655 + mamba_group_ids={1},

656 )

657 db = ChunkedTokenDatabase(

658 KeyMetadata("test-model", 0, 0, 0, 0),

@@ -659,10 +661,18 @@ def _make_partial_tail_send_thread(

661 )

662 db.set_kv_caches_base_addr([0x1000])

663 db.set_block_len([256])

664 + # Group 1 models the mamba "align" group the hand-offs reference.

665 + db_mamba = ChunkedTokenDatabase(

666 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),

667 + block_size=16,

668 + hash_block_size=4,

669 + )

670 + db_mamba.set_kv_caches_base_addr([0x2000])

671 + db_mamba.set_block_len([256])

672 return _make_store_sending_thread(

673 store,

674 coord=coord,

665 - token_databases=[db],

675 + token_databases=[db, db_mamba],

676 replicate_config=replicate_config,

677 enable_group_semantics=enable_group_semantics,

678 supports_group_ids=supports_group_ids,

@@ -673,20 +683,20 @@ def _make_partial_tail_req(block_ids: list[int]) -> ReqMeta:

683 return ReqMeta(

684 req_id="req-a",

685 token_len_chunk=0,

676 - block_ids=(block_ids,),

686 + block_ids=(block_ids, [1]),

687 block_hashes=[b"a0", b"a1", b"a2"],

688 can_save=True,

679 - partial_tail_offloads=[(1, 7, 12)],

689 + boundary_state_offloads=[(1, 7, 12)],

690 )

691

692

693 def test_partial_tail_offload_skips_null_source_blocks():

694 store = MagicMock()

695 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

686 - store.batch_put_from_multi_buffers.return_value = [256, 256]

696 + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)

697 thread = _make_partial_tail_send_thread(store)

698

689 - assert thread._maybe_offload_partial_tail(_make_partial_tail_req([0, 2, 3]))

699 + assert thread._maybe_offload_boundary_states(_make_partial_tail_req([0, 2, 3]))

700

701 keys, addrs, _sizes, _replicate_config = (

702 store.batch_put_from_multi_buffers.call_args.args

@@ -694,15 +704,21 @@ def test_partial_tail_offload_skips_null_source_blocks():

704 assert keys == [

705 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",

706 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",

707 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6132",

708 + ]

709 + # FA reads blocks 2 and 3 (block 0 is null and skipped); the mamba

710 + # boundary block reads the core-provided CoW block 7, not block_ids.

711 + assert addrs == [

712 + [0x1000 + 2 * 256],

713 + [0x1000 + 3 * 256],

714 + [0x2000 + 7 * 256],

715 ]

698 - assert addrs == [[0x1000 + 2 * 256], [0x1000 + 3 * 256]]

716

717

718 def test_store_sending_thread_skips_null_sparse_group_blocks():

719 from vllm.v1.kv_cache_interface import (

720 FullAttentionSpec,

721 KVCacheGroupSpec,

705 - MambaSpec,

722 )

723

724 store = MagicMock()

@@ -711,16 +727,11 @@ def test_store_sending_thread_skips_null_sparse_group_blocks():

727 lambda keys, addrs, sizes, replicate_config: [256] * len(keys)

728 )

729 full = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

714 - mamba = MambaSpec(

715 - block_size=16,

716 - shapes=((1, 1),),

717 - dtypes=(torch.float32,),

718 - mamba_cache_mode="align",

719 - )

730 + sparse = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=64, dtype=None)

731 coord = mooncake_store_worker.MooncakeStoreCoordinator(

732 [

733 KVCacheGroupSpec(["full"], full),

723 - KVCacheGroupSpec(["mamba"], mamba),

734 + KVCacheGroupSpec(["sparse"], sparse),

735 ],

736 scheduler_block_size=16,

737 hash_block_size=16,

@@ -732,17 +743,17 @@ def test_store_sending_thread_skips_null_sparse_group_blocks():

743 )

744 db_full.set_kv_caches_base_addr([0x1000])

745 db_full.set_block_len([256])

735 - db_mamba = ChunkedTokenDatabase(

746 + db_sparse = ChunkedTokenDatabase(

747 KeyMetadata("test-model", 0, 0, 0, 0, group_id=1),

748 block_size=16,

749 )

739 - db_mamba.set_kv_caches_base_addr([0x2000])

740 - db_mamba.set_block_len([512])

750 + db_sparse.set_kv_caches_base_addr([0x2000])

751 + db_sparse.set_block_len([512])

752

753 thread = _make_store_sending_thread(

754 store,

755 coord=coord,

745 - token_databases=[db_full, db_mamba],

756 + token_databases=[db_full, db_sparse],

757 )

758 _run_store_req(

759 thread,

@@ -756,15 +767,86 @@ def test_store_sending_thread_skips_null_sparse_group_blocks():

767 )

768

769 keys, addrs, _, _ = store.batch_put_from_multi_buffers.call_args.args

759 - mamba_keys = [key for key in keys if "@group:1" in key]

760 - assert [key.rsplit("@", 1)[-1] for key in mamba_keys] == [b"a1".hex()]

770 + sparse_keys = [key for key in keys if "@group:1" in key]

771 + assert [key.rsplit("@", 1)[-1] for key in sparse_keys] == [b"a1".hex()]

772 assert all(addr[0] >= 0x2000 for key, addr in zip(keys, addrs) if "@group:1" in key)

773

774

775 +def test_partial_tail_offload_skips_cap_omitted_mamba_group():

776 + """A Mamba group omitted by the handoff cap must not fall back to the

777 + connector's positional block table."""

778 + from vllm.v1.kv_cache_interface import (

779 + FullAttentionSpec,

780 + KVCacheGroupSpec,

781 + MambaSpec,

782 + )

783 +

784 + store = MagicMock()

785 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

786 + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)

787 +

788 + full_spec = FullAttentionSpec(

789 + block_size=4, num_kv_heads=8, head_size=64, dtype=torch.float32

790 + )

791 + mamba_spec = MambaSpec(

792 + block_size=16,

793 + shapes=((1, 1),),

794 + dtypes=(torch.float32,),

795 + mamba_cache_mode="align",

796 + )

797 + coord = mooncake_store_worker.MooncakeStoreCoordinator(

798 + [

799 + KVCacheGroupSpec(["full"], full_spec),

800 + KVCacheGroupSpec(["mamba-1"], mamba_spec),

801 + KVCacheGroupSpec(["mamba-2"], mamba_spec),

802 + ],

803 + scheduler_block_size=16,

804 + hash_block_size=4,

805 + )

806 +

807 + def make_db(group_id: int, block_size: int, base_addr: int):

808 + db = ChunkedTokenDatabase(

809 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=group_id),

810 + block_size=block_size,

811 + hash_block_size=4,

812 + )

813 + db.set_kv_caches_base_addr([base_addr])

814 + db.set_block_len([256])

815 + return db

816 +

817 + db_full = make_db(0, 4, 0x1000)

818 + db_mamba_1 = make_db(1, 16, 0x2000)

819 + db_mamba_2 = make_db(2, 16, 0x3000)

820 + thread = _make_store_sending_thread(

821 + store,

822 + coord=coord,

823 + token_databases=[db_full, db_mamba_1, db_mamba_2],

824 + )

825 +

826 + req = ReqMeta(

827 + req_id="req-a",

828 + token_len_chunk=0,

829 + block_ids=([1, 2, 3], [5], [9]),

830 + block_hashes=[b"a0", b"a1", b"a2"],

831 + can_save=True,

832 + # The scheduler accepted group 1 and omitted group 2 at boundary 12.

833 + boundary_state_offloads=[(1, 7, 12)],

834 + )

835 + assert thread._maybe_offload_boundary_states(req)

836 +

837 + keys, addrs, _sizes, _replicate_config = (

838 + store.batch_put_from_multi_buffers.call_args.args

839 + )

840 + puts = list(zip(keys, addrs, strict=True))

841 + boundary_hash = BlockHash(b"a2")

842 + assert (db_mamba_1.key_for(boundary_hash), [0x2000 + 7 * 256]) in puts

843 + assert (db_mamba_2.key_for(boundary_hash), [0x3000 + 9 * 256]) not in puts

844 +

845 +

846 def test_partial_tail_offload_replaces_stale_group_ids_after_filtering():

847 store = MagicMock()

766 - store.batch_is_exist.return_value = [1, 0, 0]

767 - store.batch_put_from_multi_buffers.return_value = [256, 256]

848 + store.batch_is_exist.return_value = [1, 0, 0, 0]

849 + store.batch_put_from_multi_buffers.return_value = [256, 256, 256]

850 replicate_config = SimpleNamespace(group_ids=["stale"])

851 thread = _make_partial_tail_send_thread(

852 store,

@@ -773,17 +855,19 @@ def test_partial_tail_offload_replaces_stale_group_ids_after_filtering():

855 supports_group_ids=True,

856 )

857

776 - assert thread._maybe_offload_partial_tail(_make_partial_tail_req([1, 2, 3]))

858 + assert thread._maybe_offload_boundary_states(_make_partial_tail_req([1, 2, 3]))

859

860 keys, _addrs, _sizes, config = store.batch_put_from_multi_buffers.call_args.args

861 assert keys == [

862 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6131",

863 "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:0@6132",

864 + "test-model@tp_rank:0@pcp0@dcp0@pp_rank:0@group:1@6132",

865 ]

866 assert config is replicate_config

867 assert config.group_ids == [

868 "vllm-mooncake-store:test-model@6131",

869 "vllm-mooncake-store:test-model@6132",

870 + "vllm-mooncake-store:test-model@6132",

871 ]

872

873

@@ -815,6 +899,165 @@ def test_partial_tail_put_failure_activates_pressure_gate():

899 assert store.batch_put_from_multi_buffers.call_count == 1

900

901

902 +def test_normal_save_excludes_mamba_group_and_null_blocks():

903 + """The positional normal save must never cover mamba chunks (align-mode

904 + block tables are not append-only) and must skip chunks whose block is the

905 + reserved null block, for any group."""

906 + from vllm.v1.kv_cache_interface import (

907 + FullAttentionSpec,

908 + KVCacheGroupSpec,

909 + MambaSpec,

910 + )

911 +

912 + store = MagicMock()

913 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

914 + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)

915 +

916 + full_spec = FullAttentionSpec(

917 + block_size=16, num_kv_heads=8, head_size=64, dtype=None

918 + )

919 + mamba_spec = MambaSpec(

920 + block_size=16,

921 + shapes=((1, 1),),

922 + dtypes=(torch.float32,),

923 + mamba_cache_mode="align",

924 + )

925 + coord = mooncake_store_worker.MooncakeStoreCoordinator(

926 + [KVCacheGroupSpec(["L0"], full_spec), KVCacheGroupSpec(["L1"], mamba_spec)],

927 + scheduler_block_size=16,

928 + hash_block_size=16,

929 + )

930 + db_full = ChunkedTokenDatabase(

931 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=0), block_size=16

932 + )

933 + db_full.set_kv_caches_base_addr([0x1000])

934 + db_full.set_block_len([256])

935 + db_mamba = ChunkedTokenDatabase(

936 + KeyMetadata("test-model", 0, 0, 0, 0, group_id=1), block_size=16

937 + )

938 + db_mamba.set_kv_caches_base_addr([0x2000])

939 + db_mamba.set_block_len([256])

940 + thread = _make_store_sending_thread(

941 + store, coord=coord, token_databases=[db_full, db_mamba]

942 + )

943 +

944 + _run_store_req(

945 + thread,

946 + ReqMeta(

947 + req_id="r0",

948 + token_len_chunk=64,

949 + # FA has a null at chunk 1; the mamba table is align-mode shaped

950 + # (interior nulls + one state block) and must not be read at all.

951 + block_ids=([1, 0, 3, 4], [0, 0, 0, 9]),

952 + block_hashes=[b"a0", b"a1", b"a2", b"a3"],

953 + can_save=True,

954 + ),

955 + )

956 +

957 + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args

958 + assert all("@group:0@" in k for k in keys)

959 + assert [k.rsplit("@", 1)[-1] for k in keys] == ["6130", "6132", "6133"]

960 + assert addrs == [[0x1000 + 256], [0x1000 + 3 * 256], [0x1000 + 4 * 256]]

961 +

962 +

963 +def test_block_aligned_snapshot_offload_uses_provided_block():

964 + """A block-aligned hand-off (sparse-retention mamba checkpoint) uploads

965 + exactly the core-pinned block under the boundary-end hash key; the

966 + positional block table entry for that chunk must be ignored."""

967 + store = MagicMock()

968 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

969 + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)

970 + thread = _make_partial_tail_send_thread(store)

971 +

972 + hs = [bytes([i + 1]) * 4 for i in range(8)] # 8 hash units = 32 tokens

973 + req = ReqMeta(

974 + req_id="req-a",

975 + token_len_chunk=0,

976 + block_ids=([1, 2, 3], [5]),

977 + block_hashes=hs,

978 + can_save=True,

979 + boundary_state_offloads=[(1, 7, 32)],

980 + )

981 + assert thread._maybe_offload_boundary_states(req)

982 +

983 + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args

984 + # boundary 32 is block-aligned for the mamba group (block 16): one key,

985 + # keyed by hs[32 // 4 - 1], read from the handed-off block 7 — not from

986 + # block_ids[1] and with no FA gap coverage.

987 + assert keys == [thread.token_databases[1].key_for(BlockHash(hs[7]))]

988 + assert addrs == [[0x2000 + 7 * 256]]

989 +

990 +

991 +def test_mixed_snapshot_and_sub_block_offloads():

992 + """A retention snapshot and the prompt-end sub-block CoW tail can arrive

993 + in one hand-off; the sub-block path covers FA gap blocks but reads the

994 + mamba boundary only from the CoW block, never positionally."""

995 + store = MagicMock()

996 + store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

997 + store.batch_put_from_multi_buffers.side_effect = lambda keys, *a: [256] * len(keys)

998 + thread = _make_partial_tail_send_thread(store)

999 +

1000 + hs = [bytes([i + 1]) * 4 for i in range(11)] # 11 hash units = 44 tokens

1001 + req = ReqMeta(

1002 + req_id="req-a",

1003 + token_len_chunk=0,

1004 + block_ids=([1, 2, 3], [0, 0, 0]),

1005 + block_hashes=hs,

1006 + can_save=True,

1007 + boundary_state_offloads=[(1, 9, 32), (1, 7, 44)],

1008 + )

1009 + assert thread._maybe_offload_boundary_states(req)

1010 +

1011 + keys, addrs, _sizes, _ = store.batch_put_from_multi_buffers.call_args.args

1012 + db_full, db_mamba = thread.token_databases

1013 + assert keys == [

1014 + # Aligned snapshot: boundary 32 from the handed-off block 9.

1015 + db_mamba.key_for(BlockHash(hs[7])),

1016 + # Sub-block boundary 44: FA gap blocks ending at 4, 8, 12.

1017 + db_full.key_for(BlockHash(hs[0])),

1018 + db_full.key_for(BlockHash(hs[1])),

1019 + db_full.key_for(BlockHash(hs[2])),

1020 + # Mamba boundary block from the CoW hand-off (block 7).

1021 + db_mamba.key_for(BlockHash(hs[10])),

1022 + ]

1023 + assert addrs == [

1024 + [0x2000 + 9 * 256],

1025 + [0x1000 + 1 * 256],

1026 + [0x1000 + 2 * 256],

1027 + [0x1000 + 3 * 256],

1028 + [0x2000 + 7 * 256],

1029 + ]

1030 +

1031 +

1032 +def test_snapshot_offload_skips_null_handoff_block():

1033 + """A hand-off the core could not materialize (null block) carries no

1034 + committed state; persisting it would poison the boundary key."""

1035 + store = MagicMock()

1036 + thread = _make_partial_tail_send_thread(store)

1037 +

1038 + hs = [bytes([i + 1]) * 4 for i in range(8)]

1039 + assert thread._maybe_offload_boundary_states(

1040 + ReqMeta(

1041 + req_id="req-a",

1042 + token_len_chunk=0,

1043 + block_ids=([1, 2, 3], [5]),

1044 + block_hashes=hs,

1045 + can_save=True,

1046 + boundary_state_offloads=[(1, NULL_BLOCK_ID, 32)],

1047 + )

1048 + )

1049 + store.batch_is_exist.assert_not_called()

1050 + store.batch_put_from_multi_buffers.assert_not_called()

1051 +

1052 +

1053 +def test_worker_meta_is_absent_before_any_completed_job():

1054 + send_thread = _make_partial_tail_send_thread(MagicMock())

1055 + store_worker = object.__new__(mooncake_store_worker.MooncakeStoreWorker)

1056 + store_worker.kv_send_thread = send_thread

1057 +

1058 + assert store_worker.build_connector_worker_meta() is None

1059 +

1060 +

1061 def test_store_sending_thread_delta_start_rank_saves_second_local_chunk():

1062 store = MagicMock()

1063 store.batch_is_exist.side_effect = lambda keys: [0] * len(keys)

tests/v1/kv_connector/unit/test_offloading_connector.py

+14 -5

@@ -212,12 +212,21 @@ def _accuracy_test(llm: LLM, subscriber: MockSubscriber | None):

212 if subscriber is not None:

213 subscriber.get_new_cpu_stored_events()

214

215 - # Pad prompt so its token count is a multiple of cpu_block_size.

216 - # Use the tokenizer directly to avoid expensive llm.generate() calls.

215 + # Align the reusable prefix; the final token must be recomputed for logits.

216 tokenizer = llm.get_tokenizer()

218 - prompt = "Let's count to 10. One, two, three, four,"

219 - while len(tokenizer.encode(prompt)) % cpu_block_size != 0:

220 - prompt = ". " + prompt

217 + prompt_token_ids = tokenizer.encode("Let's count to 10. One, two, three, four,")

218 + reusable_prefix = prompt_token_ids[:-1]

219 + final_token = prompt_token_ids[-1:]

220 +

221 + [padding_token_id] = tokenizer.encode(".", add_special_tokens=False)

222 + padding = [padding_token_id] * (-len(reusable_prefix) % cpu_block_size)

223 + insert_at = int(reusable_prefix[0] == tokenizer.bos_token_id)

224 + reusable_prefix[insert_at:insert_at] = padding

225 + assert len(reusable_prefix) % cpu_block_size == 0

226 +

227 + prompt = TokensPrompt(

228 + prompt_token_ids=reusable_prefix + final_token,

229 + )

230

231 # Seed the CPU cache with the prompt.

232 llm.generate(prompt, sampling_params, use_tqdm=False)

tests/v1/kv_connector/unit/test_scheduler_kv_connector_override.py

+16 -1

@@ -22,8 +22,13 @@ from vllm.v1.request import Request

24 class DummyConnectorMetadata(KVConnectorMetadata):

25 - def __init__(self, block_hashes_by_req: dict[str, list[BlockHash]]):

25 + def __init__(

26 + self,

27 + block_hashes_by_req: dict[str, list[BlockHash]],

28 + block_ids_by_req: dict[str, tuple[list[int], ...]],

29 + ):

30 self.block_hashes_by_req = block_hashes_by_req

31 + self.block_ids_by_req = block_ids_by_req

34 class DummyKVConnector(KVConnectorBase_V1):

@@ -47,8 +52,12 @@ class DummyKVConnector(KVConnectorBase_V1):

52 assert block_hashes_by_req is not None, (

53 "DummyKVConnector expected 'block_hashes_by_req' on scheduler_output"

54 )

55 + block_state = scheduler_output.kv_connector_block_state

56 + assert block_state is not None

57 + block_ids_by_req = block_state.block_ids

58 return DummyConnectorMetadata(

59 block_hashes_by_req=block_hashes_by_req,

60 + block_ids_by_req=block_ids_by_req,

61 )

63 def start_load_kv(self, kv_caches, finished_req_ids):

@@ -129,3 +138,9 @@ def test_connector_receives_block_hashes(_load_plugin):

138 num_tokens // block_size

139 )

140 assert meta.block_hashes_by_req[req.request_id] == req.block_hashes

141 +

142 + expected_block_ids = {

143 + req.req_id: req.block_ids for req in output.scheduled_new_reqs

144 + }

145 + assert meta.block_ids_by_req == expected_block_ids

146 + assert output.kv_connector_block_state is None

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

+20

@@ -87,6 +87,11 @@ class MooncakeStoreCoordinator:

87 for g in kv_cache_groups

88 ), "scheduler_block_size must be a multiple of each group's block_size"

89 self.kv_cache_groups = kv_cache_groups

90 + self.mamba_group_ids = {

91 + group_id

92 + for group_id, group in enumerate(kv_cache_groups)

93 + if isinstance(_unwrap_spec(group.kv_cache_spec), MambaSpec)

94 + }

95 self.hash_block_size = hash_block_size

96 self.lcm_block_size = scheduler_block_size

97 self.enable_partial_hash_hits = partial_hash_hits_enabled(

@@ -205,12 +210,23 @@ class MooncakeStoreCoordinator:

210

211 Reuses the engine's ``SingleTypeKVCacheManager.reachable_block_mask``

212 so the store retains exactly the blocks the local prefix cache would.

213 +

214 + Mamba groups are always all-False: the normal save resolves blocks

215 + positionally from the connector's append-only block-ID snapshot, but

216 + an align-mode mamba block table is not append-only (interior state

217 + blocks are nulled/freed, and speculative decoding relocates the spec

218 + blocks in place), so a positional read may hit a null, freed, or live

219 + speculative-state block and persist wrong bytes under a valid prefix

220 + hash. Mamba state is persisted only through the connector-pinned exact

221 + block hand-off path

222 + (``SchedulerOutput.kv_connector_block_state.boundary_state_offloads``).

223 """

224 return self._reachable_masks(

225 aligned_token_len,

226 start_token,

227 retention_interval=self.retention_interval,

228 num_prompt_tokens=num_prompt_tokens,

229 + exclude_mamba=True,

230 )

231

232 def lookup_mask(

@@ -237,6 +253,7 @@ class MooncakeStoreCoordinator:

253 *,

254 retention_interval: int | None,

255 num_prompt_tokens: int | None,

256 + exclude_mamba: bool = False,

257 ) -> tuple[list[bool] | None, ...]:

258 mask_alignment = (

259 self.hash_block_size

@@ -252,6 +269,9 @@ class MooncakeStoreCoordinator:

269 spec = _unwrap_spec(g.kv_cache_spec)

270 end_chunk = aligned_token_len // spec.block_size

271 start_chunk = min(end_chunk, max(0, cdiv(start_token, spec.block_size)))

272 + if exclude_mamba and isinstance(spec, MambaSpec):

273 + masks.append([False] * (end_chunk - start_chunk))

274 + continue

275 manager_cls = KVCacheSpecRegistry.get_manager_class(spec)

276 assert manager_cls is not None

277 use_eagle = g_idx in self.eagle_group_ids

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/data.py

+5 -5

@@ -406,11 +406,11 @@ class ReqMeta:

406 # serve that purpose: it is reused once a preempted request resumes, so it

407 # would release the wrong job's blocks.

408 store_job_id: int | None = None

409 - # Core-provided per-mamba-group

410 - # (group_id, cow_block_id, boundary_tokens) for this request's partial tail.

411 - # Present only on the producer's CoW step; drives the connector's offload

412 - # (the FA group's block is derived from block_ids and boundary_tokens).

413 - partial_tail_offloads: list[tuple[int, int, int]] | None = None

409 + # Core-provided (group_id, block_id, boundary_tokens) mamba "align"

410 + # boundary states. A block-aligned entry is a committed boundary snapshot;

411 + # a non-aligned entry is the sub-block CoW tail. The store-job reference

412 + # keeps each exact block alive until every worker rank finishes the job.

413 + boundary_state_offloads: list[tuple[int, int, int]] | None = None

414

415 @staticmethod

416 def from_request_tracker(

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

+121 -52

@@ -23,11 +23,12 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.worker import (

23 LookupKeyClient,

24 )

25 from vllm.logger import init_logger

26 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

27 from vllm.v1.core.block_pool import BlockPool

28 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

29 from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes

30 from vllm.v1.core.sched.output import NewRequestData, SchedulerOutput

30 -from vllm.v1.kv_cache_interface import KVCacheConfig

31 +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec

32 from vllm.v1.outputs import KVConnectorOutput

33 from vllm.v1.request import Request

@@ -80,12 +81,15 @@ class MooncakeStoreScheduler:

81 self._hash_block_size,

82 vllm_config.parallel_config.decode_context_parallel_size,

83 )

83 -

84 - # Per-request state

85 - self.load_specs: dict[str, LoadSpec] = {} # to be loaded

86 - self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests

87 - self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}

88 - self._unfinished_request_ids: set[str] = set()

84 + mamba_groups = {

85 + group_id: group.kv_cache_spec

86 + for group_id, group in enumerate(kv_cache_config.kv_cache_groups)

87 + if isinstance(group.kv_cache_spec, MambaSpec)

88 + }

89 + assert all(

90 + spec.mamba_cache_mode == "align" for spec in mamba_groups.values()

91 + ), "MooncakeStoreScheduler requires mamba_cache_mode='align'"

92 + self._boundary_state_group_ids = frozenset(mamba_groups)

94 self._gpu_block_pool: BlockPool | None = None

95 self._num_workers = vllm_config.parallel_config.world_size

@@ -93,6 +97,12 @@ class MooncakeStoreScheduler:

97 # store_job_id -> (referenced block ids, ranks yet to report completion)

98 self._pinned_saves: dict[int, tuple[list[int], int]] = {}

100 + # Per-request state

101 + self.load_specs: dict[str, LoadSpec] = {} # to be loaded

102 + self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests

103 + self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}

104 + self._unfinished_request_ids: set[str] = set()

105 +

106 def bind_gpu_block_pool(self, gpu_block_pool: BlockPool) -> None:

107 self._gpu_block_pool = gpu_block_pool

108

@@ -381,47 +391,41 @@ class MooncakeStoreScheduler:

391 if req_meta is not None:

392 meta.add_request(req_meta)

393

384 - # Flush partial-tail offloads in the step they arrive: the CoW copy is

385 - # enqueued before the connector event records, so this step's event

386 - # fences the cow block. Ride the request's save meta when present, else

387 - # emit an offload-only ReqMeta (token_len_chunk=0 skips the normal

388 - # save; can_save=True takes the normal enqueue path).

389 - step_partial_tails = getattr(scheduler_output, "partial_tail_offloads", None)

390 - if step_partial_tails and not is_consumer:

391 - pending = dict(step_partial_tails)

392 - for req_meta in meta.requests:

393 - if req_meta.can_save:

394 - groups = pending.pop(req_meta.req_id, None)

395 - if groups:

396 - req_meta.partial_tail_offloads = groups

397 - tracker = self._request_trackers.get(req_meta.req_id)

398 - if tracker is not None:

399 - tracker.has_pending_offload = True

400 - for req_id, groups in pending.items():

401 - tracker = self._request_trackers.get(req_id)

402 - req_tuple = self._unfinished_requests.get(req_id)

403 - if tracker is None or req_tuple is None:

404 - # Request finished/preempted within this step; its blocks

405 - # are going away, so the offload is conservatively dropped.

406 - logger.debug("Dropping partial-tail offload for request %s", req_id)

407 - continue

408 - assert len({boundary for _, _, boundary in groups}) == 1

409 - tracker.has_pending_offload = True

410 - meta.add_request(

411 - ReqMeta(

412 - req_id=req_id,

413 - token_len_chunk=0,

414 - block_ids=tracker.allocated_block_ids,

415 - block_hashes=req_tuple[0].block_hashes,

416 - can_save=True,

417 - num_prompt_tokens=tracker.prefill_end_tokens,

418 - partial_tail_offloads=groups,

419 - )

420 - )

394 + block_state = getattr(scheduler_output, "kv_connector_block_state", None)

395 + if (

396 + block_state is not None

397 + and block_state.boundary_state_offloads

398 + and not is_consumer

399 + ):

400 + self._handle_boundary_state_offloads(

401 + block_state.boundary_state_offloads, meta

402 + )

403

404 + self._apply_current_save_block_ids(meta, scheduler_output)

405 self._reference_save_blocks(meta)

406 return meta

407

408 + def _apply_current_save_block_ids(

409 + self,

410 + meta: MooncakeStoreConnectorMetadata,

411 + scheduler_output: SchedulerOutput,

412 + ) -> None:

413 + """Replace append-only mirrors with the core's current block tables."""

414 + save_metas = [req_meta for req_meta in meta.requests if req_meta.can_save]

415 + if not save_metas:

416 + return

417 +

418 + block_state = scheduler_output.kv_connector_block_state

419 + assert block_state is not None, (

420 + "Current block tables are required for Mooncake store jobs"

421 + )

422 + for req_meta in save_metas:

423 + block_ids = block_state.block_ids.get(req_meta.req_id)

424 + assert block_ids is not None, (

425 + f"Missing current block table for store request {req_meta.req_id}"

426 + )

427 + req_meta.block_ids = block_ids

428 +

429 def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None:

430 """Take a GPU block reference for every store job this step emits.

431

@@ -439,21 +443,86 @@ class MooncakeStoreScheduler:

443 req_meta.store_job_id = store_job_id = self._next_store_job_id

444 self._next_store_job_id += 1

445 block_ids: list[int] = []

442 - if req_meta.partial_tail_offloads:

443 - # A partial-tail CoW block is deliberately kept out of the

444 - # request's block table, so it is absent from `block_ids` even

445 - # though the worker DMAs out of it just as asynchronously.

446 - # It leads the list, as in `pop_blocks_for_free`.

447 - block_ids += [bid for _, bid, _ in req_meta.partial_tail_offloads]

446 + if req_meta.boundary_state_offloads:

447 + block_ids.extend(

448 + block_id for _, block_id, _ in req_meta.boundary_state_offloads

449 + )

450 + assert NULL_BLOCK_ID not in block_ids, (

451 + "A null block cannot back a boundary-state offload"

452 + )

453 # Every allocated block is referenced, not just the ones covering

454 # this job's token range: a rank resumes from its own last

455 # successful offset, which lags the scheduler's whenever a save was

456 # skipped or failed, so it may read anywhere below the range.

452 - block_ids += [bid for group in req_meta.block_ids for bid in group]

457 + block_ids.extend(

458 + block_id

459 + for group_id, group in enumerate(req_meta.block_ids)

460 + if group_id not in self._boundary_state_group_ids

461 + for block_id in group

462 + if block_id != NULL_BLOCK_ID

463 + )

464 + # An aligned boundary block may also be present in the request's

465 + # block table. Take and release exactly one reference per block.

466 + block_ids = list(dict.fromkeys(block_ids))

467 + assert NULL_BLOCK_ID not in block_ids

468 if not block_ids:

469 continue

470 self._pinned_saves[store_job_id] = (block_ids, self._num_workers)

456 - pool.touch([pool.blocks[bid] for bid in block_ids])

471 + pool.touch([pool.blocks[block_id] for block_id in block_ids])

472 +

473 + def _handle_boundary_state_offloads(

474 + self,

475 + offloads: dict[str, list[tuple[int, int, int]]],

476 + meta: MooncakeStoreConnectorMetadata,

477 + ) -> None:

478 + """Attach exact boundary-state blocks to store jobs for this step.

479 +

480 + Flushed in the step they arrive: the CoW copy is enqueued before the

481 + connector event records, so this step's event fences the exact block.

482 + Entries ride the request's save meta when present, else an

483 + offload-only ReqMeta (``token_len_chunk=0`` skips the normal save,

484 + ``can_save=True`` takes the normal enqueue and store-job pinning path).

485 + """

486 + save_metas = {m.req_id: m for m in meta.requests if m.can_save}

487 + for req_id, entries in offloads.items():

488 + tracker = self._request_trackers.get(req_id)

489 + req_tuple = self._unfinished_requests.get(req_id)

490 + if tracker is None or req_tuple is None:

491 + # Request finished/preempted within this step; its blocks are

492 + # going away, so the offload is conservatively dropped.

493 + logger.debug("Dropping boundary-state offload for request %s", req_id)

494 + continue

495 + accepted: list[tuple[int, int, int]] = []

496 + for group_id, block_id, boundary_tokens in entries:

497 + # Every other group stops saving at the end of this prefill, so

498 + # a mamba-only key past it can never complete a joint hybrid

499 + # hit. `prefill_end_tokens` — not the original prompt length —

500 + # is the boundary: a resumed request re-prefills and re-saves

501 + # its previously generated tokens for every group.

502 + if boundary_tokens > tracker.prefill_end_tokens:

503 + continue

504 + if block_id == NULL_BLOCK_ID:

505 + continue

506 + if group_id not in self._boundary_state_group_ids:

507 + continue

508 + accepted.append((group_id, block_id, boundary_tokens))

509 + if not accepted:

510 + continue

511 + tracker.has_pending_offload = True

512 + if (req_meta := save_metas.get(req_id)) is not None:

513 + req_meta.boundary_state_offloads = accepted

514 + continue

515 + meta.add_request(

516 + ReqMeta(

517 + req_id=req_id,

518 + token_len_chunk=0,

519 + block_ids=tracker.allocated_block_ids,

520 + block_hashes=req_tuple[0].block_hashes,

521 + can_save=True,

522 + num_prompt_tokens=tracker.prefill_end_tokens,

523 + boundary_state_offloads=accepted,

524 + )

525 + )

526

527 def update_connector_output(self, connector_output: KVConnectorOutput) -> None:

528 """Drop the block references of store jobs every rank has finished."""

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

+134 -62

@@ -648,48 +648,76 @@ class KVCacheStoreSendingThread(KVTransferThread):

648 self._skip_store_requests.clear()

649 return True

650

651 - def _maybe_offload_partial_tail(self, req_meta: ReqMeta) -> bool:

652 - """Offload the request's sub-block partial tail (its last prompt hash

653 - boundary) so a later request can hit the sub-block prefix.

654 -

655 - Covers every block from the normal save's lcm floor to the boundary:

656 - the normal save floors to ``lcm_block_size``, so a smaller-block

657 - group's full blocks in that gap are never persisted elsewhere, and

658 - the consumer's lookup needs every group at every probed boundary.

659 - Full blocks are keyed by their block-end hash, the partial boundary

660 - block by the boundary sub-hash; the mamba "align" boundary block is

661 - the core-provided CoW block. All keys are deduped against the store.

662 -

663 - Returns:

664 - True when no put is needed or every put succeeds, False otherwise.

651 + def _boundary_snapshot_puts(

652 + self, req_meta: ReqMeta, entries: list[tuple[int, int, int]]

653 + ) -> list[tuple[str, list[int], list[int], KeyMetadata]]:

654 + """Puts for committed mamba "align" boundary-state snapshots.

655 +

656 + These are block-aligned boundaries, i.e. exactly what the normal save

657 + would key — but ``store_mask`` masks mamba groups out of it entirely, so

658 + this is their *only* writer. The exclusion is not an optimization: the

659 + normal save resolves a chunk's address as

660 + ``req_meta.block_ids[g][start // block_size]``, and ``block_ids`` is the

661 + connector's append-only mirror of the core's per-group table. An

662 + align-mode table is mutated in place (a superseded state block is freed

663 + and nulled; speculative blocks relocate), and the connector is never

664 + told, so a stale mirror entry is indistinguishable from a live one — a

665 + retry of a failed or pressure-skipped chunk would read a block that now

666 + belongs to another request.

667 +

668 + Each entry's handed-off block *is* the boundary state and is pinned by

669 + the core, so it is uploaded under its boundary-end hash key and never

670 + resolved positionally.

671 """

666 - if not self.coord.enable_partial_hash_hits or not req_meta.block_hashes:

667 - return True

668 - partial_tail_offloads = req_meta.partial_tail_offloads

669 - if not partial_tail_offloads:

670 - return True

672 hash_block_size = self.coord.hash_block_size

672 - boundaries = {boundary for _, _, boundary in partial_tail_offloads}

673 + puts: list[tuple[str, list[int], list[int], KeyMetadata]] = []

674 + for group_id, block_id, boundary in entries:

675 + if boundary == 0 or block_id == NULL_BLOCK_ID:

676 + continue

677 + hash_idx = boundary // hash_block_size - 1

678 + if hash_idx >= len(req_meta.block_hashes):

679 + continue

680 + db = self.token_databases[group_id]

681 + # Distribute across ranks by the same rule as normal chunks.

682 + put_step = self.group_put_steps[group_id]

683 + put_step_rank = (self.tp_rank + group_id) % put_step

684 + if (boundary // db.block_size - 1) % put_step != put_step_rank:

685 + continue

686 + addr, size = db.prepare_value_for_block(block_id)

687 + puts.append(

688 + (db.key_for(req_meta.block_hashes[hash_idx]), addr, size, db.metadata)

689 + )

690 + return puts

691 +

692 + def _sub_block_tail_puts(

693 + self, req_meta: ReqMeta, entries: list[tuple[int, int, int]]

694 + ) -> list[tuple[str, list[int], list[int], KeyMetadata]]:

695 + """Puts for the request's sub-block partial tail (its last prompt hash

696 + boundary), so a later request can hit the sub-block prefix.

697 +

698 + Covers every group's blocks from the normal save's lcm floor to the

699 + boundary: the normal save floors to ``lcm_block_size``, so a

700 + smaller-block group's full blocks in that gap are never persisted

701 + elsewhere, and the consumer's lookup needs every group at every probed

702 + boundary. Full blocks are keyed by their block-end hash and the partial

703 + boundary block by the boundary sub-hash; a mamba "align" group

704 + contributes only its boundary block, from the core-provided CoW block.

705 + """

706 + boundaries = {boundary for _, _, boundary in entries}

707 if len(boundaries) != 1:

708 raise ValueError(

675 - "Partial-tail offloads for one request must share a boundary"

709 + "Sub-block partial-tail offloads for one request must share a boundary"

710 )

711 boundary = boundaries.pop()

678 - if boundary == 0:

679 - return True

680 - if boundary // hash_block_size - 1 >= len(req_meta.block_hashes):

681 - return True

682 - mamba_offloads = {

683 - group_id: block_id for group_id, block_id, _ in partial_tail_offloads

684 - }

712 + hash_block_size = self.coord.hash_block_size

713 + if boundary == 0 or boundary // hash_block_size - 1 >= len(

714 + req_meta.block_hashes

715 + ):

716 + return []

717

686 - keys: list[str] = []

687 - addrs: list[list[int]] = []

688 - sizes: list[list[int]] = []

689 - group_ids: list[str] | None = (

690 - [] if self.enable_group_semantics and self.supports_group_ids else None

691 - )

718 + mamba_offloads = {group_id: block_id for group_id, block_id, _ in entries}

719 saved = self._saved_offset.get(req_meta.req_id, 0)

720 + puts: list[tuple[str, list[int], list[int], KeyMetadata]] = []

721 for g_idx, db in enumerate(self.token_databases):

722 group_blocks = req_meta.block_ids[g_idx]

723 # Distribute across ranks by the same rule as normal chunks.

@@ -705,16 +733,20 @@ class KVCacheStoreSendingThread(KVTransferThread):

733 continue

734 valid_end = min((block_idx + 1) * db.block_size, boundary)

735 key_hash = req_meta.block_hashes[valid_end // hash_block_size - 1]

708 - if (

709 - g_idx in mamba_offloads

710 - and valid_end == boundary

711 - and boundary % db.block_size != 0

712 - ):

713 - block_id = mamba_offloads[g_idx]

714 - else:

715 - if block_idx >= len(group_blocks):

736 + if g_idx in mamba_offloads:

737 + if valid_end != boundary:

738 + # Interior align-mode state positions are null or

739 + # stale (the block table is not append-only) and never

740 + # valid gap content; only the boundary block is

741 + # persisted, from the core-provided hand-off.

742 continue

743 + block_id = mamba_offloads[g_idx]

744 + elif g_idx in self.coord.mamba_group_ids:

745 + continue

746 + elif block_idx < len(group_blocks):

747 block_id = group_blocks[block_idx]

748 + else:

749 + continue

750 if block_id == NULL_BLOCK_ID:

751 logger.debug(

752 "Skipping unavailable partial-tail source block "

@@ -725,20 +757,58 @@ class KVCacheStoreSendingThread(KVTransferThread):

757 )

758 continue

759 addr, size = db.prepare_value_for_block(block_id)

728 - key = db.key_for(key_hash)

729 - keys.append(key)

730 - addrs.append(addr)

731 - sizes.append(size)

732 - if group_ids is not None:

733 - group_ids.append(

734 - _make_mooncake_group_id(

735 - db.metadata,

736 - key.rsplit("@", 1)[-1],

737 - )

738 - )

760 + puts.append((db.key_for(key_hash), addr, size, db.metadata))

761 + return puts

762 +

763 + def _maybe_offload_boundary_states(self, req_meta: ReqMeta) -> bool:

764 + """Persist connector-pinned mamba "align" boundary states handed off

765 + for this request, deduped against the store.

766

740 - if not keys:

767 + This is every mamba key the connector writes — ``store_mask`` excludes

768 + mamba groups from the positional normal save, aligned boundaries

769 + included (see :meth:`_boundary_snapshot_puts`).

770 +

771 + The two entry kinds are keyed and sourced differently, so they are

772 + prepared separately and put in one batch:

773 +

774 + - block-aligned for its group: a committed boundary-state snapshot,

775 + the handed-off block itself;

776 + - not block-aligned: the sub-block CoW partial tail, which also has to

777 + cover the other groups' blocks in the normal save's lcm gap.

778 +

779 + Returns:

780 + True when no put is needed or every put succeeds, False otherwise.

781 + """

782 + offloads = req_meta.boundary_state_offloads

783 + if not offloads or not req_meta.block_hashes:

784 + return True

785 +

786 + snapshots: list[tuple[int, int, int]] = []

787 + sub_block: list[tuple[int, int, int]] = []

788 + for group_id, block_id, boundary in offloads:

789 + entry = (group_id, block_id, boundary)

790 + if boundary % self.token_databases[group_id].block_size == 0:

791 + snapshots.append(entry)

792 + else:

793 + sub_block.append(entry)

794 +

795 + puts = self._boundary_snapshot_puts(req_meta, snapshots)

796 + if sub_block and self.coord.enable_partial_hash_hits:

797 + puts.extend(self._sub_block_tail_puts(req_meta, sub_block))

798 +

799 + if not puts:

800 return True

801 + keys = [key for key, _, _, _ in puts]

802 + addrs = [addr for _, addr, _, _ in puts]

803 + sizes = [size for _, _, size, _ in puts]

804 + group_ids: list[str] | None = (

805 + [

806 + _make_mooncake_group_id(metadata, key.rsplit("@", 1)[-1])

807 + for key, _, _, metadata in puts

808 + ]

809 + if self.enable_group_semantics and self.supports_group_ids

810 + else None

811 + )

812 exists_start = time.perf_counter()

813 try:

814 exists = self.store.batch_is_exist(keys)

@@ -751,7 +821,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

821 num_failed_keys=len(keys),

822 )

823 logger.error(

754 - "Failed to check partial-tail keys for request %s: %s",

824 + "Failed to check boundary-state keys for request %s: %s",

825 req_meta.req_id,

826 e,

827 )

@@ -787,7 +857,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

857 num_failed_keys=len(keys),

858 )

859 logger.error(

790 - "Failed to put partial-tail keys for request %s: %s",

860 + "Failed to put boundary-state keys for request %s: %s",

861 req_meta.req_id,

862 e,

863 )

@@ -805,7 +875,8 @@ class KVCacheStoreSendingThread(KVTransferThread):

875 if failed:

876 failed_codes = {res[i] for i in failed}

877 logger.warning(

808 - "Partial-tail put failed for request %s: %d/%d keys failed (codes=%s)",

878 + "Boundary-state put failed for request %s: %d/%d keys failed "

879 + "(codes=%s)",

880 req_meta.req_id,

881 len(failed),

882 len(keys),

@@ -818,7 +889,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

889 if self._clear_store_pressure():

890 logger.info(

891 "Mooncake CPU/disk offloading pressure cleared after a "

821 - "successful partial-tail batch"

892 + "successful boundary-state batch"

893 )

894 return True

895

@@ -857,10 +928,10 @@ class KVCacheStoreSendingThread(KVTransferThread):

928 )

929 return

930

860 - # Offload the sub-block partial tail (independent of the normal

861 - # block-aligned save, which may be skipped this step).

862 - if req_meta.partial_tail_offloads is not None and not (

863 - self._maybe_offload_partial_tail(req_meta)

931 + # Offload the handed-off mamba boundary states (independent of the

932 + # normal positional save, which may be skipped this step).

933 + if req_meta.boundary_state_offloads is not None and not (

934 + self._maybe_offload_boundary_states(req_meta)

935 ):

936 return

937

@@ -887,6 +958,7 @@ class KVCacheStoreSendingThread(KVTransferThread):

958 # Rotate the stride phase per group to balance load across ranks.

959 put_step = self.group_put_steps[g_idx]

960 put_step_rank = (self.tp_rank + g_idx) % put_step

961 + group_blocks = block_ids_per_group[g_idx]

962 for start, end, block_hash in db.process_tokens(

963 token_len,

964 req_meta.block_hashes,

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+81 -4

@@ -1130,6 +1130,9 @@ class OffloadingConnectorScheduler:

1130 if bid != 0:

1131 self._current_batch_allocated_block_ids.add(bid)

1132

1133 + for copy in scheduler_output.kv_cache_block_copies or ():

1134 + self._current_batch_allocated_block_ids.add(copy.dst_block_id)

1135 +

1136 # Zero out stale block_ids in sliding window groups' pending-store

1137 # positions. Only sliding window groups can have stale entries (blocks

1138 # freed by remove_skipped_blocks then reallocated). Only positions in

@@ -1151,18 +1154,89 @@ class OffloadingConnectorScheduler:

1154 ):

1155 group_state.block_ids[j] = 0

1156

1157 + def _build_aligned_boundary_store_jobs(

1158 + self, handoffs: dict[str, list[tuple[int, int, int]]]

1159 + ) -> dict[int, TransferJob]:

1160 + store_jobs: dict[int, TransferJob] = {}

1161 + num_groups = len(self.config.kv_group_configs)

1162 + for req_id, entries in handoffs.items():

1163 + req_status = self._req_status.get(req_id)

1164 + if req_status is None:

1165 + continue

1166 + req = req_status.req

1167 + max_boundary = self._calc_num_offloadable_tokens(req_status, req.num_tokens)

1168 + for group_idx, block_id, boundary in entries:

1169 + group_config = self.config.kv_group_configs[group_idx]

1170 + if (

1171 + block_id == 0

1172 + or boundary > max_boundary

1173 + or boundary % group_config.tokens_per_chunk != 0

1174 + ):

1175 + continue

1176 +

1177 + key = self._make_boundary_key(req, group_idx, boundary)

1178 + store_output = self.manager.prepare_store([key], req_status.req_context)

1179 + if store_output is None:

1180 + self._connector_stats.increase_counter(

1181 + _ConnectorMetricName.ALLOCATION_FAILURE

1182 + )

1183 + continue

1184 + if not store_output.keys_to_store:

1185 + continue

1186 +

1187 + job_id = self._generate_job_id()

1188 + req_status.transfer_jobs.add(job_id)

1189 + self._block_id_to_pending_jobs.setdefault(block_id, set()).add(job_id)

1190 + self._jobs[job_id] = TransferJobStatus(

1191 + req_id=req_id,

1192 + pending_count=self.config.num_workers,

1193 + keys={key},

1194 + is_store=True,

1195 + fenced_block_ids=[block_id],

1196 + )

1197 + group_sizes = [0] * num_groups

1198 + group_sizes[group_idx] = 1

1199 + block_indices = [0] * num_groups

1200 + block_indices[group_idx] = boundary // group_config.tokens_per_block - 1

1201 + store_jobs[job_id] = TransferJob(

1202 + req_id=req_id,

1203 + src_spec=GPULoadStoreSpec(

1204 + [block_id],

1205 + group_sizes=group_sizes,

1206 + block_indices=block_indices,

1207 + ),

1208 + dst_spec=store_output.store_spec,

1209 + )

1210 + self._events_tracker.record_store(

1211 + req,

1212 + group_config,

1213 + boundary // group_config.tokens_per_chunk - 1,

1214 + key,

1215 + )

1216 + return store_jobs

1217 +

1218 def _build_partial_tail_store_jobs(

1219 self, scheduler_output: SchedulerOutput

1220 ) -> dict[int, TransferJob]:

1157 - handoffs = scheduler_output.partial_tail_offloads

1158 - if not self.config.supports_partial_tail or not handoffs:

1221 + block_state = scheduler_output.kv_connector_block_state

1222 + handoffs = block_state.boundary_state_offloads if block_state else None

1223 + if not handoffs:

1224 return {}

1225

1161 - store_jobs: dict[int, TransferJob] = {}

1226 + store_jobs = self._build_aligned_boundary_store_jobs(handoffs)

1227 + if not self.config.supports_partial_tail:

1228 + return store_jobs

1229 +

1230 for req_id, entries in handoffs.items():

1231 + entries = [

1232 + entry

1233 + for entry in entries

1234 + if entry[2] % self._partial_tail_block_size != 0

1235 + ]

1236 + if not entries:

1237 + continue

1238 req_status = self._req_status.get(req_id)

1239 assert req_status is not None

1165 - assert entries

1240 boundaries = {boundary for _, _, boundary in entries}

1241 assert len(boundaries) == 1

1242 boundary = boundaries.pop()

@@ -1282,6 +1356,9 @@ class OffloadingConnectorScheduler:

1356 group_config, group_state, num_offloadable_tokens

1357 )

1358

1359 + if group_config.requires_cow_source:

1360 + continue

1361 +

1362 start_chunk_idx = group_state.next_stored_chunk_idx

1363 if num_chunks <= start_chunk_idx:

1364 continue

vllm/v1/core/block_pool.py

@@ -716,6 +716,10 @@ class BlockPool:

716 if self.metrics_collector:

717 self.metrics_collector.on_block_accessed(block)

718

719 + def is_block_writable(self, block: KVCacheBlock) -> bool:

720 + """Return whether a block can be mutated by its sole owner."""

721 + return not block.is_null and block.ref_cnt == 1 and block.block_hash is None

722 +

723 def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None:

724 """Free a list of blocks. The blocks should be ordered by their

725 eviction priority, where the first block will be evicted first.

vllm/v1/core/kv_cache_manager.py

+12 -28

@@ -189,10 +189,6 @@ class KVCacheManager:

189 tuple(() for _ in range(self.num_kv_cache_groups))

190 )

191

192 - # Off-table cow blocks handed to a KV connector for partial-tail

193 - # offload; pinned until the request's blocks are freed.

194 - self._partial_tail_pins: dict[str, list[KVCacheBlock]] = {}

195 -

192 @property

193 def usage(self) -> float:

194 """Get the KV cache usage.

@@ -575,9 +571,6 @@ class KVCacheManager:

571 Args:

572 request: The request to free the blocks.

573 """

578 - pins = self._partial_tail_pins.pop(request.request_id, None)

579 - if pins:

580 - self.block_pool.free_blocks(pins)

574 self.coordinator.free(request.request_id)

575

576 def remove_skipped_blocks(

@@ -610,14 +603,7 @@ class KVCacheManager:

603 Returns:

604 The request's blocks in allocation order.

605 """

613 - blocks = self.coordinator.pop_blocks_for_free(request.request_id)

614 - # Pins ride the same (possibly deferred) free as the request blocks.

615 - # Preemption may release a pin under a still-queued offload — the same

616 - # exposure normal saves of table blocks already have.

617 - pins = self._partial_tail_pins.pop(request.request_id, None)

618 - if pins:

619 - blocks = pins + blocks

620 - return blocks

606 + return self.coordinator.pop_blocks_for_free(request.request_id)

607

608 def evict_blocks(self, block_ids: set[int]) -> None:

609 """evict blocks from the prefix cache by their block IDs.

@@ -854,18 +840,18 @@ class KVCacheManager:

840 retained_blocks = [block for pair in pending_copies for block in pair]

841 return copies, retained_blocks

842

857 - def take_partial_tail_offloads(self) -> dict[str, list[tuple[int, int, int]]]:

858 - """Drain producer partial-tail offload hand-offs per request.

843 + def take_boundary_state_offloads(

844 + self,

845 + ) -> dict[str, list[tuple[int, int, int]]]:

846 + """Drain this step's boundary-state hand-offs for a KV connector.

847

848 Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}``

861 - for the durable boundary blocks of producers' last-prompt-boundary

862 - partial tails. Only mamba "align" groups contribute; empty otherwise.

863 - A KV connector reads the referenced blocks and offloads them so a later

864 - request can hit the sub-block prefix.

865 -

866 - Each handed-off block lives off the request block table, so it is

867 - pinned here and unpinned when the request's blocks are freed — for a

868 - producer with saved tokens, after the connector reports sends done.

849 + for mamba "align" boundary states: the

850 + request's committed boundary-state snapshots and, on a sub-block partial

851 + hit, the CoW copy of its last-prompt-boundary state. Only mamba "align"

852 + groups contribute; empty otherwise. A connector reads the referenced

853 + blocks — never resolving them positionally — and offloads them so a

854 + later request can hit that prefix.

855 """

856 offloads: dict[str, list[tuple[int, int, int]]] = {}

857 for mgr in self.coordinator.single_type_managers:

@@ -874,9 +860,7 @@ class KVCacheManager:

860 group_id,

861 block,

862 boundary_tokens,

877 - ) in mgr.take_pending_partial_tail_offloads():

878 - self.block_pool.touch((block,))

879 - self._partial_tail_pins.setdefault(req_id, []).append(block)

863 + ) in mgr.take_pending_boundary_state_offloads():

864 offloads.setdefault(req_id, []).append(

865 (group_id, block.block_id, boundary_tokens)

866 )

vllm/v1/core/sched/output.py

+12 -5

@@ -205,6 +205,16 @@ class ScheduledEncoderInputStats:

205 output_tokens: int = 0

206

207

208 +@dataclass

209 +class KVConnectorBlockState:

210 + """Scheduler-local block state offered to a producer-side KV connector."""

211 +

212 + # Authoritative current block-table snapshots.

213 + block_ids: dict[str, tuple[list[int], ...]]

214 + # Exact Mamba "align" boundary-state hand-offs.

215 + boundary_state_offloads: dict[str, list[tuple[int, int, int]]]

216 +

217 +

218 @dataclass

219 class SchedulerOutput:

220 # list of the requests that are scheduled for the first time.

@@ -278,11 +288,8 @@ class SchedulerOutput:

288 # CoW copies to apply after zeroing new blocks and before forward.

289 kv_cache_block_copies: list[KVCacheBlockCopy] | None = None

290

281 - # Producer partial-tail offload hand-off for external KV connectors:

282 - # {request_id: [(group_id, block_id, boundary_tokens), ...]} pointing at

283 - # the durable boundary block of a producer's last-prompt-boundary partial

284 - # tail (mamba "align" CoW target). None unless partial hash hits are active.

285 - partial_tail_offloads: dict[str, list[tuple[int, int, int]]] | None = None

291 + # Scheduler-local; always None by the time this reaches a worker.

292 + kv_connector_block_state: KVConnectorBlockState | None = None

293

294 # Dynamic speculative decoding: optimal K chosen by scheduler.

295 # Number of spec tokens to schedule for the next step.

vllm/v1/core/sched/scheduler.py

+31 -15

@@ -42,6 +42,7 @@ from vllm.v1.core.sched.interface import PauseState, SchedulerInterface

42 from vllm.v1.core.sched.output import (

43 CachedRequestData,

44 GrammarOutput,

45 + KVConnectorBlockState,

46 NewRequestData,

47 ScheduledEncoderInputStats,

48 SchedulerOutput,

@@ -1263,20 +1264,32 @@ class Scheduler(SchedulerInterface):

1264 self.prev_step_scheduled_req_ids.clear()

1265 self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())

1266

1266 - # Producer partial-tail hand-off for external KV connectors. Drained

1267 - # before the CoW retentions are released below, so the pin lands while

1268 - # the cow block still holds a retention ref. Without a producer-side

1269 - # connector nothing consumes the hand-off, so skip the drain (and its

1270 - # pin); the manager drops stale entries when the request's blocks are

1271 - # popped for free.

1272 - pending_partial_tail_offloads = None

1273 - if (

1274 - self.connector is not None

1275 - and self.vllm_config.kv_transfer_config is not None

1276 - and self.vllm_config.kv_transfer_config.is_kv_producer

1277 - ):

1278 - pending_partial_tail_offloads = (

1279 - self.kv_cache_manager.take_partial_tail_offloads() or None

1267 + # Mamba "align" boundary states must be handed off with exact block ids;

1268 + # they cannot be reconstructed from a connector's append-only block

1269 + # table. Drained every step so stale offers cannot accumulate.

1270 + boundary_state_offloads = self.kv_cache_manager.take_boundary_state_offloads()

1271 +

1272 + kv_connector_block_state = None

1273 + if self.connector is not None:

1274 + snapshot_req_ids = {req.req_id for req in new_reqs_data}

1275 + snapshot_req_ids.update(

1276 + req_id

1277 + for req_id, block_ids in zip(

1278 + cached_reqs_data.req_ids,

1279 + cached_reqs_data.new_block_ids,

1280 + strict=True,

1281 + )

1282 + if block_ids

1283 + )

1284 + snapshot_req_ids.update(

1285 + req_id for req_id in boundary_state_offloads if req_id in self.requests

1286 + )

1287 + kv_connector_block_state = KVConnectorBlockState(

1288 + block_ids={

1289 + req_id: self.kv_cache_manager.get_block_ids(req_id)

1290 + for req_id in snapshot_req_ids

1291 + },

1292 + boundary_state_offloads=boundary_state_offloads,

1293 )

1294

1295 kv_cache_block_copies, cow_retained_blocks = (

@@ -1325,7 +1338,7 @@ class Scheduler(SchedulerInterface):

1338 new_block_ids_to_zero=self._get_new_block_ids_to_zero(),

1339 has_sync_kv_loads=has_sync_kv_loads,

1340 kv_cache_block_copies=pending_kv_cache_block_copies,

1328 - partial_tail_offloads=pending_partial_tail_offloads,

1341 + kv_connector_block_state=kv_connector_block_state,

1342 num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,

1343 ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(),

1344 )

@@ -1345,6 +1358,9 @@ class Scheduler(SchedulerInterface):

1358 )

1359 scheduler_output.ec_connector_metadata = ec_meta

1360

1361 + # Connector-only block state must not be dispatched to workers.

1362 + scheduler_output.kv_connector_block_state = None

1363 +

1364 # Advance the fence only for non-empty steps (those that actually

1365 # write KV and have their output processed later in update_from_output).

1366 if self.defer_block_free and total_num_scheduled_tokens > 0:

vllm/v1/core/single_type_kv_cache_manager.py

+52 -26

@@ -112,13 +112,14 @@ class SingleTypeKVCacheManager(ABC):

112 # managers (full attention, mamba "align"); harmlessly empty elsewhere.

113 self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {}

114 self._pending_cow_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = []

115 - # Partial-tail offload hand-off for external KV connectors: when a

116 - # producer registers its last-prompt-boundary partial tail and the

117 - # durable boundary block is not on the append-only request block table

118 - # (mamba "align" CoW target), record the request, group, block, and

119 - # exact token boundary so a connector can offload it under the right

120 - # hash. Populated only by mamba "align".

121 - self._pending_partial_tail_offloads: list[

115 + # Boundary-state offload hand-off for external KV connectors. A mamba

116 + # "align" block table is not append-only (interior states are

117 + # nulled/freed and speculative blocks relocate in place), so a

118 + # connector cannot resolve its state blocks positionally. Record

119 + # (request, group, block, exact token boundary) for each committed

120 + # boundary state so a connector can offload the right block under the

121 + # right hash. Populated only by mamba "align".

122 + self._pending_boundary_state_offloads: list[

123 tuple[str, int, KVCacheBlock, int]

124 ] = []

125

@@ -385,19 +386,19 @@ class SingleTypeKVCacheManager(ABC):

386 self._pending_cow_copies = []

387 return pending_copies

388

388 - def take_pending_partial_tail_offloads(

389 + def take_pending_boundary_state_offloads(

390 self,

391 ) -> list[tuple[str, int, KVCacheBlock, int]]:

391 - """Drain producer partial-tail hand-offs.

392 + """Drain producer boundary-state hand-offs.

393

394 Entries are ``(req_id, group_id, block, boundary_tokens)``.

395

395 - Only mamba "align" populates this. The block lives off the request

396 - block table, so the caller must pin it until the connector has read

397 - it — nothing else keeps it alive once the CoW retention is released.

396 + Only mamba "align" populates this. The blocks are not kept alive by

397 + the request block table for the whole request, so a caller that reads

398 + them asynchronously must pin them first.

399 """

399 - pending = self._pending_partial_tail_offloads

400 - self._pending_partial_tail_offloads = []

400 + pending = self._pending_boundary_state_offloads

401 + self._pending_boundary_state_offloads = []

402 return pending

403

404 def _apply_cow(

@@ -1292,7 +1293,7 @@ class MambaManager(SingleTypeKVCacheManager):

1293 # Requests that registered their own last-prompt-boundary partial

1294 # tail (producers). On the next step's CoW the boundary state moves

1295 # into a private cow_block; we record that block for connector

1295 - # offload (see _pending_partial_tail_offloads).

1296 + # offload (see _pending_boundary_state_offloads).

1297 self._producer_partial_tail_reqs: dict[str, int] = {}

1298

1299 @classmethod

@@ -1625,13 +1626,12 @@ class MambaManager(SingleTypeKVCacheManager):

1626 )

1627

1628 if blocks_allocated:

1628 - # reuse previous speculative blocks in this step

1629 + # Relocate exclusively owned speculative scratch blocks.

1630 for block_idx in range(

1631 prev_block_len - self.num_speculative_blocks, prev_block_len

1632 ):

1633 if block_idx < num_skipped_blocks:

1633 - req_blocks.append(req_blocks[block_idx])

1634 - req_blocks[block_idx] = self._null_block

1634 + self._relocate_speculative_block(req_blocks, block_idx)

1635 else:

1636 break

1637 num_new_blocks = num_required_blocks - len(req_blocks)

@@ -1664,7 +1664,7 @@ class MambaManager(SingleTypeKVCacheManager):

1664 # This CoW preserved a producer's own boundary

1665 # state in cow_block; hand it to the connector for

1666 # partial-tail offload once the copy has run.

1667 - self._pending_partial_tail_offloads.append(

1667 + self._pending_boundary_state_offloads.append(

1668 (

1669 request_id,

1670 self.kv_cache_group_id,

@@ -1685,17 +1685,31 @@ class MambaManager(SingleTypeKVCacheManager):

1685 returned_blocks.extend(new_blocks)

1686 return returned_blocks

1687

1688 + def _relocate_speculative_block(

1689 + self, req_blocks: list[KVCacheBlock], block_idx: int

1690 + ) -> None:

1691 + block = req_blocks[block_idx]

1692 + assert self.block_pool.is_block_writable(block), (

1693 + "Speculative Mamba blocks must be exclusively owned and unhashed "

1694 + "before relocation"

1695 + )

1696 + req_blocks.append(block)

1697 + req_blocks[block_idx] = self._null_block

1698 +

1699 def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:

1700 if self.mamba_cache_mode == "align":

1701 self._allocated_block_reqs.discard(request_id)

1702 self.last_state_block_idx.pop(request_id, None)

1703 self._num_checkpoint_blocks.pop(request_id, None)

1704 self._producer_partial_tail_reqs.pop(request_id, None)

1694 - # A hand-off whose request died in this same scheduling pass must

1695 - # not reach the connector: its unpin hook (free) has already run.

1696 - self._pending_partial_tail_offloads = [

1705 + # An offer is only guaranteed to hold committed bytes until the end

1706 + # of the pass that made it. This request's blocks are going back to

1707 + # the pool now, so drop its not-yet-offered hand-offs rather than

1708 + # let a connector claim a block another request may already have

1709 + # been handed.

1710 + self._pending_boundary_state_offloads = [

1711 entry

1698 - for entry in self._pending_partial_tail_offloads

1712 + for entry in self._pending_boundary_state_offloads

1713 if entry[0] != request_id

1714 ]

1715 return super().pop_blocks_for_free(request_id)

@@ -1722,9 +1736,9 @@ class MambaManager(SingleTypeKVCacheManager):

1736 if partial_hash is not None:

1737 self.cached_blocks_this_step.add(partial_hash)

1738 if num_cached_blocks_after > num_cached_blocks_before:

1725 - for block in self.req_to_blocks[request.request_id][

1726 - num_cached_blocks_before:num_cached_blocks_after

1727 - ]:

1739 + blocks = self.req_to_blocks[request.request_id]

1740 + for idx in range(num_cached_blocks_before, num_cached_blocks_after):

1741 + block = blocks[idx]

1742 # Skip null blocks (align-mode skipped states) and blocks that

1743 # were not cached this step — with sparse retention

1744 # (reachable_block_mask) the intermediate state snapshots carry

@@ -1732,6 +1746,18 @@ class MambaManager(SingleTypeKVCacheManager):

1746 if block.is_null or block.block_hash is None:

1747 continue

1748 self.cached_blocks_this_step.add(block.block_hash)

1749 + if self.mamba_cache_mode == "align":

1750 + # Offer every retained boundary with its exact block.

1751 + # The connector filters against its save window, which may

1752 + # extend past the original prompt during resumed prefill.

1753 + self._pending_boundary_state_offloads.append(

1754 + (

1755 + request.request_id,

1756 + self.kv_cache_group_id,

1757 + block,

1758 + (idx + 1) * self.block_size,

1759 + )

1760 + )

1761

1762 def new_step_starts(self) -> None:

1763 self.cached_blocks_this_step.clear()