diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..4c743dfbe5a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -342,6 +342,31 @@ def __init__( KVCacheManagerV2) self._draft_config = draft_config self._skip_est = skip_est + self._maybe_enable_fabric_memory_for_python_transceiver() + + def _maybe_enable_fabric_memory_for_python_transceiver(self): + """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)") def _get_model_kv_cache_manager_cls( self, diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9c82b55a4a17..53736a99c780 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: max_tokens_in_buffer: 1024 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 @@ -88,4 +89,5 @@ worker_config: max_tokens_in_buffer: 1024 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9f4b7086060d..ffa8cdfa4522 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -75,6 +75,7 @@ worker_config: max_tokens_in_buffer: 9216 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -103,6 +104,7 @@ worker_config: max_tokens_in_buffer: 9216 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index c9f73573fcf9..3f00a0063523 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 @@ -88,4 +89,5 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9610d4a5bdf9..edca66a96d87 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -63,6 +63,7 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 @@ -88,4 +89,5 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true diff --git a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9f4b7086060d..ffa8cdfa4522 100644 --- a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -75,6 +75,7 @@ worker_config: max_tokens_in_buffer: 9216 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -103,6 +104,7 @@ worker_config: max_tokens_in_buffer: 9216 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 diff --git a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9ed015cd6cda..afa141b6ae82 100644 --- a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -74,6 +74,7 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 @@ -99,4 +100,5 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true diff --git a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 5476fae34c0e..c8dffadebfa7 100644 --- a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -74,6 +74,7 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 @@ -99,4 +100,5 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true diff --git a/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml index 0701596a0bad..9e9dd0877c90 100644 --- a/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml @@ -80,6 +80,7 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true speculative_config: &id001 decoding_type: Eagle @@ -110,5 +111,6 @@ worker_config: max_tokens_in_buffer: 8448 backend: NIXL transceiver_runtime: PYTHON + kv_cache_bounce_size_mb: 512 disable_overlap_scheduler: true speculative_config: *id001