From e527ffa7b8f20ebed25a2fd443932c4b025e8adf Mon Sep 17 00:00:00 2001 From: Fergus Finn Date: Sun, 7 Jun 2026 01:08:09 +0000 Subject: [PATCH 1/7] Fix UCCL EP device selection under node-local launchers --- ep/bench/buffer.py | 83 +++++++++++++++++++++++++++++++++++++++------- ep/bench/utils.py | 17 ++-------- 2 files changed, 74 insertions(+), 26 deletions(-) diff --git a/ep/bench/buffer.py b/ep/bench/buffer.py index 0f23c8ca3..ae84568a0 100644 --- a/ep/bench/buffer.py +++ b/ep/bench/buffer.py @@ -73,6 +73,42 @@ class Buffer: # TODO(MaoZiming): Reduce SMs. UCCL Proxy should reduce the usage of SMs. num_sms: int = 20 + @staticmethod + def _env_int(name: str, default: int = -1) -> int: + try: + return int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + + @staticmethod + def _infer_device_index(rank: int) -> int: + device_count = torch.cuda.device_count() + if device_count <= 0: + return torch.cuda.current_device() + if device_count == 1: + return 0 + + local_rank = Buffer._env_int("LOCAL_RANK") + local_world_size = Buffer._env_int("LOCAL_WORLD_SIZE") + if local_world_size == device_count and 0 <= local_rank < device_count: + return local_rank + + # In one-container-per-node launchers (for example vLLM mp workers under + # a single Slurm task), every local worker can inherit LOCAL_RANK=0 while + # still seeing all local GPUs. EP ranks are laid out contiguously per node, + # so rank modulo visible devices is the local CUDA device. + return int(rank) % device_count + + @staticmethod + def _infer_local_world_size() -> int: + device_count = torch.cuda.device_count() + local_world_size = Buffer._env_int("LOCAL_WORLD_SIZE") + if local_world_size <= 0 or ( + device_count > 1 and local_world_size < device_count + ): + return device_count + return local_world_size + def __init__( self, group: dist.ProcessGroup, @@ -106,10 +142,12 @@ def __init__( is_intranode: whether to force intranode-only proxy mode. If set to `None`, infer it from the process-group topology automatically. Explicit `True` is rejected when the group spans multiple nodes. """ - if "LOCAL_RANK" in os.environ: - device_index = int(os.environ["LOCAL_RANK"]) - else: - device_index = torch.cuda.current_device() + self.rank = group.rank() + self.group_size = group.size() + self.group = group + self.device_index = self._infer_device_index(self.rank) + torch.cuda.set_device(self.device_index) + device_index = self.device_index if hasattr(ep, "get_rdma_buffer"): # Allocate outside PyTorch's CUDA allocator so RDMA/IPC sees a raw @@ -150,12 +188,12 @@ def __init__( ) rdma_buffer_ptr = self.scratch.data_ptr() - _local_world = int(os.environ.get("LOCAL_WORLD_SIZE", -1)) + _local_world = self._infer_local_world_size() self.proxies, self.workers = initialize_uccl( rdma_buffer_ptr, num_rdma_bytes, - group.rank(), - dist.get_world_size(group), + self.rank, + self.group_size, group, use_normal_mode=not low_latency_mode, is_intranode=is_intranode, @@ -164,14 +202,12 @@ def __init__( check_nvlink_connections(group) # Initialize the CPP runtime - self.rank = group.rank() - self.group_size = group.size() - self.group = group self.num_nvl_bytes = num_nvl_bytes self.num_rdma_bytes = num_rdma_bytes self.low_latency_mode = low_latency_mode self.explicitly_destroy = explicitly_destroy self._next_low_latency_combine_buffer = None + torch.cuda.set_device(self.device_index) self.runtime = ep.Buffer( self.rank, self.group_size, @@ -221,11 +257,23 @@ def __init__( for proxy in self.proxies: proxy.set_atomic_buffer_ptr(self.proxies[0].get_atomic_buffer_ptr()) + def _set_current_device(self, device: Optional[torch.device] = None) -> int: + if device is None: + device_index = self.device_index + else: + device = torch.device(device) + if device.type != "cuda": + return self.device_index + device_index = self.device_index if device.index is None else device.index + torch.cuda.set_device(device_index) + return int(device_index) + def _ll_compute_stream_ptr(self, device: torch.device): """ Return the current CUDA stream pointer for low-latency runtime calls. """ - current = torch.cuda.current_stream(device=device) + device_index = self._set_current_device(device) + current = torch.cuda.current_stream(device=device_index) return int(current.cuda_stream) def reset_rdma_buffer(self): @@ -345,6 +393,7 @@ def low_latency_dispatch( event: the event after executing the kernel (valid only if `async_finish` is set). hook: the receiving hook function (valid only if `return_recv_hook` is set). """ + self._set_current_device(x.device) for proxy in self.proxies: proxy.notify_proxy_thread_adaptive_sleeper() proxy.calculate_and_set_dispatch_recv_data_offset( @@ -514,6 +563,7 @@ def low_latency_combine( event: the event after executing the kernel (valid only if `async_finish` is set). hook: the receiving hook function (valid only if `return_recv_hook` is set). """ + self._set_current_device(x.device) if overlap: raise NotImplementedError( "low_latency_combine(overlap=True) is not implemented yet. " @@ -611,10 +661,11 @@ def get_next_low_latency_combine_buffer(self, handle: object): num_ranks = self.group.size() num_local_experts = num_experts // num_ranks num_recv_tokens = num_ranks * num_max_dispatch_tokens_per_rank + self._set_current_device(src_info.device) self._next_low_latency_combine_buffer = torch.empty( (num_local_experts, num_recv_tokens, hidden), dtype=torch.bfloat16, - device="cuda", + device=src_info.device, ) return self._next_low_latency_combine_buffer @@ -648,8 +699,10 @@ def get_comm_stream(self) -> torch.Stream: Returns: stream: the communication stream. """ + self._set_current_device() ts = self.runtime.get_comm_stream() if isinstance(ts, torch.Stream): + torch.cuda.set_device(ts.device_index) return torch.cuda.Stream( stream_id=ts.stream_id, device_index=ts.device_index, @@ -819,6 +872,7 @@ def get_dispatch_layout( if allocate_on_comm_stream: assert previous_event is not None and async_finish + self._set_current_device(topk_idx.device) alloc_ctx = ( torch.cuda.stream(self.get_comm_stream()) if allocate_on_comm_stream @@ -953,6 +1007,8 @@ def dispatch( """ # Default config config = self.get_dispatch_config(self.group_size) if config is None else config + x_for_device = x[0] if isinstance(x, tuple) else x + self._set_current_device(x_for_device.device) # Internode if self.runtime.get_num_rdma_ranks() > 1: @@ -1278,6 +1334,7 @@ def combine( """ # Default config config = self.get_combine_config(self.group_size) if config is None else config + self._set_current_device(x.device) # Internode if self.runtime.get_num_rdma_ranks() > 1: @@ -1421,6 +1478,7 @@ def internode_dispatch( assert config is not None x, x_scales = x if isinstance(x, tuple) else (x, None) + self._set_current_device(x.device) num_scales = ( 0 if x_scales is None else (1 if x_scales.dim() == 1 else x_scales.size(1)) ) @@ -1777,6 +1835,7 @@ def internode_combine( Normally, you should not directly call this function. """ assert config is not None + self._set_current_device(x.device) # Unpack handle and bias ( diff --git a/ep/bench/utils.py b/ep/bench/utils.py index 461cc1ea7..0504e5b0e 100644 --- a/ep/bench/utils.py +++ b/ep/bench/utils.py @@ -111,10 +111,7 @@ def detect_group_topology(group: dist.ProcessGroup) -> Tuple[int, int, int, bool num_nodes: number of distinct nodes spanned by the group. is_intranode: whether all ranks in the group are on the same node. """ - if "LOCAL_RANK" in os.environ: - local_rank = int(os.environ["LOCAL_RANK"]) - else: - local_rank = torch.cuda.current_device() + local_rank = torch.cuda.current_device() node_token = ( os.environ.get("NODE_RANK") @@ -160,11 +157,7 @@ def get_cpu_proxies_meta(proxies, rank, scratch_ptr, scratch_bytes, num_ranks, g "listen_ports": [proxy.get_listen_port() for proxy in proxies], } all_meta = [None] * num_ranks - # Use current device or fallback to LOCAL_RANK or 0 - if "LOCAL_RANK" in os.environ: - device_index = int(os.environ["LOCAL_RANK"]) - else: - device_index = torch.cuda.current_device() + device_index = torch.cuda.current_device() torch.cuda.set_device(device_index) dist.all_gather_object(all_meta, meta, group=group) rank2meta = {m["rank"]: m for m in all_meta} @@ -636,11 +629,7 @@ def initialize_uccl( def destroy_uccl(proxies, workers): - # Use current device or fallback to LOCAL_RANK - if "LOCAL_RANK" in os.environ: - device_index = int(os.environ["LOCAL_RANK"]) - else: - device_index = torch.cuda.current_device() + device_index = torch.cuda.current_device() if workers is not None: try: From 4379a420b1d8d404a55318f8577b1ca6f066f4ba Mon Sep 17 00:00:00 2001 From: Fergus Finn Date: Sun, 7 Jun 2026 05:14:26 +0000 Subject: [PATCH 2/7] Align disagg proxy with vLLM NIXL request contract --- ep/bench/vllm/disagg_proxy.py | 48 ++++++++++++++++++++++++++++++----- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/ep/bench/vllm/disagg_proxy.py b/ep/bench/vllm/disagg_proxy.py index bd008351c..d1884bcc7 100644 --- a/ep/bench/vllm/disagg_proxy.py +++ b/ep/bench/vllm/disagg_proxy.py @@ -18,10 +18,12 @@ import argparse import json import sys +import uuid +from urllib.parse import urlparse import aiohttp from fastapi import FastAPI, Request -from fastapi.responses import StreamingResponse +from fastapi.responses import JSONResponse, Response, StreamingResponse import uvicorn app = FastAPI() @@ -34,6 +36,17 @@ @app.post("/v1/chat/completions") async def chat_completions(request: Request): body = await request.json() + request_id = request.headers.get("X-Request-Id") or str(uuid.uuid4()) + + def request_headers() -> dict[str, str]: + headers = { + "Content-Type": "application/json", + "X-Request-Id": request_id, + } + auth = request.headers.get("Authorization") + if auth: + headers["Authorization"] = auth + return headers # Step 1: Send to prefill with max_tokens=1 to populate KV cache. # do_remote_decode=True tells prefill's NixlConnector that a remote @@ -42,16 +55,25 @@ async def chat_completions(request: Request): # block IDs, engine ID, and side channel address for the decode node. prefill_body = dict(body) prefill_body["max_tokens"] = 1 + if "max_completion_tokens" in prefill_body: + prefill_body["max_completion_tokens"] = 1 prefill_body["stream"] = False prefill_body.pop("stream_options", None) - prefill_body["kv_transfer_params"] = {"do_remote_decode": True} + prefill_body["kv_transfer_params"] = { + "do_remote_decode": True, + "do_remote_prefill": False, + "remote_engine_id": None, + "remote_block_ids": None, + "remote_host": None, + "remote_port": None, + } async with aiohttp.ClientSession() as session: # Prefill request async with session.post( f"{PREFILL_URL}/v1/chat/completions", json=prefill_body, - headers={"Content-Type": "application/json"}, + headers=request_headers(), ) as prefill_resp: if prefill_resp.status != 200: error = await prefill_resp.text() @@ -87,6 +109,9 @@ async def chat_completions(request: Request): # Step 3: Forward to decode with kv_transfer_params decode_body = dict(body) if kv_transfer_params: + kv_transfer_params = dict(kv_transfer_params) + if not kv_transfer_params.get("remote_host"): + kv_transfer_params["remote_host"] = urlparse(PREFILL_URL).hostname decode_body["kv_transfer_params"] = kv_transfer_params is_stream = body.get("stream", False) @@ -98,7 +123,7 @@ async def stream_decode(): async with s.post( f"{DECODE_URL}/v1/chat/completions", json=decode_body, - headers={"Content-Type": "application/json"}, + headers=request_headers(), ) as resp: async for chunk in resp.content.iter_any(): yield chunk @@ -108,9 +133,20 @@ async def stream_decode(): async with session.post( f"{DECODE_URL}/v1/chat/completions", json=decode_body, - headers={"Content-Type": "application/json"}, + headers=request_headers(), ) as decode_resp: - return await decode_resp.json() + body_bytes = await decode_resp.read() + content_type = decode_resp.headers.get("content-type", "") + if "application/json" in content_type: + return JSONResponse( + content=json.loads(body_bytes), + status_code=decode_resp.status, + ) + return Response( + content=body_bytes, + status_code=decode_resp.status, + media_type=content_type or None, + ) @app.get("/health") From dae27f62bd5b4dd5919ddc09569ffece6923a705 Mon Sep 17 00:00:00 2001 From: fergus barratt Date: Sun, 7 Jun 2026 07:09:42 +0100 Subject: [PATCH 3/7] Propagate disagg proxy upstream errors --- ep/bench/vllm/disagg_proxy.py | 54 ++++++++++++++++++++++++++--------- 1 file changed, 41 insertions(+), 13 deletions(-) diff --git a/ep/bench/vllm/disagg_proxy.py b/ep/bench/vllm/disagg_proxy.py index d1884bcc7..e3f887fe7 100644 --- a/ep/bench/vllm/disagg_proxy.py +++ b/ep/bench/vllm/disagg_proxy.py @@ -76,11 +76,18 @@ def request_headers() -> dict[str, str]: headers=request_headers(), ) as prefill_resp: if prefill_resp.status != 200: - error = await prefill_resp.text() + error = await prefill_resp.read() + content_type = prefill_resp.headers.get("content-type", "") print( - f"[ERROR] Prefill {prefill_resp.status}: {error}", file=sys.stderr + f"[ERROR] Prefill {prefill_resp.status}: " + f"{error.decode(errors='replace')}", + file=sys.stderr, + ) + return Response( + content=error, + status_code=prefill_resp.status, + media_type=content_type or None, ) - return {"error": f"Prefill failed: {error}"} prefill_result = await prefill_resp.json() # Step 2: Extract kv_transfer_params from prefill response. @@ -117,18 +124,39 @@ def request_headers() -> dict[str, str]: is_stream = body.get("stream", False) if is_stream: + decode_session = aiohttp.ClientSession() + decode_resp = await decode_session.post( + f"{DECODE_URL}/v1/chat/completions", + json=decode_body, + headers=request_headers(), + ) + + if decode_resp.status != 200: + body_bytes = await decode_resp.read() + content_type = decode_resp.headers.get("content-type", "") + decode_resp.release() + await decode_session.close() + if "application/json" in content_type: + return JSONResponse( + content=json.loads(body_bytes), + status_code=decode_resp.status, + ) + return Response( + content=body_bytes, + status_code=decode_resp.status, + media_type=content_type or None, + ) async def stream_decode(): - async with aiohttp.ClientSession() as s: - async with s.post( - f"{DECODE_URL}/v1/chat/completions", - json=decode_body, - headers=request_headers(), - ) as resp: - async for chunk in resp.content.iter_any(): - yield chunk - - return StreamingResponse(stream_decode(), media_type="text/event-stream") + try: + async for chunk in decode_resp.content.iter_any(): + yield chunk + finally: + decode_resp.release() + await decode_session.close() + + media_type = decode_resp.headers.get("content-type") or "text/event-stream" + return StreamingResponse(stream_decode(), media_type=media_type) else: async with session.post( f"{DECODE_URL}/v1/chat/completions", From 35d633ccb5c6602cbdb7994817af2a7eab225cb3 Mon Sep 17 00:00:00 2001 From: fergus barratt Date: Sun, 7 Jun 2026 07:16:42 +0100 Subject: [PATCH 4/7] Allow internode test SM override --- ep/bench/test_internode.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/ep/bench/test_internode.py b/ep/bench/test_internode.py index d23472512..bda39ecac 100644 --- a/ep/bench/test_internode.py +++ b/ep/bench/test_internode.py @@ -581,7 +581,9 @@ def test_loop( if args.test_ll_compatibility: ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9 - if torch.version.cuda: + if args.num_sms is not None: + num_sms = args.num_sms + elif torch.version.cuda: num_sms = 24 elif torch.version.hip: num_sms = 64 if num_nodes < 4 else 32 @@ -699,6 +701,12 @@ def test_loop( parser.add_argument( "--num-experts", type=int, default=256, help="Number of experts (default: 256" ) + parser.add_argument( + "--num-sms", + type=int, + default=None, + help="Override the number of SMs used by high-throughput kernels", + ) parser.add_argument( "--test-ll-compatibility", action="store_true", From a8b4fd4bd8502c1c44923c8d221c817c1932727f Mon Sep 17 00:00:00 2001 From: fergus barratt Date: Sun, 7 Jun 2026 07:48:33 +0100 Subject: [PATCH 5/7] Add skewed internode EP test controls --- ep/bench/test_internode.py | 48 ++++++++++++++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 5 deletions(-) diff --git a/ep/bench/test_internode.py b/ep/bench/test_internode.py index bda39ecac..b6f41f2a3 100644 --- a/ep/bench/test_internode.py +++ b/ep/bench/test_internode.py @@ -91,6 +91,14 @@ def test_main( ): # Settings num_tokens, hidden = args.num_tokens, args.hidden + if args.rank_num_tokens: + rank_num_tokens = [int(value) for value in args.rank_num_tokens.split(",")] + if len(rank_num_tokens) != num_ranks: + raise ValueError( + f"--rank-num-tokens has {len(rank_num_tokens)} entries, " + f"but world size is {num_ranks}" + ) + num_tokens = rank_num_tokens[rank] num_topk_groups, num_topk, num_experts = ( args.num_topk_groups, args.num_topk, @@ -103,6 +111,8 @@ def test_main( f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}", flush=True, ) + if args.rank_num_tokens: + print(f"[config] rank_num_tokens={rank_num_tokens}", flush=True) # Random data x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank @@ -193,7 +203,16 @@ def test_main( rdma_buffer_size, nvl_buffer_size = 512, (720 if num_ranks in (144, 160) else 512) if num_ranks == 24: nvl_buffer_size = 540 - config = Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size) + dispatch_config = ( + Buffer.get_dispatch_config(num_ranks) + if args.use_default_configs + else Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size) + ) + combine_config = ( + Buffer.get_combine_config(num_ranks) + if args.use_default_configs + else dispatch_config + ) # Test dispatch # noinspection PyShadowingNames @@ -220,7 +239,7 @@ def check_data(check_x, recv_gbl_rank_prefix_sum): "num_tokens_per_rdma_rank": num_tokens_per_rdma_rank, "is_token_in_rank": is_token_in_rank, "num_tokens_per_expert": num_tokens_per_expert, - "config": config, + "config": dispatch_config, "async_finish": async_mode, } if with_topk: @@ -293,7 +312,7 @@ def check_data(check_x, recv_gbl_rank_prefix_sum): check_data(recv_topk_weights, recv_gbl_rank_prefix_sum) # Test `num_worst_tokens != 0` - if with_topk: + if with_topk and not args.skip_worst_tokens: num_worst_tokens = num_tokens * num_ranks dispatch_args.update({"num_worst_tokens": num_worst_tokens}) ( @@ -331,7 +350,7 @@ def check_data(check_x, recv_gbl_rank_prefix_sum): dispatch_args = { "x": current_x, "handle": handle, - "config": config, + "config": dispatch_config, "async_finish": async_mode, } if previous_mode: @@ -357,7 +376,7 @@ def check_data(check_x, recv_gbl_rank_prefix_sum): "x": recv_x, "bias": (bias_0, bias_1), "handle": handle, - "config": config, + "config": combine_config, "async_finish": async_mode, } if with_topk: @@ -680,6 +699,15 @@ def test_loop( parser.add_argument( "--num-tokens", type=int, default=4096, help="Number of tokens (default: 4096)" ) + parser.add_argument( + "--rank-num-tokens", + type=str, + default=None, + help=( + "Comma-separated token counts by global rank. Overrides " + "--num-tokens for skewed per-rank dispatch tests." + ), + ) parser.add_argument( "--hidden", type=int, default=7168, help="Hidden dimension size (default: 7168)" ) @@ -717,6 +745,16 @@ def test_loop( action="store_true", help="run only the first dispatch/combine correctness variant", ) + parser.add_argument( + "--skip-worst-tokens", + action="store_true", + help="skip the num_worst_tokens stress subcase", + ) + parser.add_argument( + "--use-default-configs", + action="store_true", + help="use Buffer.get_dispatch_config/get_combine_config for correctness", + ) parser.add_argument( "--fixed-dispatch-nvl-chunk", type=int, From 56609d8617bbca16cacef946459fd25942efbe03 Mon Sep 17 00:00:00 2001 From: fergus barratt Date: Wed, 10 Jun 2026 14:50:18 +0100 Subject: [PATCH 6/7] Revert "Fix UCCL EP device selection under node-local launchers" This reverts commit 4e0bd792913a751123c3b50a1a98f535e3857465. --- ep/bench/buffer.py | 83 +++++++--------------------------------------- ep/bench/utils.py | 17 ++++++++-- 2 files changed, 26 insertions(+), 74 deletions(-) diff --git a/ep/bench/buffer.py b/ep/bench/buffer.py index ae84568a0..0f23c8ca3 100644 --- a/ep/bench/buffer.py +++ b/ep/bench/buffer.py @@ -73,42 +73,6 @@ class Buffer: # TODO(MaoZiming): Reduce SMs. UCCL Proxy should reduce the usage of SMs. num_sms: int = 20 - @staticmethod - def _env_int(name: str, default: int = -1) -> int: - try: - return int(os.environ.get(name, default)) - except (TypeError, ValueError): - return default - - @staticmethod - def _infer_device_index(rank: int) -> int: - device_count = torch.cuda.device_count() - if device_count <= 0: - return torch.cuda.current_device() - if device_count == 1: - return 0 - - local_rank = Buffer._env_int("LOCAL_RANK") - local_world_size = Buffer._env_int("LOCAL_WORLD_SIZE") - if local_world_size == device_count and 0 <= local_rank < device_count: - return local_rank - - # In one-container-per-node launchers (for example vLLM mp workers under - # a single Slurm task), every local worker can inherit LOCAL_RANK=0 while - # still seeing all local GPUs. EP ranks are laid out contiguously per node, - # so rank modulo visible devices is the local CUDA device. - return int(rank) % device_count - - @staticmethod - def _infer_local_world_size() -> int: - device_count = torch.cuda.device_count() - local_world_size = Buffer._env_int("LOCAL_WORLD_SIZE") - if local_world_size <= 0 or ( - device_count > 1 and local_world_size < device_count - ): - return device_count - return local_world_size - def __init__( self, group: dist.ProcessGroup, @@ -142,12 +106,10 @@ def __init__( is_intranode: whether to force intranode-only proxy mode. If set to `None`, infer it from the process-group topology automatically. Explicit `True` is rejected when the group spans multiple nodes. """ - self.rank = group.rank() - self.group_size = group.size() - self.group = group - self.device_index = self._infer_device_index(self.rank) - torch.cuda.set_device(self.device_index) - device_index = self.device_index + if "LOCAL_RANK" in os.environ: + device_index = int(os.environ["LOCAL_RANK"]) + else: + device_index = torch.cuda.current_device() if hasattr(ep, "get_rdma_buffer"): # Allocate outside PyTorch's CUDA allocator so RDMA/IPC sees a raw @@ -188,12 +150,12 @@ def __init__( ) rdma_buffer_ptr = self.scratch.data_ptr() - _local_world = self._infer_local_world_size() + _local_world = int(os.environ.get("LOCAL_WORLD_SIZE", -1)) self.proxies, self.workers = initialize_uccl( rdma_buffer_ptr, num_rdma_bytes, - self.rank, - self.group_size, + group.rank(), + dist.get_world_size(group), group, use_normal_mode=not low_latency_mode, is_intranode=is_intranode, @@ -202,12 +164,14 @@ def __init__( check_nvlink_connections(group) # Initialize the CPP runtime + self.rank = group.rank() + self.group_size = group.size() + self.group = group self.num_nvl_bytes = num_nvl_bytes self.num_rdma_bytes = num_rdma_bytes self.low_latency_mode = low_latency_mode self.explicitly_destroy = explicitly_destroy self._next_low_latency_combine_buffer = None - torch.cuda.set_device(self.device_index) self.runtime = ep.Buffer( self.rank, self.group_size, @@ -257,23 +221,11 @@ def __init__( for proxy in self.proxies: proxy.set_atomic_buffer_ptr(self.proxies[0].get_atomic_buffer_ptr()) - def _set_current_device(self, device: Optional[torch.device] = None) -> int: - if device is None: - device_index = self.device_index - else: - device = torch.device(device) - if device.type != "cuda": - return self.device_index - device_index = self.device_index if device.index is None else device.index - torch.cuda.set_device(device_index) - return int(device_index) - def _ll_compute_stream_ptr(self, device: torch.device): """ Return the current CUDA stream pointer for low-latency runtime calls. """ - device_index = self._set_current_device(device) - current = torch.cuda.current_stream(device=device_index) + current = torch.cuda.current_stream(device=device) return int(current.cuda_stream) def reset_rdma_buffer(self): @@ -393,7 +345,6 @@ def low_latency_dispatch( event: the event after executing the kernel (valid only if `async_finish` is set). hook: the receiving hook function (valid only if `return_recv_hook` is set). """ - self._set_current_device(x.device) for proxy in self.proxies: proxy.notify_proxy_thread_adaptive_sleeper() proxy.calculate_and_set_dispatch_recv_data_offset( @@ -563,7 +514,6 @@ def low_latency_combine( event: the event after executing the kernel (valid only if `async_finish` is set). hook: the receiving hook function (valid only if `return_recv_hook` is set). """ - self._set_current_device(x.device) if overlap: raise NotImplementedError( "low_latency_combine(overlap=True) is not implemented yet. " @@ -661,11 +611,10 @@ def get_next_low_latency_combine_buffer(self, handle: object): num_ranks = self.group.size() num_local_experts = num_experts // num_ranks num_recv_tokens = num_ranks * num_max_dispatch_tokens_per_rank - self._set_current_device(src_info.device) self._next_low_latency_combine_buffer = torch.empty( (num_local_experts, num_recv_tokens, hidden), dtype=torch.bfloat16, - device=src_info.device, + device="cuda", ) return self._next_low_latency_combine_buffer @@ -699,10 +648,8 @@ def get_comm_stream(self) -> torch.Stream: Returns: stream: the communication stream. """ - self._set_current_device() ts = self.runtime.get_comm_stream() if isinstance(ts, torch.Stream): - torch.cuda.set_device(ts.device_index) return torch.cuda.Stream( stream_id=ts.stream_id, device_index=ts.device_index, @@ -872,7 +819,6 @@ def get_dispatch_layout( if allocate_on_comm_stream: assert previous_event is not None and async_finish - self._set_current_device(topk_idx.device) alloc_ctx = ( torch.cuda.stream(self.get_comm_stream()) if allocate_on_comm_stream @@ -1007,8 +953,6 @@ def dispatch( """ # Default config config = self.get_dispatch_config(self.group_size) if config is None else config - x_for_device = x[0] if isinstance(x, tuple) else x - self._set_current_device(x_for_device.device) # Internode if self.runtime.get_num_rdma_ranks() > 1: @@ -1334,7 +1278,6 @@ def combine( """ # Default config config = self.get_combine_config(self.group_size) if config is None else config - self._set_current_device(x.device) # Internode if self.runtime.get_num_rdma_ranks() > 1: @@ -1478,7 +1421,6 @@ def internode_dispatch( assert config is not None x, x_scales = x if isinstance(x, tuple) else (x, None) - self._set_current_device(x.device) num_scales = ( 0 if x_scales is None else (1 if x_scales.dim() == 1 else x_scales.size(1)) ) @@ -1835,7 +1777,6 @@ def internode_combine( Normally, you should not directly call this function. """ assert config is not None - self._set_current_device(x.device) # Unpack handle and bias ( diff --git a/ep/bench/utils.py b/ep/bench/utils.py index 0504e5b0e..461cc1ea7 100644 --- a/ep/bench/utils.py +++ b/ep/bench/utils.py @@ -111,7 +111,10 @@ def detect_group_topology(group: dist.ProcessGroup) -> Tuple[int, int, int, bool num_nodes: number of distinct nodes spanned by the group. is_intranode: whether all ranks in the group are on the same node. """ - local_rank = torch.cuda.current_device() + if "LOCAL_RANK" in os.environ: + local_rank = int(os.environ["LOCAL_RANK"]) + else: + local_rank = torch.cuda.current_device() node_token = ( os.environ.get("NODE_RANK") @@ -157,7 +160,11 @@ def get_cpu_proxies_meta(proxies, rank, scratch_ptr, scratch_bytes, num_ranks, g "listen_ports": [proxy.get_listen_port() for proxy in proxies], } all_meta = [None] * num_ranks - device_index = torch.cuda.current_device() + # Use current device or fallback to LOCAL_RANK or 0 + if "LOCAL_RANK" in os.environ: + device_index = int(os.environ["LOCAL_RANK"]) + else: + device_index = torch.cuda.current_device() torch.cuda.set_device(device_index) dist.all_gather_object(all_meta, meta, group=group) rank2meta = {m["rank"]: m for m in all_meta} @@ -629,7 +636,11 @@ def initialize_uccl( def destroy_uccl(proxies, workers): - device_index = torch.cuda.current_device() + # Use current device or fallback to LOCAL_RANK + if "LOCAL_RANK" in os.environ: + device_index = int(os.environ["LOCAL_RANK"]) + else: + device_index = torch.cuda.current_device() if workers is not None: try: From c02fb08ad81bb25aa975f29afe77812d79afabba Mon Sep 17 00:00:00 2001 From: fergus barratt Date: Wed, 10 Jun 2026 14:50:41 +0100 Subject: [PATCH 7/7] ep: use the current CUDA device; never read LOCAL_RANK In one-container-per-node launchers (vLLM mp workers under a single Slurm task) every local worker can inherit LOCAL_RANK=0 while seeing all GPUs, so reading it for device or proxy-metadata decisions binds everything to cuda:0. The launcher owns placement and sets the current device before constructing the Buffer (ours also unsets the inherited Slurm env in launch.sh); use torch.cuda.current_device() everywhere, matching upstream DeepEP, which never reads LOCAL_RANK. --- ep/bench/buffer.py | 5 +---- ep/bench/utils.py | 17 +++-------------- 2 files changed, 4 insertions(+), 18 deletions(-) diff --git a/ep/bench/buffer.py b/ep/bench/buffer.py index 0f23c8ca3..6e65bcd2e 100644 --- a/ep/bench/buffer.py +++ b/ep/bench/buffer.py @@ -106,10 +106,7 @@ def __init__( is_intranode: whether to force intranode-only proxy mode. If set to `None`, infer it from the process-group topology automatically. Explicit `True` is rejected when the group spans multiple nodes. """ - if "LOCAL_RANK" in os.environ: - device_index = int(os.environ["LOCAL_RANK"]) - else: - device_index = torch.cuda.current_device() + device_index = torch.cuda.current_device() if hasattr(ep, "get_rdma_buffer"): # Allocate outside PyTorch's CUDA allocator so RDMA/IPC sees a raw diff --git a/ep/bench/utils.py b/ep/bench/utils.py index 461cc1ea7..0504e5b0e 100644 --- a/ep/bench/utils.py +++ b/ep/bench/utils.py @@ -111,10 +111,7 @@ def detect_group_topology(group: dist.ProcessGroup) -> Tuple[int, int, int, bool num_nodes: number of distinct nodes spanned by the group. is_intranode: whether all ranks in the group are on the same node. """ - if "LOCAL_RANK" in os.environ: - local_rank = int(os.environ["LOCAL_RANK"]) - else: - local_rank = torch.cuda.current_device() + local_rank = torch.cuda.current_device() node_token = ( os.environ.get("NODE_RANK") @@ -160,11 +157,7 @@ def get_cpu_proxies_meta(proxies, rank, scratch_ptr, scratch_bytes, num_ranks, g "listen_ports": [proxy.get_listen_port() for proxy in proxies], } all_meta = [None] * num_ranks - # Use current device or fallback to LOCAL_RANK or 0 - if "LOCAL_RANK" in os.environ: - device_index = int(os.environ["LOCAL_RANK"]) - else: - device_index = torch.cuda.current_device() + device_index = torch.cuda.current_device() torch.cuda.set_device(device_index) dist.all_gather_object(all_meta, meta, group=group) rank2meta = {m["rank"]: m for m in all_meta} @@ -636,11 +629,7 @@ def initialize_uccl( def destroy_uccl(proxies, workers): - # Use current device or fallback to LOCAL_RANK - if "LOCAL_RANK" in os.environ: - device_index = int(os.environ["LOCAL_RANK"]) - else: - device_index = torch.cuda.current_device() + device_index = torch.cuda.current_device() if workers is not None: try: