[pranav-nvidia]

Dev Engineer Review

MiniMaxM3KVCacheManagerV2 can now host explicit draft-layer specs and route draft buffers and page indices through their own pool lifecycle. It excludes draft layers from sparse target selection and rejects unsupported dtypes, layouts, and target/draft pool coalescing.

KVCacheManagerV2 adds support for subclass-provided layer buffers and window-based target/draft pool sizing when the required quota and lifecycle conditions are met. These changes add new cache allocation and indexing paths; verify those paths with the drafter configuration that will enable unified KV hosting. The supplied objectives state that current main does not activate this path.

No current review findings or test execution results were supplied. The reported test totals in the objectives are author-reported, not independently verified here.

QA Engineer Review

The new unit test file covers draft-layer ownership, dtype and buffer layout, page indices, request lifecycle, sparse-layer selection, unsupported configurations, and window-sized pool ratios. CUDA-dependent cases are marked to skip without CUDA; the pool-ratio test is CPU-only. The available output shows assertions and test declarations, but no test run results.

The test file did not appear in the searched tests/integration/test_lists/ paths. It is a unit test, so no CI or manual-QA integration-list entry was identified. Coverage verdict: needs follow-up to confirm test execution in the relevant CUDA environment.

Per-File QA Perspective

• tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py — Verify explicit draft-layer buffer access, index conversion, dummy-request handling, and rejection of unsupported layouts and dtype combinations.
• tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py — Verify explicit layer-buffer setup and window-based pool-ratio calculation, including fallback conditions and context-chunk quota boundaries.
• tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py — Adds CUDA-gated cache integration tests and CPU-only pool-ratio tests. No integration test-list entry was found.

Description

Groundwork for running MiniMax-M3.1's DSpark drafter on the unified target/draft KV cache from #20486 and #20664. On its own this PR changes nothing at runtime: nothing on main builds MiniMaxM3KVCacheManagerV2 with layer specs yet, and the unified-KV gate still keeps MiniMax-M3 on the separate draft pool. A follow-up lets the gate admit it.

KVCacheManagerV2 already hosts a drafter's layers as explicit layer_specs with their own lifecycle, but refuses any subclass that adds extra per-layer buffers, which rules out the MiniMax-M3 manager and its index-K buffers. This PR teaches the MiniMax-M3 manager to host them. The core of it is by Zheyu Fu.

• MiniMaxM3KVCacheManagerV2 opts in with hosts_unified_draft_layers. With an NVFP4 cache, dense target layer specs become FP8 to keep the hybrid sparse-NVFP4 / dense-FP8 layout. Draft layers use the base view in their own dtype and keep attention-op pools, because the drafter reads their block offsets; their page indices take the base conversion that matches that view, while target layers keep raw slot ids. Draft layers are never sparse or native target layers, and the global target count comes from the checkpoint, so a PP rank's subset of specs keeps the checkpoint's sparse split.
• The manager refuses layouts it cannot serve: NVFP4 draft layers (block scales belong to the sparse target layers), target specs outside the cache dtype, draft layers sharing a lifecycle (and so a coalesced pool) with target layers, and disaggregated draft layers under a non-HND K/V layout (drafters write head-major K/V).
• With explicit draft layers the draft view is the manager itself; add_dummy_requests drops it as a draft manager, so CUDA-graph padding does not create each dummy request twice.
• With layer specs, KVCacheManagerV2 appends a subclass's extra buffers (index-K) to the spec-built layer configs. A subclass that addresses its pools by size can keep NVFP4 data and scales in one coalescing group (_separate_nvfp4_coalescing); the MiniMax-M3 manager does, so the target's pools are unchanged. get_layer_bytes_per_token(Role.INDEX_KEY) answers for a manager with layer specs, whose specs do not carry index-K.
• KVCacheManagerV2._windowed_draft_pool_ratio: when the draft layers form one windowed lifecycle and neither a pool ratio nor avg_seq_len is set, the hot tier gives the draft lifecycle every admissible sequence's window plus one context chunk, over the resume watermark and capped at an equal-token share. The backend's default split models a single 2048-token request and says nothing about how many sequences hold a draft window at once. Lifecycles are keyed as the backend keys them, and the ratio comes with a one-context-chunk allocation floor, so rounding cannot leave the draft pool short of a full chunk. When the quota cannot hold that floor, the backend keeps its own split.
• KVCacheManagerV2.layer_specs defaults to None on the class, for managers built without __init__.

Test Coverage

tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py (new; GPU cases need CUDA storage only, no Blackwell kernels):

• test_draft_layer_specs_join_the_hybrid_nvfp4_manager: layer dtypes, target page strides and index-K shapes match a manager without draft layers, index-K bytes per token, the draft lifecycle and its base view, the block-offset rows the drafter reads, draft-layer page indices, the window-sized ratio, and dummy requests with the draft view passed back in.
• test_page_indices_follow_the_layer_owner, test_specs_decide_the_sparse_target_layers (default and checkpoint splits, with and without local draft specs), test_target_specs_need_the_cache_dtype, test_nvfp4_draft_layers_are_rejected, test_draft_layers_need_their_own_lifecycle, test_disagg_draft_layers_need_the_hnd_layout.
• test_window_sized_draft_pool_admits_a_full_context_chunk: a separately pooled draft lifecycle admits a full context chunk.
• test_window_sized_draft_ratio (CPU, aggregated and disaggregated, PP 1 and 2): ratio arithmetic, the resume watermark, the equal-token cap, the quota guard, and None without a chunk bound, without a windowed draft lifecycle, or when sinks put target and draft in one lifecycle.

Every test in the file fails on main. The existing MiniMax-M3 MSA, KV cache manager, MiniMax-M3 KV transfer and model suites show no new failures (3092 passed; the 2 failures are present on main as well).

PR Checklist

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

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[pranav-nvidia]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

No actionable comments were generated in the recent review.

Recent review info

Run configuration

• Source mode: Auto
• Repository: NVIDIA/TensorRT-LLM
• Branch: main
• File: .coderabbit.yaml
• Commit: 7969af8
• Layers: Repository YAML
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 7206c206-abf1-41a8-906c-bae0836839bd

Commits

Reviewing files that changed from the base of the PR and between 7969af8 and 5679974.

Files selected for processing: 3

• tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
• tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
• tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Walkthrough

MiniMax-M3 now supports explicit draft KV-cache layer specs alongside target layers. KVCacheManagerV2 configures per-layer buffers and can size a windowed draft pool. Tests cover configuration constraints, pool sizing, cache views, and index mappings.

Changes

Unified draft KV-cache layers

Layer / File(s)	Summary
Explicit layer buffers and pool sizing tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py, tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py	KVCacheManagerV2 supports subclass-provided buffers for explicit layer specs and can calculate a windowed draft-pool ratio when pool ratio and typical step are unset. Tests cover ratio calculation and fallback conditions.
MiniMax-M3 draft-layer configuration tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py, tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py	MiniMaxM3KVCacheManagerV2 identifies draft layers, applies target and draft spec constraints, and excludes draft layers from target sparse/native selection. Tests cover layer selection and rejected configurations.
Draft cache views and index routing tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py, tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_unified_draft_kv.py	Draft-layer buffer and cache-index access uses the draft pool’s base conversions. Integration tests cover cache views, index mappings, request cleanup, and allocation across target and draft pools.

Estimated code review effort: 4 (Complex) | ~35 minutes

Change: Feature

Suggested reviewers: schetlur-nv, lowsfer, lfr-0531

Merge Risk: Minimal · up to 56799

No actionable merge-blocking issue remains after the reviewed cache-routing and pool-sizing checks.

Pre-merge checks: 5

Check name	Status	Explanation
Title check		The title follows the repository template and clearly identifies the main change: hosting MiniMax-M3 drafter layers in the hybrid KV cache manager.
Description check		The description explains the motivation, implementation details, rejected configurations, test coverage, reported results, and checklist status. It is complete and relevant to the changeset.
Docstring Coverage		Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentati…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 80.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 20 functions across 3 files.

Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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