From 58e200aac2c465f62f2d7b92b6472d7e4db5e725 Mon Sep 17 00:00:00 2001 From: Bo Deng Date: Tue, 26 May 2026 03:56:05 +0000 Subject: [PATCH 1/4] [TRTLLM-12958][feat] Enable gen-only spec dec Signed-off-by: Bo Deng --- tensorrt_llm/_torch/disaggregation/native/peer.py | 11 +++++++---- tensorrt_llm/_torch/disaggregation/native/transfer.py | 5 +++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 4 ++++ 3 files changed, 16 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 0d305cee40df..0a917548b463 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -97,11 +97,14 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: self_layers = sum(self._ri.layer_num_per_pp) peer_layers = sum(peer_ri.layer_num_per_pp) if self_layers != peer_layers: - logger.warning( - "PeerRegistrar: total layer count mismatch " - f"(local={self_layers}, peer={peer_layers})." + # Allow mismatch when one side has speculative (e.g. MTP) layers + # that the other side doesn't. The pool_mapping logic will only + # transfer layers that exist on both sides. + logger.info( + "PeerRegistrar: layer count differs " + f"(local={self_layers}, peer={peer_layers}), " + "allowing partial layer transfer." ) - return False return True diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 75398ac38971..89b51a02ec4b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1717,6 +1717,11 @@ def unpack_aux(self, request: LlmRequest) -> None: self._aux_buffer.get_slot_data(self.aux_slot) ) request.py_first_gen_tokens = first_gen_tokens # type: ignore[attr-defined] + # When CTX has no MTP but GEN does, draft_tokens will be empty. + # Pad with dummy zeros so GEN's MTP forward sees uniform draft_len + # across the batch. The dummy tokens will be rejected on first verify. + if not draft_tokens and self._aux_buffer._max_draft_len > 0: + draft_tokens = [0] * self._aux_buffer._max_draft_len request.py_draft_tokens = draft_tokens # type: ignore[attr-defined] if request.py_disaggregated_params is not None: request.py_disaggregated_params.ctx_usage = { diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 7dcfd154c075..daa63df562d6 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4214,6 +4214,10 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): req.py_kv_transfer_timed_out = False first_gen_tokens = req.context_phase_params.first_gen_tokens ctx_draft_tokens = req.context_phase_params.draft_tokens + if not ctx_draft_tokens and self.model_engine.enable_spec_decode: + # CTX has no MTP — fill dummy draft tokens so GEN's MTP + # forward sees uniform draft_len. Dummies will be rejected. + ctx_draft_tokens = [0] * self.model_engine.max_draft_len req.py_draft_tokens = [] if ctx_draft_tokens is None else ctx_draft_tokens beam_width = req.py_beam_width for beam in range(0, beam_width): From 124cc70676f2c79103f9d3e13487a9e3dd932819 Mon Sep 17 00:00:00 2001 From: Bo Deng Date: Tue, 26 May 2026 10:53:24 +0000 Subject: [PATCH 2/4] clean codes && add tests Signed-off-by: Bo Deng --- .../_torch/disaggregation/native/peer.py | 2 +- .../_torch/disaggregation/native/transfer.py | 5 -- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++- .../accuracy/test_disaggregated_serving.py | 90 +++++++++++++++++++ .../test_lists/test-db/l0_dgx_h100.yml | 2 + 5 files changed, 100 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/peer.py b/tensorrt_llm/_torch/disaggregation/native/peer.py index 0a917548b463..80eaa30bb59e 100644 --- a/tensorrt_llm/_torch/disaggregation/native/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/peer.py @@ -100,7 +100,7 @@ def _check_peer_compatible(self, peer_ri: RankInfo) -> bool: # Allow mismatch when one side has speculative (e.g. MTP) layers # that the other side doesn't. The pool_mapping logic will only # transfer layers that exist on both sides. - logger.info( + logger.warning( "PeerRegistrar: layer count differs " f"(local={self_layers}, peer={peer_layers}), " "allowing partial layer transfer." diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 89b51a02ec4b..75398ac38971 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -1717,11 +1717,6 @@ def unpack_aux(self, request: LlmRequest) -> None: self._aux_buffer.get_slot_data(self.aux_slot) ) request.py_first_gen_tokens = first_gen_tokens # type: ignore[attr-defined] - # When CTX has no MTP but GEN does, draft_tokens will be empty. - # Pad with dummy zeros so GEN's MTP forward sees uniform draft_len - # across the batch. The dummy tokens will be rejected on first verify. - if not draft_tokens and self._aux_buffer._max_draft_len > 0: - draft_tokens = [0] * self._aux_buffer._max_draft_len request.py_draft_tokens = draft_tokens # type: ignore[attr-defined] if request.py_disaggregated_params is not None: request.py_disaggregated_params.ctx_usage = { diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index daa63df562d6..f16a1f81ed4b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4215,9 +4215,13 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): first_gen_tokens = req.context_phase_params.first_gen_tokens ctx_draft_tokens = req.context_phase_params.draft_tokens if not ctx_draft_tokens and self.model_engine.enable_spec_decode: - # CTX has no MTP — fill dummy draft tokens so GEN's MTP - # forward sees uniform draft_len. Dummies will be rejected. - ctx_draft_tokens = [0] * self.model_engine.max_draft_len + # CTX has no speculative decoding — fill dummy draft tokens + # so model_engine builds the correct input shape (1 + draft_len + # tokens per gen request). Dummies will be rejected on verify, + # and the draft model will produce real tokens for next step. + ctx_draft_tokens = [ + 0 + ] * self.model_engine.max_total_draft_tokens req.py_draft_tokens = [] if ctx_draft_tokens is None else ctx_draft_tokens beam_width = req.py_beam_width for beam in range(0, beam_width): diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 74911add4cfd..2782bfb8db8c 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -780,6 +780,63 @@ def test_eagle3(self, overlap_scheduler, eagle3_one_model): self.MODEL_PATH) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.skip_less_device(2) + @skip_pre_hopper + def test_gen_only_spec_dec(self): + speculative_decoding_config = { + "decoding_type": "Eagle", + "max_draft_len": 4, + "speculative_model": + f"{llm_models_root()}/EAGLE3-LLaMA3.1-Instruct-8B", + "eagle3_one_model": True, + } + ctx_server_config = { + "disable_overlap_scheduler": + True, # BS=1 does not need overlap scheduling + "kv_cache_config": { + "free_gpu_memory_fraction": 0.5, + "enable_block_reuse": True # reuse on context requests + }, + "max_num_tokens": 13393 * 2, + "max_batch_size": 1, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 4096, + }, + "cuda_graph_config": None, + } + gen_server_config = { + "disable_overlap_scheduler": False, + "speculative_config": speculative_decoding_config, + "kv_cache_config": { + "free_gpu_memory_fraction": 0.5, + "enable_block_reuse": False + }, + "max_num_tokens": 13393 * 2, + "max_batch_size": 16, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 4096, + }, + "cuda_graph_config": None, + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + } + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, gen_server_config, + self.MODEL_PATH) as llm: + run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.skip_less_device(2) @pytest.mark.skip_less_device_memory(32000) @pytest.mark.parametrize("backend", ["xgrammar", "llguidance"]) @@ -1001,6 +1058,39 @@ def test_gen_only_sync(self): ) as llm: run_accuracy_test(llm, self.MODEL_NAME, ["GSM8K"]) + @pytest.mark.skip_less_device(8) + @skip_pre_hopper + def test_gen_only_spec_dec(self): + ctx_server_config = {"disable_overlap_scheduler": True} + gen_server_config = {"disable_overlap_scheduler": False} + cache_transceiver_config = { + "backend": "NIXL", + "max_tokens_in_buffer": 4096, + "transceiver_runtime": "PYTHON", + } + ctx_server_config["cache_transceiver_config"] = cache_transceiver_config + gen_server_config["cache_transceiver_config"] = cache_transceiver_config + gen_server_config["speculative_config"] = { + "decoding_type": "MTP", + "max_draft_len": 2 + } + disaggregated_server_config = { + "hostname": "localhost", + "backend": "pytorch", + "context_servers": { + "num_instances": 1 + }, + "generation_servers": { + "num_instances": 1 + } + } + with launch_disaggregated_llm(disaggregated_server_config, + ctx_server_config, + gen_server_config, + self.MODEL_PATH, + tensor_parallel_size=4) as llm: + run_accuracy_test(llm, self.MODEL_NAME, ["MMLU", "GSM8K"]) + @pytest.mark.skip_less_device(8) @parametrize_with_ids("overlap_scheduler", [True, False]) @parametrize_with_ids("mtp_nextn", [0, 2]) diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 7a516247f1d2..110d0951459c 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -56,6 +56,8 @@ l0_dgx_h100: - disaggregated/test_disaggregated.py::test_disaggregated_cancel_large_context_requests[DeepSeek-V3-Lite-bf16] # llmapi - unittest/llmapi/test_mpi_session.py::test_llmapi_launch_multiple_tasks + - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_gen_only_spec_dec + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_spec_dec - condition: ranges: system_gpu_count: From e35a9198425830b68304949b25587837859d6cbd Mon Sep 17 00:00:00 2001 From: Bo Deng Date: Sun, 7 Jun 2026 11:33:09 +0000 Subject: [PATCH 3/4] fix eagle3 acc Signed-off-by: Bo Deng --- tensorrt_llm/_torch/disaggregation/transceiver.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index e6da14561cb6..5b0b872e322a 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -178,6 +178,12 @@ def _create_kv_slice( groups.append(np.array([], dtype=np.int64)) continue block_ids = adapter.get_block_ids(req, idx, lg) + # Limit to prompt_len blocks, matching C++ cacheFormatter behavior. + # Extra blocks from num_extra_kv_tokens (speculative decoding) have + # uninitialized KV data and must not be transferred. + prompt_blocks = (req.prompt_len + tpb - 1) // tpb + if block_ids.size > prompt_blocks: + block_ids = block_ids[:prompt_blocks] window_size = lg.sliding_window_size if window_size is not None: From 3ee277ee1cee031e8910f047d838e9668970b630 Mon Sep 17 00:00:00 2001 From: Bo Deng Date: Wed, 10 Jun 2026 02:18:19 +0000 Subject: [PATCH 4/4] fix Signed-off-by: Bo Deng --- tensorrt_llm/_torch/disaggregation/transceiver.py | 7 +++---- tests/integration/test_lists/test-db/l0_dgx_h100.yml | 1 - tests/integration/test_lists/test-db/l0_dgx_h200.yml | 1 + 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 5b0b872e322a..d1344ec6495e 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -181,14 +181,13 @@ def _create_kv_slice( # Limit to prompt_len blocks, matching C++ cacheFormatter behavior. # Extra blocks from num_extra_kv_tokens (speculative decoding) have # uninitialized KV data and must not be transferred. - prompt_blocks = (req.prompt_len + tpb - 1) // tpb - if block_ids.size > prompt_blocks: - block_ids = block_ids[:prompt_blocks] + total_blocks = (req.prompt_len + tpb - 1) // tpb + if block_ids.size > total_blocks: + block_ids = block_ids[:total_blocks] window_size = lg.sliding_window_size if window_size is not None: # Drop stale blocks the manager may still expose (V1 pre-eviction). - total_blocks = (req.prompt_len + tpb - 1) // tpb stale_end = max(0, (req.prompt_len + 1 - window_size) // tpb) expected_valid = max(0, total_blocks - stale_end) if block_ids.size > expected_valid: diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 110d0951459c..3ce96aa19f3b 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -57,7 +57,6 @@ l0_dgx_h100: # llmapi - unittest/llmapi/test_mpi_session.py::test_llmapi_launch_multiple_tasks - accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_gen_only_spec_dec - - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_spec_dec - condition: ranges: system_gpu_count: diff --git a/tests/integration/test_lists/test-db/l0_dgx_h200.yml b/tests/integration/test_lists/test-db/l0_dgx_h200.yml index bea7ec084b82..aa468157ea3a 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h200.yml @@ -41,6 +41,7 @@ l0_dgx_h200: - disaggregated/test_disaggregated.py::test_disaggregated_ctxpp4_genpp4[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_deepseek_v3_lite_fp8_ctxtp2ep2pp2_gentp4_one_mtp_block_reuse[DeepSeek-V3-Lite-fp8] - unittest/llmapi/test_llm_pytorch.py::test_nemotron_nas_lora + - accuracy/test_disaggregated_serving.py::TestDeepSeekV3Lite::test_gen_only_spec_dec - condition: ranges: system_gpu_count: