[TRTLLM-12123][feat] Add per-iteration request-aggregate counters to InflightBatchingStats (#13199)

nv-yna · web-flow · commit 734a1460f67e · 2026-04-27T09:17:24.000-04:00
Signed-off-by: Yuewei Na &lt;nv-yna@users.noreply.github.com&gt;
Co-authored-by: Yuewei Na &lt;nv-yna@users.noreply.github.com&gt;
diff --git a/cpp/include/tensorrt_llm/executor/types.h b/cpp/include/tensorrt_llm/executor/types.h
@@ -321,6 +321,35 @@ struct InflightBatchingStats
     SizeType32 microBatchId;
     /// @brief Average number of tokens decoded per request per iteration
     float avgNumDecodedTokensPerIter;
+    /// @brief Context tokens for scheduled context requests that are read from
+    /// KV cache rather than computed this iteration. Covers prefix-cache hits
+    /// and previously-chunked tokens for chunked-prefill continuations.
+    /// Complements @ref numCtxTokens (tokens computed this iteration).
+    SizeType32 numCtxKvTokens;
+    /// @brief Total KV context length (prompt + generated-so-far) summed
+    /// across scheduled generation (decode) requests.
+    SizeType32 numGenKvTokens;
+    /// @brief Number of context (prefill) requests waiting in the executor
+    /// request queue — submitted but not yet scheduled. Excludes non-normal
+    /// control items (shutdown/cancel) and requests without a payload.
+    SizeType32 numQueuedContextRequests;
+    /// @brief Sum of prompt-token counts across queued context requests (the
+    /// requests counted in @ref numQueuedContextRequests).
+    SizeType32 numQueuedCtxTokens;
+    /// @brief Number of generation-only requests waiting in the executor
+    /// request queue. On a disaggregated-decode engine these are requests
+    /// that have completed prefill elsewhere and are awaiting KV-cache
+    /// transfer before they can start decoding. Always 0 on a
+    /// non-disaggregated or disaggregated-prefill engine.
+    SizeType32 numQueuedGenRequests;
+    /// @brief Sum of prompt-token counts across queued generation-only
+    /// requests (the requests counted in @ref numQueuedGenRequests). Acts
+    /// as the KV-budget these requests will need once their KV transfer
+    /// completes.
+    SizeType32 numQueuedGenKvTokens;
+    /// @brief Total KV context length summed across paused (preempted-decode)
+    /// requests. Complements @ref numPausedRequests (count).
+    SizeType32 numPausedKvTokens;
 };
 
 /// @brief Struct that holds speculative decoding stats
diff --git a/cpp/tensorrt_llm/executor/jsonSerialization.cpp b/cpp/tensorrt_llm/executor/jsonSerialization.cpp
@@ -30,7 +30,8 @@ NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(KvCacheStats, maxNumBlocks, freeNumBlocks, us
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(
     StaticBatchingStats, numScheduledRequests, numContextRequests, numCtxTokens, numGenTokens, emptyGenSlots);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(InflightBatchingStats, numScheduledRequests, numContextRequests, numGenRequests,
-    numPausedRequests, numCtxTokens, microBatchId, avgNumDecodedTokensPerIter);
+    numPausedRequests, numCtxTokens, microBatchId, avgNumDecodedTokensPerIter, numCtxKvTokens, numGenKvTokens,
+    numQueuedContextRequests, numQueuedCtxTokens, numQueuedGenRequests, numQueuedGenKvTokens, numPausedKvTokens);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(SpecDecodingStats, numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens,
     acceptanceLength, iterLatencyMS, draftOverhead);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(IterationStats, timestamp, iter, iterLatencyMS, newActiveRequestsQueueLatencyMS,
diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp
@@ -1897,8 +1897,16 @@ InflightBatchingStats Serialization::deserializeInflightBatchingStats(std::istre
     auto numCtxTokens = su::deserialize<SizeType32>(is);
     auto microBatchId = su::deserialize<SizeType32>(is);
     auto avgNumDecodedTokensPerIter = su::deserialize<float>(is);
+    auto numCtxKvTokens = su::deserialize<SizeType32>(is);
+    auto numGenKvTokens = su::deserialize<SizeType32>(is);
+    auto numQueuedContextRequests = su::deserialize<SizeType32>(is);
+    auto numQueuedCtxTokens = su::deserialize<SizeType32>(is);
+    auto numQueuedGenRequests = su::deserialize<SizeType32>(is);
+    auto numQueuedGenKvTokens = su::deserialize<SizeType32>(is);
+    auto numPausedKvTokens = su::deserialize<SizeType32>(is);
     return InflightBatchingStats{numScheduledRequests, numContextRequests, numGenRequests, numPausedRequests,
-        numCtxTokens, microBatchId, avgNumDecodedTokensPerIter};
+        numCtxTokens, microBatchId, avgNumDecodedTokensPerIter, numCtxKvTokens, numGenKvTokens,
+        numQueuedContextRequests, numQueuedCtxTokens, numQueuedGenRequests, numQueuedGenKvTokens, numPausedKvTokens};
 }
 
 void Serialization::serialize(InflightBatchingStats const& inflightBatchingStats, std::ostream& os)
