[TRTLLM-12112][feat] Support v2 KV cache stats

cpp/include/tensorrt_llm/batch_manager/llmRequest.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2025, NVIDIA CORPORATION.  All rights reserved.
+ * Copyright (c) 2022-2026, NVIDIA CORPORATION.  All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -1835,6 +1835,15 @@ class GenericLlmRequest
         return mPerfMetrics.kvCacheMetrics.numNewAllocatedBlocks;
     }
 
+    void updateKvCachePerfMetrics(
+        SizeType32 allocTotalBlocks, SizeType32 allocNewBlocks, SizeType32 reusedBlocks, SizeType32 missedBlocks)
+    {
+        updateAllocTotalBlocksPerRequest(allocTotalBlocks);
+        updateAllocNewBlocksPerRequest(allocNewBlocks);
+        updateReusedBlocksPerRequest(reusedBlocks);
+        updateMissedBlocksPerRequest(missedBlocks);
+    }
+
     void updateReusedBlocksPerRequest(SizeType32 reusedBlocksPerRequest)
     {
         mPerfMetrics.kvCacheMetrics.numReusedBlocks += reusedBlocksPerRequest;

cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -191,11 +191,19 @@ void initBindings(nb::module_& m)
         .def_prop_ro("is_disagg_context_complete_state", &GenLlmReq::isDisaggContextCompleteState)
         .def_prop_ro("stage", &GenLlmReq::getRequestStage)
         .def_prop_ro("kv_cache_transfer_time_ms", &GenLlmReq::getKvCacheTransferTimeMS)
+        .def_prop_ro("kv_cache_transfer_start", &GenLlmReq::getKvCacheTransferStart)
+        .def_prop_ro("kv_cache_transfer_end", &GenLlmReq::getKvCacheTransferEnd)
         .def_prop_ro("kv_cache_size", &GenLlmReq::getKvCacheSize)
+        .def("set_kv_cache_transfer_start", &GenLlmReq::setKvCacheTransferStart, nb::arg("time"))
+        .def("set_kv_cache_transfer_end", &GenLlmReq::setKvCacheTransferEnd, nb::arg("time"))
+        .def("set_kv_cache_size", &GenLlmReq::setKvCacheSize, nb::arg("target_buffer_size"))
+        .def("update_kv_cache_size", &GenLlmReq::updateKvCacheSize, nb::arg("target_buffer_size"))
         .def_prop_ro("avg_decoded_tokens_per_iter", &GenLlmReq::getAvgDecodedTokensPerIter)
         .def_prop_ro("alloc_total_blocks", &GenLlmReq::getAllocTotalBlocksPerRequest)
         .def_prop_ro("alloc_new_blocks", &GenLlmReq::getAllocNewBlocksPerRequest)
         .def("alloc_context_logits", &GenLlmReq::allocContextLogitsHost, nb::arg("vocab_size"), nb::arg("logit_dtype"))
+        .def("update_kv_cache_perf_metrics", &GenLlmReq::updateKvCachePerfMetrics, nb::arg("alloc_total_blocks"),
+            nb::arg("alloc_new_blocks"), nb::arg("reused_blocks"), nb::arg("missed_blocks"))
         .def_prop_ro("reused_blocks", &GenLlmReq::getReusedBlocksPerRequest)
         .def_prop_ro("missed_blocks", &GenLlmReq::getMissedBlocksPerRequest)
         .def_prop_ro("kv_cache_hit_rate", &GenLlmReq::getKVCacheHitRatePerRequest)

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py

@@ -1,3 +1,18 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 from collections import defaultdict
 from typing import Dict, List, Optional, Tuple
 
@@ -21,6 +36,7 @@
     AttentionLayerConfig,
     BatchDesc,
     BufferConfig,
+    DataRole,
     GpuCacheTierConfig,
     HostCacheTierConfig,
     KVCacheDesc,
@@ -229,6 +245,22 @@ def __init__(
             device="cpu",
         )
 
+    def _format_kv_cache_pool_lifecycle_entry(self, layer_id: LayerId, role: DataRole) -> str:
+        layer_semantics = self._manager_layer_id_to_layer_attn.get(layer_id)
+        if layer_semantics is None:
+            return super()._format_kv_cache_pool_lifecycle_entry(layer_id, role)
+
+        model_layer_idx, attn_type = layer_semantics
+        attr = self.impl._storage.get_buffer_attr(layer_id, role)
+        pool_group_id = self.impl._storage.get_pool_group_index(attr.life_cycle_id)
+        lifecycle = self.impl._life_cycles.get_life_cycle(attr.life_cycle_id)
+        return (
+            f"deepseek_role={attn_type.name}, "
+            f"compress_ratio={self._compress_ratios[model_layer_idx]}, "
+            f"pool_group_id={int(pool_group_id)}, "
+            f"lifecycle_id={int(attr.life_cycle_id)}, lifecycle={lifecycle}"
+        )
+
     def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> torch.Tensor:
         """
         Get the buffers for a specific layer and attention type.
@@ -365,6 +397,7 @@ def _build_cache_config(
         """
         layers: List[AttentionLayerConfig] = []
         layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {}
+        manager_layer_id_to_layer_attn: Dict[LayerId, Tuple[int, DeepseekV4AttentionType]] = {}
 
         def _add_layer(
             layer_idx: int, attn_type: DeepseekV4AttentionType, sliding_window_size: int | None
@@ -373,6 +406,7 @@ def _add_layer(
             layer_id = LayerId(len(layers))
             # update the mapping from layer index and attention type to layer id
             layer_attn_to_layer_id[layer_idx, attn_type] = layer_id
+            manager_layer_id_to_layer_attn[layer_id] = (layer_idx, attn_type)
             # add the layer to the layers list
             layer_config = AttentionLayerConfig(
                 layer_id=layer_id,
@@ -433,6 +467,7 @@ def _add_layer(
                 )
         # the mapping from layer index and attention type to layer id
         self._layer_attn_to_layer_id = layer_attn_to_layer_id
+        self._manager_layer_id_to_layer_attn = manager_layer_id_to_layer_attn
         # number of layers in the KVCacheManagerPy
         self._num_manager_layers = len(layers)
 
@@ -476,6 +511,7 @@ def _add_layer(
             vocab_size=vocab_size,
             cache_tiers=cache_tiers,
             max_util_for_resume=kv_cache_config.max_util_for_resume,
+            enable_stats=self.enable_stats,
             layers=layers,
             typical_step=typical_step,
             constraints=constraints,

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -274,6 +274,8 @@ def create_draft_kv_cache_manager_maybe(
         max_beam_width=ad_config.max_beam_width,
         kv_connector_manager=None,  # KV connector manager not used in AutoDeploy (no disagg support)
         estimating_kv_cache=False,
+        enable_kv_cache_stats=ad_config.enable_iter_perf_stats
+        or getattr(ad_config, "return_perf_metrics", False),
     )
 
 

tensorrt_llm/_torch/disaggregation/native/transfer.py

@@ -470,11 +470,13 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta):
                     str(write_meta.slice_id).encode("ascii"),
                     b"True",  # is_last_slice — ensures receiver resolves its task future
                     AgentResult.FAILED.value.encode("ascii"),
+                    b"0",
                 ]
             )
             return
 
         agent_result = AgentResult.SUCCESS
+        transferred_bytes = int(write_meta.sizes.sum()) if write_meta.sizes.size > 0 else 0
         if write_meta.src_ptrs.size > 0:
             request = Sender._make_agent_request(write_meta, device_id=self._device_id)
             if timer:
@@ -510,9 +512,14 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta):
                 str(write_meta.slice_id).encode("ascii"),
                 str(write_meta.is_last_slice).encode("ascii"),
                 agent_result.value.encode("ascii"),
+                str(transferred_bytes if agent_result == AgentResult.SUCCESS else 0).encode(
+                    "ascii"
+                ),
             ]
         )
 
