From 223df8c629660862697d4ddd7c10dada3f2eebd5 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Sat, 18 Jul 2026 17:34:04 -0700 Subject: [PATCH 1/2] fix: exclude ADP/cuda-graph dummy requests from spec-decode acceptance stats The per-iteration speculative-decoding stats loop in _append_iter_stats counted attention-DP (and cuda-graph padding) dummy generation requests as real drafted requests. On a disagg-GEN worker under attention-DP, an idle rank is padded with a *generation* dummy that carries synthetic py_draft_tokens=[1]*k (resource_manager.add_dummy_requests is_gen branch) and is assigned a spec-sampler accepted count, so it entered the AR/AL aggregate as a mostly-rejected drafted request and depressed the reported acceptance rate/length. The parallel speculation-gate loop already filters is_dummy; mirror that here. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3ab3b8d4f6d9..b8eaa9cf600a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2017,6 +2017,16 @@ def _update_iter_stats( # Aggregate stats from all generation requests for req in scheduled_batch.generation_requests: + # Skip attention-DP / cuda-graph padding dummy requests: they + # carry synthetic py_draft_tokens=[1]*k (resource_manager + # add_dummy_requests is_gen branch) and a spec-sampler-assigned + # accepted count, which otherwise pollute the specdec AR/AL + # metric (the parallel speculation-gate loop already filters + # is_dummy). Without this, idle-rank gen dummies inflate the + # denominator with mostly-rejected rows and depress AR/AL, + # especially on disagg-GEN workers under attention-DP. + if getattr(req, 'is_dummy', False): + continue draft_len = getattr(req, 'num_draft_tokens', 0) py_draft_tokens = getattr(req, 'py_draft_tokens', None) py_num_accepted = getattr(req, 'py_num_accepted_draft_tokens', From 1b974706695b4f917bd55f7296a28e7b609820e3 Mon Sep 17 00:00:00 2001 From: Xiao Wang <24860335+xwang233@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:55:50 -0700 Subject: [PATCH 2/2] [TRTLLM-14417][test] Add regression test for specdec dummy-request filtering Address review feedback on the specdec AR/AL dummy-request fix: - Add focused unit tests in test_iter_stats_populate.py covering the speculative-decoding aggregation branch of _update_iter_stats. The existing dummy-request tests leave stats.specdec_stats as None, so that branch was previously uncovered. New tests allocate specdec_stats and verify that attention-DP / CUDA-graph dummy generation requests are excluded from draft/accepted token counts and acceptance length. - Collapse the verbose filter comment to a concise one-liner per reviewer suggestion. Signed-off-by: Xiao Wang <24860335+xwang233@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 9 +- .../executor/test_iter_stats_populate.py | 111 +++++++++++++++++- 2 files changed, 111 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9c0e0a744c28..a49ec713dad0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -2012,14 +2012,7 @@ def _update_iter_stats( # Aggregate stats from all generation requests for req in scheduled_batch.generation_requests: - # Skip attention-DP / cuda-graph padding dummy requests: they - # carry synthetic py_draft_tokens=[1]*k (resource_manager - # add_dummy_requests is_gen branch) and a spec-sampler-assigned - # accepted count, which otherwise pollute the specdec AR/AL - # metric (the parallel speculation-gate loop already filters - # is_dummy). Without this, idle-rank gen dummies inflate the - # denominator with mostly-rejected rows and depress AR/AL, - # especially on disagg-GEN workers under attention-DP. + # exclude attention dp dummy / CUDA Graph padding requests from AL calculation if getattr(req, 'is_dummy', False): continue draft_len = getattr(req, 'num_draft_tokens', 0) diff --git a/tests/unittest/_torch/executor/test_iter_stats_populate.py b/tests/unittest/_torch/executor/test_iter_stats_populate.py index 051047c4442b..1aad42de9499 100644 --- a/tests/unittest/_torch/executor/test_iter_stats_populate.py +++ b/tests/unittest/_torch/executor/test_iter_stats_populate.py @@ -47,7 +47,7 @@ ADPIterStatsBuffer, ) from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import RankIterStatsPayload, RankState -from tensorrt_llm.bindings.executor import InflightBatchingStats, IterationStats +from tensorrt_llm.bindings.executor import InflightBatchingStats, IterationStats, SpecDecodingStats class _StubRequest: @@ -73,6 +73,9 @@ def __init__( is_attention_dp_dummy: bool = False, is_cuda_graph_dummy: bool = False, is_dummy_request: bool = False, + num_draft_tokens: int = 0, + py_draft_tokens=None, + py_num_accepted_draft_tokens=None, ): self.context_current_position = context_current_position # py_last_context_chunk = (begin_compute, end_compute). For a fresh @@ -88,6 +91,14 @@ def __init__( self.is_attention_dp_dummy = is_attention_dp_dummy self.is_cuda_graph_dummy = is_cuda_graph_dummy self.is_dummy_request = is_dummy_request + # Speculative-decoding accessors read by the specdec-stats aggregation + # branch of ``_update_iter_stats``. Defaults match a non-drafted request + # (draft_len resolves to 0, so it is skipped by the ``draft_len > 0`` + # gate). A dummy gen request mirrors resource_manager.add_dummy_requests' + # is_gen branch: py_draft_tokens = [1] * k plus an accepted count. + self.num_draft_tokens = num_draft_tokens + self.py_draft_tokens = py_draft_tokens + self.py_num_accepted_draft_tokens = py_num_accepted_draft_tokens @property def is_dummy(self) -> bool: @@ -189,6 +200,7 @@ def _invoke_update_iter_stats( num_ctx_tokens, enable_attention_dp=False, scheduled_batch_stats=None, + specdec=False, ): """Call real ``PyExecutor._update_iter_stats`` unbound; return the stats. @@ -207,6 +219,11 @@ def _invoke_update_iter_stats( Explicit scheduled-batch counters passed to ``_update_iter_stats``. When None, defaults to the same partial ``num_ctx_tokens`` payload used by older tests. + specdec : bool + Allocate ``stats.specdec_stats`` so the speculative-decoding + aggregation branch runs. On a fresh ``IterationStats`` this member is + None (production sets it in ``_get_init_iter_stats`` only when spec + decode is enabled), and the branch is skipped otherwise. """ from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, ScheduledBatchStats @@ -226,6 +243,8 @@ def _invoke_update_iter_stats( # 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() + if specdec: + stats.specdec_stats = SpecDecodingStats() with patch( "tensorrt_llm._torch.pyexecutor.py_executor.torch.cuda.mem_get_info", @@ -587,6 +606,96 @@ def test_num_gen_kv_tokens_uses_scheduled_batch_stats(): assert ifb.num_gen_kv_tokens == 1024 +# --------------------------------------------------------------------------- +# Speculative-decoding aggregation tests: the specdec branch of +# ``_update_iter_stats`` sums draft/accepted tokens across generation +# requests and derives the acceptance length. Attention-DP / CUDA-graph +# padding adds *generation* dummies carrying synthetic +# ``py_draft_tokens = [1] * k`` and a spec-sampler accepted count +# (resource_manager.add_dummy_requests is_gen branch); they must not enter +# the aggregate or they depress the reported acceptance rate/length. This +# path is distinct from the KV-token dummy filtering above and is only +# reached when ``stats.specdec_stats`` is allocated (spec decode enabled). +# --------------------------------------------------------------------------- + + +def test_specdec_excludes_dummy_generation_requests(): + """Regression guard: dummy gen requests must not pollute specdec AR/AL. + + One real drafted request (k=4, 3 accepted) alongside an attention-DP + dummy and a CUDA-graph dummy, each mirroring the is_gen dummy branch + (``num_draft_tokens == 0`` with ``py_draft_tokens = [1] * 4``). The + aggregate must reflect only the real request: 4 draft tokens, 3 + accepted, one drafted request, and acceptance_length = (3 + 1) / 1. + + Without the ``is_dummy`` filter each dummy would contribute draft_len=4 + (counted from the non-zero ``py_draft_tokens``) and 0 accepted, giving + num_requests_with_draft=3 and a depressed acceptance_length of + (3 + 3) / 3 = 2.0 instead of 4.0. + """ + real = _StubRequest(num_draft_tokens=4, py_num_accepted_draft_tokens=3) + adp_dummy = _StubRequest( + py_draft_tokens=[1, 1, 1, 1], + py_num_accepted_draft_tokens=0, + is_attention_dp_dummy=True, + ) + cuda_graph_dummy = _StubRequest( + py_draft_tokens=[1, 1, 1, 1], + py_num_accepted_draft_tokens=0, + is_cuda_graph_dummy=True, + ) + stats = _invoke_update_iter_stats( + _StubScheduledBatch(gen_reqs=[real, adp_dummy, cuda_graph_dummy]), + [], + num_ctx_tokens=0, + specdec=True, + ) + sd = stats.specdec_stats + assert sd.num_draft_tokens == 4 + assert sd.num_accepted_tokens == 3 + assert sd.num_requests_with_draft_tokens == 1 + assert sd.acceptance_length == 4.0 + + +def test_specdec_all_dummy_generation_requests_yield_zero(): + """An iteration of only dummy gen requests reports no drafted work.""" + dummies = [ + _StubRequest(py_draft_tokens=[1, 1, 1, 1], is_attention_dp_dummy=True), + _StubRequest(py_draft_tokens=[1, 1, 1, 1], is_cuda_graph_dummy=True), + ] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(gen_reqs=dummies), + [], + num_ctx_tokens=0, + specdec=True, + ) + sd = stats.specdec_stats + assert sd.num_draft_tokens == 0 + assert sd.num_accepted_tokens == 0 + assert sd.num_requests_with_draft_tokens == 0 + assert sd.acceptance_length == 0.0 + + +def test_specdec_aggregates_multiple_real_requests(): + """Baseline: with no dummies, every real drafted request is aggregated.""" + reqs = [ + _StubRequest(num_draft_tokens=4, py_num_accepted_draft_tokens=3), + _StubRequest(num_draft_tokens=4, py_num_accepted_draft_tokens=1), + ] + stats = _invoke_update_iter_stats( + _StubScheduledBatch(gen_reqs=reqs), + [], + num_ctx_tokens=0, + specdec=True, + ) + sd = stats.specdec_stats + assert sd.num_draft_tokens == 8 + assert sd.num_accepted_tokens == 4 + assert sd.num_requests_with_draft_tokens == 2 + # (total_accepted + num_requests_with_draft) / num_requests_with_draft + assert sd.acceptance_length == (4 + 2) / 2 + + # --------------------------------------------------------------------------- # Attention-DP fanout tests: completed rank-local payloads are carried by the # next ADP allgather, then rank 0 appends one row per ADP rank.