Skip to content

[None][fix] Account for PARD draft KV capacity in cache manager V2 #18932

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
yizhang-nv wants to merge 3 commits into from
Closed

[None][fix] Account for PARD draft KV capacity in cache manager V2 #18932

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fa5a84a
[None][fix] Fix KVCM V2 PARD cache accounting
yizhang-nv Sep 9, 2026
1de090c
[None][refactor] Share V2 KV cache capacity accounting
yizhang-nv Sep 9, 2026
b8169da
[None][test] Run PARD accuracy cases with KV cache manager V2
yizhang-nv Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 35 additions & 16 deletions tensorrt_llm/_torch/pyexecutor/_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -762,7 +762,7 @@ def _per_manager_cache_cost(self,
tokens_per_block=self._tokens_per_block,
max_seq_len=self._max_seq_len,
max_batch_size=self._max_batch_size,
max_num_tokens=self._max_num_tokens if is_draft else 0,
max_num_tokens=getattr(self, "_max_num_tokens", 0),

@zhaoyangwang-nvidia zhaoyangwang-nvidia Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_max_num_tokens is set unconditionally in __init__ (L631) and read directly everywhere else in this file (L879, L1470, L1671, L2058), so this fallback is only needed by the hand-built object.__new__(KvCacheCreator) double in test_kv_cache_budget_split._make_creator, which never sets it. On main the conditional expression self._max_num_tokens if is_draft else 0 short-circuited and never touched the attribute for is_draft=False; reading it unconditionally is what makes those ~15 _split_kv_cache_budget_for_draft tests need the fallback.

Please set _max_num_tokens in that fixture (as test_kv_cache_estimation.py L169 and test_dual_pool_kv_cache.py L179 already do) and use self._max_num_tokens here -- otherwise a genuinely missing attribute silently estimates a zero fixed cost instead of failing.

kv_cache_config=kv_cache_config,
spec_config=self._speculative_config,
is_draft=is_draft,
Expand Down Expand Up @@ -797,14 +797,29 @@ def _get_kv_size_per_token(self,
use_separate_draft_kv_cache=use_separate_draft_kv_cache)
if self._is_encoder_decoder():
total += CacheCost.from_raw(self._get_cross_kv_size_per_token())
draft_cost = self._get_draft_cache_cost(
kv_cache_config,
use_separate_draft_kv_cache=use_separate_draft_kv_cache,
)
if draft_cost is not None:
total += draft_cost
return total

def _get_draft_cache_cost(
self,
kv_cache_config: KvCacheConfig,
*,
use_separate_draft_kv_cache: bool,
) -> Optional[CacheCost]:
"""Return the draft manager's standalone cache cost, if it has one."""
if self._draft_model_engine is not None:
draft_model_config = self._draft_model_engine.model.model_config
draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls(
self._draft_model_engine, kv_cache_config)
total += self._per_manager_cache_cost(draft_kv_cache_manager_cls,
draft_model_config,
kv_cache_config)
elif use_separate_draft_kv_cache:
return self._per_manager_cache_cost(draft_kv_cache_manager_cls,
draft_model_config,
kv_cache_config)
if use_separate_draft_kv_cache:
# One-model draft with separate KV cache layout.
# Pass num_layers explicitly since the HF config may report a
# different layer count than what is actually used at runtime
Expand All @@ -822,20 +837,19 @@ def _get_kv_size_per_token(self,
effective_draft_config,
draft_kv_cache_config,
is_disagg=self._is_disagg)
total += self._per_manager_cache_cost(
draft_kv_cache_manager_cls,
effective_draft_config,
draft_kv_cache_config,
is_draft=True)
return self._per_manager_cache_cost(draft_kv_cache_manager_cls,
effective_draft_config,
draft_kv_cache_config,
is_draft=True)
elif self._mapping.is_last_pp_rank():
# EAGLE3/MTP: draft layers only on last PP rank
total += self._per_manager_cache_cost(
return self._per_manager_cache_cost(
self._kv_cache_manager_cls,
effective_draft_config,
draft_kv_cache_config,
num_layers=self._get_num_draft_layers(),
is_draft=True)
return total
return None

def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
allocated_bytes: int) -> int:
Expand Down Expand Up @@ -1679,16 +1693,21 @@ def _get_target_and_draft_cache_costs(
"""Per-manager KV cache costs for target and draft layers."""
target_kv_cache_config = (kv_cache_config if kv_cache_config is not None
else self._kv_cache_config)
total_kv = self._get_kv_size_per_token(target_kv_cache_config)
use_separate_draft_kv_cache = (
self._should_create_separate_draft_kv_cache())
target_kv = self._per_manager_cache_cost(
self._kv_cache_manager_cls,
self._model_engine.model.model_config,
target_kv_cache_config,
use_separate_draft_kv_cache=use_separate_draft_kv_cache)
draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope,
intercept=total_kv.intercept - target_kv.intercept)
# Estimate the draft component directly so its independently modelled
# affine intercept is preserved exactly.
draft_kv = self._get_draft_cache_cost(
target_kv_cache_config,
use_separate_draft_kv_cache=use_separate_draft_kv_cache,
)
if draft_kv is None:
return None
costs = (target_kv, draft_kv)
if any(cost.slope < 0 or cost.intercept < 0 or (
cost.slope == 0 and cost.intercept == 0) for cost in costs):
Expand Down Expand Up @@ -1790,7 +1809,7 @@ def _split_kv_cache_budget_for_draft(
raise ValueError(
f"KV cache GPU budget ({total_budget / GB:.2f} GiB) is "
f"insufficient after the combined fixed cost "
f"({intercept_total / GB:.2f} GiB, e.g. mamba SSM state) "
f"({intercept_total / GB:.2f} GiB, e.g. SWA or mamba state) "
f"for target+draft. Increase free_gpu_memory_fraction or "
f"max_gpu_total_bytes, or reduce max_batch_size (the fixed "
f"cost scales with batch size).")
Expand Down
187 changes: 105 additions & 82 deletions tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,11 @@
"iter_partial_reused_blocks",
"iter_missed_blocks",
)

# Every generation step reserves the golden/base token in addition to any
# speculative draft tokens. Keep this shared with the capacity-growth paths so
# cache estimation and runtime allocation cannot drift apart.
BASE_GENERATION_TOKEN_COUNT = 1
KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS = tuple(
field_name
for field_name in KV_CACHE_ITERATION_STATS_DELTA_FIELDS
Expand Down Expand Up @@ -331,23 +336,20 @@ def _estimate_swa_cache_size(
*,
context: bool,
scratch: bool,
generation_capacity_headroom: Optional[int] = None,
generation_capacity_headroom: int = BASE_GENERATION_TOKEN_COUNT,
) -> tuple[int, int]:
tokens_per_block = int(tokens_per_block)
size_per_token = 0
size_per_request = 0
scratch_keys = set()
for layer_size, window_size in zip(layer_sizes, attention_windows):
if window_size is not None and window_size > 0:
if generation_capacity_headroom is None:
window_blocks = math.ceil(window_size / tokens_per_block)
else:
# Match DFlash's retained boundary page and capacity reserved
# ahead of committed history for the next draft step.
window_blocks = (
math.ceil((window_size + generation_capacity_headroom - 1) / tokens_per_block)
+ 1
)
# Match AttnLifeCycle.get_stale_range(): the live interval contains
# window_size + generation_capacity_headroom - 1 tokens. Across all
# page offsets, that interval touches at most the count below.
window_blocks = (
math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1
)
window_tokens = window_blocks * tokens_per_block
if not context:
size_per_request += window_tokens * layer_size
Expand All @@ -363,16 +365,63 @@ def _estimate_swa_cache_size(
return size_per_token, size_per_request


def _get_dflash_generation_kv_capacity_headroom(spec_config) -> Optional[int]:
"""DFlash KV capacity reserved ahead of committed history."""
from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode
def _get_kv_reserve_draft_tokens(spec_config, *, is_draft: bool) -> int:
"""Return the same speculative KV reserve used by runtime resize paths."""
if spec_config is None:
return 0
reserve = spec_config.max_total_draft_tokens
if (
is_draft
and getattr(spec_config, "use_dynamic_tree", False)
and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0
):
draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len
reserve = max(reserve, draft_loop_tokens)
return reserve


def _estimate_cache_size_components(
layer_sizes: Sequence[int],
attention_windows: Sequence[int | None],
tokens_per_block: int,
*,
scratch: bool,
generation_capacity_headroom: int,
) -> tuple[int, int, int]:
"""Return context/generation bytes per token and generation bytes per request.

Static profiling and runtime quota conversion must charge the same SWA
retention pages and context scratch space. Resume-watermark normalization
is separate from these usable-capacity costs.
"""
full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows)
context_swa_size, _ = _estimate_swa_cache_size(
layer_sizes, attention_windows, tokens_per_block, context=True, scratch=scratch
)
generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size(
layer_sizes,
attention_windows,
tokens_per_block,
context=False,
scratch=False,
generation_capacity_headroom=generation_capacity_headroom,
)
return (
full_attn_size + context_swa_size,
full_attn_size + generation_swa_size,
generation_swa_per_request,
)

if spec_config is None or spec_config.spec_dec_mode != SpeculativeDecodingMode.DFLASH:
return None

def _get_generation_kv_capacity_headroom(spec_config, *, is_draft: bool) -> int:
"""Maximum capacity lead over history used by generation KV allocation."""
from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens

return get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step
if spec_config is None:
return BASE_GENERATION_TOKEN_COUNT
reserve = _get_kv_reserve_draft_tokens(spec_config, is_draft=is_draft)
dynamic_reserve = reserve - spec_config.max_total_draft_tokens
return get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step + dynamic_reserve


def _get_single_swa_pool_slot_bytes(
Expand Down Expand Up @@ -1029,15 +1078,12 @@ def __init__(
if not self._supports_reuse_match_backoff:
self.reuse_match_backoff = 0
# Mirror V1's KV reserve sizing (see V1 __init__ for rationale).
self._kv_reserve_draft_tokens = self.max_total_draft_tokens
if (
self.is_draft
and spec_config is not None
and getattr(spec_config, "use_dynamic_tree", False)
and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0
):
draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len
self._kv_reserve_draft_tokens = max(self.max_total_draft_tokens, draft_loop_tokens)
self._kv_reserve_draft_tokens = _get_kv_reserve_draft_tokens(
spec_config, is_draft=self.is_draft
)
self._generation_kv_capacity_headroom = _get_generation_kv_capacity_headroom(
spec_config, is_draft=self.is_draft
)

self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
self.enable_stats = enable_stats
Expand Down Expand Up @@ -1417,11 +1463,8 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
self.max_seq_len = int(max_num_tokens)

# Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel).
# Account for max single-sequence capacity = seq_len + extra KV tokens +
# _kv_reserve_draft_tokens (see __init__) + 1 base decode token.
max_seq_capacity = (
self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1
)
# Include the same maximum generation lead used by allocation and sizing.
max_seq_capacity = self.max_seq_len + self._generation_kv_capacity_headroom
self.max_blocks_per_seq = (
max_seq_capacity + self._ledger_tokens_per_block - 1
) // self._ledger_tokens_per_block
Expand Down Expand Up @@ -1707,33 +1750,26 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:

def _get_max_tokens_from_quota_impl(self, quota: int) -> float:
layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components()
full_attn_size_per_token = _estimate_full_attn_size_per_token(
layer_sizes, attention_windows
)
context_swa_size_per_token, _ = _estimate_swa_cache_size(
(
context_size_per_token,
generation_size_per_token,
generation_swa_size_per_request,
) = _estimate_cache_size_components(
layer_sizes,
attention_windows,
self.tokens_per_block,
context=True,
scratch=self.enable_swa_scratch_reuse,
)
(
generation_swa_size_per_token,
generation_swa_size_per_request,
) = _estimate_swa_cache_size(
layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False
generation_capacity_headroom=self._generation_kv_capacity_headroom,
)
size_per_batch = self.max_batch_size * generation_swa_size_per_request
if quota < size_per_batch:
return 0
context_size_per_token = full_attn_size_per_token + context_swa_size_per_token
context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch
if quota <= context_limit_quota:
if context_size_per_token <= 0:
return float("inf")
return (quota - size_per_batch) / context_size_per_token

generation_size_per_token = full_attn_size_per_token + generation_swa_size_per_token
if generation_size_per_token <= 0:
return float("inf")
return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token
Expand All @@ -1750,34 +1786,24 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int:

def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int:
layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components()
full_attn_size_per_token = _estimate_full_attn_size_per_token(
layer_sizes, attention_windows
)
(
context_swa_size_per_token,
_,
) = _estimate_swa_cache_size(
context_size_per_token,
generation_size_per_token,
generation_swa_size_per_request,
) = _estimate_cache_size_components(
layer_sizes,
attention_windows,
self.tokens_per_block,
context=True,
scratch=self.enable_swa_scratch_reuse,
)
(
generation_swa_size_per_token,
generation_swa_size_per_request,
) = _estimate_swa_cache_size(
layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False
generation_capacity_headroom=self._generation_kv_capacity_headroom,
)
context_tokens = min(max_tokens, self.max_num_tokens)
generation_tokens = max_tokens - context_tokens
generation_quota = (
max_tokens * full_attn_size_per_token
+ generation_tokens * generation_swa_size_per_token
return int(
context_tokens * context_size_per_token
+ generation_tokens * generation_size_per_token
+ self.max_batch_size * generation_swa_size_per_request
)
context_extra_quota = context_tokens * context_swa_size_per_token
return int(generation_quota + context_extra_quota)

def _get_event_num_blocks_per_cache_level(
self,
Expand Down Expand Up @@ -2546,7 +2572,7 @@ def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int:

Grows *current_capacity* by 1 + draft tokens.
"""
return current_capacity + 1 + self._generation_draft_slots(req)
return current_capacity + BASE_GENERATION_TOKEN_COUNT + self._generation_draft_slots(req)

def _generation_draft_slots(self, req: LlmRequest) -> int:
"""Physical draft width reserved for one iteration. Dynamic-tree draft pools
Expand Down Expand Up @@ -4162,37 +4188,34 @@ def get_cache_size_per_token(
backoff,
max_seq_len,
)
full_attn_size_per_token = _estimate_full_attn_size_per_token(
layer_sizes, attention_windows
generation_capacity_headroom = _get_generation_kv_capacity_headroom(
spec_config, is_draft=is_draft
)
dflash_headroom = _get_dflash_generation_kv_capacity_headroom(spec_config)
is_dflash_draft = is_draft and dflash_headroom is not None
generation_capacity_headroom = dflash_headroom if is_dflash_draft else None
swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size(
(
context_size_per_token,
cache_size_per_token,
swa_size_per_request,
) = _estimate_cache_size_components(
layer_sizes,
attention_windows,
tokens_per_block,
context=False,
scratch=False,
scratch=bool(kwargs.get("enable_swa_scratch_reuse", False)),

@zhaoyangwang-nvidia zhaoyangwang-nvidia Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

kwargs.get("enable_swa_scratch_reuse") is never supplied by the only production caller: KvCacheCreator._per_manager_cache_cost forwards just use_separate_draft_kv_cache / num_layers (all four call sites at _util.py L793/819/840/846/1698), so this is always False here. The runtime uses kv_cache_config.enable_swa_scratch_reuse and not is_draft (__init__ L1013) for the same estimate in _get_quota_from_max_tokens_impl/_get_max_tokens_from_quota_impl. On a 30-SWA-layer / 2048 B-per-token-per-layer config with tokens_per_block=64 and max_num_tokens=8192, that is a ~464 MiB divergence between the static intercept and the runtime quota whenever the flag is on.

This classmethod already receives kv_cache_config and is_draft, so it can derive the flag with the same expression __init__ uses. As written, test_v2_static_and_runtime_cache_costs_agree passes the kwarg explicitly and therefore cannot catch the drift.

generation_capacity_headroom=generation_capacity_headroom,
)
context_swa_size_per_token = 0
if is_dflash_draft:
context_swa_size_per_token, _ = _estimate_swa_cache_size(
layer_sizes,
attention_windows,
tokens_per_block,
context=True,
scratch=False,
)
# The affine slope covers all tokens; context additionally retains SWA
# pages for the current token batch beyond the generation windows.
fixed_cost = (
swa_size_per_request * max_batch_size + context_swa_size_per_token * max_num_tokens
swa_size_per_request * max_batch_size
+ (context_size_per_token - cache_size_per_token) * max_num_tokens

@zhaoyangwang-nvidia zhaoyangwang-nvidia Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This term is a real fix -- the runtime quota conversion has always charged context_tokens * context_swa_size_per_token, while the static estimate for the target manager charged zero (max_num_tokens was passed as 0 for is_draft=False, and context_swa_size_per_token was only computed on the DFlash-draft path). But it now lands on every SWA model, not just PARD.

On a Gemma3-27B-shaped config (30 SWA layers @ w=1024, 10 full-attn, 2048 B/token/layer, tpb=64, bs=8, max_num_tokens=8192) the target intercept goes from ~0.47 GiB to ~0.97 GiB: ~469 MiB from this context term and ~30 MiB from the +1 page that the reworked window_blocks formula adds at headroom=1. Could the description and the test coverage call this out and include an SWA model, so the available-capacity drop is a known consequence rather than a surprise?

(For what it is worth, I brute-forced the new ceil((w + headroom - 2) / tpb) + 1 against AttnLifeCycle.get_stale_range() over tpb in {16,32,64}, w in [1,300), headroom in {1,2,5,8,11,20} -- it is exactly the tight maximum, so both the old ceil(w/tpb) under-count and the old DFlash -1 over-count are corrected.)

)
cache_size_per_token = full_attn_size_per_token + swa_size_per_token
bytes_per_slot = _get_single_swa_pool_slot_bytes(
layer_sizes, attention_windows, tokens_per_block
)
if is_dflash_draft and bytes_per_slot is not None:
if is_draft and fixed_cost > 0 and bytes_per_slot is not None:

@zhaoyangwang-nvidia zhaoyangwang-nvidia Sep 10, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Minor robustness point, not currently reachable: kv_cache_config is Optional[KvCacheConfig] = None in this signature but is dereferenced unguarded two lines below (kv_cache_config.max_util_for_resume). Today the only caller that omits it (_get_cross_kv_size_per_token, _util.py L1948) leaves is_draft=False so the guard short-circuits, but this PR widens the condition from "DFlash draft" to any draft manager with a single SWA pool. Either handle None explicitly or make the parameter required.

# The affine intercept is the configured quota needed to preserve
# the fixed usable capacity at V2's resume watermark. Keep the
# allocator's page rounding inside the manager estimator rather
# than extending the generic CacheCost model with pool geometry.
required_slots = math.ceil(fixed_cost / bytes_per_slot)
resume_util = float(np.float32(kv_cache_config.max_util_for_resume))
fixed_cost = math.ceil(required_slots / resume_util) * bytes_per_slot
Expand Down
Loading
Loading