[None][fix] Fix gen-only async kvtransfer hang

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -2594,8 +2594,11 @@ def _executor_loop_pp(self):
                         DisaggScheduleStyle.GENERATION_FIRST
                         for req in self.active_requests)
                     self._check_disagg_transfer_progress_when_idle(
-                        num_fitting_reqs, fitting_disagg_gen_init_requests,
-                        wait_for_disagg_gen_transfer_progress, all_gen_first)
+                        num_fitting_reqs,
+                        fitting_disagg_gen_init_requests,
+                        wait_for_disagg_gen_transfer_progress,
+                        all_gen_first,
+                        is_idle=scheduled_batch.batch_size == 0)
 
                 self.num_scheduled_requests = scheduled_batch.batch_size
 
@@ -3507,20 +3510,27 @@ def _sync_disagg_ctx_status_entry(self, local_need_check: bool) -> int:
         return int(local_need_check)
 
     def _check_disagg_transfer_progress_when_idle(
-            self, num_fitting_reqs: int,
+            self,
+            num_fitting_reqs: int,
             fitting_disagg_gen_init_requests: List[LlmRequest],
             wait_for_disagg_gen_transfer_progress: bool,
-            all_gen_first: bool) -> None:
-        local_need_check = (num_fitting_reqs == 0
-                            and not fitting_disagg_gen_init_requests)
+            all_gen_first: bool,
+            is_idle: bool = False) -> None:
+        local_needs_progress = (num_fitting_reqs == 0
+                                and not fitting_disagg_gen_init_requests)
+
+        uses_async_gen_transfer = self._uses_async_disagg_gen_transfer()
 
         # A synchronous GEN receive is rank-local and blocking. One rank can
         # still be receiving while another is idle, so entering either the
-        # generation or context progress collective here is unsafe.
-        if not self._uses_async_disagg_gen_transfer():
+        # generation or context progress collective here is unsafe. The
+        # gen-only-no-context benchmark skips KV transfer entirely, so its
+        # ranks remain aligned and may safely poll context progress.
+        if (not uses_async_gen_transfer
+                and not self._is_disagg_gen_only_no_context_benchmark()):
             return
 
-        local_need_gen_check = (local_need_check
+        local_need_gen_check = (uses_async_gen_transfer and local_needs_progress
                                 and wait_for_disagg_gen_transfer_progress)
 
         any_need_gen_check = self._sync_disagg_gen_status_entry(
@@ -3533,12 +3543,15 @@ def _check_disagg_transfer_progress_when_idle(
             self._check_disagg_gen_cache_transfer_status(1)
             return
 
-        any_need_check = self._sync_disagg_ctx_status_entry(local_need_check)
+        local_need_ctx_check = is_idle or (uses_async_gen_transfer
+                                           and local_needs_progress)
+        any_need_check = self._sync_disagg_ctx_status_entry(
+            local_need_ctx_check)
         if any_need_check > 0:
-            if local_need_check and not all_gen_first:
+            if local_need_ctx_check and not all_gen_first:
                 logger.warning(
-                    "num_fitting_reqs=0 and fitting_disagg_gen_init_requests is empty, may not have enough kvCache"
-                )
+                    "Executor is idle or no disaggregated generation request "
+                    "fits; waiting for context KV cache transfer progress")
                 # Local conditions warrant a blocking wait for at least one
                 # in-flight transfer to complete so KV blocks can be freed.
                 self._check_disagg_ctx_cache_transfer_status(1)
@@ -3699,8 +3712,11 @@ def _prepare_and_schedule_batch(self):
                 schedule_style == DisaggScheduleStyle.GENERATION_FIRST
                 for req in self.active_requests)
             self._check_disagg_transfer_progress_when_idle(
-                num_fitting_reqs, admitted_disagg_gen_init_requests,
-                wait_for_disagg_gen_transfer_progress, all_gen_first)
+                num_fitting_reqs,
+                admitted_disagg_gen_init_requests,
+                wait_for_disagg_gen_transfer_progress,
+                all_gen_first,
+                is_idle=scheduled_batch.batch_size == 0)
 
             # In gen-only benchmark mode, all requests must fit in KV cache
             # simultaneously. If some requests are stuck in INIT state and the

