[https://nvbugs/6541356][fix] align cache transceiver harness setup

examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py

@@ -55,7 +55,10 @@
 from tensorrt_llm._torch.distributed import Distributed
 from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
-from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver
+from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import (
+    create_kv_cache_transceiver,
+    maybe_enable_fabric_memory_for_python_transceiver,
+)
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, LlmRequestType
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
 from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, CacheTransceiverConfig
@@ -551,6 +554,21 @@ def reset_sock():
     zmq_sock = open_sock()
 
     cases = build_cases(cfg)
+    # The C++ fabric-memory env getter is cached on its first KV pool
+    # allocation. Enable the default before any matrix case builds a pool, even
+    # when a C++ transceiver case appears before Python+V1 in the matrix.
+    python_v1_case = next(
+        (case for case in cases if case["runtime"] == "PYTHON" and case["cache_manager"] == "V1"),
+        None,
+    )
+    if python_v1_case is not None:
+        maybe_enable_fabric_memory_for_python_transceiver(
+            CacheTransceiverConfig(
+                backend=python_v1_case["backend"],
+                transceiver_runtime="PYTHON",
+            ),
+            KVCacheManager,
+        )
     req_lens = cfg["test_matrix"]["request_lengths"]
     warmup = cfg["test_matrix"]["warmup_requests"]
     num_req = cfg["test_matrix"]["num_requests_per_length"]

jenkins/scripts/perf/cluster_env.py

@@ -44,15 +44,12 @@
         "rocep198s0:1,rocep199s0:1,rocep205s0:1,rocep206s0:1"
         " UCX_IB_GID_INDEX=auto UCX_IB_TRAFFIC_CLASS=52 UCX_IB_SL=0",
     ),
-    # oci-aga: avoid transports that fail on this VF fabric, disable DEVX to
-    # avoid UAR allocation failures, and pin the GPU-connected rail VFs.
+    # oci-aga: use TCP over IPv4 alongside the local CUDA/shared-memory
+    # transports.
     (
         "oci-aga*",
         "*",
-        "export UCX_TLS=^tcp,rc_gda,gga UCX_IB_MLX5_DEVX=n "
-        "UCX_NET_DEVICES="
-        "rdma_vf_rail0:1,rdma_vf_rail1:1,rdma_vf_rail2:1,rdma_vf_rail3:1 "
-        "UCX_IB_TRAFFIC_CLASS=96 TRTLLM_NIXL_NUM_THREADS=1",
+        "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp UCX_TCP_AF_PRIO=inet",
     ),
     # nsc-svg: UCX picks wrong RDMA devices; pin the usable mlx5 ports.
     (

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,
             )
         )
 

jenkins/scripts/perf/submit.py

@@ -569,6 +569,31 @@ def _is_output_file_part(part):
     return ("", worker_pytest_command, disagg_server_pytest_command, benchmark_pytest_command)
 
 
