diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 5f14cc348eab..fbbcad881f58 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/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( diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 529851947843..9e21f634a13a 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/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):