From efe60576c2fba2b5d7521d2a307d562b266166a2 Mon Sep 17 00:00:00 2001
From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Date: Tue, 4 Aug 2026 17:02:09 -0700
Subject: [PATCH 1/2] [https://nvbugs/6480621][fix] Preserve KV ownership in
 disaggregated precheck

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
---
 jenkins/scripts/perf/local/submit.py          |   2 +
 jenkins/scripts/perf/submit.py                |  31 ++
 .../_torch/disaggregation/native/transfer.py  |  47 ++-
 .../_torch/disaggregation/transceiver.py      |   5 +-
 .../cache_transceiver_precheck/README.md      |  10 +-
 .../precheck_config.py                        |  17 +-
 .../run_precheck.py                           | 133 +++++---
 .../test_cache_transceiver_precheck_e2e.py    |  21 +-
 .../test_transceiver_bounded_polling.py       | 134 +++++++-
 .../test_cache_transceiver_precheck_config.py |  93 +++++-
 .../test_cache_transceiver_precheck_run.py    | 305 ++++++++++++++++++
 tests/unittest/scripts/test_perf_submit.py    |  43 ++-
 12 files changed, 768 insertions(+), 73 deletions(-)

diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py
index 8e571eb2f890..94fee1b89ca3 100755
--- a/jenkins/scripts/perf/local/submit.py
+++ b/jenkins/scripts/perf/local/submit.py
@@ -967,6 +967,7 @@ def main():
                     hardware_config.get("gpus_per_ctx_server", 0) or 0,
                     hardware_config.get("gpus_per_gen_server", 0) or 0,
                 ),
+                llm_models_root=args.llm_models_root,
             )
         )
 
@@ -974,6 +975,7 @@ def main():
         srun_args_lines.extend(
             [
                 "--container-env=DISAGG_SERVING_TYPE",
+                "--container-env=LLM_MODELS_ROOT",
                 "--container-env=pytestCommand",
             ]
         )
diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py
index 72b176df3a95..c36db498b560 100755
--- a/jenkins/scripts/perf/submit.py
+++ b/jenkins/scripts/perf/submit.py
@@ -588,6 +588,34 @@ def get_test_output_dir(script_prefix_lines, test_case_name):
     return os.path.join(output_dir, test_case_name) if test_case_name else output_dir
 
 
+def extract_pytest_command_env(script_prefix_lines, name):
+    """Read a leading environment assignment from the exported pytest command."""
+    line = next((ln for ln in script_prefix_lines if "export pytestCommand=" in ln), None)
+    if line is None:
+        raise ValueError("launch prefix does not export pytestCommand")
+    try:
+        outer_tokens = shlex.split(line)
+    except ValueError as e:
+        raise ValueError(f"cannot parse exported pytestCommand: {e}") from e
+    command_assignment = next(
+        (token for token in outer_tokens if token.startswith("pytestCommand=")), None
+    )
+    if command_assignment is None:
+        raise ValueError("launch prefix has a malformed pytestCommand export")
+    command = command_assignment.partition("=")[2]
+    try:
+        command_tokens = shlex.split(command)
+    except ValueError as e:
+        raise ValueError(f"cannot parse pytestCommand payload: {e}") from e
+    for token in command_tokens:
+        if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*=", token):
+            break
+        key, value = token.split("=", 1)
+        if key == name:
+            return value
+    raise ValueError(f"pytestCommand does not set leading environment variable {name}")
+
+
 def remove_whitespace_lines(lines):
     return [line.strip() for line in lines if line.strip()]
 
@@ -791,6 +819,7 @@ def main():
         # Enable/kill-switch policy and timeouts live in precheck_config
         # (single owner, shared with the local flow).
         pcfg = _import_precheck_config(args.llm_src)
+        llm_models_root = extract_pytest_command_env(script_prefix_lines, "LLM_MODELS_ROOT")
         script_prefix_lines.extend(
             pcfg.precheck_prefix_lines(
                 config,
@@ -801,12 +830,14 @@ def main():
                     hardware_config["gpus_per_ctx_server"],
                     hardware_config["gpus_per_gen_server"],
                 ),
+                llm_models_root=llm_models_root,
                 stage_name=args.stage_name,
             )
         )
         srun_args_lines.extend(
             [
                 "--container-env=DISAGG_SERVING_TYPE",
+                "--container-env=LLM_MODELS_ROOT",
                 "--container-env=pytestCommand",
             ]
         )
diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py
index 6f6bb151c318..ddd7729f7182 100644
--- a/tensorrt_llm/_torch/disaggregation/native/transfer.py
+++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py
@@ -79,6 +79,10 @@
 # Number of worker threads for KV transfer queues (default: 1)
 KV_TRANSFER_NUM_THREADS = int(os.environ.get("TRTLLM_KV_TRANSFER_NUM_THREADS", "1"))
 
+# Keep standalone TxSession waits responsive to cancellation even when callers
+# do not configure a sender-future wait slice.
+_FALLBACK_TX_WAIT_SLICE_S = 1.0
+
 
 @dataclass
 class RecvReqInfo:
@@ -1324,11 +1328,13 @@ def has_transferring_tasks(self) -> bool:
     def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]:
         """Poll or block until KV (and optionally aux) transfer finishes.
 
-        With blocking=True (default): waits up to _timeout_s for each task.
+        With blocking=True (default): retries bounded wait slices until a
+        successful transfer finishes. Errors and cancellation remain terminal;
+        callers must not interpret them as proof that peer writes quiesced.
         With blocking=False: polls non-blockingly; returns None if any KV task
         or aux is not yet done.
         """
-        if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED):
+        if self.has_failed():
             return WaitResult.FAILED
         if not self.kv_tasks:
             return None
@@ -1350,17 +1356,42 @@ def wait_complete(self, blocking: bool = True) -> Optional[WaitResult]:
                     return None
             return WaitResult.COMPLETED
 
+        # ``_timeout_s`` bounds one scheduler wait slice; it is not a transfer
+        # deadline. A successful blockAll must not return merely because one
+        # slice expired: NIXL may still be reading the request's KV pages.
+        wait_slice_s = self._timeout_s
+        if wait_slice_s is None or wait_slice_s <= 0:
+            wait_slice_s = _FALLBACK_TX_WAIT_SLICE_S
+        # The loop preserves standalone block-until-terminal behavior while a
+        # bounded slice keeps cancellation and sibling failure observable.
         for task in self.kv_tasks:
-            if not task.wait(timeout=self._timeout_s):
-                return WaitResult.TIMEOUT
+            while not task.wait(timeout=wait_slice_s):
+                # cancel() leaves a TRANSFERRING task's event unset until the
+                # physical writer finishes. Preserve the bounded-slice control
+                # point so blockAll can still report terminal cancellation.
+                # Check every task: a sibling slice can fail while this one is
+                # still pending without setting the session terminal status.
+                if self.has_failed():
+                    return WaitResult.FAILED
             if task.status == TaskStatus.ERROR:
                 return WaitResult.FAILED
