clean codes && add tests

Signed-off-by: Bo Deng <deemod@nvidia.com>

Author: bo-nv

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."

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 = {

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -3742,9 +3742,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):

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"])
@@ -1041,6 +1098,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])

tests/integration/test_lists/test-db/l0_dgx_h100.yml

@@ -58,6 +58,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: