Skip to content

[https://nvbugs/6438658][fix] Fix KV cache estimation capacity #16545

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

Merged
jiaganc merged 3 commits into from
Jul 22, 2026
Merged

[https://nvbugs/6438658][fix] Fix KV cache estimation capacity #16545

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
2e26722
[None][fix] Fix KV cache estimation capacity
jiaganc Jul 20, 2026
8604e52
Remove duplicate DeepSeek V4 cache config tests
jiaganc Jul 20, 2026
a119c9b
fix: forward max num tokens in DeepSeek V4 cache manager
jiaganc Jul 21, 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
117 changes: 10 additions & 107 deletions tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,16 +14,13 @@
# limitations under the License.

from collections import defaultdict
from dataclasses import replace
from typing import Dict, List, Optional, Tuple

import torch

from tensorrt_llm._torch.pyexecutor import llm_request
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import (
GPU_LEVEL,
BlockReusePolicy,
KVCacheManagerV2,
)
from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import GPU_LEVEL, KVCacheManagerV2
from tensorrt_llm._utils import (
TensorWrapper,
convert_to_torch_tensor,
Expand All @@ -39,16 +36,11 @@
from tensorrt_llm.runtime import ModelConfig
from tensorrt_llm.runtime.kv_cache_manager_v2 import (
AttentionLayerConfig,
BatchDesc,
BufferConfig,
DataRole,
GpuCacheTierConfig,
HostCacheTierConfig,
KVCacheDesc,
LayerId,
PageIndexMode,
ScratchDesc,
SwaScratchReuseConfig,
)
from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy
from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
Expand Down Expand Up @@ -254,7 +246,6 @@ def __init__(

self._init_indexer_dtype(sparse_attn_config)

# _build_cache_config() needs them to build constraints
self._max_input_len = max_input_len
self._max_num_tokens = max_num_tokens

Expand All @@ -271,6 +262,7 @@ def __init__(
vocab_size=vocab_size,
mapping=mapping,
dtype=dtype,
max_num_tokens=max_num_tokens,
**kwargs,
)
self.is_vswa = True # DeepSeek-V4 must has VSWA
Expand Down Expand Up @@ -768,16 +760,9 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:
return float("inf")
return self._max_num_tokens + (quota - context_limit_quota) / generation_size_per_token

def _build_cache_config(
self,
kv_cache_config: KvCacheConfig,
*,
tokens_per_block: int,
vocab_size: int | None,
cache_tiers: List[GpuCacheTierConfig | HostCacheTierConfig],
) -> KVCacheManagerConfigPy:
def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy:
"""
Create the cache manager config for DeepSeek-V4.
Add DeepSeek-V4 layers to the cache config.
"""
layers: List[AttentionLayerConfig] = []
layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {}
Expand Down Expand Up @@ -864,91 +849,9 @@ def _add_layer(
# number of layers in the KVCacheManagerPy
self._num_manager_layers = len(layers)

# Build constraints and typical_step for better pool ratio.
max_batch_size = self.max_batch_size
max_seq_len = self.max_seq_len
max_num_tokens = self._max_num_tokens
max_draft_len = self._max_draft_len
typical_step = None
constraints = []
if kv_cache_config.pool_ratio is None:
typical_seq_len = (
kv_cache_config.avg_seq_len
if kv_cache_config.avg_seq_len is not None
else max_seq_len
)
if typical_seq_len > max_seq_len:
raise ValueError(
f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or "
f"equal to max_seq_len ({max_seq_len})"
)

# For aggregated serving in large batch size:
# Use 1 context request + (max_batch_size - 1) generation requests as
# the typical step. An all-generation typical_step over-provisions the
# compressed-cache pool at the expense of the SWA pool, starving the
# SWA pool and artificially capping the achievable batch size.
ctx_capacity = (
max_num_tokens if max_num_tokens is not None else typical_seq_len
) + self.num_extra_kv_tokens
generation_history_length = max(0, typical_seq_len - max_draft_len - 1)
typical_step = BatchDesc(
kv_caches=[
KVCacheDesc(capacity=ctx_capacity, history_length=0),
]
+ [
KVCacheDesc(
capacity=typical_seq_len,
history_length=generation_history_length,
)
]
* (max_batch_size - 1),
)

# Constraint 1: cuda graph generation warmup — one decode request that has
# accumulated to the tail of max_seq_len. Using history_length=max_seq_len-1
# (instead of 0) lets SWA / SSM pools collapse to their windowed working set,
# while full-cache pools still need max_seq_len/tokens_per_block blocks
# because they don't age.
constraints.append(
BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1)])
)

# Constraint 2: general / chunked-prefill warmup — one fresh context request
# at max_num_tokens (the per-iteration token budget).
if max_num_tokens is not None:
constraints.append(
BatchDesc(
[
KVCacheDesc(
capacity=max_num_tokens + self.num_extra_kv_tokens, history_length=0
)
]
)
)

scratch_reuse_config = None
if self.enable_swa_scratch_reuse:
# Context requests will allocate num_extra_kv_tokens tokens for spec decoding.
# Cache manager should not take them into account when calculating scratch range.
# Therefore set max_rewind_len to num_extra_kv_tokens.
scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)

return KVCacheManagerConfigPy(
tokens_per_block=tokens_per_block,
cache_tiers=cache_tiers,
max_util_for_resume=kv_cache_config.max_util_for_resume,
enable_partial_reuse=kv_cache_config.enable_partial_reuse,
swa_scratch_reuse=scratch_reuse_config,
commit_min_snapshot=(
kv_cache_config.enable_block_reuse
and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE
),
return replace(
config,
layers=layers,
typical_step=typical_step,
constraints=constraints,
enable_stats=self.enable_stats,
initial_pool_ratio=kv_cache_config.pool_ratio,
)

def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None:
Expand Down Expand Up @@ -1163,9 +1066,9 @@ def get_layer_bytes_per_token(
local_layer_idx: int,
data_role: DataRole,
) -> int:
raise NotImplementedError(
"DeepSeek-V4 doesn't support get_layer_bytes_per_token, use _get_attn_bytes_per_block"
)
# The generic layers in the base config are replaced by
# _build_cache_config, so their buffer sizes are only placeholders.
return 1

def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor:
"""
Expand Down
4 changes: 2 additions & 2 deletions tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ def __init__(
disable_index_value_layer_ids = list(sparse_layer_ids)

# Must be set BEFORE super().__init__ — the base
# ``_build_cache_config`` invokes ``_extra_buffers_per_layer``
# ``_build_base_config`` invokes ``_extra_buffers_per_layer``
# which reads these attributes.
self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids)
self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids)
Expand Down Expand Up @@ -198,7 +198,7 @@ def _extra_buffers_per_layer(self, *, tokens_per_block):

``size`` is bytes per **block**: ``1 * sparse_index_dim *
elem_bytes * tokens_per_block``. Keyed by **local** layer id —
the base ``_build_cache_config`` iterates local ids, so keying
the base ``_build_base_config`` iterates local ids, so keying
by global ids would silently skip registration on non-trivial
PP ranks.
"""
Expand Down
Loading
Loading