+def _get_pytest_command_env_var(script_prefix_lines, env_name):
+    """Extract an environment assignment from the inbound pytest command."""
+    pytest_command_line = next(
+        (ln for ln in script_prefix_lines if "export pytestCommand=" in ln), None
+    )
+    if pytest_command_line is None:
+        return None
+
+    command_value = pytest_command_line.split("=", 1)[1]
+    try:
+        parts = shlex.split(command_value)
+        # A fully quoted export value is parsed as one token on the first
+        # pass. Parse that token once more to recover its command words.
+        if len(parts) == 1 and any(char.isspace() for char in parts[0]):
+            parts = shlex.split(parts[0])
+    except ValueError as error:
+        raise ValueError(f"Invalid inbound pytestCommand: {error}") from error
+
+    prefix = f"{env_name}="
+    for part in parts:
+        if part.startswith(prefix):
+            return part[len(prefix) :]
+    return None
+
+
 def get_test_output_dir(script_prefix_lines, test_case_name):
     """Build the per-test output directory from the inbound pytestCommand.
 
@@ -748,6 +773,12 @@ def main():
         ucx_tls_server_cmd = ucx_tls_cmd
 
         pytest_common_vars = ""
+        llm_models_root = _get_pytest_command_env_var(script_prefix_lines, "LLM_MODELS_ROOT")
+        if not llm_models_root:
+            raise ValueError(
+                "LLM_MODELS_ROOT is missing from the inbound pytestCommand; "
+                "the cache-transceiver precheck cannot resolve model defaults"
+            )
         script_prefix_lines.extend(
             [
                 worker_pytest_command,
@@ -802,6 +833,7 @@ def main():
                     hardware_config["gpus_per_gen_server"],
                 ),
                 stage_name=args.stage_name,
+                llm_models_root=llm_models_root,
             )
         )
         srun_args_lines.extend(

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -54,7 +54,9 @@
 from .dwdp import DwdpManager
 from .guided_decoder import GuidedDecoder
 from .kv_cache_manager_v2 import KVCacheManagerV2
-from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver
+from .kv_cache_transceiver import (
+    AttentionTypeCpp, create_kv_cache_transceiver,
+    maybe_enable_fabric_memory_for_python_transceiver)
 from .llm_request import ExecutorResponse, LlmRequestState
 from .mamba_cache_manager import (BaseMambaCacheManager,
                                   CppMambaHybridCacheManager,
@@ -528,28 +530,8 @@ def __init__(
         self._maybe_enable_fabric_memory_for_python_transceiver()
 
     def _maybe_enable_fabric_memory_for_python_transceiver(self) -> None:
-        """Default TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 for the Python
-        transceiver on the C++ V1 KV cache manager.
-
-        The Python transceiver (KvCacheTransceiverV2) transfers KV blocks
-        directly out of the C++ pool, so the pool should be allocated with
-        fabric memory to enable MNNVL transfers. This must run before any
-        pool allocation because the C++ env getter caches the value on first
-        read. Explicit user settings are respected, and platforms without
-        fabric memory support fall back to standard allocation in C++.
-        """
-        if (self._cache_transceiver_config is None
-                or self._cache_transceiver_config.backend is None or
-                self._cache_transceiver_config.transceiver_runtime != "PYTHON"):
-            return
-        if not issubclass(self._kv_cache_manager_cls, KVCacheManager):
-            return
-        if os.environ.get("TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY") is None:
-            os.environ["TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"] = "1"
-            logger.info(
-                "Python cache transceiver with C++ KV cache manager detected; "
-                "defaulting TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 (set it "
-                "to 0 explicitly to disable)")
+        maybe_enable_fabric_memory_for_python_transceiver(
+            self._cache_transceiver_config, self._kv_cache_manager_cls)
 
     def _get_model_kv_cache_manager_cls(
         self,

tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py

@@ -2,7 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from abc import ABC, abstractmethod
-from os import getenv
+from os import environ, getenv
 from typing import Any, Dict, List, Optional
 
 import tensorrt_llm
@@ -30,10 +30,40 @@
 _DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV = "TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP"
 _DISAGG_LAYERWISE_ENV = "TRTLLM_DISAGG_LAYERWISE"
 _TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV = "TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER"
+_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV = "TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"
 _SUPPORTED_INFLIGHT_CANCEL_NIXL_BACKEND = "UCX"
 _disagg_inflight_cancel_enabled_cache: Optional[bool] = None
 
 
+def maybe_enable_fabric_memory_for_python_transceiver(
+        cache_transceiver_config: Optional[CacheTransceiverConfig],
+        kv_cache_manager_cls: type) -> None:
+    """Default the C++ V1 KV pool to fabric memory for the Python transceiver.
+
+    This must run before any KV pool allocation because the C++ environment
+    getter caches the value on first read. Explicit user settings are always
+    respected.
+
+    Args:
+        cache_transceiver_config: Configuration used to select the cache
+            transceiver runtime and backend.
+        kv_cache_manager_cls: KV-cache manager class to check for C++ V1 pool
+            allocation.
+    """
+    if (cache_transceiver_config is None
+            or cache_transceiver_config.backend is None
+            or cache_transceiver_config.transceiver_runtime != "PYTHON"):
+        return
+    if not issubclass(kv_cache_manager_cls, KVCacheManager):
+        return
+    if getenv(_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV) is None:
+        environ[_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV] = "1"
+        logger.info(
+            "Python cache transceiver with C++ KV cache manager detected; "
+            f"defaulting {_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV}=1 (set it "
+            "to 0 explicitly to disable)")
+
+
 def is_disagg_inflight_cancel_enabled() -> bool:
     """Return whether disaggregated in-flight KV transfer cancellation is enabled."""
     global _disagg_inflight_cancel_enabled_cache

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

@@ -29,6 +29,7 @@ l0_sanity_check:
     - llmapi/test_llm_examples.py::test_llmapi_sampling
     - llmapi/test_llm_examples.py::test_llmapi_runtime
     - examples/test_llm_api_with_mpi.py::test_llm_api_single_gpu_with_mpirun[TinyLlama-1.1B-Chat-v1.0] ISOLATION
+    - unittest/others/test_kv_cache_transceiver.py::test_maybe_enable_fabric_memory_for_python_transceiver
     - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[NIXL-mha-ctx_fp16_gen_fp16]
     - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[UCX-mha-ctx_fp16_gen_fp16]
     - unittest/others/test_kv_cache_transceiver.py::test_cpp_nixl_sync_transfer_stress

tests/integration/test_lists/waives.txt

@@ -377,8 +377,6 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-r1-fp4_128k8
 perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-v32-fp4_32k4k_con2048_ctx1_dep4_gen1_dep32_eplb288_mtp1_ccb-NIXL] SKIP (https://nvbugs/6550133)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6550133)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6490049)
-perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-b200_deepseek-r1-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6561566)
-perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-b200_deepseek-r1-fp4_8k1k_con256_ctx1_dep4_gen1_dep8_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6561566)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_8k1k_con1024_ctx1_dep4_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6561566)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_8k1k_con1_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6561566)
 perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6550133)

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 = {
@@ -128,7 +129,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,
+    stage_name="",
+    llm_models_root=None,
 ):
     """Launch-script export lines wiring the precheck gate.
 