tests/integration/defs/disaggregated/test_configs/disagg_config_mamba_bs1_concurrency2.yaml

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+hostname: localhost
+model: NVIDIA-Nemotron-Nano-9B-v2
+backend: pytorch
+context_servers:
+  num_instances: 1
+  max_batch_size: 1
+  max_num_tokens: 1024
+  max_seq_len: 1024
+  trust_remote_code: true
+  kv_cache_config:
+    enable_block_reuse: false
+    use_kv_cache_manager_v2: false
+    mamba_ssm_cache_dtype: float16
+    free_gpu_memory_fraction: 0.2
+  cache_transceiver_config:
+    backend: NIXL
+    transceiver_runtime: PYTHON
+    max_tokens_in_buffer: 1024
+generation_servers:
+  num_instances: 1
+  max_batch_size: 1
+  max_num_tokens: 1024
+  max_seq_len: 1024
+  trust_remote_code: true
+  kv_cache_config:
+    enable_block_reuse: false
+    use_kv_cache_manager_v2: false
+    mamba_ssm_cache_dtype: float16
+    free_gpu_memory_fraction: 0.2
+  cache_transceiver_config:
+    backend: NIXL
+    transceiver_runtime: PYTHON
+    max_tokens_in_buffer: 1024

tests/integration/defs/disaggregated/test_disaggregated.py

@@ -47,6 +47,8 @@
 from tensorrt_llm._utils import mpi_disabled
 from tensorrt_llm.logger import logger
 
+MAMBA_BS1_CONCURRENCY2_MODEL = "NVIDIA-Nemotron-Nano-9B-v2"
+
 
 @dataclass
 class TestConfig:
@@ -365,6 +367,8 @@ def get_test_config(test_desc, example_dir, test_root):
         f"{test_configs_root}/disagg_config_ctxtp2_gentp2_llama31_8b_ucx.yaml",
         "mamba_conc_greater_than_mbs":
         f"{test_configs_root}/disagg_config_mamba_conc_greater_than_mbs.yaml",
+        "mamba_bs1_concurrency2":
+        f"{test_configs_root}/disagg_config_mamba_bs1_concurrency2.yaml",
     }
 
     if test_desc not in config_map:
@@ -657,6 +661,9 @@ def setup_disagg_cluster(
     startup_callback=None,
     startup_tick: int = 30,
     perf_metrics_output_dir: str | None = None,
+    ctx_env: dict[str, str] | None = None,
+    gen_env: dict[str, str] | None = None,
+    share_gpu: bool = False,
 ) -> tuple[dict[str, Any], list[ProcessWrapper], list[ProcessWrapper],
            ProcessWrapper, int, str]:
     """Load config, launch workers + disagg server, wait for ready.
@@ -757,7 +764,7 @@ def setup_disagg_cluster(
                                work_dir,
                                port=0,
                                device=device_ids,
-                               env=env,
+                               env=ctx_env or env,
                                save_log=save_log,
                                worker_index=i)
             ctx_workers.append(w)
@@ -767,6 +774,8 @@ def setup_disagg_cluster(
             )
             next_device += gpus_per_ctx
 
+        if share_gpu:
+            next_device = 0
         for i in range(num_gen_instances):
             device_ids = ",".join(
                 str(d) for d in dict.fromkeys((next_device + j) % num_gpus
@@ -776,7 +785,7 @@ def setup_disagg_cluster(
                                work_dir,
                                port=0,
                                device=device_ids,
-                               env=env,
+                               env=gen_env or env,
                                save_log=save_log,
                                worker_index=i)
             gen_workers.append(w)
@@ -936,7 +945,11 @@ def run_disaggregated_test(example_dir,
                            disagg_schedule_style=None,
                            post_client_test=None,
                            assert_gen_log_contains=None,
-                           perf_metrics_output_dir=None):
+                           perf_metrics_output_dir=None,
+                           ctx_env=None,
+                           gen_env=None,
+                           share_gpu=False,
+                           server_start_timeout=300):
     """Run disaggregated test using service discovery instead of MPI.
 
     If assert_gen_log_contains is set, the generation-worker logs are captured and, after the
@@ -950,14 +963,23 @@ def run_disaggregated_test(example_dir,
 
     run_env = env.copy() if env else os.environ.copy()
     run_env["UCX_TLS"] = get_ucx_tls()
+    ctx_run_env = run_env.copy()
+    if ctx_env:
+        ctx_run_env.update(ctx_env)
+    gen_run_env = run_env.copy()
+    if gen_env:
+        gen_run_env.update(gen_env)
 
     config_file = get_test_config(test_desc, example_dir,
                                   os.path.dirname(__file__))
     config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \
         setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd,
+                             server_start_timeout=server_start_timeout,
                              schedule_style=disagg_schedule_style,
                              save_log=assert_gen_log_contains is not None,
-                             perf_metrics_output_dir=perf_metrics_output_dir)
+                             perf_metrics_output_dir=perf_metrics_output_dir,
+                             ctx_env=ctx_run_env, gen_env=gen_run_env,
+                             share_gpu=share_gpu)
 
     server_host = config.get("hostname", "localhost")
 
@@ -1043,6 +1065,66 @@ def test_disaggregated_single_gpu(disaggregated_test_root,
                            cwd=llm_venv.get_working_directory())
 
 
+def _verify_mamba_bs1_concurrency2(server_url: str) -> None:
+
+    async def run() -> None:
+        timeout = aiohttp.ClientTimeout(total=120)
+        prompts = (
+            "Write one sentence about disaggregated inference.",
+            "Write one sentence about recurrent-state transfer.",
+        )
+
+        async def send(session: aiohttp.ClientSession, prompt: str) -> None:
+            payload = {
+                "model": MAMBA_BS1_CONCURRENCY2_MODEL,
+                "prompt": prompt,
+                "max_tokens": 16,
+                "temperature": 0,
+                "ignore_eos": True,
+            }
+            async with session.post(f"{server_url}/v1/completions",
+                                    json=payload,
+                                    timeout=timeout) as response:
+                body = await response.json()
+                assert response.status == 200, body
+                assert body.get("choices"), body
+
+        async with aiohttp.ClientSession() as session:
+            await asyncio.gather(*(send(session, prompt) for prompt in prompts))
+
+    asyncio.run(run())
+
+
+@skip_pre_blackwell
+@pytest.mark.timeout(900)
+def test_disaggregated_mamba_bs1_concurrency2(disaggregated_example_root,
+                                              llm_venv):
+    model_path = f"{llm_models_root()}/{MAMBA_BS1_CONCURRENCY2_MODEL}"
+    env = llm_venv._new_env.copy()
+    repo_root = os.path.abspath(
+        os.path.join(os.path.dirname(__file__), "../../../.."))
+    env["LLM_ROOT"] = repo_root
+    env["PYTHONPATH"] = os.pathsep.join(path for path in (repo_root,
+                                                          env.get("PYTHONPATH"))
+                                        if path)
+    env.pop("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", None)
+    env["TRTLLM_NIXL_NUM_THREADS"] = "1"
+    worker_env = {"TRTLLM_DISAGG_BENCHMARK_GEN_ONLY": "1"}
+    run_disaggregated_test(
+        disaggregated_example_root,
+        "mamba_bs1_concurrency2",
+        num_iters=0,
+        env=env,
+        model_path=model_path,
+        cwd=llm_venv.get_working_directory(),
+        post_client_test=_verify_mamba_bs1_concurrency2,
+        ctx_env=worker_env,
+        gen_env=worker_env,
+        share_gpu=True,
+        server_start_timeout=600,
+    )
+
+
 @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'],
                          indirect=True)
 def test_disaggregated_tinyllama_multi_orchestrator(disaggregated_test_root,

tests/integration/test_lists/qa/llm_function_core.txt

@@ -881,6 +881,7 @@ disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt
 disaggregated/test_disaggregated.py::test_disaggregated_kv_cache_time_output[TinyLlama-1.1B-Chat-v1.0]
 disaggregated/test_disaggregated.py::test_disaggregated_load_balance[TinyLlama-1.1B-Chat-v1.0]
 disaggregated/test_disaggregated.py::test_disaggregated_logprobs_serving[llama-3.1-8b-instruct]
+disaggregated/test_disaggregated.py::test_disaggregated_mamba_bs1_concurrency2
 disaggregated/test_disaggregated.py::test_disaggregated_mixed[TinyLlama-1.1B-Chat-v1.0]
 disaggregated/test_disaggregated.py::test_disaggregated_multi_gpu[TinyLlama-1.1B-Chat-v1.0]
 disaggregated/test_disaggregated.py::test_disaggregated_ngram[TinyLlama-1.1B-Chat-v1.0]

tests/unittest/_torch/executor/test_benchmark_disagg.py

@@ -1151,7 +1151,7 @@ def test_healthy_fill_phase_does_not_kill(self):
         )
         ex._handle_errors.assert_not_called()
 
-    def test_partial_transfer_admission_uses_only_admitted_requests(self):
+    def test_partial_transfer_admission_uses_only_admitted_requests(self) -> None:
         """The admitted subset is prepared and passed to the idle check."""
         admitted_req = _make_active_request(in_init=True)
         deferred_req = _make_active_request(in_init=True)
@@ -1166,7 +1166,7 @@ def test_partial_transfer_admission_uses_only_admitted_requests(self):
         ex._apply_disagg_transfer_admission.assert_called_once_with(candidates)
         ex._prepare_disagg_gen_init.assert_called_once_with([admitted_req])
         ex._check_disagg_transfer_progress_when_idle.assert_called_once_with(
-            0, [admitted_req], False, False
+            0, [admitted_req], False, False, is_idle=True
         )
         ex._handle_errors.assert_not_called()