[None][fix] improve BCG + ADP trigger time in agg mode

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -3949,11 +3949,28 @@ def get_padded_prefill_tokens(tokens: int) -> int:
                 and self._prefill_cuda_graph_num_tokens):
             max_captured_num_tokens = self._prefill_cuda_graph_num_tokens[-1]
             if attn_all_rank_num_tokens is not None:
-                has_ctx_requests = num_ctx_requests != 0 or (
-                    all_rank_ctx_requests is not None
-                    and any(ctx_requests != 0
-                            for ctx_requests in all_rank_ctx_requests))
-                can_run_prefill_cuda_graph = (has_ctx_requests
+                # Every rank is padded to bucket(max(all_rank_num_tokens))
+                # below, so a rank with no context work of its own still
+                # replays the prefill graph at the busiest rank's token count.
+                # Requiring context work everywhere keeps that amplification
+                # off generation-only ranks; previously one context request
+                # anywhere was enough, which in aggregate serving turned a
+                # one-token generation step into a full prefill replay.
+                #
+                # Disaggregated serving still qualifies: an otherwise idle
+                # rank receives a context dummy from
+                # `PyExecutor._pad_attention_dp_dummy_request`, which runs
+                # before `_schedule()` and so is counted here. Aggregate
+                # serving keeps a generation dummy instead
+                # (`_update_adp_dummy_role` returns early with no cache
+                # transceiver), which is what separates the two paths.
+                if all_rank_ctx_requests is not None:
+                    every_rank_has_ctx_requests = all(
+                        ctx_requests != 0
+                        for ctx_requests in all_rank_ctx_requests)
+                else:
+                    every_rank_has_ctx_requests = num_ctx_requests != 0
+                can_run_prefill_cuda_graph = (every_rank_has_ctx_requests
                                               and max(attn_all_rank_num_tokens)
                                               <= max_captured_num_tokens)
                 all_ranks_can_run_prefill_cuda_graph = list(

tests/unittest/llmapi/test_llm_args.py

@@ -2167,17 +2167,52 @@ def tp_allgather(self, value):
         engine.enable_attention_dp = True
         engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.BREAKABLE
         engine._prefill_cuda_graph_num_tokens = [128, 256, 512]
-        engine._get_all_rank_ctx_requests = lambda _: [0, 1, 0, 0]
+        engine._get_all_rank_ctx_requests = lambda _: [1, 1, 1, 1]
 
-        all_rank_num_tokens = [1, 129, 1, 1]
+        all_rank_num_tokens = [129, 129, 129, 129]
         engine.dist = FakeDist([True, True, True, True])
-        assert engine._get_padding_params(1, 0,
+        assert engine._get_padding_params(129, 1,
                                           all_rank_num_tokens) == (256, True,
                                                                    [256] * 4)
 
         engine.dist = FakeDist([True, False, True, True])
         assert engine._get_padding_params(
-            1, 0, all_rank_num_tokens) == (1, False, all_rank_num_tokens)
+            129, 1, all_rank_num_tokens) == (129, False, all_rank_num_tokens)
+
+    def test_attention_dp_prefill_graph_skips_generation_only_ranks(self):
+        from tensorrt_llm._torch.pyexecutor.model_engine import \
+            PyTorchModelEngine
+
+        class EchoDist:
+            """Reports the same decision every rank would reach on its own.
+
+            The gate is derived from the gathered context-request counts,
+            which are identical everywhere, so echoing the local answer is
+            what the real all-gather returns.
+            """
+
+            def tp_allgather(self, value):
+                return [value] * 4
+
+        engine = object.__new__(PyTorchModelEngine)
+        engine.enable_attention_dp = True
+        engine.prefill_cuda_graph_backend = PrefillCudaGraphBackend.BREAKABLE
+        engine._prefill_cuda_graph_num_tokens = [128, 256, 512]
+        engine.dist = EchoDist()
+
+        # Context work on every rank: the group shares one padded bucket.
+        engine._get_all_rank_ctx_requests = lambda _: [1, 1, 1, 1]
+        busy = [129, 129, 129, 129]
+        assert engine._get_padding_params(129, 1,
+                                          busy) == (256, True, [256] * 4)
+
+        # One generation-only rank is enough to keep the whole group out of
+        # the graph. Padding would lift its single token to the busiest
+        # rank's bucket, so it would replay a full prefill to produce one
+        # token -- the aggregate-serving regression this gate prevents.
+        engine._get_all_rank_ctx_requests = lambda _: [0, 1, 1, 1]
+        mixed = [1, 129, 129, 129]
+        assert engine._get_padding_params(1, 0, mixed) == (1, False, mixed)
 
     def test_torch_compile_config_does_not_populate_legacy_capture_buckets(
             self):