Skip to content

ep: vLLM serving integration fixes (device selection, disagg proxy, test controls) #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
fergusfinn wants to merge 7 commits into from
Closed

ep: vLLM serving integration fixes (device selection, disagg proxy, test controls) #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
e527ffa
Fix UCCL EP device selection under node-local launchers
fergusfinn Jun 7, 2026
4379a42
Align disagg proxy with vLLM NIXL request contract
fergusfinn Jun 7, 2026
dae27f6
Propagate disagg proxy upstream errors
fergusfinn Jun 7, 2026
35d633c
Allow internode test SM override
fergusfinn Jun 7, 2026
a8b4fd4
Add skewed internode EP test controls
fergusfinn Jun 7, 2026
56609d8
Revert "Fix UCCL EP device selection under node-local launchers"
fergusfinn Jun 10, 2026
c02fb08
ep: use the current CUDA device; never read LOCAL_RANK
fergusfinn Jun 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 1 addition & 4 deletions ep/bench/buffer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: Missing rationale for device selection change.

Why it matters: This line replaces the previous LOCAL_RANK environment variable read. Future maintainers might wonder why LOCAL_RANK was removed, especially since it's still commonly used in other distributed training frameworks. Without context, someone might reintroduce the old logic when debugging launcher issues.

Suggested fix: Add a brief comment explaining the motivation:

# Use current CUDA device rather than LOCAL_RANK: in one-container-per-node
# launchers (e.g., vLLM mp workers under Slurm), all workers may inherit
# LOCAL_RANK=0 while seeing all GPUs. The launcher sets the current device
# before constructing the Buffer, matching upstream DeepEP behavior.
device_index = torch.cuda.current_device()

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: Device selection now relies on torch.cuda.current_device() instead of LOCAL_RANK env var.

Why it matters: Per PyTorch docs, current_device() returns the currently selected device (default 0 if set_device() wasn't called). This is correct for one-container-per-node launchers where the launcher sets the device before construction. However, it introduces an ordering dependency: init_dist() must be called before Buffer.__init__().

Suggested fix: Add a comment documenting this requirement, or add a defensive check:

device_index = torch.cuda.current_device()
# Note: caller must ensure torch.cuda.set_device() was called first

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: Good fix for device selection robustness.

Why it matters: Reading LOCAL_RANK from the environment fails under launch mechanisms that don't set this variable (e.g., some srun configurations or custom launchers). Using torch.cuda.current_device() defers to PyTorch's device management, which is set correctly by torch.distributed.run via the device_id parameter.

Suggested fix: This change is correct. Consider applying the same pattern to other files that read LOCAL_RANK (e.g., utils.py:detect_group_topology already uses torch.cuda.current_device()).


if hasattr(ep, "get_rdma_buffer"):
# Allocate outside PyTorch's CUDA allocator so RDMA/IPC sees a raw
Expand Down
58 changes: 52 additions & 6 deletions ep/bench/test_internode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,14 @@ def test_main(
):
# Settings
num_tokens, hidden = args.num_tokens, args.hidden
if args.rank_num_tokens:

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: Consider validating rank_num_tokens entries are non-negative.

Why it matters: While the test is primarily for internal validation, negative token counts could cause confusing downstream errors in tensor allocation or kernel launches. A simple validation would provide clearer feedback.

Suggested fix: Add after line 96:

if any(n < 0 for n in rank_num_tokens):
    raise ValueError("--rank-num-tokens values must be non-negative")

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,
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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})
(
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand Down Expand Up @@ -581,7 +600,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
Expand Down Expand Up @@ -678,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)"
)
Expand All @@ -699,6 +729,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",
Expand All @@ -709,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,
Expand Down
17 changes: 3 additions & 14 deletions ep/bench/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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}
Expand Down Expand Up @@ -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:
Expand Down
100 changes: 82 additions & 18 deletions ep/bench/vllm/disagg_proxy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand All @@ -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
Expand All @@ -42,23 +55,39 @@ 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"] = {

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: Expanded kv_transfer_params structure aligns with vLLM NIXL connector expectations.

Why it matters: The additional fields (do_remote_prefill, remote_engine_id, remote_block_ids, remote_host, remote_port) provide explicit control over KV cache transfer behavior. Setting them to None initially allows the prefill node to populate them.

Suggested fix: None required — this matches the expected vLLM request contract.

"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()
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.
Expand Down Expand Up @@ -87,30 +116,65 @@ 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)

if is_stream:
decode_session = aiohttp.ClientSession()

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: Manual session management creates a resource leak risk.

Why it matters: If decode_session.post() throws an exception before returning decode_resp, the finally block in stream_decode() never executes, leaving the session unclosed. This can exhaust connection pool resources under sustained error conditions.

Suggested fix: Use async with for automatic cleanup:

async with aiohttp.ClientSession() as decode_session:
    decode_resp = await decode_session.post(...)
    # ... rest of streaming logic

Or wrap the entire streaming block in a try/finally that ensures await decode_session.close() is called.

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():

@doubleword-code doubleword-code Bot Jun 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: Streaming generator cleanup is correct but complex.

Why it matters: The try/finally in stream_decode() properly releases the response and closes the session. However, this pattern is error-prone — if future modifications forget the finally block, resources leak.

Suggested fix: Consider using asynccontextmanager to make the lifecycle more explicit, or consolidate with the non-streaming path to reduce duplication.

async with aiohttp.ClientSession() as s:
async with s.post(
f"{DECODE_URL}/v1/chat/completions",
json=decode_body,
headers={"Content-Type": "application/json"},
) 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",
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")
Expand Down