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/test_internode.py b/ep/bench/test_internode.py index d23472512..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: @@ -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 @@ -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)" ) @@ -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", @@ -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, 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: diff --git a/ep/bench/vllm/disagg_proxy.py b/ep/bench/vllm/disagg_proxy.py index bd008351c..e3f887fe7 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,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"] = { + "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. @@ -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() + 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={"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")