[None][feat] Support V2 Mamba disaggregated serving - #9
Conversation
94db6dc to
6ffdf55
Compare
00620eb to
a8b2d43
Compare
45e3b73 to
8c85708
Compare
| layer_offsets: Dict[int, int], | ||
| overlapping_layers: List[int], | ||
| slot: int, | ||
| layer_slot0_addresses: Optional[Dict[int, int]] = None, |
There was a problem hiding this comment.
layer_slot0_addresses looks like it could be derived rather than stored. In the single contiguous pool it's equivalent to base_address + lid*num_slots*stride, so it may be redundant. And physical_slot_stride_bytes captures a genuinely necessary concept, but perhaps a single field would suffice since the conv/ssm pair is always equal (and derivable from the two slot_bytes summed). block_stride_bytes might also read a bit more clearly.
There was a problem hiding this comment.
In KVCacheManagerV2 the states are arranged in slot-first order, so we can't reuse the calculation of V1.
| ssm_states: Optional[PhysicalPool] = None | ||
| conv_section_bytes: Optional[List[int]] = None | ||
| ssm_bytes_per_head: Optional[int] = None | ||
| conv_layer_slot0_addresses: Optional[Dict[int, int]] = None |
There was a problem hiding this comment.
Would it be worth collapsing these four fields into a single optional slot_stride_bytes on PhysicalPool (defaulting to slot_bytes)? The per-layer slot0_addresses look derivable from base_address + lid*num_slots*stride, and the conv/ssm stride pair is always equal, so one stride field would let V1 and V2 share the same pointer formula (V1 just falls back to slot_bytes, unchanged). It also reads a bit more naturally, since "a slot whose stride exceeds its payload" is really a property of the pool.
There was a problem hiding this comment.
conv/ssm stride pair is always equal
This is not true in V2.
0ae65cc to
935fd99
Compare
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
935fd99 to
6558b78
Compare
22ce579
into
agent/v2-mamba-snapshot-reuse-core
@coderabbitai summary
Stack
agent/v2-mamba-snapshot-reuse-coreDescription
This PR adds disaggregated context-to-generation state transfer for the V2 Mamba hybrid cache manager.
The changes:
auto/default, explicit C++, and UCX routes remain on the C++ manager unless model runtime resolution selects Python/NIXL.This PR does not depend on the save-last snapshot-reuse follow-up.
Test Coverage
64 passed)2 passed)2 passed)18 passed)pre-commit run --from-ref 0ff3a34e47 --to-ref HEADThe history squash was tree-preserving; no code changed during the squash.
PR Checklist