-        if self._need_aux and self.aux_task is not None:
-            if not self.aux_task.wait(timeout=self._timeout_s):
-                return WaitResult.TIMEOUT
+        if self._need_aux:
+            if self.aux_task is None:
+                return (
+                    WaitResult.FAILED
+                    if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED)
+                    else None
+                )
+            while not self.aux_task.wait(timeout=wait_slice_s):
+                if self.has_failed():
+                    return WaitResult.FAILED
             if self.aux_task.status == TaskStatus.ERROR:
                 return WaitResult.FAILED
-        return WaitResult.COMPLETED
+        return (
+            WaitResult.FAILED
+            if self.status in (SessionStatus.ERROR, SessionStatus.CANCELLED)
+            else WaitResult.COMPLETED
+        )
 
     def set_exception(self, reason: str = ""):
         msg = f"TxSession {self.disagg_request_id} exception"
diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py
index 6968159a90d7..8a6314d02206 100644
--- a/tensorrt_llm/_torch/disaggregation/transceiver.py
+++ b/tensorrt_llm/_torch/disaggregation/transceiver.py
@@ -703,8 +703,9 @@ def check_context_transfer_status(
             elif result is None:
                 continue
             elif result == WaitResult.TIMEOUT:
-                logger.warning(
-                    f"TxSession rid={session.disagg_request_id} timed out after {self._sender_future_timeout_ms}ms"
+                logger.debug(
+                    f"TxSession rid={session.disagg_request_id} not ready after "
+                    f"{self._sender_future_timeout_ms}ms wait slice; keeping it in progress"
                 )
                 timed_out.append(rid)
             else:
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
index 2588c1362157..48bd794290dd 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md
@@ -16,7 +16,7 @@ starts.
 | Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. |
 | Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. |
 | Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). |
-| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves against the model class's `get_model_defaults()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. |
+| Same KV cache manager version + transceiver runtime | The launch generator forwards the real test's `LLM_MODELS_ROOT` into every precheck process. Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and requires a registered model class, then resolves against its optional `get_model_defaults()` hook (no hook means empty defaults, matching serving). `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. An unresolved manager-version `auto` setting fails with INIT_ERROR instead of silently assuming V1. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination also fails fast. |
 
 Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are
 supported: data is seeded per (request, **global** layer) and constant along
@@ -77,13 +77,13 @@ csv/ctx_<i>/<uuid>_<rank>.csv       # Python transceiver per-task perf
 
 ```bash
 # Inspect what a yaml resolves to (no GPU needed):
-python3 run_precheck.py --role gen --server-idx 0 --dry-run \
+LLM_MODELS_ROOT='<models>' python3 run_precheck.py --role gen --server-idx 0 --dry-run \
     --config ../disaggregated/<test>.yaml --work-dir /tmp/ct --llm-src <repo>
 
 # On a SLURM allocation: one srun per instance, e.g. ctx dep4 + gen dep8:
-srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py --role ctx --server-idx 0 \
-    --config <yaml> --work-dir <shared-dir> --llm-src <repo> &
-srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \
+LLM_MODELS_ROOT='<models>' srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py \
+    --role ctx --server-idx 0 --config <yaml> --work-dir <shared-dir> --llm-src <repo> &
+LLM_MODELS_ROOT='<models>' srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \
     --config <yaml> --work-dir <shared-dir> --llm-src <repo> &
 wait
 ```
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
index 916f2c4678a3..adc28b45e245 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py
@@ -26,6 +26,7 @@
 
 import json
 import os
+import shlex
 
 # Optional per-yaml overrides live under a `cache_transceiver_precheck:` block.
 PRECHECK_DEFAULTS = {
@@ -61,8 +62,9 @@
     "verify_data": True,
 }
 
-# Fallback KV shape when the model directory cannot be resolved: the precheck
-# still exercises the exact network path, just with a synthetic cache shape.
+# Fallback KV shape for dry-runs and explicitly selected manager versions when
+# the model directory cannot be resolved. Manager-version "auto" resolution
+# fails fast instead of silently pairing this shape with V1.
 FALLBACK_KV_SHAPE = {
     "num_layers": 32,
     "num_kv_heads": 8,
@@ -128,7 +130,13 @@ def default_step_timeout_s(max_world):
 
 
 def precheck_prefix_lines(
-    cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, max_world, stage_name=""
+    cfg,
+    benchmark_mode,
+    config_path_expr,
+    ucx_tls_cmd,
+    max_world,
+    llm_models_root,
+    stage_name="",
 ):
     """Launch-script export lines wiring the precheck gate.
 
@@ -164,6 +172,9 @@ def precheck_prefix_lines(
         f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode"
     )
     lines = [
+        # Keep this as a top-level assignment. shlex.quote() is not safe when
+        # nested inside the double-quoted pytestCommand exports below.
+        f"export LLM_MODELS_ROOT={shlex.quote(llm_models_root)}",
         f"export ctPrecheckEnabled={int(enabled)}",
         # The external srun timeout must cover the driver's first-rep NIXL
         # wire-up allowance; the default derives from the same formula the
diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
index 19c6bf70c856..5f6753f20bea 100644
--- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
+++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
@@ -272,28 +272,33 @@ def _pattern_like(shape, dtype, device, seed):
     return rnd.to(dtype).to(device).expand(nb, kv, heads, tok, dim)
 
 
-def _request_block_views(kvm, rid):
-    """Yield (global_layer, buffer, valid_block_indices) for this rank."""
+def _request_block_views(kvm, rid, prompt_len):
+    """Yield the prompt blocks transferred for this request on this rank."""
+    num_prompt_blocks = (prompt_len + kvm.tokens_per_block - 1) // kvm.tokens_per_block
     for global_layer in kvm.pp_layers:
         blocks = kvm.get_batch_cache_indices([rid], layer_idx=global_layer)[0]
-        valid = [b for b in blocks if b >= 0]
+        # V2 may reserve extra KV tokens for speculative decoding. At an exact
+        # block boundary those tokens allocate an additional page, but the
+        # transceiver intentionally trims its slice to prompt_len blocks.
+        # Verify the same payload range instead of the untransferred page.
+        valid = [b for b in blocks if b >= 0][:num_prompt_blocks]
         if not valid:
             continue
         buf = kvm.get_buffers(global_layer, kv_layout="HND")
         yield global_layer, buf, valid
 
 
-def fill_request(kvm, rid):
-    for global_layer, buf, valid in _request_block_views(kvm, rid):
+def fill_request(kvm, rid, prompt_len):
+    for global_layer, buf, valid in _request_block_views(kvm, rid, prompt_len):
         shape = (len(valid), *buf.shape[1:])
         buf[valid] = _pattern_like(shape, buf.dtype, buf.device, seed_for(rid, global_layer))
 
 
-def verify_request(kvm, rid):
+def verify_request(kvm, rid, prompt_len):
     """Returns (ok, detail) comparing received blocks to the expected pattern."""
     import torch
 
-    for global_layer, buf, valid in _request_block_views(kvm, rid):
+    for global_layer, buf, valid in _request_block_views(kvm, rid, prompt_len):
         recv = buf[valid]
         exp = _pattern_like(recv.shape, recv.dtype, recv.device, seed_for(rid, global_layer))
         recv_f, exp_f = recv.float(), exp.float()  # fp8 lacks direct compare ops
@@ -320,8 +325,8 @@ def _lookup_model_cls(model_dir):
 def resolve_model_prefs(model_dir, side, cache_cfg):
     """Mirror serving's model-preference resolution (PR #15823 semantics).
 
-    - use_kv_cache_manager_v2 == "auto" (yaml absent): adopt the model
-      class's get_model_defaults() value, default False
+    - use_kv_cache_manager_v2 == "auto" (yaml absent): require the model
+      class and adopt its get_model_defaults() value when the hook exists
       (llm_utils._resolve_kv_cache_manager_v2_auto).
     - cache_cfg.transceiver_runtime == "auto": adopt
       model_cls.get_preferred_transceiver_runtime(), NIXL-gated, via the
@@ -333,6 +338,12 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
 
     api = load_internal_apis()
     model_cls, hf_view = _lookup_model_cls(model_dir)
+    setting = side["use_kv_cache_manager_v2"]
+    if setting == "auto" and model_cls is None:
+        raise RuntimeError(
+            "use_kv_cache_manager_v2 is 'auto', but the precheck could not resolve "
+            f"a registered model class from model_dir={model_dir!r}; refusing to assume V1"
+        )
 
     # Runtime BEFORE V2, like serving: the V2 resolver's disagg gating reads
     # cache_cfg.transceiver_runtime and treats an unresolved "auto" as non-PYTHON.
@@ -347,17 +358,15 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
                 flush=True,
             )
 
-    setting = side["use_kv_cache_manager_v2"]
     if setting == "auto":
         defaults = {}
-        if model_cls is not None:
+        if hasattr(model_cls, "get_model_defaults"):
             try:
                 defaults = model_cls.get_model_defaults(None) or {}
-            except Exception as e:  # noqa: BLE001 - model hooks may need llm_args
-                print(
-                    f"[precheck] WARNING: get_model_defaults failed ({e!r}); assuming V1",
-                    flush=True,
-                )
+            except Exception as e:  # noqa: BLE001 - model hooks are third-party extension points
+                raise RuntimeError(
+                    f"get_model_defaults failed for {model_cls.__name__}; refusing to assume V1"
+                ) from e
         try:
             # The REAL serving resolver, via the same shim pattern as the
             # runtime resolution below -- one owner for the 'auto' semantics.
@@ -368,11 +377,8 @@ def resolve_model_prefs(model_dir, side, cache_cfg):
                 cache_transceiver_config=cache_cfg,
             )
             use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, defaults))
