[https://nvbugs/6417488][fix] Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and…

[trtllm-agent]

Summary

• Root cause: PR [None][feat] DSv4 follow-up: runtime KV and cache foundations #15633 renamed base hook _build_pool_mapping_tensors → _prepare_page_table_tensor; MiniMaxM3KVCacheManagerV2 override still used the old name, so it never ran and the base's exact_div asserted on M3's INDEX_KEY-coalesced 3x pool stride.
• Fix: Rename M3 override to _prepare_page_table_tensor(index_mapper_capacity) and populate the base's full state contract (kv_cache_pool_pointers, kv_cache_pool_mapping, index_scales, kv_offset, host_kv_cache_block_offsets); compute per-layer offset from layer_grouping position instead of via the broken exact_div.
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6417488

Summary by CodeRabbit

• Bug Fixes
  • Improved MiniMax-M3 sparse KV cache page-table preparation to correctly handle INDEX_KEY coalescing in the V2 cache layout, improving reliability when multiple buffers share the same pool.
  • Fixed per-pool/per-layer cache offset derivation and stride normalization to ensure accurate cache indexing.
  • Corrected cache initialization sizing and host KV block offset computation for all cache types, and updated per-request page-index capping to honor the configured block limits.

[coderabbitai]

📝 Walkthrough

Walkthrough

The MiniMax-M3 sparse KV cache manager now prepares page-table tensors from layer_grouping, derives pool offsets with the updated layout rules, and can cap per-request batch cache indices using num_blocks_per_seq.

Changes

MiniMax-M3 Page Table Preparation

Layer / File(s)	Summary
Import updates and override rationale tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py	Imports BufferConfig and exact_div, removes the prior LayerId/typed_range import, and updates the override docstring.
Page table and pool mapping computation tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py	Reworks _prepare_page_table_tensor(index_mapper_capacity) to derive kv_cache_pool_mapping from layer_grouping, compute per-pool kv_offset, and initialize host_kv_cache_block_offsets from the mapper capacity.
Batch cache index capping tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py	Updates _get_batch_cache_indices_by_pool_id to slice base page indices with an optional per-request num_blocks_per_seq limit.

Estimated code review effort: 3 (Moderate) | ~25 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is tied to the bug fix and uses the required ticket/type format, even if it is a bit verbose.
Description check	✅ Passed	The description explains the issue, fix, and test plan, but it does not follow the template's exact section headings.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[liji-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58453 [ run ] triggered by Bot. Commit: b9d6414 Link to invocation

[liji-nv]

My local try shows this change does not solve all the problem.

MiniMaxM3KVCacheManagerV2._get_batch_cache_indices_by_pool_id()
got an unexpected keyword argument 'num_blocks_per_seq'

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py`:
- Around line 435-443: The MinimaxM3 cache manager override is bypassing the
base SWA scratch-copy setup, so the scratch-reuse tensors are never initialized
when enable_swa_scratch_reuse is enabled. In CacheManager’s override where
host_kv_cache_block_offsets is allocated, make sure to preserve the base hook
behavior by calling _prepare_swa_scratch_copy_tensors(index_mapper_capacity) on
the enabled path, or otherwise delegate to the base initialization before
returning, so later scratch-reuse code in CacheManager and related SWA paths can
find the expected tensors.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 9dce3103-99a0-4840-98a3-6ed2447e3d1a

📥 Commits

Reviewing files that changed from the base of the PR and between d8c3ef4 and fa43f8d.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py

[liji-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58461 [ run ] triggered by Bot. Commit: 59133cb Link to invocation

[tensorrt-cicd]

PR_Github #58453 [ run ] completed with state ABORTED. Commit: b9d6414

Link to invocation

[tensorrt-cicd]

PR_Github #58461 [ run ] completed with state SUCCESS. Commit: 59133cb
/LLM/main/L0_MergeRequest_PR pipeline #47072 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

[liji-nv]

Waiting #16218 for a more proper fix.

[liji-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59844 [ run ] triggered by Bot. Commit: 89ffcfd Link to invocation

[tensorrt-cicd]

PR_Github #59844 [ run ] completed with state SUCCESS. Commit: 89ffcfd
/LLM/main/L0_MergeRequest_PR pipeline #48250 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

[BowenFu]

LGTM — the M3-specific override rename fixes a silently-dead hook (exact_div assert); the base kv_cache_manager_v2 change swaps try/except-KeyError for an all_buffer_ids membership check that is behavior-preserving — default non-sparse models still return None. Covered by a new unit test.

[nvpohanh]

[by Codex] @yizhang-nv Could you please review this PR? Thank you!