@@ -1910,6 +1918,13 @@ void Serialization::serialize(InflightBatchingStats const& inflightBatchingStats
     su::serialize(inflightBatchingStats.numCtxTokens, os);
     su::serialize(inflightBatchingStats.microBatchId, os);
     su::serialize(inflightBatchingStats.avgNumDecodedTokensPerIter, os);
+    su::serialize(inflightBatchingStats.numCtxKvTokens, os);
+    su::serialize(inflightBatchingStats.numGenKvTokens, os);
+    su::serialize(inflightBatchingStats.numQueuedContextRequests, os);
+    su::serialize(inflightBatchingStats.numQueuedCtxTokens, os);
+    su::serialize(inflightBatchingStats.numQueuedGenRequests, os);
+    su::serialize(inflightBatchingStats.numQueuedGenKvTokens, os);
+    su::serialize(inflightBatchingStats.numPausedKvTokens, os);
 }
 
 size_t Serialization::serializedSize(InflightBatchingStats const& inflightBatchingStats)
@@ -1922,6 +1937,13 @@ size_t Serialization::serializedSize(InflightBatchingStats const& inflightBatchi
     totalSize += su::serializedSize(inflightBatchingStats.numCtxTokens);
     totalSize += su::serializedSize(inflightBatchingStats.microBatchId);
     totalSize += su::serializedSize(inflightBatchingStats.avgNumDecodedTokensPerIter);
+    totalSize += su::serializedSize(inflightBatchingStats.numCtxKvTokens);
+    totalSize += su::serializedSize(inflightBatchingStats.numGenKvTokens);
+    totalSize += su::serializedSize(inflightBatchingStats.numQueuedContextRequests);
+    totalSize += su::serializedSize(inflightBatchingStats.numQueuedCtxTokens);
+    totalSize += su::serializedSize(inflightBatchingStats.numQueuedGenRequests);
+    totalSize += su::serializedSize(inflightBatchingStats.numQueuedGenKvTokens);
+    totalSize += su::serializedSize(inflightBatchingStats.numPausedKvTokens);
     return totalSize;
 }
 
diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp
@@ -131,7 +131,14 @@ void initBindings(nb::module_& m)
         .def_rw("num_paused_requests", &tle::InflightBatchingStats::numPausedRequests)
         .def_rw("num_ctx_tokens", &tle::InflightBatchingStats::numCtxTokens)
         .def_rw("micro_batch_id", &tle::InflightBatchingStats::microBatchId)
