[TRTLLM-6444] Add some UCX trouble shooting docs and print UCX related logs (NVIDIA#6085)

Signed-off-by: Lizhi Zhou <[email protected]>

Author: reasonsolo

docs/source/advanced/disaggregated-service.md

@@ -32,6 +32,10 @@ TRT-LLM uses some environment variables to control the behavior of disaggregated
 
 * `TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM`: The maximum number of concurrent KV cache sends. The default value is `4`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0.
 
+There are some other useful environment variables that may help when encountering failures or performance issues.
+
+* `NCCL_GRAPH_MIXING_SUPPORT`: With the default value `1`, the CUDA driver may create too many CUDA streams while working with one CUDA graph, leading to performance drop. Setting it to `0` will reduce the number of CUDA streams, but please make sure there are no other NCCL ops outside the one CUDA graph, otherwise it's unsafe.
+
 ## Troubleshooting and FAQ
 
 ### General FAQs
@@ -80,3 +84,13 @@ A. Yes, TRT-LLM supports using GPU direct RDMA for inter-node KV cache transfer.
 *Q. What causes the substantial bandwidth fluctuations in kvCache transfers, especially during the first few requests following service initialization?*
 
 A. The communication for kvCache transfer between executors are established dynamically. The connection establishment process incurs significant overhead, which explains the apparently lower kvCache transfer bandwidth observed during the initial requests after service startup. This lower bandwidth reflects the inclusion of connection establishment overhead. When conducting benchmarks, it is recommended to perform a warm-up phase to ensure accurate performance measurements.
+
+*Q. When my servers are running on different NVLink domains, some servers hang or have a lower performance. How to fix that?
+
+A. NVLink domain can be found with `nvidia-smi -q` in the `Fabric.ClusterUUID` field. A few UCX environment variables can be adjusted when your servers have different NVLink domains:
+
+* `UCX_CUDA_IPC_ENABLE_MNNVL`: Set to `n`. This also can reduce UCX timeout error messages like `UCX  ERROR   cuMemImportFromShareableHandle failed: invalid resource handle`, although these errors don't necessarily cause your trtllm-serve to fail.
+
+* `UCX_NET_DEVICES`: Check if this is set correctly, or unset this variable to allow UCX to use all possible devices.
+
+* `UCX_RNDV_SCHEME`: Set to `get_zcopy` or `put_zcopy` on GB200 for better performance. The default value is `auto`.

tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py

@@ -31,29 +31,36 @@ def create_kv_cache_transceiver(
         mapping: Mapping, kv_cache_manager: KVCacheManager,
         attention_type: AttentionTypeCpp,
         cache_transceiver_config: CacheTransceiverConfig):
-    if cache_transceiver_config is None or (cache_transceiver_config.backend
-                                            is None):
+    if cache_transceiver_config is None or cache_transceiver_config.backend is None:
         logger.info("cache_transceiver is disabled")
         return None
-    if (cache_transceiver_config.backend == BackendTypeCpp.DEFAULT):
-
-        backend_type = BackendTypeCpp.UCX
-        if getenv("TRTLLM_USE_UCX_KVCACHE"):
-            backend_type = BackendTypeCpp.UCX
-        elif getenv("TRTLLM_USE_NIXL_KVCACHE"):
-            backend_type = BackendTypeCpp.NIXL
-        elif getenv("TRTLLM_USE_MPI_KVCACHE"):
-            backend_type = BackendTypeCpp.MPI
-        cache_transceiver_config.backend = backend_type
-
-    if (cache_transceiver_config.backend == BackendTypeCpp.MPI):
+
+    if cache_transceiver_config.backend == BackendTypeCpp.DEFAULT:
+        # When cache_transceiver_config.backend is not set, fallback to env_vars settings
+        # UCX is the default backend
+        cache_transceiver_config.backend = BackendTypeCpp.UCX
+        # Ordered by priority
+        env_vars = [("TRTLLM_USE_NIXL_KVCACHE", BackendTypeCpp.NIXL),
+                    ("TRTLLM_USE_MPI_KVCACHE", BackendTypeCpp.MPI)]
+        for env_var, be_type in env_vars:
+            if getenv(env_var) == "1":
+                logger.warning(
+                    f"{env_var}=1 is set, but it's recommended to set cache_transceiver_config.backend in yaml config"
+                )
+                cache_transceiver_config.backend = be_type
+                break
+
+    if cache_transceiver_config.backend == BackendTypeCpp.MPI:
         logger.warning(
             "MPI CacheTransceiver is deprecated, UCX or NIXL is recommended")
-    cache_transceiver = BindKvCacheTransceiver(mapping, kv_cache_manager,
-                                               attention_type,
-                                               cache_transceiver_config)
-
-    return cache_transceiver
+    elif cache_transceiver_config.backend == BackendTypeCpp.UCX:
+        logger.info(
+            f"Using UCX kv-cache transceiver. If your devices are not in the same domain, please consider setting "
+            f"UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server "
+            f"hangs or lower-than-expected performance.")
+
+    return BindKvCacheTransceiver(mapping, kv_cache_manager, attention_type,
+                                  cache_transceiver_config)
 
 
 class KvCacheTransceiver(ABC):