fix: use current CUDA device instead of tp_rank for SymmDeviceMemory allocation

flashinfer/comm/trtllm_ar.py

@@ -597,11 +597,14 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion(
         if not use_symm_dev_mem:
             ipc_handles.append(create_shared_buffer(aligned_size, group))
         else:
+            # Use torch.cuda.current_device() instead of tp_rank to support
+            # base_gpu_id != 0 scenarios where the actual CUDA device index
+            # differs from the TP rank.
             symm_mem = SymmDeviceMemory(
                 aligned_size,
                 tp_size,
                 tp_rank,
-                torch.device("cuda", tp_rank).index,
+                torch.cuda.current_device(),
                 comm_backend,
                 enable_multicast=False,
                 allocate_signal_pads=False,

flashinfer/comm/trtllm_mnnvl_ar.py

@@ -132,11 +132,14 @@ def __init__(
         )
 
         # Allocate the workspace
+        # Use torch.cuda.current_device() instead of mapping.local_rank to
+        # support base_gpu_id != 0 scenarios where the actual CUDA device
+        # index differs from the TP rank / local_rank.
         self.mcast_buffer_handle = McastGPUBuffer(
             requested_workspace_size,
             mapping.tp_size,
             mapping.tp_rank,
-            torch.device("cuda", mapping.local_rank),
+            torch.device("cuda", torch.cuda.current_device()),
             comm_backend,
         )
 
@@ -167,7 +170,7 @@ def __init__(
         self.buffer_flags = torch.tensor(
             [0, 2, self.buffer_size_bytes, 0, *num_bytes_to_clear, 0],
             dtype=torch.uint32,
-            device=torch.device("cuda", mapping.local_rank),
+            device=torch.device("cuda", torch.cuda.current_device()),
         )
 
         self.uc_ptrs_dev = self.mcast_buffer_handle.get_buffer_ptrs_dev()

tests/comm/test_trtllm_allreduce_fusion.py

@@ -25,9 +25,15 @@
 
 
 def _run_correctness_worker(
-    world_size, rank, dtype, hidden_dim, distributed_init_port, legacy_api=True
+    world_size,
+    rank,
+    dtype,
+    hidden_dim,
+    distributed_init_port,
+    legacy_api=True,
+    gpu_offset=0,
 ):
-    device = torch.device(f"cuda:{rank}")
+    device = torch.device(f"cuda:{rank + gpu_offset}")
     torch.cuda.set_device(device)
     distributed_init_method = f"tcp://localhost:{distributed_init_port}"
     dist.init_process_group(
@@ -39,7 +45,6 @@ def _run_correctness_worker(
     group = dist.group.WORLD
 
     try:
-        device = torch.device(f"cuda:{rank}")
         token_nums = [1, 128, 1024, 2048]
         pattern_codes = [
             comm.AllReduceFusionPattern.kAllReduce,
@@ -422,19 +427,24 @@ def multi_process_parallel(
     hidden_dim: int,
     test_target: Any,
     target_args: tuple = (),
+    gpu_offset: int = 0,
 ) -> None:
     mp.set_start_method("spawn", force=True)
 
     procs = []
     distributed_init_port = get_open_port()
     for i in range(world_size):
         proc_args = (
-            world_size,
-            i,
-            dtype,
-            hidden_dim,
-            distributed_init_port,
-        ) + target_args
+            (
+                world_size,
+                i,
+                dtype,
+                hidden_dim,
+                distributed_init_port,
+            )
+            + target_args
+            + (gpu_offset,)
+        )
         proc = mp.Process(target=test_target, args=proc_args, name=f"Worker-{i}")
         proc.start()
         procs.append(proc)
@@ -473,6 +483,43 @@ def test_trtllm_allreduce_fusion(world_size, dtype, hidden_dim, legacy_api):
     print(f"allreduce fusion tp = {world_size} ({api_str} API): OK")
 
 
+@pytest.mark.parametrize("world_size", [2, 4])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("legacy_api", [True, False])
+def test_trtllm_allreduce_fusion_gpu_offset(world_size, dtype, legacy_api):
+    """Test allreduce fusion when CUDA device index != TP rank (base_gpu_id > 0).
+
+    Simulates sglang colocate mode where inference engines run on non-zero
+    base GPUs (e.g. GPUs 4-7 with TP ranks 0-3).
+    See: https://github.com/flashinfer-ai/flashinfer/pull/2662
+    """
+    np.random.seed(42)
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+    available_gpus = torch.cuda.device_count()
+    gpu_offset = available_gpus - world_size
+    if gpu_offset <= 0:
+        pytest.skip(
+            f"Need more than {world_size} GPUs to test gpu_offset>0 "
+            f"(have {available_gpus})"
+        )
+    api_str = "legacy" if legacy_api else "unified"
+    print(
+        f"Running gpu_offset test: world_size={world_size}, gpu_offset={gpu_offset}, "
+        f"{api_str} API (GPUs {gpu_offset}..{gpu_offset + world_size - 1})"
+    )
+
+    multi_process_parallel(
+        world_size,
+        dtype,
+        1024,
+        _run_correctness_worker,
+        target_args=(legacy_api,),
+        gpu_offset=gpu_offset,
+    )
+    print(f"gpu_offset allreduce fusion tp={world_size} ({api_str} API): OK")
+
+
 if __name__ == "__main__":
     # Test both legacy and unified APIs
     print("Testing legacy API...")