[TRTLLM-12123][feat] Add per-iteration request-aggregate counters to InflightBatchingStats

cpp/include/tensorrt_llm/executor/types.h

@@ -321,6 +321,24 @@ 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 Total KV context length summed across paused (preempted-decode)
+    /// requests. Complements @ref numPausedRequests (count).
+    SizeType32 numPausedKvTokens;
 };
 
 /// @brief Struct that holds speculative decoding stats

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, numPausedKvTokens);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(SpecDecodingStats, numDraftTokens, numAcceptedTokens, numRequestsWithDraftTokens,
     acceptanceLength, iterLatencyMS, draftOverhead);
 NLOHMANN_DEFINE_TYPE_NON_INTRUSIVE(IterationStats, timestamp, iter, iterLatencyMS, newActiveRequestsQueueLatencyMS,

cpp/tensorrt_llm/executor/serialization.cpp

@@ -1897,8 +1897,14 @@ 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 numPausedKvTokens = su::deserialize<SizeType32>(is);
     return InflightBatchingStats{numScheduledRequests, numContextRequests, numGenRequests, numPausedRequests,
-        numCtxTokens, microBatchId, avgNumDecodedTokensPerIter};
+        numCtxTokens, microBatchId, avgNumDecodedTokensPerIter, numCtxKvTokens, numGenKvTokens,
+        numQueuedContextRequests, numQueuedCtxTokens, numPausedKvTokens};
 }
 
 void Serialization::serialize(InflightBatchingStats const& inflightBatchingStats, std::ostream& os)
@@ -1910,6 +1916,11 @@ 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.numPausedKvTokens, os);
 }
 
 size_t Serialization::serializedSize(InflightBatchingStats const& inflightBatchingStats)
@@ -1922,6 +1933,11 @@ 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.numPausedKvTokens);
     return totalSize;
 }
 

cpp/tensorrt_llm/nanobind/executor/bindings.cpp

@@ -131,7 +131,12 @@ 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_paused_kv_tokens", &tle::InflightBatchingStats::numPausedKvTokens);
 
     nb::class_<tle::SpecDecodingStats>(m, "SpecDecodingStats")
         .def(nb::init<>())

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -1183,6 +1183,87 @@ 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 requests whose payload is missing (e.g.
+        # disagg generation-only) so downstream consumers see only real
+        # backlog. Each queued item is a RequestQueueItem wrapping an
+        # ExecutorRequest (tle::Request).
+        num_queued_context_requests = 0
+        num_queued_ctx_tokens = 0
+        for item in list(self.executor_request_queue.get_request_queue().queue):
+            if not getattr(item, "is_normal_request", False):
+                continue
+            if item.request is None:
+                continue
+            num_queued_context_requests += 1
+            try:
+                num_queued_ctx_tokens += len(item.request.input_token_ids)
+            except Exception:
+                # input_token_ids unavailable for unusual request shapes;
+                # skip the token count for this item.
+                pass
+
+        # 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_paused_kv_tokens = num_paused_kv_tokens
+
         return stats
 
     def _append_iter_stats(self,

tensorrt_llm/executor/base_worker.py

@@ -664,7 +664,13 @@ 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.
+        stats_dict.setdefault("attentionDpRank", 0)
 
         if req_stats is not None and len(req_stats) > 0:
             stats_dict["requestStats"] = []

tests/unittest/llmapi/test_llm.py

@@ -2127,6 +2127,43 @@ def validate_stats(
         if pytorch_backend:
             assert result["numCompletedRequests"] == expected_num_completed
 
+            # Per-iteration request-aggregate fields populated by
+            # PyExecutor._update_iter_stats inside inflightBatchingStats.
+            # Assert presence (a missing key indicates a serializer or
+            # RPC-path regression) and sane per-iteration values (a
+            # zero-under-load value indicates a mis-wired populate block).
+            new_aggregate_keys = (
+                "numCtxKvTokens",
+                "numGenKvTokens",
+                "numQueuedContextRequests",
+                "numQueuedCtxTokens",
+                "numPausedKvTokens",
+            )
+            for k in new_aggregate_keys:
+                assert k in ifbStats, f"iter {iter}: missing ifbStats key {k}"
+                assert isinstance(
+                    ifbStats[k],
+                    int), (f"iter {iter}: ifbStats key {k} not int "
+                           f"(got {type(ifbStats[k])})")
+                assert ifbStats[
+                    k] >= 0, f"iter {iter}: ifbStats key {k} negative"
+
+            if iter < context_iterations:
+                # Prefill iteration: at least one scheduled context request
+                # and nonzero numCtxTokens. numCtxTokens is sourced from
+                # model_engine.iter_states after _forward_step for this
+                # batch, so it is overlap-safe under every scheduler
+                # configuration.
+                assert ifbStats["numContextRequests"] >= 1, f"iter: {iter}"
+                assert ifbStats["numGenRequests"] == 0, f"iter: {iter}"
+                assert ifbStats["numCtxTokens"] > 0, f"iter: {iter}"
+            else:
+                # Generation iteration: at least one decode request with
+                # nonzero total KV context length.
+                assert ifbStats["numGenRequests"] >= 1, f"iter: {iter}"
+                assert ifbStats["numGenKvTokens"] > 0, f"iter: {iter}"
+                assert ifbStats["numContextRequests"] == 0, f"iter: {iter}"
+
 
 def llm_get_stats_test_harness(tp_size: int = 1,
                                pp_size: int = 1,