Skip to content

[None][feat] Support V2 Mamba disaggregated serving#16472

Closed
VALLIS-NERIA wants to merge 9 commits into
NVIDIA:mainfrom
VALLIS-NERIA:agent/v2-mamba-disaggregated-serving
Closed

[None][feat] Support V2 Mamba disaggregated serving#16472
VALLIS-NERIA wants to merge 9 commits into
NVIDIA:mainfrom
VALLIS-NERIA:agent/v2-mamba-disaggregated-serving

Conversation

@VALLIS-NERIA

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

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 depends on #16471. It does not depend on the save-last snapshot-reuse PR. Because an upstream PR cannot use a fork-only branch as its base, this draft temporarily shows the prerequisite changes until the core PR merges and this branch is rebased.

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 TP (1 passed)
  • V2-to-V2 Python/NIXL transfer with TP mismatch (1 passed)
  • Mixed/V2 native transfer and physical-pool registration regressions (10 passed)
  • Effective backend/runtime routing matrix (8 passed)
  • pre-commit run --from-ref origin/main --to-ref HEAD

PR Checklist

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

@VALLIS-NERIA

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59617 [ run ] triggered by Bot. Commit: d71c1ee Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59617 [ run ] completed with state SUCCESS. Commit: d71c1ee
/LLM/main/L0_MergeRequest_PR pipeline #48054 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
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 d71c1ee to 9dac8ed Compare July 16, 2026 08:04
@VALLIS-NERIA

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59671 [ run ] triggered by Bot. Commit: 9dac8ed Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59671 [ run ] completed with state FAILURE. Commit: 9dac8ed
/LLM/main/L0_MergeRequest_PR pipeline #48104 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
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 9dac8ed to 00620eb Compare July 16, 2026 09:09

Copy link
Copy Markdown
Collaborator Author

Reopened as a stacked draft at VALLIS-NERIA#9 so it can target the V2 Mamba core branch and show only its single squashed commit. Closing this upstream draft as superseded.

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