diff --git a/cpp/include/tensorrt_llm/executor/types.h b/cpp/include/tensorrt_llm/executor/types.h index 77f910455c57..0800865df7f1 100644 --- 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 index 81f128ffbff5..45b716f68b81 100644 --- 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 index 12291854f5cd..60eb094810c8 100644 --- 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(is); auto microBatchId = su::deserialize(is); auto avgNumDecodedTokensPerIter = su::deserialize(is); + auto numCtxKvTokens = su::deserialize(is); + auto numGenKvTokens = su::deserialize(is); + auto numQueuedContextRequests = su::deserialize(is); + auto numQueuedCtxTokens = su::deserialize(is); + auto numQueuedGenRequests = su::deserialize(is); + auto numQueuedGenKvTokens = su::deserialize(is); + auto numPausedKvTokens = su::deserialize(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 index 78c90a86ca37..fbec513de3a1 100644 --- 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_(m, "SpecDecodingStats") .def(nb::init<>()) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6564be1ba449..7b1deb50d665 100644 --- 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) @@ -1183,6 +1183,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 index ba71f339ba8a..076e18c07d71 100644 --- 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 index c683c7f76694..db6873d04cc3 100644 --- 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 index c33502a8ccb7..9730b3324df3 100644 --- a/tests/unittest/llmapi/test_llm.py +++ b/tests/unittest/llmapi/test_llm.py @@ -2127,6 +2127,45 @@ 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", + "numQueuedGenRequests", + "numQueuedGenKvTokens", + "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, diff --git a/tests/unittest/pyexecutor/test_iter_stats_populate.py b/tests/unittest/pyexecutor/test_iter_stats_populate.py new file mode 100644 index 000000000000..f5496361879c --- /dev/null +++ b/tests/unittest/pyexecutor/test_iter_stats_populate.py @@ -0,0 +1,511 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the per-iteration request-aggregate fields on InflightBatchingStats. + +``PyExecutor._update_iter_stats`` populates the following members on +``stats.inflight_batching_stats`` in addition to the pre-existing +``num_context_requests`` / ``num_gen_requests`` / ``num_paused_requests`` / +``num_scheduled_requests`` / ``micro_batch_id``: + + * ``num_ctx_kv_tokens`` — tokens read from prior state + (prefix-cache hits + previously-chunked tokens) summed across + scheduled context requests; dummy-filtered. + * ``num_gen_kv_tokens`` — total KV context length summed across + scheduled generation requests; dummy-filtered. + * ``num_queued_context_requests`` — normal context-type requests + (``REQUEST_TYPE_CONTEXT_AND_GENERATION`` / ``REQUEST_TYPE_CONTEXT_ONLY``) + sitting in the executor_request_queue; excludes shutdown/cancel + control items and requests without a payload. + * ``num_queued_ctx_tokens`` — prompt-token sum across the above. + * ``num_queued_gen_requests`` — generation-only queued requests + (``REQUEST_TYPE_GENERATION_ONLY``), typically disagg-decode items + awaiting KV transfer before decoding. + * ``num_queued_gen_kv_tokens`` — prompt-token sum across the above; + acts as the KV budget each request will occupy post-transfer. + * ``num_paused_kv_tokens`` — total KV context length summed across + paused (preempted-decode) requests; dummy-filtered. + +These tests invoke the real ``PyExecutor._update_iter_stats`` as an +unbound call against a minimal fake ``self``. Exercising the production +function directly (rather than a duplicated shim) catches drift from +refactors, renamed attributes, or silent unit changes in the populate +block. +""" + +from __future__ import annotations + +import types +from unittest.mock import MagicMock, patch + +from tensorrt_llm.bindings.executor import InflightBatchingStats, IterationStats + + +class _StubRequest: + """Stub LlmRequest exposing only the accessors ``_update_iter_stats`` reads. + + Attributes: + py_last_context_chunk: tuple (start, end) — the (begin_compute, + begin_compute + chunk_size) pair cached by ``_update_request_states`` + before state mutation. ``start`` is the primary source of + ``num_ctx_kv_tokens``. Set to None for decode/paused reqs. + context_current_position: fallback source consulted when + ``py_last_context_chunk`` is None. + _num_tokens: return value for ``get_num_tokens()``, used for decode + and paused (preempted-decode) requests. + """ + + def __init__( + self, + *, + context_chunk_size: int = 0, + context_current_position: int = 0, + num_tokens: int = 0, + is_attention_dp_dummy: bool = False, + ): + self.context_current_position = context_current_position + # py_last_context_chunk = (begin_compute, end_compute). For a fresh + # prefill, begin_compute == 0; for continuations, == prev position. + if context_chunk_size > 0: + self.py_last_context_chunk = ( + context_current_position, + context_current_position + context_chunk_size, + ) + else: + self.py_last_context_chunk = None + self._num_tokens = num_tokens + self.is_attention_dp_dummy = is_attention_dp_dummy + + def get_num_tokens(self, beam: int = 0) -> int: + return self._num_tokens + + +class _StubScheduledBatch: + def __init__(self, context_reqs=None, gen_reqs=None, paused_reqs=None): + self.context_requests = list(context_reqs or []) + self.generation_requests = list(gen_reqs or []) + self.paused_requests = list(paused_reqs or []) + + @property + def num_context_requests(self): + return len(self.context_requests) + + @property + def num_generation_requests(self): + return len(self.generation_requests) + + +class _StubQueueItem: + _next_id = 0 + + def __init__( + self, + input_token_ids, + *, + is_normal_request: bool = True, + request_type=None, + ): + # request_type defaults to REQUEST_TYPE_CONTEXT_AND_GENERATION, + # matching the executor::Request constructor default. Tests that + # exercise disagg routing (REQUEST_TYPE_CONTEXT_ONLY / + # REQUEST_TYPE_GENERATION_ONLY) pass it explicitly. + if request_type is None: + from tensorrt_llm.bindings.executor import RequestType + + request_type = RequestType.REQUEST_TYPE_CONTEXT_AND_GENERATION + self.request = types.SimpleNamespace( + input_token_ids=input_token_ids, + request_type=request_type, + ) + self.is_normal_request = is_normal_request + # Stable id so the populate block's drift-warning log has a value + # to reference even if a future branch exercises item.id. + _StubQueueItem._next_id += 1 + self.id = _StubQueueItem._next_id + + +def _build_fake_self(queued_items, iter_states): + """Minimal 'self' for ``PyExecutor._update_iter_stats(self, ...)``. + + Stubs only the ``self.*`` attributes the method actually reads: + + Setup reads (run before per-request aggregation): + * ``max_num_active_requests``, ``iter_counter`` — scalars + * ``executor_request_queue.get_request_queue_size()`` — for + top-level ``num_queued_requests`` (not one of the fields + exercised here) + * ``resource_manager.resource_managers.get(...)`` — returns None so + the KV-cache-stats block is skipped entirely + * ``drafter`` — None (spec-decode block no-ops when + ``stats.specdec_stats`` is None, which is the default on a fresh + ``IterationStats``) + + Per-request aggregation reads: + * ``executor_request_queue.get_request_queue().queue`` — source for + ``num_queued_context_requests`` / ``num_queued_ctx_tokens`` + * ``model_engine.iter_states`` — stubbed but not read by the + request-aggregate fields under test here; the regression test + ``test_num_ctx_kv_tokens_ignores_iter_states_side_channel`` + verifies the populate block does not read it + """ + fake = MagicMock() + fake.max_num_active_requests = 64 + fake.iter_counter = 1 + fake.executor_request_queue.get_request_queue_size.return_value = len(queued_items) + fake.executor_request_queue.get_request_queue.return_value.queue = queued_items + fake.resource_manager.resource_managers.get.return_value = None + fake.drafter = None + fake.model_engine = types.SimpleNamespace(iter_states=iter_states) + return fake + + +def _invoke_update_iter_stats(scheduled_batch, queued_items, *, num_ctx_tokens): + """Call real ``PyExecutor._update_iter_stats`` unbound; return the stats. + + Patches ``torch.cuda.mem_get_info`` so the method can run on hosts + without a live CUDA context (the call is unconditional for + ``gpu_mem_usage`` but the value is not consumed by the fields + under test). + + Parameters + ---------- + scheduled_batch : _StubScheduledBatch + queued_items : list[_StubQueueItem] + num_ctx_tokens : int | None + If int, wired into ``model_engine.iter_states = {"num_ctx_tokens": num_ctx_tokens}``. + If None, ``iter_states`` is set to None. The populate block under + test does not consume this value; it is plumbed so regression + tests can verify the side channel remains unread. + """ + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + iter_states = None if num_ctx_tokens is None else {"num_ctx_tokens": num_ctx_tokens} + + fake_self = _build_fake_self(queued_items, iter_states) + + stats = IterationStats() + # The method reads ``stats.inflight_batching_stats.*`` unconditionally; + # the default on a fresh IterationStats is None, so we allocate one. + stats.inflight_batching_stats = InflightBatchingStats() + + with patch( + "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.mem_get_info", + return_value=(1 << 30, 1 << 30), + ): + PyExecutor._update_iter_stats( + fake_self, + stats, + iter_latency_ms=10.0, + num_completed_requests=0, + scheduled_batch=scheduled_batch, + micro_batch_id=0, + ) + return stats + + +# --------------------------------------------------------------------------- +# Populate tests: call the real ``_update_iter_stats`` and assert on the +# inflight_batching_stats request-aggregate fields it populates. +# --------------------------------------------------------------------------- + + +def test_empty_iteration(): + stats = _invoke_update_iter_stats(_StubScheduledBatch(), [], num_ctx_tokens=0) + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 0 + assert ifb.num_gen_requests == 0 + assert ifb.num_paused_requests == 0 + assert ifb.num_ctx_kv_tokens == 0 + assert ifb.num_gen_kv_tokens == 0 + assert ifb.num_queued_context_requests == 0 + assert ifb.num_queued_ctx_tokens == 0 + assert ifb.num_queued_gen_requests == 0 + assert ifb.num_queued_gen_kv_tokens == 0 + assert ifb.num_paused_kv_tokens == 0 + + +def test_prefill_only_no_prefix_cache(): + # Two fresh prefill requests: prompts of 100 and 200 tokens, chunk size + # == full prompt (no chunked prefill). No prefix cache hits. + reqs = [ + _StubRequest(context_chunk_size=100, context_current_position=0), + _StubRequest(context_chunk_size=200, context_current_position=0), + ] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=reqs), [], num_ctx_tokens=300 + ) + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 2 + assert ifb.num_ctx_kv_tokens == 0 # py_last_context_chunk[0] == 0 for both + + +def test_prefill_with_prefix_cache_hit(): + # Prompt 1000 tokens; 256 already in prefix cache (prepopulatedPromptLen). + # Chunk size = remaining = 744. py_last_context_chunk = (256, 1000); + # start=256 is the precomputed-tokens count. + reqs = [ + _StubRequest(context_chunk_size=744, context_current_position=256), + ] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=reqs), [], num_ctx_tokens=744 + ) + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 1 + assert ifb.num_ctx_kv_tokens == 256 + + +def test_chunked_prefill_continuation(): + # Chunked prefill: 3-chunk request, each chunk 512. This is step 2: + # chunk size 512, previously computed 512 (== context_current_position). + # py_last_context_chunk = (512, 1024); start=512. + reqs = [ + _StubRequest(context_chunk_size=512, context_current_position=512), + ] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=reqs), [], num_ctx_tokens=512 + ) + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 1 + assert ifb.num_ctx_kv_tokens == 512 + + +def test_decode_only(): + # Two decode requests: 1024 total context and 2048 total context. + reqs = [ + _StubRequest(num_tokens=1024), + _StubRequest(num_tokens=2048), + ] + stats = _invoke_update_iter_stats(_StubScheduledBatch(gen_reqs=reqs), [], num_ctx_tokens=0) + ifb = stats.inflight_batching_stats + assert ifb.num_gen_requests == 2 + assert ifb.num_gen_kv_tokens == 3072 + assert ifb.num_context_requests == 0 + assert ifb.num_ctx_kv_tokens == 0 + + +def test_mixed_prefill_and_decode(): + ctx = [_StubRequest(context_chunk_size=128, context_current_position=0)] + gen = [_StubRequest(num_tokens=500), _StubRequest(num_tokens=700)] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=ctx, gen_reqs=gen), [], num_ctx_tokens=128 + ) + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 1 + assert ifb.num_gen_requests == 2 + assert ifb.num_gen_kv_tokens == 1200 + + +def test_queued_context_requests_from_request_queue(): + items = [ + _StubQueueItem(input_token_ids=list(range(256))), + _StubQueueItem(input_token_ids=list(range(1024))), + ] + stats = _invoke_update_iter_stats(_StubScheduledBatch(), items, num_ctx_tokens=0) + ifb = stats.inflight_batching_stats + assert ifb.num_queued_context_requests == 2 + assert ifb.num_queued_ctx_tokens == 1280 + + +def test_queued_filters_non_normal_requests(): + # Shutdown / cancel / control items should be ignored. + items = [ + _StubQueueItem(input_token_ids=list(range(100)), is_normal_request=False), + _StubQueueItem(input_token_ids=list(range(50))), + ] + stats = _invoke_update_iter_stats(_StubScheduledBatch(), items, num_ctx_tokens=0) + ifb = stats.inflight_batching_stats + assert ifb.num_queued_context_requests == 1 + assert ifb.num_queued_ctx_tokens == 50 + + +def test_queued_routes_by_request_type(): + # Disaggregated serving: a decode engine receives + # REQUEST_TYPE_GENERATION_ONLY items that await KV transfer from a + # prefill engine before starting decode. They belong in the + # queued-gen counters, not the queued-context counters. A prefill + # engine sees REQUEST_TYPE_CONTEXT_ONLY items, which are real + # queued-prefill work and count in the context counters alongside + # default-type (CONTEXT_AND_GENERATION) items. + from tensorrt_llm.bindings.executor import RequestType + + items = [ + # Non-disagg default: queued-context work, len=100. + _StubQueueItem( + input_token_ids=list(range(100)), + request_type=RequestType.REQUEST_TYPE_CONTEXT_AND_GENERATION, + ), + # Disagg prefill side: queued-context work, len=200. + _StubQueueItem( + input_token_ids=list(range(200)), + request_type=RequestType.REQUEST_TYPE_CONTEXT_ONLY, + ), + # Disagg decode side: queued-gen work, len=512. + _StubQueueItem( + input_token_ids=list(range(512)), + request_type=RequestType.REQUEST_TYPE_GENERATION_ONLY, + ), + # Disagg decode side: queued-gen work, len=1024. + _StubQueueItem( + input_token_ids=list(range(1024)), + request_type=RequestType.REQUEST_TYPE_GENERATION_ONLY, + ), + ] + stats = _invoke_update_iter_stats(_StubScheduledBatch(), items, num_ctx_tokens=0) + ifb = stats.inflight_batching_stats + # Two context-flavoured items -> queued-context counters. + assert ifb.num_queued_context_requests == 2 + assert ifb.num_queued_ctx_tokens == 300 # 100 + 200 + # Two generation-only items -> queued-gen counters. + assert ifb.num_queued_gen_requests == 2 + assert ifb.num_queued_gen_kv_tokens == 1536 # 512 + 1024 + + +def test_paused_decode_requests(): + paused = [ + _StubRequest(num_tokens=300), + _StubRequest(num_tokens=800), + ] + stats = _invoke_update_iter_stats(_StubScheduledBatch(paused_reqs=paused), [], num_ctx_tokens=0) + ifb = stats.inflight_batching_stats + assert ifb.num_paused_requests == 2 + assert ifb.num_paused_kv_tokens == 1100 + + +def test_attention_dp_dummy_filtering_on_kv_token_fields(): + # Dummy-padding added by ``_pad_attention_dp_dummy_request`` must not + # contribute to the KV-token-weighted fields under test + # (num_ctx_kv_tokens, num_gen_kv_tokens, num_paused_kv_tokens). + # The existing count fields (num_context_requests / num_gen_requests / + # num_paused_requests) are set directly from ``scheduled_batch`` + # properties earlier in _update_iter_stats, so they DO include dummies. + # This asymmetry is by design for the new KV-token fields: Dynamo + # wants "real work only" for its autoscaling signal. + ctx = [ + _StubRequest( + context_chunk_size=100, + context_current_position=0, + is_attention_dp_dummy=True, + ), + _StubRequest(context_chunk_size=200, context_current_position=50), + ] + gen = [_StubRequest(num_tokens=1024, is_attention_dp_dummy=True)] + paused = [_StubRequest(num_tokens=500, is_attention_dp_dummy=True)] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=ctx, gen_reqs=gen, paused_reqs=paused), + [], + num_ctx_tokens=300, + ) + ifb = stats.inflight_batching_stats + # Count fields include dummies (populated directly from scheduled_batch). + assert ifb.num_context_requests == 2 + assert ifb.num_gen_requests == 1 + assert ifb.num_paused_requests == 1 + # KV-token-weighted new fields filter dummies. + assert ifb.num_ctx_kv_tokens == 50 # only the non-dummy's start + assert ifb.num_gen_kv_tokens == 0 # dummy gen filtered + assert ifb.num_paused_kv_tokens == 0 # dummy paused filtered + + +def test_full_mixed_iteration(): + # Realistic scenario: 3 prefill (1 fresh, 2 continuing chunks), 4 decode, + # 2 preempted, 3 queued. + ctx = [ + _StubRequest(context_chunk_size=1024, context_current_position=0), + _StubRequest(context_chunk_size=512, context_current_position=1024), + _StubRequest(context_chunk_size=256, context_current_position=768), + ] + gen = [_StubRequest(num_tokens=n) for n in (500, 1500, 2500, 3500)] + paused = [_StubRequest(num_tokens=n) for n in (400, 900)] + qitems = [_StubQueueItem(input_token_ids=list(range(n))) for n in (256, 512, 1024)] + + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=ctx, gen_reqs=gen, paused_reqs=paused), + qitems, + num_ctx_tokens=1024 + 512 + 256, + ) + + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 3 + # py_last_context_chunk[0] per req = 0, 1024, 768 + assert ifb.num_ctx_kv_tokens == 0 + 1024 + 768 + assert ifb.num_gen_requests == 4 + assert ifb.num_gen_kv_tokens == 500 + 1500 + 2500 + 3500 + assert ifb.num_queued_context_requests == 3 + assert ifb.num_queued_ctx_tokens == 256 + 512 + 1024 + assert ifb.num_paused_requests == 2 + assert ifb.num_paused_kv_tokens == 400 + 900 + + +def test_num_ctx_kv_tokens_ignores_iter_states_side_channel(): + """Regression guard: num_ctx_kv_tokens must not read iter_states. + + Under overlap scheduling the ``model_engine.iter_states["num_ctx_tokens"]`` + side channel is rewritten by the current iteration's forward step + before the previous batch's stats are emitted, so consuming it would + report the wrong iteration's count. This test pushes a deliberately + wrong value into ``iter_states`` and confirms the field is still + computed from per-request ``py_last_context_chunk``. If someone + re-introduces the side-channel read, this test fails. + """ + ctx = [_StubRequest(context_chunk_size=744, context_current_position=256)] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(context_reqs=ctx), [], num_ctx_tokens=99999 + ) + # py_last_context_chunk = (256, 1000); precomputed = start = 256. + # If the side channel were consulted, we'd see 99999. + ifb = stats.inflight_batching_stats + assert ifb.num_context_requests == 1 + assert ifb.num_ctx_kv_tokens == 256 + + +# --------------------------------------------------------------------------- +# Serializer test: confirm the C++ NLOHMANN serializer exposes every new +# InflightBatchingStats member under its expected camelCase key, nested +# inside the ``inflightBatchingStats`` object. Distinct values per field +# so a cross-wiring bug (e.g. swapped keys) fails the assertion rather +# than round-tripping silently. +# --------------------------------------------------------------------------- + + +def test_to_json_str_roundtrip_includes_new_inflight_batching_stats_fields(): + """Confirm the C++ NLOHMANN serializer emits every new InflightBatchingStats key. + + Every new request-aggregate field must appear in the JSON dict under + its expected camelCase key, nested inside ``inflightBatchingStats``. + """ + import json as _json + + ifb = InflightBatchingStats() + ifb.num_scheduled_requests = 12 + ifb.num_context_requests = 5 + ifb.num_gen_requests = 7 + ifb.num_paused_requests = 3 + ifb.num_ctx_tokens = 2048 + ifb.micro_batch_id = 4 + ifb.avg_num_decoded_tokens_per_iter = 1.25 + ifb.num_ctx_kv_tokens = 256 + ifb.num_gen_kv_tokens = 9000 + ifb.num_queued_context_requests = 11 + ifb.num_queued_ctx_tokens = 4096 + ifb.num_queued_gen_requests = 6 + ifb.num_queued_gen_kv_tokens = 2345 + ifb.num_paused_kv_tokens = 1500 + + stats = IterationStats() + stats.inflight_batching_stats = ifb + + d = _json.loads(stats.to_json_str()) + ifb_d = d["inflightBatchingStats"] + # Existing keys still round-trip. + assert ifb_d["numScheduledRequests"] == 12 + assert ifb_d["numContextRequests"] == 5 + assert ifb_d["numGenRequests"] == 7 + assert ifb_d["numPausedRequests"] == 3 + assert ifb_d["numCtxTokens"] == 2048 + # New keys round-trip under the expected camelCase. + assert ifb_d["numCtxKvTokens"] == 256 + assert ifb_d["numGenKvTokens"] == 9000 + assert ifb_d["numQueuedContextRequests"] == 11 + assert ifb_d["numQueuedCtxTokens"] == 4096 + assert ifb_d["numQueuedGenRequests"] == 6 + assert ifb_d["numQueuedGenKvTokens"] == 2345 + assert ifb_d["numPausedKvTokens"] == 1500