Skip to content

[None][feat] Native V2 KV cache event publishing #17023

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
tanmayv25 wants to merge 8 commits into
base: main
Choose a base branch
from
Draft

[None][feat] Native V2 KV cache event publishing #17023

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
92b8fd3
feat: publish native v2 kv cache events
alec-flowers Jul 25, 2026
195fd1d
perf: optimize native v2 kv event production
alec-flowers Jul 26, 2026
7c58c0d
refactor: streamline native KV event manager
tanmayv25 Jul 29, 2026
5a66ee1
refactor: nest KV events config under KvCacheConfig
tanmayv25 Jul 29, 2026
92f92de
fix: address independent review of native KV events
tanmayv25 Jul 29, 2026
902c2ee
refactor: collapse KVEventAdapter into NativeKVCacheEventManager
tanmayv25 Jul 29, 2026
bee8413
fix: address xhigh code-review findings for native KV events
tanmayv25 Jul 29, 2026
8f3474b
fix: address native KV events review comments
tanmayv25 Aug 5, 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
18 changes: 14 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 @@ -1142,6 +1142,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_cache_config.kv_events_config,
)

if not self._skip_est:
Expand Down Expand Up @@ -1858,7 +1861,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 @@ -1989,6 +1993,12 @@ 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
elif kv_events_config is not None and kv_events_config.enable_kv_cache_events:
logger.warning(
"kv_cache_config.kv_events_config is set but native KV event "
"publishing requires KV cache manager V2; events will not be "
f"published for {kv_cache_manager_cls.__name__}.")
if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2):
manager_extra_kwargs["is_disagg"] = is_disagg

Expand Down
Loading