Skip to content

[None][perf] optimize native V2 KV event production #16876

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

Draft
alec-flowers wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

[None][perf] optimize native V2 KV event production #16876

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
98436f7
feat: publish native v2 kv cache events
alec-flowers Jul 25, 2026
22e6f17
perf: optimize native v2 kv event production
alec-flowers Jul 26, 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
13 changes: 9 additions & 4 deletions tensorrt_llm/_torch/pyexecutor/_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,9 @@
# isort: off
from tensorrt_llm.llmapi.llm_args import (
CacheTransceiverConfig, CapacitySchedulerPolicy, EagleDecodingConfig,
KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig, PeftCacheConfig,
SamplerType, SchedulerConfig, SparseAttentionConfig, SpeculativeConfig,
TorchLlmArgs, WaitingQueuePolicy)
KVEventsConfig, KvCacheCompressionConfig, KvCacheConfig, MTPDecodingConfig,
PeftCacheConfig, SamplerType, SchedulerConfig, SparseAttentionConfig,
SpeculativeConfig, TorchLlmArgs, WaitingQueuePolicy)
# isort: on
from tensorrt_llm.logger import logger
from tensorrt_llm.lora_helper import (LoraConfig,
Expand Down Expand Up @@ -1012,6 +1012,9 @@ def _create_kv_cache_manager(
execution_stream=self._execution_stream,
layer_mask=spec_dec_layer_mask,
is_disagg=self._is_disagg,
kv_events_config=None
if estimating_kv_cache or model_engine.is_draft_model else
self._llm_args.kv_events_config,
)

if not self._skip_est:
Expand Down Expand Up @@ -1721,7 +1724,8 @@ def _create_kv_cache_manager(
num_kv_heads: Optional[Union[int, List[int]]] = None,
head_dim: Optional[int] = None,
kv_cache_type=None,
is_disagg: bool = False) -> KVCacheManager:
is_disagg: bool = False,
kv_events_config: Optional[KVEventsConfig] = None) -> KVCacheManager:
"""
Returns:
A KVCacheManager instance for the given model engine or model config
Expand Down Expand Up @@ -1839,6 +1843,7 @@ def _create_kv_cache_manager(
manager_extra_kwargs = {}
if issubclass(kv_cache_manager_cls, KVCacheManagerV2):
manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats
manager_extra_kwargs["kv_events_config"] = kv_events_config

if is_mla(config):
kv_cache_manager = kv_cache_manager_cls(
Expand Down
Loading
Loading