diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index e262fd825ccc..79a953d782ee 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2987,14 +2987,20 @@ def __init__( @property def disagg_request_id(self) -> int: + """The key this session goes on the wire under, so it must be the id the + ctx TxSession registered with: every REQUEST_DATA, result and cancel + message is matched by it on both sides. + + ``ctx_request_id`` is that id: the ctx server sets it from its own + TxSession key when it answers. ``disagg_request_id`` only agrees with it + when the orchestrator handed both servers the same id, so it is a + fallback rather than the first choice. + """ params = self._base_args.params - if params.disagg_request_id is not None: - return params.disagg_request_id - # ctx_request_id is set on gen-side requests to the ctx server's request ID, - # which matches the key the ctx TxSession registered under. Fall back to - # the local request_id only when neither field is available. if params.ctx_request_id is not None: return params.ctx_request_id + if params.disagg_request_id is not None: + return params.disagg_request_id return self.request_id @contextmanager diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py index 55721b44aceb..b8ab00cd256c 100644 --- a/tensorrt_llm/_torch/speculative/eagle3.py +++ b/tensorrt_llm/_torch/speculative/eagle3.py @@ -120,7 +120,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): self.slot_ids = [] for req in context_batch: if req.is_first_context_chunk: - slot_id = self.slot_manager.add_slot(req.request_id) + # A padding dummy (e.g. the attention-DP idle-rank dummy) may + # already hold a slot from add_dummy_requests -- reuse it + # instead of re-adding, which SlotManager.add_slot only + # tolerates for the CUDA-graph dummy id. + slot_id = self.slot_manager.get_slot(req.request_id) + if slot_id is None: + slot_id = self.slot_manager.add_slot(req.request_id) self.slot_ids.append(slot_id) if self.use_relaxed_acceptance_for_thinking: self.relaxed_delta_pool[slot_id].fill_(0) diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 7363ace81da9..2107fbed26c8 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -80,7 +80,13 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): # allocate hidden state tensors for req in context_batch: if req.is_first_context_chunk: - slot_id = self.slot_manager.add_slot(req.request_id) + # A padding dummy (e.g. the attention-DP idle-rank dummy) may + # already hold a slot from add_dummy_requests -- reuse it + # instead of re-adding, which SlotManager.add_slot only + # tolerates for the CUDA-graph dummy id. + slot_id = self.slot_manager.get_slot(req.request_id) + if slot_id is None: + slot_id = self.slot_manager.add_slot(req.request_id) if self.use_relaxed_acceptance_for_thinking: self.mtp_relaxed_delta_pool[slot_id].copy_( 0, non_blocking=True) diff --git a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py index a3d36cf9bdf6..15f5d4cdab31 100644 --- a/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py +++ b/tests/unittest/_torch/speculative/test_spec_slot_pool_sizing.py @@ -166,6 +166,46 @@ def test_mtp_slot_pool_survives_a_full_overlap_turnover(): assert all(0 <= slot < POOL + 1 for slot in retiring + incoming) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="MTP hidden-state pools are CUDA tensors") +def test_mtp_prepare_resources_reuses_a_preallocated_dummy_slot(): + """A slot registered by ``add_dummy_requests`` must be reused, not re-added. + + The attention-DP idle-rank dummy registers its slot via + ``add_dummy_requests``, then the same request flows through + ``prepare_resources`` as an ordinary first-context-chunk request. + ``SlotManager.add_slot`` only tolerates a duplicate add for the + CUDA-graph dummy id, so a regression that drops the reuse guard in + ``prepare_resources`` crashes on any other preallocated id. + """ + mgr = MTPHiddenStatesManager( + _mtp_config(), torch.float16, hidden_size=8, max_num_requests=R, num_seq_slots=POOL + ) + + dummy_id = 0 + mgr.add_dummy_requests([dummy_id]) + preallocated_slot = mgr.slot_manager.get_slot(dummy_id) + + scheduled_batch = types.SimpleNamespace( + context_requests=[ + types.SimpleNamespace(request_id=dummy_id, is_first_context_chunk=True), + ] + ) + mgr.prepare_resources(scheduled_batch) + + assert mgr.slot_manager.get_slot(dummy_id) == preallocated_slot + + other_id = 1 + scheduled_batch = types.SimpleNamespace( + context_requests=[ + types.SimpleNamespace(request_id=other_id, is_first_context_chunk=True), + ] + ) + mgr.prepare_resources(scheduled_batch) + + assert mgr.slot_manager.get_slot(other_id) is not None + assert mgr.slot_manager.get_slot(other_id) != preallocated_slot + + # --------------------------------------------------------------------------- # Resource managers. The plumbing is per-branch, so the AST guard below makes # forgetting a branch a test failure rather than a runtime IndexError. @@ -355,6 +395,52 @@ def test_eagle3_keeps_the_max_seq_len_floor(): assert mgr.slot_manager.max_num_requests == 1024 + 1 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Eagle3 hidden states are CUDA tensors") +def test_eagle3_prepare_resources_reuses_a_preallocated_dummy_slot(): + """A slot registered by ``add_dummy_requests`` must be reused, not re-added. + + The attention-DP idle-rank dummy registers its slot via + ``add_dummy_requests``, then the same request flows through + ``prepare_resources`` as an ordinary first-context-chunk request. + ``SlotManager.add_slot`` only tolerates a duplicate add for the + CUDA-graph dummy id, so a regression that drops the reuse guard in + ``prepare_resources`` crashes on any other preallocated id. + """ + mgr = Eagle3ResourceManager( + _eagle_config(), + torch.float16, + hidden_size=8, + max_num_requests=R, + max_seq_len=4, + max_num_tokens=64, + num_seq_slots=POOL, + ) + + dummy_id = 0 + mgr.add_dummy_requests([dummy_id]) + preallocated_slot = mgr.slot_manager.get_slot(dummy_id) + + scheduled_batch = types.SimpleNamespace( + context_requests=[ + types.SimpleNamespace(request_id=dummy_id, is_first_context_chunk=True), + ] + ) + mgr.prepare_resources(scheduled_batch) + + assert mgr.slot_manager.get_slot(dummy_id) == preallocated_slot + + other_id = 1 + scheduled_batch = types.SimpleNamespace( + context_requests=[ + types.SimpleNamespace(request_id=other_id, is_first_context_chunk=True), + ] + ) + mgr.prepare_resources(scheduled_batch) + + assert mgr.slot_manager.get_slot(other_id) is not None + assert mgr.slot_manager.get_slot(other_id) != preallocated_slot + + def _sa_manager(num_seq_slots, **config_kwargs): config = SAConfig(max_seq_len=1024, max_slots=R, **config_kwargs) return SuffixAutomatonManager(config, R, 1024, num_seq_slots=num_seq_slots) diff --git a/tests/unittest/disaggregated/kv_transfer_harness.py b/tests/unittest/disaggregated/kv_transfer_harness.py index 305938ed40c8..3a9ae090f67a 100644 --- a/tests/unittest/disaggregated/kv_transfer_harness.py +++ b/tests/unittest/disaggregated/kv_transfer_harness.py @@ -396,7 +396,10 @@ def run_kv_transfer_test( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, + # Mirrors a real ctx response: ContextPhaseParams.req_id + # resolves disagg-first, so it equals the id the ctx TxSession + # registered under, not the ctx server's local request id. + ctx_request_id=unique_rid, ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index e140f53eb11d..608eecf826e0 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -600,12 +600,13 @@ def _init_pool_data(managers, tp, is_mla, use_v2, fill_random=True, seed_base=10 # Add sequence to manager # --------------------------------------------------------------------------- def _make_gen_request( - gen_rid, req_len, unique_rid, ctx_rid, ctx_dp_rank, ctx_info_endpoint, sampling_params + gen_rid, req_len, unique_rid, ctx_dp_rank, ctx_info_endpoint, sampling_params ): """Build one GENERATION_ONLY request wired to its context peer. Used to mint a distinct request per helix CP rank (each needs its own - rank-local prompt_len); mirrors the inline construction in run_transfer_test. + rank-local prompt_len); mirrors the inline construction in run_transfer_test, + including ctx_request_id carrying the key the ctx TxSession registered under. """ req = LlmRequest( request_id=gen_rid, @@ -618,7 +619,7 @@ def _make_gen_request( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) req.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, + ctx_request_id=unique_rid, ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, @@ -1354,7 +1355,7 @@ def run_transfer_test( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, + ctx_request_id=unique_rid, ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, @@ -1389,7 +1390,6 @@ def run_transfer_test( gen_rid, req_len, unique_rid, - ctx_rid, ctx_dp_rank, ctx_info_endpoint, sampling_params, diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 9494849f23fd..f6f807ed962e 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -607,7 +607,7 @@ def _gen_ranks(req_idx: int) -> List[int]: llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_req.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, + ctx_request_id=unique_rid, ctx_dp_rank=_ctx_ranks(req_idx)[0] if enable_attention_dp else 0, ctx_info_endpoint=ctx_endpoint, disagg_request_id=unique_rid, diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 92b7a836032f..d2969c7d244b 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -976,7 +976,7 @@ def add_and_verify_request( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_request.py_request_id, + ctx_request_id=unique_rid, ctx_dp_rank=ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, @@ -1456,7 +1456,7 @@ def test_transfer_with_gen_prefix_offset(use_v2, chunk_size_blocks): llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=0, + ctx_request_id=unique_rid, ctx_dp_rank=0, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, @@ -1712,7 +1712,7 @@ def _setup_chunked_request(setup, ctx_request_id, gen_request_id, request_len): llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_request.py_request_id, + ctx_request_id=unique_rid, ctx_dp_rank=0, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, @@ -1903,7 +1903,7 @@ def test_session_has_transferring_tasks_false(): ) gen_request.py_disaggregated_params = DisaggregatedParams( disagg_request_id=unique_rid, - ctx_request_id=ctx_request.py_request_id, + ctx_request_id=unique_rid, schedule_style=1, ) @@ -1975,7 +1975,7 @@ def make_gen_request(request_id): llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) req.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=request_id, + ctx_request_id=rid, ctx_dp_rank=0, ctx_info_endpoint=bad_endpoint, disagg_request_id=rid, diff --git a/tests/unittest/disaggregated/test_kv_transfer_mp.py b/tests/unittest/disaggregated/test_kv_transfer_mp.py index 539a6039523e..375012b118e5 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_mp.py +++ b/tests/unittest/disaggregated/test_kv_transfer_mp.py @@ -391,7 +391,7 @@ def process_and_verify_request( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_request_id, + ctx_request_id=unique_rid, ctx_dp_rank=0, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid, diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index f31a4ca7a17c..0c4576c435d5 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -757,7 +757,7 @@ def run_mamba_transfer_test( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) gen_req.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_rid, + ctx_request_id=unique_rid, ctx_dp_rank=0, ctx_info_endpoint=ctx_endpoint, disagg_request_id=unique_rid, diff --git a/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py b/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py index d414d18fb48b..667abfe1aefa 100644 --- a/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py +++ b/tests/unittest/disaggregated/test_py_cache_transceiver_mp.py @@ -527,7 +527,7 @@ def create_request( llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, ) request.py_disaggregated_params = DisaggregatedParams( - ctx_request_id=ctx_request_id, + ctx_request_id=unique_rid, ctx_dp_rank=actual_ctx_dp_rank, ctx_info_endpoint=ctx_info_endpoint, disagg_request_id=unique_rid,