Skip to content

[None][feat] Support V2 Mamba disaggregated serving - #9

Merged
VALLIS-NERIA merged 1 commit into
agent/v2-mamba-snapshot-reuse-corefrom
agent/v2-mamba-disaggregated-serving
Jul 21, 2026
Merged

[None][feat] Support V2 Mamba disaggregated serving#9
VALLIS-NERIA merged 1 commit into
agent/v2-mamba-snapshot-reuse-corefrom
agent/v2-mamba-disaggregated-serving

Conversation

@VALLIS-NERIA

Copy link
Copy Markdown
Owner

@coderabbitai summary

Stack

  • Base branch: agent/v2-mamba-snapshot-reuse-core
  • Prerequisite: VALLIS-NERIA/TensorRT-LLM#8
  • This PR contains exactly one commit relative to its base.

Description

This PR adds disaggregated context-to-generation state transfer for the V2 Mamba hybrid cache manager.

The changes:

  • Represent a Mamba layer group in transfer metadata with its state-pool address, byte stride, and layer count, and preserve those fields through serialization.
  • Extract V2 convolution and SSM pages from the actual KVCMv2 physical pools, including coalesced equal-size state pools, and transfer them through the Python/NIXL transceiver while preserving legacy layer-major registration.
  • Reconstruct peer-side Mamba state pages for both matching and mismatched tensor-parallel layouts.
  • Select the V2 manager only after resolving the effective backend and runtime: V2 requires NIXL + Python; auto/default, explicit C++, and UCX routes remain on the C++ manager unless model runtime resolution selects Python/NIXL.
  • Add model-level preferred-runtime selection for the supported hybrid architectures.
  • Cover runtime/backend routing, metadata round trips, page extraction, same-TP transfer, TP-mismatch transfer, and inflight-cancellation gating.

This PR does not depend on the save-last snapshot-reuse follow-up.

Test Coverage

  • CPU routing, transceiver, serialization, and pointer tests (64 passed)
  • CUDA V2 page-table tests (2 passed)
  • V2-to-V2 Python/NIXL transfer with matching and mismatched TP (2 passed)
  • Mixed/V2 native transfer, physical-pool registration, and routing regressions (18 passed)
  • pre-commit run --from-ref 0ff3a34e47 --to-ref HEAD

The history squash was tree-preserving; no code changed during the squash.

PR Checklist

  • Please check this after reviewing the above items as appropriate for this PR.

@VALLIS-NERIA
VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch 3 times, most recently from 94db6dc to 6ffdf55 Compare July 20, 2026 10:19
@VALLIS-NERIA
VALLIS-NERIA force-pushed the agent/v2-mamba-disaggregated-serving branch from 00620eb to a8b2d43 Compare July 21, 2026 02:40
@VALLIS-NERIA
VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch 3 times, most recently from 45e3b73 to 8c85708 Compare July 21, 2026 05:41
layer_offsets: Dict[int, int],
overlapping_layers: List[int],
slot: int,
layer_slot0_addresses: Optional[Dict[int, int]] = None,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

conv/ssm stride pair is always equal

This is not true in V2.

@VALLIS-NERIA
VALLIS-NERIA force-pushed the agent/v2-mamba-disaggregated-serving branch 3 times, most recently from 0ae65cc to 935fd99 Compare July 21, 2026 09:45
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
@VALLIS-NERIA
VALLIS-NERIA force-pushed the agent/v2-mamba-disaggregated-serving branch from 935fd99 to 6558b78 Compare July 21, 2026 10:19
@VALLIS-NERIA
VALLIS-NERIA marked this pull request as ready for review July 21, 2026 11:13
@VALLIS-NERIA
VALLIS-NERIA merged commit 22ce579 into agent/v2-mamba-snapshot-reuse-core Jul 21, 2026
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants