[None][fix] Fix KV cache V2 OOM with separate draft KV cache (EAGLE3/MTP)

[yizhang-nv]

@coderabbitai summary

Description

Fix OOM failures when KV cache manager V2 is used with one-model speculative decoding (EAGLE3/MTP) that requires a separate draft KV cache.

Root cause: When V2 creates both a target and a draft KV cache manager, both managers independently claim the full max_gpu_total_bytes budget, leading to OOM during prepare_resources.

Changes:

• Budget splitting (_util.py): Split max_gpu_total_bytes proportionally between target and draft KV cache managers based on per-token KV cache sizes. Only applies to final manager creation (not estimation).
• Layer count fix (_util.py, resource_manager.py): Add num_layers parameter to get_cache_size_per_token for draft models where the HF config layer count differs from runtime (e.g. EAGLE3: config says 1 layer, runtime uses 4). Prevents pp_layers() from giving wrong results for draft layers.
• PP rank handling (_util.py): Only count draft KV cache cost on the last PP rank where EAGLE3/MTP draft layers reside.
• Sparse attention config (_util.py): Pass sparse_attention_config from effective_draft_config to draft KV cache manager for MoE models (e.g. DeepSeek V3 with DSA).
• BAD_PAGE_INDEX fix (kvCacheManagerV2Utils.cu): Use 0 instead of BAD_PAGE_INDEX for padding in copyBatchBlockOffsetsToDeviceKernel to avoid out-of-bounds memory access in attention kernels.
• Graceful resize failure (resource_manager.py): Handle resize failures in dummy request preparation by returning None instead of raising ValueError, letting the scheduler retry.

Test Coverage

TBD — V2 is not yet enabled by default; will be covered by the enable-v2-by-default PR.

PR Checklist

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

[coderabbitai]

📝 Walkthrough

Walkthrough

Updates KV-cache management for draft models including budget splitting logic, explicit layer count propagation, and resource reset methods. Modifies KV-cache size calculation signatures to accept optional layer counts, affecting both Python executor utilities and resource manager classes.

Changes

Cohort / File(s)	Summary
KV Cache Sentinel Handling cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu	Modified BAD_PAGE_INDEX handling in copyBatchBlockOffsetsToDeviceKernel; when sentinel value encountered, maps to 0 instead of propagating the sentinel for both key and value caches.
Draft KV Cache Budget Management tensorrt_llm/_torch/pyexecutor/_util.py	Added budget splitting logic for separate draft KV caches with conditional size calculation per pipeline rank. Introduced _get_num_draft_layers, _split_kv_cache_budget_for_draft helpers, and extended _create_one_model_draft_kv_cache_manager to accept kv_cache_config_override parameter.
KV Cache Manager API Updates tensorrt_llm/_torch/pyexecutor/resource_manager.py	Extended get_cache_size_per_token signatures in both KVCacheManager and KVCacheManagerV2 to accept optional num_layers parameter. Added reset_reuse_state methods to both classes for clearing reusable blocks state.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 53.85% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly describes the main fix: resolving OOM issues with KV cache V2 when using separate draft KV cache for EAGLE3/MTP.
Description check	✅ Passed	PR description comprehensively explains the root cause, lists all specific changes across multiple files, and provides clear rationale.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

📝 Coding Plan

• Generate coding plan for human review comments

---

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

Tip

You can customize the high-level summary generated by CodeRabbit.

Configure the reviews.high_level_summary_instructions setting to provide custom instructions for generating the high-level summary.

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/pyexecutor/_util.py (1)

754-779: ⚠️ Potential issue | 🟠 Major

Draft KV budget can remain unsplit while a draft manager is still created on non-last PP ranks.

At Line 754-779, draft_kv_cache_config may stay None (when draft_kv <= 0), but one-model draft manager creation still proceeds with the unsplit target config. Combined with the fallback non-empty PP layer behavior in get_pp_layers (tensorrt_llm/_torch/pyexecutor/resource_manager.py, Line 147-150), this can still reserve extra KV budget and reintroduce OOM pressure.

💡 Proposed fix

@@
-        draft_kv_cache_config = None
-        if (not estimating_kv_cache
-                and self._should_create_separate_draft_kv_cache()
-                and issubclass(self._kv_cache_manager_cls, KVCacheManagerV2)):
+        draft_kv_cache_config = None
+        should_host_one_model_draft_on_this_rank = (
+            self._speculative_config.spec_dec_mode.is_external_drafter()
+            or self._mapping.is_last_pp_rank())
+        if (not estimating_kv_cache
+                and self._should_create_separate_draft_kv_cache()
+                and issubclass(self._kv_cache_manager_cls, KVCacheManagerV2)
+                and should_host_one_model_draft_on_this_rank):
             draft_kv_cache_config = self._split_kv_cache_budget_for_draft()
@@
-        elif self._should_create_separate_draft_kv_cache():
-            draft_kv_cache_manager = self._create_one_model_draft_kv_cache_manager(
-                estimating_kv_cache,
-                kv_cache_config_override=draft_kv_cache_config)
+        elif self._should_create_separate_draft_kv_cache():
+            if (issubclass(self._kv_cache_manager_cls, KVCacheManagerV2)
+                    and not should_host_one_model_draft_on_this_rank):
+                draft_kv_cache_manager = None
+            else:
+                draft_kv_cache_manager = self._create_one_model_draft_kv_cache_manager(
+                    estimating_kv_cache,
+                    kv_cache_config_override=draft_kv_cache_config)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 754 - 779, The code may
create a one-model draft KV manager even when draft_kv_cache_config is None
(e.g., draft_kv <= 0), causing unintentional KV reservations; update the branch
that calls _create_one_model_draft_kv_cache_manager so it only creates
draft_kv_cache_manager when draft_kv_cache_config is not None (or >0) — i.e.,
add an explicit guard checking draft_kv_cache_config before invoking
_create_one_model_draft_kv_cache_manager (and similarly ensure
_split_kv_cache_budget_for_draft returns None for no-split cases); reference
symbols: draft_kv_cache_config, _split_kv_cache_budget_for_draft,
_create_one_model_draft_kv_cache_manager, _create_kv_cache_manager, and consider
interplay with get_pp_layers behavior to avoid reserving extra KV on non-last PP
ranks.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 2121-2122: The failure path after creating the draft KV cache
leaks the per-request draft resources because release_resources(req) does not
call draft_kv_cache_manager.free_resources(current_request); update the error
handling where kv_cache.resize(new_capacity) can fail (the block that calls
release_resources(req) and returns None) to explicitly call
draft_kv_cache_manager.free_resources(current_request) (or the equivalent
cleanup method) before calling release_resources(req) and returning, ensuring
draft_kv resources for current_request are freed even on resize failure.

---

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 754-779: The code may create a one-model draft KV manager even
when draft_kv_cache_config is None (e.g., draft_kv <= 0), causing unintentional
KV reservations; update the branch that calls
_create_one_model_draft_kv_cache_manager so it only creates
draft_kv_cache_manager when draft_kv_cache_config is not None (or >0) — i.e.,
add an explicit guard checking draft_kv_cache_config before invoking
_create_one_model_draft_kv_cache_manager (and similarly ensure
_split_kv_cache_budget_for_draft returns None for no-split cases); reference
symbols: draft_kv_cache_config, _split_kv_cache_budget_for_draft,
_create_one_model_draft_kv_cache_manager, _create_kv_cache_manager, and consider
interplay with get_pp_layers behavior to avoid reserving extra KV on non-last PP
ranks.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 5ed31748-802c-44d3-953b-174d66262207

📥 Commits

Reviewing files that changed from the base of the PR and between 60091ff and 3e982f1.

📒 Files selected for processing (3)

• cpp/tensorrt_llm/batch_manager/kvCacheManagerV2Utils.cu
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #38972 [ run ] triggered by Bot. Commit: b945ae1 Link to invocation

[tensorrt-cicd]

PR_Github #38972 [ run ] completed with state SUCCESS. Commit: b945ae1
/LLM/main/L0_MergeRequest_PR pipeline #30254 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

Link to invocation

[yizhang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #39006 [ run ] triggered by Bot. Commit: b945ae1 Link to invocation

[tensorrt-cicd]

PR_Github #39006 [ run ] completed with state SUCCESS. Commit: b945ae1
/LLM/main/L0_MergeRequest_PR pipeline #30285 completed with status: 'SUCCESS'

CI Report

Link to invocation

[HuiGao-NV]

LGTM