-        except Exception as e:  # noqa: BLE001 - fall back like a missing model
-            print(
-                f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); assuming V1", flush=True
-            )
-            use_v2 = False
+        except Exception as e:  # noqa: BLE001 - resolver spans model extension hooks
+            raise RuntimeError("V2 'auto' resolution failed; refusing to assume V1") from e
     else:
         use_v2 = bool(setting)
     return use_v2
@@ -508,7 +514,8 @@ def free_sequence(kvm, req, use_v2):
 def _wait_gen_complete(xcvr, req, runtime, llm_request_state):
     """Block until this gen request's receive finishes (or errors).
 
-    PYTHON transceiver: check_gen_transfer_status(None) blocks for all. C++:
+    The Python transceiver is handled once per wave in gen_run_wave(), where
+    its returned request IDs can be checked before releasing KV pages. For C++,
     the int API can return before THIS request completes on a cold link, so
     poll for a terminal state (bounded by signal.alarm + hang detector).
     Logs periodic progress so a stalled transfer shows WHICH request is stuck
@@ -516,8 +523,7 @@ def _wait_gen_complete(xcvr, req, runtime, llm_request_state):
     "RDMA write never completed").
     """
     if runtime == "PYTHON":
-        xcvr.check_gen_transfer_status(None)
-        return
+        raise ValueError("Python generation waves must be checked as a batch")
     terminal = (
         llm_request_state.DISAGG_GENERATION_TRANS_COMPLETE,
         llm_request_state.DISAGG_TRANS_ERROR,
@@ -924,12 +930,15 @@ def ctx_run_wave(self, peer_idx, li, req_len, rep, wave):
                 rid = self._pair_rid(peer_idx, li, rep, pair)
                 req = make_request(True, rid, req_len, self.runtime)
                 add_sequence(self.kvm, req, req_len, self.use_v2)
-                fill_request(self.kvm, rid)
+                # Track ownership as soon as allocation succeeds. A later
+                # setup failure must retain the pages rather than free storage
+                # that an asynchronously dispatched sender may still read.
+                reqs[pair] = req
+                fill_request(self.kvm, rid, req_len)
                 tensorrt_llm.logger.info(
                     f"[ctx{self.server_idx} r{self.rank}] rid={rid} len={req_len}: send START"
                 )
                 self.xcvr.respond_and_send_async(req)
-                reqs[pair] = req
         except Exception as e:  # noqa: BLE001 - relayed to gen, then raised
             local_err = e
         reason = self._consensus_error(local_err)
@@ -963,7 +972,9 @@ def ctx_run_wave(self, peer_idx, li, req_len, rep, wave):
         reason = self.comm.bcast(reason, root=0)
 
         if reason is not None:
-            self._free_all(reqs)
+            # Send setup can fail asymmetrically after another rank dispatched
+            # work. A block-all collective is not safe from that state, and
+            # reusing the pages is worse than retaining them until teardown.
             raise _TransferError(f"ctx send setup failed: {reason}")
         return params_by_pair, reqs
 
@@ -974,16 +985,33 @@ def ctx_finish_wave(self, reqs):
         t0 = time.monotonic()
         local_err = None
         try:
-            self.xcvr.check_context_transfer_status(None)  # block-all
+            completed, failed = self.xcvr.check_context_transfer_status(None)  # block-all
+            completed_rids = set(completed)
+            failed_rids = set(failed)
+            missing = [
+                p
+                for p, req in reqs.items()
+                if req.py_request_id not in completed_rids | failed_rids
+            ]
+            if missing:
+                raise _TransferError(
+                    f"block-all returned before terminal status for pairs {missing}"
+                )
+            failed_pairs = [p for p, req in reqs.items() if req.py_request_id in failed_rids]
+            if failed_pairs:
+                raise _TransferError(f"ctx transfer failed for pairs {failed_pairs}")
             bad = [
                 p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR
             ]
             if bad:
-                local_err = _TransferError(f"ctx DISAGG_TRANS_ERROR on pairs {bad}")
+                raise _TransferError(f"ctx DISAGG_TRANS_ERROR on pairs {bad}")
+            # Only successful IDs prove that every sender finished reading its
+            # source pages. Failed/cancelled/missing waves retain ownership
+            # until process teardown because their task events need not prove
+            # physical NIXL quiescence.
+            self._free_all(reqs)
         except Exception as e:  # noqa: BLE001
             local_err = e
-        finally:
-            self._free_all(reqs)
         states = {p: str(r.state) for p, r in reqs.items()}
         tensorrt_llm.logger.info(
             f"[ctx{self.server_idx} r{self.rank}] wave sends finished in "
@@ -1012,23 +1040,40 @@ def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair):
                     False, rid, req_len, self.runtime, ctx_params=params_by_pair[pair]
                 )
                 add_sequence(self.kvm, req, req_len, self.use_v2)