+        if agent_result == AgentResult.SUCCESS:
+            session.record_kv_transfer_bytes(transferred_bytes)
         task.transferred_count += 1
         if timer:
             timer.record_task_end(write_meta.peer_rank)
@@ -934,6 +941,7 @@ def _send_failed_result_to_receiver(self, info: RecvReqInfo):
                     str(slice_id).encode("ascii"),
                     b"True",  # is_last_slice
                     AgentResult.FAILED.value.encode("ascii"),
+                    b"0",
                 ]
             )
         except Exception as e:
@@ -1059,6 +1067,7 @@ def __init__(
         self.kv_tasks = []
         self.aux_task = None
         self.lock = threading.Lock()
+        self._transferred_kv_bytes = 0
 
         self._exception: Optional[Exception] = None
         self._closed = False
@@ -1079,6 +1088,15 @@ def disagg_request_id(self) -> int:
             return params.ctx_request_id
         return self.request_id
 
+    def record_kv_transfer_bytes(self, transferred_bytes: int) -> None:
+        with self.lock:
+            self._transferred_kv_bytes += transferred_bytes
+
+    @property
+    def transferred_kv_bytes(self) -> int:
+        with self.lock:
+            return self._transferred_kv_bytes
+
     @property
     def status(self) -> SessionStatus:
         if self._terminal_status is not None:
@@ -1497,12 +1515,26 @@ def _handle_cancel_session(self, message: list[bytes]):
             session.cancel()
 
     def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]):