-        .def_rw("avg_num_decoded_tokens_per_iter", &tle::InflightBatchingStats::avgNumDecodedTokensPerIter);
+        .def_rw("avg_num_decoded_tokens_per_iter", &tle::InflightBatchingStats::avgNumDecodedTokensPerIter)
+        .def_rw("num_ctx_kv_tokens", &tle::InflightBatchingStats::numCtxKvTokens)
+        .def_rw("num_gen_kv_tokens", &tle::InflightBatchingStats::numGenKvTokens)
+        .def_rw("num_queued_context_requests", &tle::InflightBatchingStats::numQueuedContextRequests)
+        .def_rw("num_queued_ctx_tokens", &tle::InflightBatchingStats::numQueuedCtxTokens)
+        .def_rw("num_queued_gen_requests", &tle::InflightBatchingStats::numQueuedGenRequests)
+        .def_rw("num_queued_gen_kv_tokens", &tle::InflightBatchingStats::numQueuedGenKvTokens)
+        .def_rw("num_paused_kv_tokens", &tle::InflightBatchingStats::numPausedKvTokens);
 
     nb::class_<tle::SpecDecodingStats>(m, "SpecDecodingStats")
         .def(nb::init<>())
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -27,7 +27,7 @@
                                             FinishReason, InflightBatchingStats,
                                             IterationStats, KvCacheStats,
                                             RequestStage, RequestStats,
-                                            SpecDecodingStats,
+                                            RequestType, SpecDecodingStats,
                                             StaticBatchingStats)
 from tensorrt_llm.bindings.internal.batch_manager import (LlmRequestType,
                                                           ReqIdsSet)
