diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index c4ebd8c6460c..5cd45dbcbc96 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/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,28 @@ 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, + ) + print( + f"[{role} rank={rank}] PYTHON+V1 case in matrix: " + "TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=" + f"{os.environ.get('TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY')} " + "applies to every case in this run, including C++ transceiver ones", + flush=True, + ) req_lens = cfg["test_matrix"]["request_lengths"] warmup = cfg["test_matrix"]["warmup_requests"] num_req = cfg["test_matrix"]["num_requests_per_length"] diff --git a/jenkins/scripts/perf/cluster_env.py b/jenkins/scripts/perf/cluster_env.py index cbc342aa2d24..923975089d80 100644 --- a/jenkins/scripts/perf/cluster_env.py +++ b/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. ( diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 8e571eb2f890..4c911c3f0e19 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, ) ) diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index f1c54fb13690..3adf7b6211e4 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -630,6 +630,39 @@ 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].strip() + try: + parts = shlex.split(command_value) + # A fully quoted export value ("pytest ...") is parsed as one token on + # the first pass; parse that token once more to recover its command + # words. Only unwrap when the raw value is wrapped in an outer quote + # pair — a bare single token with escaped whitespace is already the + # final value and must not be re-split. + wrapped_in_quotes = ( + len(command_value) >= 2 + and command_value[0] in "\"'" + and command_value[-1] == command_value[0] + ) + if len(parts) == 1 and wrapped_in_quotes: + 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. @@ -809,6 +842,7 @@ 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") script_prefix_lines.extend( [ worker_pytest_command, @@ -852,6 +886,20 @@ 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) + # The model root is only consumed by the precheck (auto KV-cache-manager + # resolution needs the model config). Fail fast only when the precheck + # will actually run; otherwise degrade to a warning so stages whose + # pytestCommand does not carry LLM_MODELS_ROOT inline keep submitting. + if not llm_models_root: + if pcfg.precheck_enabled(config): + raise ValueError( + "LLM_MODELS_ROOT is missing from the inbound pytestCommand; " + "the cache-transceiver precheck cannot resolve model defaults" + ) + print( + "WARNING: LLM_MODELS_ROOT not found in the inbound pytestCommand; " + "cache-transceiver precheck is disabled for this config so continuing" + ) script_prefix_lines.extend( pcfg.precheck_prefix_lines( config, @@ -863,6 +911,7 @@ def main(): hardware_config["gpus_per_gen_server"], ), stage_name=args.stage_name, + llm_models_root=llm_models_root, ) ) srun_args_lines.extend( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index dc7b753a89ef..cef6b018f105 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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, @@ -553,28 +555,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, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index d4868acf0226..65f613337fff 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/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 diff --git a/tests/integration/test_lists/test-db/l0_sanity_check.yml b/tests/integration/test_lists/test-db/l0_sanity_check.yml index db5c37afffbe..d4ae779eb363 100644 --- a/tests/integration/test_lists/test-db/l0_sanity_check.yml +++ b/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 diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index f69821b2131c..4fa735bc7800 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -353,7 +353,6 @@ perf/test_perf_sanity.py::test_e2e[aggr_upload-qwen3_5_397b_fp4_blackwell-qwen3_ perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws1_1k1k] SKIP (https://nvbugs/6153575) 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_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6566777) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_32k4k_con2048_ctx1_dep4_gen1_dep32_eplb288_mtp1_ccb-NIXL] SKIP (https://nvbugs/6374872) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_deepseek-v32-fp4_8k1k_con4096_ctx1_dep4_gen1_dep32_eplb256_mtp0_ccb-NIXL] SKIP (https://nvbugs/6490049) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6490049) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index ddf2514b7a19..2588c1362157 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -26,15 +26,13 @@ layers) is read from the real model's `config.json` under `$LLM_MODELS_ROOT`. ## Enabling / disabling -**Off by default** until the gate is validated on the post-merge stages -(the precheck is a launch-script gate, not a pytest case, so it cannot be -waived in `waives.txt` — this default is the waive). To opt in: +**On by default** for every disaggregated perf-sanity test. To opt out: - per test yaml: ```yaml cache_transceiver_precheck: - enabled: true + enabled: false # optional overrides (defaults in precheck_config.PRECHECK_DEFAULTS): # request_lengths: [1024, 8192] # num_requests: 2 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 8219a2581d9c..973561e871c4 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 = { @@ -127,8 +128,37 @@ def default_step_timeout_s(max_world): return 900 + wireup_timeout_s(max_world) +def precheck_enabled(cfg): + """Resolve the precheck enable/kill-switch policy for a test config. + + On by default; yaml opts out per test; the env var (when set) overrides + the yaml either way (global kill switch). Parse the usual boolean spellings + so a well-meant TRTLLM_DISAGG_CT_PRECHECK=true force-enable is not silently + read as "off"; reject anything ambiguous instead of guessing. + """ + env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") + if env is not None: + val = env.strip().lower() + if val in ("1", "true", "on", "yes"): + return True + if val in ("0", "false", "off", "no"): + return False + raise ValueError( + "TRTLLM_DISAGG_CT_PRECHECK must be a boolean " + f"(1/0/true/false/on/off/yes/no), got {env!r}" + ) + knobs = cfg.get("cache_transceiver_precheck", {}) or {} + return bool(knobs.get("enabled", True)) + + 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. @@ -139,33 +169,14 @@ def precheck_prefix_lines( expressions ($llmSrcNode etc.), expanded at sbatch runtime. """ knobs = cfg.get("cache_transceiver_precheck", {}) or {} - # Off by default until the gate is validated on the post-merge stages - # (the precheck is a launch-script gate, not a pytest case, so it cannot - # be waived in waives.txt — this default is the waive). Yaml opts in per - # test; the env var (when set) overrides the yaml either way (global kill - # switch). Parse the usual boolean spellings so a well-meant - # TRTLLM_DISAGG_CT_PRECHECK=true force-enable is not silently read as - # "off"; reject anything ambiguous instead of guessing. - env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") - if env is not None: - val = env.strip().lower() - if val in ("1", "true", "on", "yes"): - enabled = True - elif val in ("0", "false", "off", "no"): - enabled = False - else: - raise ValueError( - "TRTLLM_DISAGG_CT_PRECHECK must be a boolean " - f"(1/0/true/false/on/off/yes/no), got {env!r}" - ) - else: - enabled = bool(knobs.get("enabled", False)) + enabled = precheck_enabled(cfg) cmd = ( "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" f"run_precheck.py --config {config_path_expr} " "--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 @@ -175,8 +186,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 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..29ac898f88c8 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/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, @@ -272,28 +278,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 @@ -883,6 +894,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 ) @@ -924,7 +937,7 @@ 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) + 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" ) @@ -1042,7 +1055,7 @@ def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair): local_err = _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}") elif 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 diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 99dc35706f03..e4bc6be40402 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/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,11 +361,102 @@ 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 test_precheck_commands_split_pytest_common_vars(monkeypatch): + # $PYTEST_COMMON_VARS is spliced unquoted on purpose: bash word splitting + # must yield separate K=V env-assignment tokens ahead of the executable. + # Values containing spaces are unsupported by design — this pins the + # expected splitting behavior. + 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="/models", + ) + shell_script = "\n".join( + [ + "CTX_WORKER_ENV_VARS=", + "GEN_WORKER_ENV_VARS=", + 'PYTEST_COMMON_VARS="FOO=1 BAR=two"', + "llmSrcNode=/repo", + "testOutputDir=/tmp/output", + "config=/tmp/config.yaml", + *lines, + "printf '%s\\n' \"$pytestCommandCTXPrecheck\"", + ] + ) + + result = subprocess.run( + ["bash"], input=shell_script, capture_output=True, check=True, text=True + ) + tokens = shlex.split(result.stdout.splitlines()[0]) + + python_index = tokens.index("python3") + for assignment in ("FOO=1", "BAR=two"): + assert tokens.index(assignment) < python_index + + 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")) +def test_precheck_enabled_helper(monkeypatch): + # submit.py consults this helper to decide whether a missing model root is + # fatal — it must mirror the policy encoded in ctPrecheckEnabled. + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + assert pcfg.precheck_enabled({}) is True + assert pcfg.precheck_enabled({"cache_transceiver_precheck": {"enabled": False}}) is False + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "0") + assert pcfg.precheck_enabled({}) is False + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "true") + assert pcfg.precheck_enabled({"cache_transceiver_precheck": {"enabled": False}}) is True + + def test_precheck_env_kill_switch_truthy(monkeypatch): """The TRTLLM_DISAGG_CT_PRECHECK kill switch parses the usual boolean spellings. @@ -373,7 +466,7 @@ def test_precheck_env_kill_switch_truthy(monkeypatch): cfg = {"cache_transceiver_precheck": {"enabled": True}} monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) assert _enabled_line(cfg).endswith("=1") # yaml opt-in - assert _enabled_line({}).endswith("=0") # off by default (waived) + assert _enabled_line({}).endswith("=1") # on by default for v in ("1", "true", "on", "YES", " True "): monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", v) assert _enabled_line(cfg).endswith("=1"), v diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py index ed060d1e9053..6962140c2a12 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 # --------------------------------------------------------------------------- # @@ -397,6 +427,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")), diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 241a4aee0d09..9a8521c65592 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -21,8 +21,10 @@ import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm from tensorrt_llm._torch.distributed import Distributed -from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import \ - create_kv_cache_transceiver +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +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) from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import \ @@ -42,6 +44,33 @@ KV_TRANSFER_COMPLETION_MARGIN_S = 10.0 +@pytest.mark.parametrize( + "runtime,manager_cls,initial_value,expected_value", + [ + ("PYTHON", KVCacheManager, None, "1"), + ("PYTHON", KVCacheManager, "0", "0"), + ("PYTHON", KVCacheManagerV2, None, None), + ("CPP", KVCacheManager, None, None), + # "auto" is resolved to PYTHON/CPP by serving before this helper runs; + # callers that construct CacheTransceiverConfig directly must resolve + # it first — the helper deliberately leaves "auto" untouched. + ("auto", KVCacheManager, None, None), + ], +) +def test_maybe_enable_fabric_memory_for_python_transceiver( + monkeypatch, runtime, manager_cls, initial_value, expected_value): + env_name = "TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY" + if initial_value is None: + monkeypatch.delenv(env_name, raising=False) + else: + monkeypatch.setenv(env_name, initial_value) + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + + maybe_enable_fabric_memory_for_python_transceiver(config, manager_cls) + + assert os.environ.get(env_name) == expected_value + + @pytest.mark.parametrize("transceiver_runtime", ["CPP", "auto"]) def test_cpp_transceiver_rejects_mixed_mamba_manager(transceiver_runtime): config = CacheTransceiverConfig(backend="NIXL", diff --git a/tests/unittest/scripts/test_cluster_env.py b/tests/unittest/scripts/test_cluster_env.py index 26905068b10f..8486854d2a78 100644 --- a/tests/unittest/scripts/test_cluster_env.py +++ b/tests/unittest/scripts/test_cluster_env.py @@ -83,10 +83,7 @@ def test_gpu_type_from_supported_gpus( ), ( "oci-aga-cs-001", - "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", ), ("aws-cmh", "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp"), ("aws-dfw-prod", "export UCX_TLS=^gdr_copy"), diff --git a/tests/unittest/scripts/test_perf_submit.py b/tests/unittest/scripts/test_perf_submit.py index 3a824627abbd..37b3ef41649a 100644 --- a/tests/unittest/scripts/test_perf_submit.py +++ b/tests/unittest/scripts/test_perf_submit.py @@ -16,6 +16,7 @@ import importlib.util import json +import shlex from pathlib import Path from types import ModuleType @@ -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(SUBMIT_PATHS[0], 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) - - def _select_ci_test_case_line( ci_submit_module: ModuleType, tmp_path: Path, @@ -273,3 +274,45 @@ def test_ci_submit_rejects_missing_pytest_split_durations( pytest_options="--splits 1 --group 1 --durations-path /remote/.test_durations", split_group=1, ) + + +@pytest.mark.parametrize("quote_command", [False, True]) +@pytest.mark.parametrize("model_root", ["/scratch/llm-models", "/models/o'hare with spaces"]) +def test_ci_extracts_model_root_from_inbound_pytest_command( + ci_submit_module, model_root, quote_command +): + command = ( + f"LLM_ROOT=/repo LLM_MODELS_ROOT={shlex.quote(model_root)} " + "/repo/tensorrt_llm/llmapi/trtllm-llmapi-launch pytest -q" + ) + command_value = shlex.quote(command) if quote_command else command + prefix_lines = [f"export pytestCommand={command_value}"] + + assert ( + ci_submit_module._get_pytest_command_env_var(prefix_lines, "LLM_MODELS_ROOT") == model_root + ) + + +def test_ci_single_token_with_escaped_whitespace_is_not_resplit(ci_submit_module): + # A bare (unwrapped) export value that parses to a single token with + # embedded whitespace is already final — re-splitting it would corrupt + # the value. Only quote-wrapped whole commands get a second parse. + prefix_lines = [r"export pytestCommand=LLM_MODELS_ROOT=/models\ with\ spaces"] + + assert ( + ci_submit_module._get_pytest_command_env_var(prefix_lines, "LLM_MODELS_ROOT") + == "/models with spaces" + ) + + +def test_ci_missing_model_root_is_detectable(ci_submit_module): + prefix_lines = ["export pytestCommand='LLM_ROOT=/repo pytest -q'"] + + assert ci_submit_module._get_pytest_command_env_var(prefix_lines, "LLM_MODELS_ROOT") is None + + +def test_ci_rejects_invalid_pytest_command(ci_submit_module): + prefix_lines = ['export pytestCommand="LLM_MODELS_ROOT=/models'] + + with pytest.raises(ValueError, match="Invalid inbound pytestCommand"): + ci_submit_module._get_pytest_command_env_var(prefix_lines, "LLM_MODELS_ROOT")