-        msg_type, peer_rank, unique_rid, slice_id_str, is_last_slice_str, status = decode_message(
-            message
-        )
+        decoded_message = decode_message(message)
+        if len(decoded_message) == 6:
+            msg_type, peer_rank, unique_rid, slice_id_str, is_last_slice_str, status = (
+                decoded_message
+            )
+            transferred_bytes = 0
+        else:
+            (
+                msg_type,
+                peer_rank,
+                unique_rid,
+                slice_id_str,
+                is_last_slice_str,
+                status,
+                transferred_bytes,
+            ) = decoded_message
         peer_rank = int(peer_rank)
         unique_rid = int(unique_rid)
         sender_slice_id = int(slice_id_str)
+        transferred_bytes = int(transferred_bytes)
         if msg_type.encode("ascii") != MessageType.KV_AGENT_RESULT:
             logger.error(
                 f"_process_kv_agent_result: unexpected msg_type={msg_type!r}, expected KV_AGENT_RESULT"
@@ -1515,7 +1547,11 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]):
             )
             return
         session.process_kv_agent_result(
-            peer_rank, sender_slice_id, is_last_slice_str == "True", AgentResult(status)
+            peer_rank,
+            sender_slice_id,
+            is_last_slice_str == "True",
+            AgentResult(status),
+            transferred_bytes,
         )
 
     def _process_aux_agent_result(self, _send_id: bytes, message: list[bytes]):
@@ -1575,6 +1611,7 @@ def __init__(
         self._aux_status: TaskStatus = TaskStatus.INIT
         self._sender_endpoints: set[str] = set()
         self.lock = threading.Lock()
+        self._transferred_kv_bytes = 0
         self._receiver.setup_session(self)
 
     @property
@@ -1589,6 +1626,11 @@ def disagg_request_id(self) -> int:
             return params.ctx_request_id
         return self.request_id
 
+    @property
+    def transferred_kv_bytes(self) -> int:
+        with self.lock:
+            return self._transferred_kv_bytes
+
     @property
     def status(self) -> SessionStatus:
         if self._terminal_status is not None:
@@ -1623,7 +1665,12 @@ def receive(self, slice: KVSlice) -> None:
         self._receiver.dispatch_task(task)
 
     def process_kv_agent_result(
-        self, peer_rank: int, sender_slice_id: int, is_last_slice: bool, status: AgentResult
+        self,
+        peer_rank: int,
+        sender_slice_id: int,
+        is_last_slice: bool,
+        status: AgentResult,
+        transferred_bytes: int = 0,
     ):
         with self.lock:
             assert sender_slice_id < len(self._kv_tasks), (
@@ -1633,6 +1680,7 @@ def process_kv_agent_result(
             )
             task = self._kv_tasks[sender_slice_id]
             if status == AgentResult.SUCCESS:
+                self._transferred_kv_bytes += transferred_bytes
                 if is_last_slice:
                     task.last_slice_count += 1
                     if task.last_slice_count == task.expected_transfers: