diff --git a/flashinfer/comm/trtllm_ar.py b/flashinfer/comm/trtllm_ar.py index 4392468e80..d66d10fb5a 100644 --- a/flashinfer/comm/trtllm_ar.py +++ b/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, diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 3f8e198146..512802e732 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/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() diff --git a/tests/comm/test_trtllm_allreduce_fusion.py b/tests/comm/test_trtllm_allreduce_fusion.py index 5ae027c4ff..538e653cd1 100644 --- a/tests/comm/test_trtllm_allreduce_fusion.py +++ b/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,6 +427,7 @@ 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) @@ -429,12 +435,16 @@ def multi_process_parallel( 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...")