+                # Track every allocated sequence before receive dispatch. On
+                # setup failure, retain all pages until process teardown.
+                reqs[pair] = req
                 tensorrt_llm.logger.info(
                     f"[gen{self.server_idx} r{self.rank}] rid={rid} len={req_len}: recv START"
                 )
                 self.xcvr.request_and_receive_async(req)
-                reqs[pair] = req
         except Exception as e:  # noqa: BLE001
             local_err = e
         reason = self._consensus_error(local_err)
         if reason is not None:
-            self._free_all(reqs)
             raise _TransferError(f"gen receive setup failed: {reason}")
 
         mismatch = ""
         t0 = time.monotonic()
+        safe_to_free = False
         try:
-            for pair, req in reqs.items():
-                _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state)
+            if self.runtime == "PYTHON":
+                completed, failed, cancelled = self.xcvr.check_gen_transfer_status(None)
+                completed_rids = set(completed)
+                failed_rids = set(failed)
+                cancelled_rids = {req.py_request_id for req in cancelled}
+                expected_rids = {req.py_request_id for req in reqs.values()}
+                missing_rids = expected_rids - completed_rids - failed_rids - cancelled_rids
+                if failed_rids or cancelled_rids or missing_rids:
+                    raise _TransferError(
+                        "Python gen block-all did not complete every request: "
+                        f"failed={sorted(failed_rids)} "
+                        f"cancelled={sorted(cancelled_rids)} "
+                        f"missing={sorted(missing_rids)}"
+                    )
+            else:
+                for req in reqs.values():
+                    _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state)
             if reqs:
                 tensorrt_llm.logger.info(
                     f"[gen{self.server_idx} r{self.rank}] wave recvs finished in "
@@ -1039,16 +1084,26 @@ def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair):
                 p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR
             ]
             if bad:
-                local_err = _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}")
-            elif self.plan["verify_data"] and rep >= self.plan["warmup_requests"]:
+                raise _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}")
+            incomplete = [
+                p
+                for p, r in reqs.items()
+                if r.state != self.llm_request_state.DISAGG_GENERATION_TRANS_COMPLETE
+            ]
+            if incomplete:
+                raise _TransferError(f"gen requests not complete for pairs {incomplete}")
+            # Remote writes and local CUDA work are now complete. Later byte
+            # verification failures do not invalidate the ownership proof.
+            safe_to_free = True
+            if self.plan["verify_data"] and rep >= self.plan["warmup_requests"]:
                 for pair, req in reqs.items():
-                    ok, detail = verify_request(self.kvm, req.py_request_id)
+                    ok, detail = verify_request(self.kvm, req.py_request_id, req_len)
                     if not ok:
                         mismatch = f"pair={pair} {detail}"
                         break
         except Exception as e:  # noqa: BLE001
             local_err = e
-        finally:
+        if safe_to_free:
             self._free_all(reqs)
         reason = self._consensus_error(local_err)
         if reason is not None:
diff --git a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py
index e8a92f99c99c..9b106a1d8d43 100644
--- a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py
+++ b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py
@@ -118,12 +118,12 @@ def _terminate_process_groups(processes):
                     pass
 
 
-def _disagg_yaml(num_ctx, num_gen, ctx_tp, gen_tp, request_lengths=(64,)):
+def _disagg_yaml(num_ctx, num_gen, ctx_tp, gen_tp, request_lengths=(64,), mtp_draft_len=0):
     """Minimal disagg perf-sanity yaml shaped like the checked-in configs."""
     tokens_per_block = 32
 
     def side(tp):