@@ -163,6 +170,7 @@ def precheck_prefix_lines(
         "--work-dir $testOutputDir/cache_transceiver_precheck "
         f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode"
     )
+    model_root_env = f"LLM_MODELS_ROOT={shlex.quote(llm_models_root)}" if llm_models_root else ""
     lines = [
         f"export ctPrecheckEnabled={int(enabled)}",
         # The external srun timeout must cover the driver's first-rep NIXL
@@ -172,8 +180,14 @@ def precheck_prefix_lines(
         f"{int(knobs.get('step_timeout_s', default_step_timeout_s(max_world)))}",
         "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/"
         "disaggregated/slurm_precheck_run.sh",
-        f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS {cmd} --role ctx"',
-        f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS {cmd} --role gen"',
+        # Quote the complete assignment as an export value. The launch script
+        # expands it into pytestCommand*Precheck, whose later eval interprets
+        # the inner shlex-quoted model path.
+        f"export ctPrecheckModelRootEnv={shlex.quote(model_root_env)}",
+        f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $ctPrecheckModelRootEnv '
+        f'$CTX_WORKER_ENV_VARS $PYTEST_COMMON_VARS {cmd} --role ctx"',
+        f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $ctPrecheckModelRootEnv '
+        f'$GEN_WORKER_ENV_VARS $PYTEST_COMMON_VARS {cmd} --role gen"',
     ]
     if stage_name:
         # Suite name for the synthetic junit xml the gate writes on failure

tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py

@@ -204,7 +204,10 @@ def load_internal_apis():
     from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING
     from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector
     from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
-    from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver
+    from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import (
+        create_kv_cache_transceiver,
+        maybe_enable_fabric_memory_for_python_transceiver,
+    )
     from tensorrt_llm._torch.pyexecutor.llm_request import (
         LlmRequest,
         LlmRequestState,
@@ -237,6 +240,9 @@ def load_internal_apis():
         KVCacheManager=KVCacheManager,
         KVCacheManagerV2=KVCacheManagerV2,
         create_kv_cache_transceiver=create_kv_cache_transceiver,
+        maybe_enable_fabric_memory_for_python_transceiver=(
+            maybe_enable_fabric_memory_for_python_transceiver
+        ),
         LlmRequest=LlmRequest,
         LlmRequestState=LlmRequestState,
         LlmRequestType=LlmRequestType,
@@ -883,6 +889,8 @@ def setup(self, kv_shape, max_req_len):
                 "(the C++ transceiver only supports the V1 manager)"
             )
 
+        manager_cls = api.KVCacheManagerV2 if self.use_v2 else api.KVCacheManager
+        api.maybe_enable_fabric_memory_for_python_transceiver(cache_cfg, manager_cls)
         self.kvm = build_kv_cache_manager(
             kv_shape, self.plan, self.side, self.mapping, max_req_len, self.use_v2
         )

tests/unittest/others/test_cache_transceiver_precheck_config.py

@@ -19,6 +19,8 @@
 
 import json
 import os
+import shlex
+import subprocess
 import sys
 
 import pytest
@@ -359,6 +361,47 @@ def test_wireup_timeout_derivation():
     assert plan["wireup_timeout_s"] == 42
 
 
+@pytest.mark.parametrize("model_root", ["/models with spaces", "/models/o'hare"])
+def test_precheck_commands_propagate_model_root(monkeypatch, model_root):
+    # CI provides the model root inside its inbound pytest command, not in the
+    # environment of the Python process that generates the launch script.
+    monkeypatch.delenv("LLM_MODELS_ROOT", raising=False)
+    monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False)
+    lines = pcfg.precheck_prefix_lines(
+        {},
+        "e2e",
+        "$config",
+        "unset UCX_TLS &&",
+        max_world=8,
+        llm_models_root=model_root,
+    )
+    shell_script = "\n".join(
+        [
+            "CTX_WORKER_ENV_VARS=",
+            "GEN_WORKER_ENV_VARS=",
+            "PYTEST_COMMON_VARS=",
+            "llmSrcNode=/repo",
+            "testOutputDir=/tmp/output",
+            "config=/tmp/config.yaml",
+            *lines,
+            "printf '%s\\n' \"$pytestCommandCTXPrecheck\"",
+            "printf '%s\\n' \"$pytestCommandGENPrecheck\"",
+        ]
+    )
+
+    result = subprocess.run(
+        ["bash"], input=shell_script, capture_output=True, check=True, text=True
+    )
+    commands = result.stdout.splitlines()
+
+    assert len(commands) == 2
+    for command in commands:
+        tokens = shlex.split(command)
+        assignment = f"LLM_MODELS_ROOT={model_root}"
+        assert assignment in tokens
+        assert tokens.index(assignment) < tokens.index("python3")
+
+
 def _enabled_line(cfg):
     lines = pcfg.precheck_prefix_lines(cfg, "e2e", "$c", "unset &&", max_world=8)
     return next(x for x in lines if x.startswith("export ctPrecheckEnabled"))

tests/unittest/others/test_cache_transceiver_precheck_run.py

@@ -397,6 +397,17 @@ def test_serving_resolvers(self, api):
         rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters
         assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3
 
+    def test_deepseek_v4_auto_selects_kv_cache_manager_v2(self, api, tmp_path):
+        model_dir = tmp_path / "deepseek-v4"
+        model_dir.mkdir()
+        (model_dir / "config.json").write_text(
+            json.dumps({"architectures": ["DeepseekV4ForCausalLM"]})
+        )
+        side = {"use_kv_cache_manager_v2": "auto"}
+        cache_cfg = api.CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON")
+
+        assert rp.resolve_model_prefs(str(model_dir), side, cache_cfg) is True
+
     def test_enum_members(self, api):
         for enum, members in (
             (api.DataType, ("FP8", "HALF", "BF16")),