Skip to content

[None][feat] Add V2 Mamba hybrid cache with snapshot reuse#16471

Closed
VALLIS-NERIA wants to merge 8 commits into
NVIDIA:mainfrom
VALLIS-NERIA:agent/v2-mamba-snapshot-reuse-core
Closed

[None][feat] Add V2 Mamba hybrid cache with snapshot reuse#16471
VALLIS-NERIA wants to merge 8 commits into
NVIDIA:mainfrom
VALLIS-NERIA:agent/v2-mamba-snapshot-reuse-core

Conversation

@VALLIS-NERIA

Copy link
Copy Markdown
Collaborator

@coderabbitai summary

Description

This PR adds the V2 Mamba hybrid-cache core and snapshot reuse on top of the generic scheduler/resource-manager contracts in #16427.

The changes:

  • Add V2MambaHybridCacheManager, which stores attention KV blocks together with Mamba convolution and SSM state in KVCMv2-managed pools.
  • Expose role-aware pool addresses and block indices so linear-attention state can use logical views over the underlying block layout.
  • Track live state and reusable snapshots separately, commit snapshots only at reusable boundaries, and advance snapshot history consistently across decode and MTP replay.
  • Keep CUDA state-index storage stable across attention-DP dummy insertion and remove request-to-index mappings whenever requests are freed.
  • Preserve explicit C++/mixed-manager routing and legacy disaggregated-serving fallbacks; this PR does not add V2 Mamba state transfer.
  • Add focused coverage for allocation, snapshot lookup/commit, replay state, dummy indices, scheduler chunk boundaries, stats, and transceiver integration.

This PR depends on #16427. GitHub cannot use a branch from a fork as the base of an upstream PR, so the draft diff temporarily includes #16427 until that PR merges and this branch is rebased. The commit-minimum integration formerly carried by this work is already available on main through #16185.

Save-last snapshot reuse and V2 Mamba disaggregated transfer are intentionally split into follow-up PRs.

Test Coverage

  • pytest -q tests/unittest/_torch/executor/test_mamba_cache_manager.py (44 passed)
  • pytest -q tests/unittest/_torch/executor/test_kv_cache_manager_v2.py (10 passed)
  • pytest -q tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py (18 passed)
  • pytest -q tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (102 passed; the one stale-config failure was fixed and its focused rerun passed)
  • Focused Python/NIXL native transceiver tests after rebuilding bindings against current main (7 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 #59619 [ run ] triggered by Bot. Commit: 1461b2d 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>
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59619 [ run ] completed with state SUCCESS. Commit: 1461b2d
/LLM/main/L0_MergeRequest_PR pipeline #48056 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>
@VALLIS-NERIA VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch from 1461b2d to 149013b Compare July 16, 2026 07:53
@VALLIS-NERIA

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59668 [ run ] triggered by Bot. Commit: 149013b Link to invocation

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
@VALLIS-NERIA VALLIS-NERIA force-pushed the agent/v2-mamba-snapshot-reuse-core branch from 149013b to 0ff3a34 Compare July 16, 2026 09:09

Copy link
Copy Markdown
Collaborator Author

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #59668 [ run ] completed with state SUCCESS. Commit: 149013b
/LLM/main/L0_MergeRequest_PR pipeline #48101 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

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