-        return {
+        config = {
             "tensor_parallel_size": tp,
             "pipeline_parallel_size": 1,
             "kv_cache_config": {
@@ -139,6 +139,12 @@ def side(tp):
                 "max_tokens_in_buffer": 512,
             },
         }
+        if mtp_draft_len:
+            config["speculative_config"] = {
+                "decoding_type": "MTP",
+                "max_draft_len": mtp_draft_len,
+            }
+        return config
 
     return {
         "metadata": {"model_dir_name": "tiny-llama"},
@@ -311,6 +317,17 @@ def test_precheck_passes(tmp_path, num_ctx, num_gen, ctx_tp, gen_tp):
         assert peers == {f"ctx_{ci}" for ci in range(num_ctx)}
 
 
+@pytest.mark.timeout(300)
+def test_precheck_passes_mtp_exact_block_boundary(tmp_path):
+    """Reserved MTP tokens must not expand the verified transfer payload."""
+    pytest.importorskip("mpi4py")
+    cfg = _disagg_yaml(1, 1, 1, 1, request_lengths=(64,), mtp_draft_len=3)
+    config_path, models_root = _write_inputs(tmp_path, cfg)
+    work_dir, launched = _launch_instances(tmp_path, _jobs(cfg, config_path), models_root)
+    _wait_all(launched)
+    _assert_all_passed(work_dir, launched)
+
+
 @pytest.mark.timeout(300)
 def test_precheck_fails_fast_on_fingerprint_mismatch(tmp_path):
     """Mismatched ctx/gen yamls must produce FAIL verdicts, not a hang.
diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py
index 76eff6adc439..ce9cad43ebeb 100644
--- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py
+++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py
@@ -82,14 +82,17 @@ def close(self) -> None:
 
 
 class _FakeTask:
-    def __init__(self, status: TaskStatus, wait_result: bool = True) -> None:
+    def __init__(self, status: TaskStatus, wait_result: bool | list[bool] = True) -> None:
         self.status = status
-        self._wait_result = wait_result
+        self._wait_results = list(wait_result) if isinstance(wait_result, list) else [wait_result]
         self.wait_calls: list[Optional[float]] = []
 
     def wait(self, timeout: Optional[float] = None) -> bool:
         self.wait_calls.append(timeout)
-        return self._wait_result
+        result = self._wait_results.pop(0) if len(self._wait_results) > 1 else self._wait_results[0]
+        if result and self.status != TaskStatus.ERROR:
+            self.status = TaskStatus.TRANSFERRED
+        return result
 
 
 def _make_transceiver(
@@ -122,9 +125,10 @@ def _make_tx_session(
     *,
     need_aux: bool = False,
     aux_task: Optional[_FakeTask] = None,
+    timeout_s: Optional[float] = 0.25,
 ) -> TxSession:
     session = object.__new__(TxSession)
-    session._timeout_s = 0.25
+    session._timeout_s = timeout_s
     session._need_aux = need_aux
     session._terminal_status = None
     session.receiver_ready = True
@@ -356,11 +360,127 @@ def test_ctx_consensus_fastpath_skips_when_idle(monkeypatch) -> None:
     transceiver._ctx_consensus.assert_called_once()
 
 
-def test_tx_session_wait_complete_defaults_to_blocking() -> None:
-    task = _FakeTask(TaskStatus.INIT, wait_result=False)
+def test_tx_session_blocking_wait_retries_wait_slices_until_complete() -> None:
+    task = _FakeTask(TaskStatus.INIT, wait_result=[False, True])
     session = _make_tx_session([task])
 
-    assert session.wait_complete() == WaitResult.TIMEOUT
+    assert session.wait_complete() == WaitResult.COMPLETED
+    assert task.wait_calls == [0.25, 0.25]
+
+
+def test_context_transfer_status_block_all_drains_wait_slices_before_close() -> None:
+    task = _FakeTask(TaskStatus.INIT, wait_result=[False, True])
+    session = _make_tx_session([task])
+    transceiver = _make_transceiver({15: session}, {15: _FakeRequest()})
+
+    completed, failed = transceiver.check_context_transfer_status(None)
+
+    assert completed == [15]
+    assert failed == []
+    assert task.wait_calls == [0.25, 0.25]
+    assert session._closed
+    assert 15 not in transceiver._send_sessions
+
+
+def test_tx_session_blocking_wait_treats_cancelled_session_as_terminal() -> None:
+    task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
+    session = _make_tx_session([task])
+    session._terminal_status = SessionStatus.CANCELLED
+
+    assert session.wait_complete(blocking=True) == WaitResult.FAILED
+    assert task.wait_calls == []
+
+
+def test_tx_session_blocking_wait_observes_cancellation_between_slices() -> None:
+    task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
+    session = _make_tx_session([task])
+    wait = task.wait
+
+    def cancel_during_wait(timeout: Optional[float] = None) -> bool:
+        result = wait(timeout)
+        session._terminal_status = SessionStatus.CANCELLED
+        return result
+
+    task.wait = cancel_during_wait
+
+    assert session.wait_complete(blocking=True) == WaitResult.FAILED
+    assert task.wait_calls == [0.25]
+
+
+@pytest.mark.parametrize("timeout_s", [None, 0.0, -1.0])
+def test_tx_session_blocking_wait_uses_fallback_without_positive_timeout(
+    timeout_s: Optional[float],
+) -> None:
+    task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
+    session = _make_tx_session([task], timeout_s=timeout_s)
+    wait = task.wait
+
+    def cancel_during_wait(timeout: Optional[float] = None) -> bool:
+        result = wait(timeout)
+        session._terminal_status = SessionStatus.CANCELLED
+        return result
+
+    task.wait = cancel_during_wait
+
+    assert session.wait_complete(blocking=True) == WaitResult.FAILED
+    assert task.wait_calls == [1.0]
+
+
+def test_tx_session_blocking_wait_treats_task_failure_as_terminal() -> None:
+    failed_task = _FakeTask(TaskStatus.ERROR)
+    pending_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=[False, True])
+    session = _make_tx_session([failed_task, pending_task])
+
+    assert session.wait_complete(blocking=True) == WaitResult.FAILED
+    assert failed_task.wait_calls == []
+    # A failed task event does not prove sibling physical writers quiesced, so
+    # precheck callers retain the wave instead of treating failure as drained.
+    assert pending_task.wait_calls == []
+
+
+def test_tx_session_blocking_wait_detects_failed_sibling_behind_pending_task() -> None:
+    pending_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
+    failed_task = _FakeTask(TaskStatus.ERROR)
+    session = _make_tx_session([pending_task, failed_task])
+
+    assert session.wait_complete(blocking=True) == WaitResult.FAILED
+    assert pending_task.wait_calls == []
+    assert failed_task.wait_calls == []
+
+
+def test_tx_session_blocking_wait_retries_aux_wait_slices() -> None:
+    kv_task = _FakeTask(TaskStatus.TRANSFERRED)
+    aux_task = _FakeTask(TaskStatus.INIT, wait_result=[False, True])
+    session = _make_tx_session([kv_task], need_aux=True, aux_task=aux_task)
+
+    assert session.wait_complete(blocking=True) == WaitResult.COMPLETED
+    assert kv_task.wait_calls == [0.25]
+    assert aux_task.wait_calls == [0.25, 0.25]
+
+
+def test_tx_session_blocking_aux_wait_observes_cancellation_between_slices() -> None:
+    kv_task = _FakeTask(TaskStatus.TRANSFERRED)
+    aux_task = _FakeTask(TaskStatus.TRANSFERRING, wait_result=False)
+    session = _make_tx_session([kv_task], need_aux=True, aux_task=aux_task)
+    wait = aux_task.wait
+
+    def cancel_during_wait(timeout: Optional[float] = None) -> bool:
+        result = wait(timeout)
+        session._terminal_status = SessionStatus.CANCELLED
+        return result
+
+    aux_task.wait = cancel_during_wait
+
+    assert session.wait_complete(blocking=True) == WaitResult.FAILED
+    assert kv_task.wait_calls == [0.25]
+    assert aux_task.wait_calls == [0.25]
+
+
+def test_tx_session_blocking_wait_keeps_missing_aux_pending() -> None:
+    task = _FakeTask(TaskStatus.TRANSFERRED)
+    session = _make_tx_session([task], need_aux=True)
+
+    assert session.wait_complete(blocking=True) is None
     assert task.wait_calls == [0.25]
 
 
diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py
index 0bcff41adb48..6d7ab301c71c 100644
--- a/tests/unittest/others/test_cache_transceiver_precheck_config.py
+++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py
@@ -19,7 +19,9 @@
 
 import json
 import os
+import subprocess
 import sys
+import types
 
 import pytest
 
@@ -298,6 +300,55 @@ def test_use_kv_cache_manager_v2_flags():
     assert pcfg.side_plan(plan, "gen")["use_kv_cache_manager_v2"] is True
 
 
+def test_resolve_model_prefs_auto_requires_registered_model(monkeypatch):
+    monkeypatch.setattr(rp, "load_internal_apis", lambda: types.SimpleNamespace())
+    monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (None, None))
+    cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON")
+
+    with pytest.raises(RuntimeError, match="refusing to assume V1"):
+        rp.resolve_model_prefs(None, {"use_kv_cache_manager_v2": "auto"}, cache_cfg)
+
+
+def test_resolve_model_prefs_auto_propagates_model_default_failure(monkeypatch):
+    class FailingModel:
+        @classmethod
+        def get_model_defaults(cls, _llm_args):
+            raise RuntimeError("model hook failed")
+
+    monkeypatch.setattr(rp, "load_internal_apis", lambda: types.SimpleNamespace())
+    monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (FailingModel, object()))
+    cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON")
+
+    with pytest.raises(RuntimeError, match="get_model_defaults failed.*refusing to assume V1"):
+        rp.resolve_model_prefs("/model", {"use_kv_cache_manager_v2": "auto"}, cache_cfg)
+
+
+def test_resolve_model_prefs_auto_propagates_resolver_failure(monkeypatch):
+    class Model:
+        @classmethod
+        def get_model_defaults(cls, _llm_args):
+            return {"use_kv_cache_manager_v2": True}
+
+    def fail_resolver(_shim, _defaults):
+        raise RuntimeError("resolver failed")
+
+    api = types.SimpleNamespace(resolve_kv_cache_manager_v2_auto=fail_resolver)
+    monkeypatch.setattr(rp, "load_internal_apis", lambda: api)
+    monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (Model, object()))
+    cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON")
+
+    with pytest.raises(RuntimeError, match="V2 'auto' resolution failed.*refusing to assume V1"):
+        rp.resolve_model_prefs("/model", {"use_kv_cache_manager_v2": "auto"}, cache_cfg)
+
+
+def test_resolve_model_prefs_explicit_v1_does_not_require_model(monkeypatch):
+    monkeypatch.setattr(rp, "load_internal_apis", lambda: types.SimpleNamespace())
+    monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (None, None))
+    cache_cfg = types.SimpleNamespace(transceiver_runtime="PYTHON")
+
+    assert not rp.resolve_model_prefs(None, {"use_kv_cache_manager_v2": False}, cache_cfg)
+
+
 def test_model_kv_shape_vocab_size(tmp_path):
     model_dir = tmp_path / "m"
     model_dir.mkdir()
@@ -360,10 +411,50 @@ def test_wireup_timeout_derivation():
 
 
 def _enabled_line(cfg):
-    lines = pcfg.precheck_prefix_lines(cfg, "e2e", "$c", "unset &&", max_world=8)
+    lines = pcfg.precheck_prefix_lines(
+        cfg,
+        "e2e",
+        "$c",
+        "unset &&",
+        max_world=8,
+        llm_models_root="/models",
+    )
     return next(x for x in lines if x.startswith("export ctPrecheckEnabled"))
 
 
+@pytest.mark.parametrize(
+    "model_root",
+    (
+        "/models with spaces",
+        "/models/it's",
+        "/models/$HOME/$(must-not-run)",
+    ),
+)
+def test_precheck_commands_export_model_root_safely(model_root):
+    lines = pcfg.precheck_prefix_lines(
+        _disagg_yaml(),
+        "e2e",
+        "$config",
+        "unset UCX_TLS &&",
+        max_world=8,
+        llm_models_root=model_root,
+    )
+
+    commands = [line for line in lines if "pytestCommand" in line]
+    assert len(commands) == 2
+    assert all("python3" in line for line in commands)
+
+    script = "\n".join(lines) + '\nprintf "%s" "$LLM_MODELS_ROOT"\n'
+    result = subprocess.run(
+        ["bash"],
+        input=script,
+        text=True,
+        capture_output=True,
+        check=True,
+    )
+    assert result.stdout == model_root
+
+
 def test_precheck_env_kill_switch_truthy(monkeypatch):
     """The TRTLLM_DISAGG_CT_PRECHECK kill switch parses the usual boolean spellings.
 
diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py
index ed060d1e9053..c42b7264acc7 100644
--- a/tests/unittest/others/test_cache_transceiver_precheck_run.py
+++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py
@@ -80,6 +80,36 @@ def test_seed_for_deterministic_and_distinct():
     assert all(0 <= s <= 0x7FFFFFFF for s in seeds)
 
 
+@pytest.mark.parametrize(("prompt_len", "expected_blocks"), ((1024, 8), (7408, 58)))
+def test_request_block_views_excludes_untransferred_speculative_page(prompt_len, expected_blocks):
+    """V2's reserved MTP tokens must not expand the verified transfer range."""
+    tokens_per_block = 128
+    num_allocated = (prompt_len + 2 + tokens_per_block - 1) // tokens_per_block
+    allocated = [-1] + list(range(num_allocated))
+    buffer = object()
+
+    def get_batch_cache_indices(request_ids, layer_idx):
+        assert request_ids == [7]
+        assert layer_idx == 4
+        return [allocated]
+
+    def get_buffers(global_layer, kv_layout):
+        assert global_layer == 4
+        assert kv_layout == "HND"
+        return buffer
+
+    kvm = types.SimpleNamespace(
+        tokens_per_block=tokens_per_block,
+        pp_layers=[4],
+        get_batch_cache_indices=get_batch_cache_indices,
+        get_buffers=get_buffers,
+    )
+
+    views = list(rp._request_block_views(kvm, rid=7, prompt_len=prompt_len))
+
+    assert views == [(4, buffer, list(range(expected_blocks)))]
+
+
 # --------------------------------------------------------------------------- #
 # HMAC control-channel wire format
 # --------------------------------------------------------------------------- #
@@ -327,6 +357,281 @@ def test_timeout_budgets():
     assert rp.wave_timeout_s(plan, 1, 0) == 180
 
 
+# --------------------------------------------------------------------------- #
+# Model preference resolution
+# --------------------------------------------------------------------------- #
+def test_resolve_model_prefs_allows_registered_class_without_defaults_hook(monkeypatch):
+    model_cls = type("ModelWithoutDefaultsHook", (), {})
+    cache_cfg = types.SimpleNamespace(transceiver_runtime="CPP")
+    calls = []
+
+    def resolve_v2(shim, defaults):
+        calls.append((shim, defaults))
+        return False
+
+    monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (model_cls, object()))
+    monkeypatch.setattr(
+        rp,
+        "load_internal_apis",
+        lambda: types.SimpleNamespace(resolve_kv_cache_manager_v2_auto=resolve_v2),
+    )
+
+    use_v2 = rp.resolve_model_prefs(
+        "/models/example", {"use_kv_cache_manager_v2": "auto"}, cache_cfg
+    )
+
+    assert use_v2 is False
+    assert len(calls) == 1
+    assert calls[0][1] == {}
+
+
+# --------------------------------------------------------------------------- #
+# Transfer ownership
+# --------------------------------------------------------------------------- #
+def _ctx_finish_runner(monkeypatch, check_status):
+    events = []
+    monkeypatch.setitem(
+        sys.modules,
+        "tensorrt_llm",
+        types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)),
+    )
+    runner = object.__new__(rp.PrecheckRunner)
+    runner.xcvr = types.SimpleNamespace(check_context_transfer_status=check_status)
+    runner.llm_request_state = types.SimpleNamespace(DISAGG_TRANS_ERROR="error")
+    runner.server_idx = 0
+    runner.rank = 0
+    runner._consensus_error = lambda err: None if err is None else repr(err)
+    runner._free_all = lambda reqs: events.append(("free", sorted(reqs)))
+    return runner, events
+
+
+def test_ctx_finish_wave_frees_only_after_block_all_returns_every_request(monkeypatch):
+    events = []
+
+    def check_status(at_least_request_num):
+        events.append(("block_all", at_least_request_num))
+        return [101, 102], []
+
+    runner, free_events = _ctx_finish_runner(monkeypatch, check_status)
+    reqs = {
+        0: types.SimpleNamespace(py_request_id=101, state="in_progress"),
+        1: types.SimpleNamespace(py_request_id=102, state="in_progress"),
+    }
+    runner._free_all = lambda owned: events.append(("free", sorted(owned)))
+
+    runner.ctx_finish_wave(reqs)
+
+    assert events == [("block_all", None), ("free", [0, 1])]
+    assert free_events == []
+
+
+def test_ctx_finish_wave_retains_pages_when_block_all_omits_request(monkeypatch):
+    runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], []))
+    reqs = {
+        0: types.SimpleNamespace(py_request_id=101, state="in_progress"),
+        1: types.SimpleNamespace(py_request_id=102, state="in_progress"),
+    }
+
+    with pytest.raises(rp._TransferError, match="block-all returned before terminal"):
+        runner.ctx_finish_wave(reqs)
+
+    assert events == []
+
+
+def test_ctx_finish_wave_retains_pages_when_block_all_raises(monkeypatch):
+    def check_status(_n):
+        raise RuntimeError("interrupted")
+
+    runner, events = _ctx_finish_runner(monkeypatch, check_status)
+    reqs = {0: types.SimpleNamespace(py_request_id=101, state="in_progress")}
+
+    with pytest.raises(rp._TransferError, match="interrupted"):
+        runner.ctx_finish_wave(reqs)
+
+    assert events == []
+
+
+def test_ctx_finish_wave_retains_pages_when_request_failed(monkeypatch):
+    runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [102]))
+    reqs = {
+        0: types.SimpleNamespace(py_request_id=101, state="in_progress"),
+        1: types.SimpleNamespace(py_request_id=102, state="error"),
+    }
+
+    with pytest.raises(rp._TransferError, match=r"ctx transfer failed for pairs \[1\]"):
+        runner.ctx_finish_wave(reqs)
+
+    assert events == []
+
+
+def test_ctx_run_wave_setup_error_retains_allocated_pages(monkeypatch):
+    monkeypatch.setitem(
+        sys.modules,
+        "tensorrt_llm",
+        types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)),
+    )
+    monkeypatch.setattr(
+        rp,
+        "make_request",
+        lambda _is_ctx, rid, _req_len, _runtime: types.SimpleNamespace(
+            py_request_id=rid, context_phase_params=None
+        ),
+    )
+    monkeypatch.setattr(rp, "add_sequence", lambda *_args: None)
+    monkeypatch.setattr(rp, "fill_request", lambda *_args: None)
+
+    calls = {"send": 0, "free": 0}
+
+    def respond(_req):
+        calls["send"] += 1
+        if calls["send"] == 2:
+            raise RuntimeError("injected setup failure")
+
+    runner = object.__new__(rp.PrecheckRunner)
+    runner.runtime = "PYTHON"
+    runner.kvm = object()
+    runner.use_v2 = True
+    runner.server_idx = 0
+    runner.rank = 0
+    runner.is_leader = True
+    runner.side = {"parallel": {"enable_attention_dp": False}}
+    runner.mapping = types.SimpleNamespace(pp_rank=0)
+    runner.xcvr = types.SimpleNamespace(respond_and_send_async=respond)
+    runner.comm = types.SimpleNamespace(
+        gather=lambda obj, root=0: [obj],
+        bcast=lambda obj, root=0: obj,
+    )
+    runner._owned = lambda _wave: [0, 1]
+    runner._pair_rid = lambda _peer, _li, _rep, pair: 101 + pair
+    runner._consensus_error = lambda err: None if err is None else repr(err)
+    runner._free_all = lambda _reqs: calls.__setitem__("free", calls["free"] + 1)
+
+    with pytest.raises(rp._TransferError, match="injected setup failure"):
+        runner.ctx_run_wave(0, 0, 64, 0, [0, 1])
+
+    assert calls == {"send": 2, "free": 0}
+
+
+def _gen_run_wave_runner(monkeypatch, outcome):
+    requests = {}
+    events = []
+    states = types.SimpleNamespace(
+        DISAGG_GENERATION_TRANS_COMPLETE="complete",
+        DISAGG_TRANS_ERROR="error",
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "torch",
+        types.SimpleNamespace(
+            cuda=types.SimpleNamespace(synchronize=lambda: events.append("cuda_sync"))
+        ),
+    )
+    monkeypatch.setitem(
+        sys.modules,
+        "tensorrt_llm",
+        types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)),
+    )
+
+    def make_request(_is_ctx, rid, _req_len, _runtime, ctx_params=None):
+        req = types.SimpleNamespace(py_request_id=rid, state="in_progress")
+        requests[rid] = req
+        return req
+
+    monkeypatch.setattr(rp, "make_request", make_request)
+    monkeypatch.setattr(rp, "add_sequence", lambda *_args: None)
+
+    def check_status(_at_least_request_num):
+        completed, failed, cancelled = outcome(requests)
+        for rid in completed:
+            requests[rid].state = states.DISAGG_GENERATION_TRANS_COMPLETE
+        for rid in failed:
+            requests[rid].state = states.DISAGG_TRANS_ERROR
+        return completed, failed, [requests[rid] for rid in cancelled]
+
+    runner = object.__new__(rp.PrecheckRunner)
+    runner.runtime = "PYTHON"
+    runner.kvm = object()
+    runner.use_v2 = True
+    runner.server_idx = 0
+    runner.rank = 0
+    runner.side = {"parallel": {"enable_attention_dp": False}}
+    runner.mapping = types.SimpleNamespace(pp_rank=0)
+    runner.llm_request_state = states
+    runner.plan = {"verify_data": False, "warmup_requests": 0}
+    runner.xcvr = types.SimpleNamespace(
+        request_and_receive_async=lambda _req: None,
+        check_gen_transfer_status=check_status,
+    )
+    runner.comm = types.SimpleNamespace(allgather=lambda value: [value])
+    runner._owned = lambda _wave: [0, 1]
+    runner._pair_rid = lambda _peer, _li, _rep, pair: 101 + pair
+    runner._consensus_error = lambda err: None if err is None else repr(err)
+    runner._free_all = lambda owned: events.append(("free", sorted(owned)))
+    return runner, events
+
+
+def test_gen_run_wave_frees_only_after_every_python_receive_completes(monkeypatch):
+    runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], []))
+
+    ok, detail = runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()})
+
+    assert ok and not detail
+    assert events == ["cuda_sync", ("free", [0, 1])]
+
+
+def test_gen_run_wave_checks_python_status_on_empty_owner_rank(monkeypatch):
+    calls = []
+
+    def outcome(requests):
+        calls.append(dict(requests))
+        return [], [], []
+
+    runner, events = _gen_run_wave_runner(monkeypatch, outcome)
+    runner._owned = lambda _wave: []
+
+    ok, detail = runner.gen_run_wave(0, 0, 64, 0, [0, 1], {})
+
+    assert ok and not detail
+    assert calls == [{}]
+    assert events == ["cuda_sync", ("free", [])]
+
+
+def test_gen_run_wave_setup_error_retains_allocated_pages(monkeypatch):
+    runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([], [], []))
+    calls = 0
+
+    def receive(_req):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            raise RuntimeError("injected setup failure")
+
+    runner.xcvr.request_and_receive_async = receive
+
+    with pytest.raises(rp._TransferError, match="injected setup failure"):
+        runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()})
+
+    assert calls == 2
+    assert events == []
+
+
+@pytest.mark.parametrize(
+    ("outcome", "message"),
+    (
+        (lambda _reqs: ([101], [102], []), r"failed=\[102\]"),
+        (lambda _reqs: ([101], [], [102]), r"cancelled=\[102\]"),
+        (lambda _reqs: ([101], [], []), r"missing=\[102\]"),
+    ),
+)
+def test_gen_run_wave_retains_pages_without_all_successes(monkeypatch, outcome, message):
+    runner, events = _gen_run_wave_runner(monkeypatch, outcome)
+
+    with pytest.raises(rp._TransferError, match=message):
+        runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()})
+
+    assert events == []
+
+
 # --------------------------------------------------------------------------- #
 # Internal-API contract (imports tensorrt_llm; no GPU work)
 # --------------------------------------------------------------------------- #