@@ -1198,6 +1198,108 @@ def _update_iter_stats(self, stats, iter_latency_ms, num_completed_requests,
             # Calculate draft overhead
             stats.specdec_stats.draft_overhead = 0.0 if iter_latency_ms <= 0.0 else float(
                 draft_latency_ms) / float(iter_latency_ms)
+
+        # Extra per-iteration request-aggregate counters attached to
+        # inflight_batching_stats. These complement the existing
+        # num_context_requests / num_gen_requests / num_ctx_tokens /
+        # num_paused_requests members with token-weighted counts and
+        # queue/paused KV accounting.
+
+        # Tokens read from prior state (prefix-cache hits and
+        # previously-chunked tokens) summed across scheduled context
+        # requests; complements num_ctx_tokens (tokens computed this
+        # iteration). Read from py_last_context_chunk, a Python-side
+        # cache set by _update_request_states before state mutation — it
+        # stays valid after the request transitions to
+        # GENERATION_IN_PROGRESS, unlike the C++ getContextChunkSize() /
+        # getContextCurrentPosition() accessors that would raise
+        # RuntimeError on a mutated request.
+        num_ctx_kv_tokens = 0
+        for req in scheduled_batch.context_requests:
+            if getattr(req, "is_attention_dp_dummy", False):
+                continue
+            last_chunk = getattr(req, "py_last_context_chunk", None)
+            if last_chunk is not None and last_chunk[0] is not None:
+                start, _end = last_chunk
+                num_ctx_kv_tokens += start
+            else:
+                try:
+                    num_ctx_kv_tokens += \
+                        req.context_current_position
+                except RuntimeError:
+                    pass
+
+        # Total KV context length (prompt + tokens generated so far)
+        # summed across scheduled generation requests.
+        num_gen_kv_tokens = 0
+        for req in scheduled_batch.generation_requests:
+            if getattr(req, "is_attention_dp_dummy", False):
+                continue
+            try:
+                num_gen_kv_tokens += req.get_num_tokens(0)
+            except RuntimeError:
+                pass
+
+        # Normal requests waiting in the executor_request_queue that have
+        # never been scheduled. Excludes non-normal control items
+        # (shutdown/cancel) and items with a missing payload. Each queued
+        # item is a RequestQueueItem wrapping an ExecutorRequest
+        # (tle::Request). Requests are routed by request_type:
+        #   - CONTEXT_AND_GENERATION (default) and CONTEXT_ONLY
+        #     (disagg-prefill side) -> queued-context counters.
+        #   - GENERATION_ONLY (disagg-decode side, awaiting KV transfer
+        #     before they can start decoding) -> queued-gen counters.
+        # On a non-disagg engine all items land in the context counters;
+        # on a disagg-decode engine all items land in the gen counters.
+        num_queued_context_requests = 0
+        num_queued_ctx_tokens = 0
+        num_queued_gen_requests = 0
+        num_queued_gen_kv_tokens = 0
+        for item in list(self.executor_request_queue.get_request_queue().queue):
+            if not item.is_normal_request:
+                continue
+            if item.request is None:
+                continue
+            try:
+                token_count = len(item.request.input_token_ids)
+            except (AttributeError, TypeError) as e:
+                # Unusual request shape with no usable token payload;
+                # exclude from all queued counters so downstream consumers
+                # see consistent per-request averages. Not expected on the
+                # current API (ExecutorRequest construction requires a
+                # non-empty input_token_ids), logged so future API drift
+                # surfaces instead of being silently dropped.
+                logger.warning(f"Excluding queued item {item.id} from queued "
+                               f"counters: input_token_ids not readable "
+                               f"({type(e).__name__})")
+                continue
+            if item.request.request_type == RequestType.REQUEST_TYPE_GENERATION_ONLY:
+                num_queued_gen_requests += 1
+                num_queued_gen_kv_tokens += token_count
+            else:
+                num_queued_context_requests += 1
+                num_queued_ctx_tokens += token_count
+
+        # Total KV context length summed across paused (preempted-decode)
+        # requests — were decoding but got evicted back to the waiting
+        # pool for this iteration.
+        num_paused_kv_tokens = 0
+        for req in scheduled_batch.paused_requests:
+            if getattr(req, "is_attention_dp_dummy", False):
+                continue
+            try:
+                num_paused_kv_tokens += req.get_num_tokens(0)
+            except RuntimeError:
+                pass
+
+        stats.inflight_batching_stats.num_ctx_kv_tokens = num_ctx_kv_tokens
+        stats.inflight_batching_stats.num_gen_kv_tokens = num_gen_kv_tokens
+        stats.inflight_batching_stats.num_queued_context_requests = num_queued_context_requests
+        stats.inflight_batching_stats.num_queued_ctx_tokens = num_queued_ctx_tokens
+        stats.inflight_batching_stats.num_queued_gen_requests = num_queued_gen_requests
+        stats.inflight_batching_stats.num_queued_gen_kv_tokens = num_queued_gen_kv_tokens
+        stats.inflight_batching_stats.num_paused_kv_tokens = num_paused_kv_tokens
+
         return stats
 
     def _append_iter_stats(self,
diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py
@@ -664,7 +664,15 @@ def get_disaggregated_params(self) -> dict:
     def _stats_serializer(stats) -> str:
         iteration_stats, req_stats = stats[0], stats[1]
         kv_iter_stats = stats[2] if len(stats) > 2 else None
+
         stats_dict = json.loads(iteration_stats.to_json_str())
+        # Tag with dp_rank=0 so Dynamo's adapter can always read
+        # stat["attentionDpRank"] without a missing-key branch. Attention-DP
+        # per-rank emission is a follow-up; today FPM only flows under
+        # non-attention-DP.
+        # TODO(https://jirasw.nvidia.com/browse/TRTLLM-12123): implement
+        # per-rank IterationStats delivery under attention-DP.
+        stats_dict.setdefault("attentionDpRank", 0)
 
         if req_stats is not None and len(req_stats) > 0:
             stats_dict["requestStats"] = []
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -191,6 +191,96 @@ def test_pad_generation_requests(self) -> None:
 
         kv_cache_manager.shutdown()
 
+    def test_pad_batch_strips_cudagraph_dummies_on_clean_exit(self) -> None:
+        # Regression guard for the invariant that CUDAGraphRunner.pad_batch's
+        # `finally` strips every is_cuda_graph_dummy=True entry from
+        # scheduled_requests.generation_requests before the `with` block
+        # exits. Downstream consumers of scheduled_batch.generation_requests
+        # — including the per-iteration stats populate block in
+        # PyExecutor._update_iter_stats — rely on never observing
+        # cudagraph dummies.
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache()
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+
+        # batch_size=5 rounds up to 8 (nearest captured graph size in the
+        # fixture config) -> padding_size=3, deterministically.
+        real_batch_size = 5
+        max_seq_len = 1
+        real_requests = [
+            _create_request(max_seq_len, i) for i in range(real_batch_size)
+        ]
+        real_ids = [id(r) for r in real_requests]
+
+        batch = ScheduledRequests()
+        batch.generation_requests = list(real_requests)
+
+        with model_engine.cuda_graph_runner.pad_batch(
+                batch, resource_manager) as padded_batch:
+            # Positive assertion that padding actually fired — guards
+            # against a vacuous pass where padding was a no-op.
+            self.assertGreater(
+                len(padded_batch.generation_requests), real_batch_size,
+                "padding did not fire; fixture config may have drifted "
+                "so that 5 no longer rounds up to 8")
+            # Every appended entry past the original count is a
+            # cudagraph-flagged dummy.
+            for req in padded_batch.generation_requests[real_batch_size:]:
+                self.assertTrue(
+                    getattr(req, "is_cuda_graph_dummy", False),
+                    "pad_batch appended a request without "
+                    "is_cuda_graph_dummy=True")
+            # Real requests' identities and order are untouched.
+            self.assertEqual([
+                id(r)
+                for r in padded_batch.generation_requests[:real_batch_size]
+            ], real_ids)
+
+        # After the with-block: finally must have sliced off the padding.
+        self.assertEqual(
+            len(batch.generation_requests), real_batch_size,
+            "pad_batch.finally did not strip cudagraph dummies — "
+            "downstream consumers of scheduled_batch.generation_requests "
+            "would observe the leaked dummies")
+        for req in batch.generation_requests:
+            self.assertFalse(
+                getattr(req, "is_cuda_graph_dummy", False),
+                "cudagraph dummy leaked out of pad_batch's finally")
+
+        kv_cache_manager.shutdown()
+
+    def test_pad_batch_strips_cudagraph_dummies_on_exception(self) -> None:
+        # The strip must fire even when the body raises. This is the
+        # critical property of `finally` vs. a plain trailing statement —
+        # it guards the invariant on the error path. A refactor that
+        # accidentally dropped the `finally` would be caught here but not
+        # by the clean-exit variant.
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache()
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+
+        real_batch_size = 5
+        real_requests = [_create_request(1, i) for i in range(real_batch_size)]
+
+        batch = ScheduledRequests()
+        batch.generation_requests = list(real_requests)
+
+        class _ForwardBoom(Exception):
+            pass
+
+        with self.assertRaises(_ForwardBoom):
+            with model_engine.cuda_graph_runner.pad_batch(
+                    batch, resource_manager) as padded_batch:
+                self.assertGreater(len(padded_batch.generation_requests),
+                                   real_batch_size)
+                raise _ForwardBoom()
+
+        self.assertEqual(len(batch.generation_requests), real_batch_size)
+        for req in batch.generation_requests:
+            self.assertFalse(getattr(req, "is_cuda_graph_dummy", False))
+
+        kv_cache_manager.shutdown()
+
     def test_position_id_preparation(self):
         model_engine, kv_cache_manager = create_model_engine_and_kvcache()
         resource_manager = ResourceManager(
diff --git a/tests/unittest/llmapi/test_llm.py b/tests/unittest/llmapi/test_llm.py
diff --git a/tests/unittest/pyexecutor/test_iter_stats_populate.py b/tests/unittest/pyexecutor/test_iter_stats_populate.py