diff --git a/tests/unittest/scripts/test_perf_submit.py b/tests/unittest/scripts/test_perf_submit.py
index ed4fb8dedea7..ede64792e5ef 100644
--- a/tests/unittest/scripts/test_perf_submit.py
+++ b/tests/unittest/scripts/test_perf_submit.py
@@ -23,8 +23,9 @@
 from pytest_split.algorithms import LeastDurationAlgorithm
 
 REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
+CI_SUBMIT_PATH = REPO_ROOT / "jenkins" / "scripts" / "perf" / "submit.py"
 SUBMIT_PATHS = (
-    REPO_ROOT / "jenkins" / "scripts" / "perf" / "submit.py",
+    CI_SUBMIT_PATH,
     REPO_ROOT / "jenkins" / "scripts" / "perf" / "local" / "submit.py",
 )
 EXAMPLE_SUBMIT_PATH = REPO_ROOT / "examples" / "disaggregated" / "slurm" / "benchmark" / "submit.py"
@@ -53,6 +54,11 @@ def submit_module(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatc
     return _load_module(request.param, monkeypatch)
 
 
+@pytest.fixture
+def ci_submit_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType:
+    return _load_module(CI_SUBMIT_PATH, monkeypatch)
+
+
 @pytest.fixture
 def example_submit_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType:
     return _load_module(EXAMPLE_SUBMIT_PATH, monkeypatch)
@@ -72,11 +78,6 @@ def test_get_benchmark_config_accepts_positive_integer(submit_module: ModuleType
     assert benchmark_config["concurrency"] == int(concurrency)
 
 
-@pytest.fixture
-def ci_submit_module(monkeypatch: pytest.MonkeyPatch) -> ModuleType:
-    return _load_module(SUBMIT_PATHS[0], monkeypatch)
-
-
 @pytest.mark.parametrize(
     "concurrency",
     (True, 1.5, [], {}, "0", 0, "-1", -1, "1.5", "not-an-integer", None),
@@ -270,3 +271,33 @@ def test_ci_submit_rejects_missing_pytest_split_durations(
             script_prefix_lines,
             split_group=1,
         )
+
+
+@pytest.mark.parametrize(
+    ("assignment", "expected"),
+    (
+        ("LLM_MODELS_ROOT=/models", "/models"),
+        ("LLM_MODELS_ROOT='/models with spaces'", "/models with spaces"),
+        ("LLM_MODELS_ROOT=/models/cache=production", "/models/cache=production"),
+    ),
+)
+def test_extract_pytest_command_env(ci_submit_module: ModuleType, assignment: str, expected: str):
+    lines = [f'export pytestCommand="LLM_ROOT=/src {assignment} COLUMNS=300 pytest -vv"']
+
+    assert ci_submit_module.extract_pytest_command_env(lines, "LLM_MODELS_ROOT") == expected
+
+
+def test_extract_pytest_command_env_rejects_missing_leading_assignment(
+    ci_submit_module: ModuleType,
+):
+    lines = ['export pytestCommand="LLM_ROOT=/src pytest LLM_MODELS_ROOT=/too-late"']
+
+    with pytest.raises(ValueError, match="does not set leading environment variable"):
+        ci_submit_module.extract_pytest_command_env(lines, "LLM_MODELS_ROOT")
+
+
+def test_extract_pytest_command_env_rejects_malformed_export(ci_submit_module: ModuleType):
+    lines = ['export pytestCommand="LLM_ROOT=/src LLM_MODELS_ROOT=/models pytest']
+
+    with pytest.raises(ValueError, match="cannot parse exported pytestCommand"):
+        ci_submit_module.extract_pytest_command_env(lines, "LLM_MODELS_ROOT")

From 1714136e4dbfbe2c4570067da6572d268a058997 Mon Sep 17 00:00:00 2001
From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
Date: Fri, 31 Jul 2026 12:20:27 -0700
Subject: [PATCH 2/2] [https://nvbugs/6480621][test] Validate 60-second KV
 transfer timeout

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>
---
 ...k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml
index 36edcb022db4..ac86e99bd425 100644
--- a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml
+++ b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL.yaml
@@ -67,7 +67,7 @@ worker_config:
       load_balancer: tests/scripts/perf-sanity/disaggregated/deepseek-v4-pro-eplb/moe_load_balancer_gen_ep32_slots384.yaml
     cache_transceiver_config:
       max_tokens_in_buffer: 8192
-      kv_transfer_timeout_ms: 600000
+      kv_transfer_timeout_ms: 60000
       backend: NIXL
       transceiver_runtime: PYTHON
     disable_overlap_scheduler: false
@@ -98,7 +98,7 @@ worker_config:
       load_balancer: tests/scripts/perf-sanity/disaggregated/deepseek-v4-pro-eplb/moe_load_balancer_ctx_ep4_384.yaml
     cache_transceiver_config:
       max_tokens_in_buffer: 8192
-      kv_transfer_timeout_ms: 600000
+      kv_transfer_timeout_ms: 60000
       backend: NIXL
       transceiver_runtime: PYTHON
     disable_overlap_scheduler: true