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

[fergusfinn]

Summary

vLLM serving integration fixes for UCCL EP, found while serving
DeepSeek-V4-Flash on Isambard (GH200/CXI) with vLLM's
deepep_high_throughput all2all backend.

Python-only and CXI-agnostic — nothing here touches the CXI transport.
The PR is based on cxi-ep only because that is the branch we serve from;
all of it should apply to the verbs/EFA paths too.

• EP device selection: the Python wrapper now uses the current CUDA
  device and never reads LOCAL_RANK. Under node-local launchers (one
  container per node, several worker processes inside) every vLLM worker can
  inherit LOCAL_RANK=0, so all local workers were initializing UCCL
  resources against cuda:0.
• Disaggregated serving proxy: align request handling with vLLM's NIXL
  contract (forward request IDs and expected remote decode metadata) and
  propagate upstream prefill/decode errors instead of converting failures
  into 200 OK zero-token responses — previously vllm bench serve could
  report misleading successes.
• test_internode.py controls for reproducing vLLM-shaped traffic:
  skewed-token routing and an SM-count override.

Note: earlier iterations of this branch also carried two C++ fixes (the
zero-token clear extent and the host atomic-buffer memset) and the
zero-layout regression guard test; those have graduated into the cxi-ep
base branch (#1) and are no longer part of this diff.

Validation

Isambard GH200/CXI build:

USE_CXI=1 UCCL_NUM_MAX_NVL_PEERS=4 TORCH_CUDA_ARCH_LIST=9.0 python3 setup.py build_ext --inplace

Two-node high-throughput internode smoke:

python3 ep/bench/test_internode.py \
  --num-tokens=128 --hidden=1024 --num-experts=64 --num-topk=4 \
  --smoke-one --skip-worst-tokens

vLLM serving validation (2 GH200 nodes, DeepSeek-V4-Flash, DP_GLOBAL=8,
DP_LOCAL=4, TP=1, expert parallel, --all2all-backend deepep_high_throughput; two endpoints each driven by vllm bench serve,
ISL/OSL 1024/1024, 2048 prompts at concurrency 2048):

• node0: 2048 successful, 0 failed, 3439.65 output tok/s
• node1: 2048 successful, 0 failed, 3443.19 output tok/s
• combined: 4096/4096 successful, 6882.84 output tok/s
  (860.36 output tok/s/GPU)
• both health endpoints OK after the benchmark; post-run log scan found no
  EngineDead, OutOfMemory, DeepEP.*timeout, worker death, or traceback
  signatures.

Remaining scope

This validates the EP8 high-concurrency shape. It does not claim every
EP16/4-node serving shape is fully validated: the latest EP16 attempt failed
during startup warmup (OOM) before any client traffic, tracked separately.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, add credits to your account and enable them for code reviews in your settings.

[doubleword-code]

Summary

This PR addresses vLLM serving integration fixes across three main areas: (1) device selection for multi-worker launchers, (2) disagg proxy error handling and request contract alignment, and (3) test controls for skewed internode EP scenarios. The changes are well-scoped and address real issues seen in production deployments.

Verdict: Needs minor changes before merge. The core fixes are sound, but there are a few edge cases and documentation gaps that should be addressed.

Research notes

• Fetched DeepEP upstream patterns for device selection: the change to use torch.cuda.current_device() instead of reading LOCAL_RANK aligns with how vLLM and other modern launchers work (they set the CUDA device before constructing worker objects).
• The zero-token layout guard regression test (test_zero_layout_guard.py) validates the fix in uccl_ep.cc:545 where num_ranks was incorrectly used instead of num_rdma_ranks when clearing the RDMA rank count buffer.
• The disagg proxy changes align with vLLM's NIXL connector expectations for kv_transfer_params structure.

Suggested next steps

1. Blocking: Verify the num_rdma_ranks variable is in scope at the cudaMemsetAsync call site in uccl_ep.cc (line 545).
2. Non-blocking: Add a comment in buffer.py explaining why torch.cuda.current_device() is preferred over LOCAL_RANK for device selection.
3. Non-blocking: Consider adding unit tests for the disagg proxy error paths (e.g., prefill returning 500, missing kv_transfer_params).
4. Nit: Add trailing newline to ep/src/proxy.cpp.

General findings

1. Device selection strategy (buffer.py, utils.py)

The change from reading os.environ["LOCAL_RANK"] to using torch.cuda.current_device() is correct for one-container-per-node launchers where vLLM sets the CUDA device before constructing the Buffer. However, this assumes the caller has already called torch.cuda.set_device() before Buffer initialization. The commit message mentions this, but adding an assertion or documentation in the code would help future maintainers.

2. Zero-token RDMA rank count clearing (uccl_ep.cc:545)

Critical bug fix: when num_tokens=0, the code was clearing num_ranks * sizeof(int) bytes instead of num_rdma_ranks * sizeof(int). This could cause memory corruption when num_rdma_ranks < num_ranks (multi-node setups). The regression test in test_zero_layout_guard.py properly validates this fix with a guard value pattern.

3. Disagg proxy error handling (disagg_proxy.py)

Significant improvement in error propagation:

• Properly forwards HTTP status codes and content types from prefill/decode back to clients
• Handles both JSON and non-JSON error responses
• Preserves Authorization headers in proxied requests
• Adds request ID tracking for debugging

One gap: streaming error handling closes the session properly, but non-streaming errors don't explicitly close the session (though Python's GC will eventually clean it up).

4. CXI transport backend integration

The USE_CXI gated changes in proxy.cpp, rdma.cpp, and new cxi_transport.* files implement a libfabric-based transport for CXI/Slingshot NICs. The code follows RAII patterns and properly handles error cases. The barrier implementation using atomic_add64 on host-pinned memory is clever but adds complexity.

5. Test control additions (test_internode.py)

Useful additions for performance tuning and skewed load testing:

• --rank-num-tokens for per-rank token distribution control
• --fixed-dispatch-* and --fixed-combine-* for targeted benchmarking
• --smoke-one for quick correctness validation
• Support for 4-GPU nodes (relaxed assertion from num_local_ranks == 8 to in (4, 8))

6. Low-latency kernel refactoring (internode_ll.cu)

The extraction of recv_via_local_or_ipc boolean improves readability and reduces code duplication in the dispatch/combine receive loops. The logic is equivalent to the original but more maintainable.

---

Overall, this is a solid set of fixes that address real production issues. The most critical change is the zero-token memset size fix, which is properly tested. The device selection change requires trust in the launcher setup but matches industry patterns.

[doubleword-code]

Summary

This PR delivers critical fixes for vLLM EP serving integration across three main areas:

1. Device selection correctness: Replaces LOCAL_RANK environment variable reads with torch.cuda.current_device(), fixing multi-worker launcher scenarios where all workers inherit identical LOCAL_RANK values.
2. Disaggregation proxy alignment: Updates disagg_proxy.py to match vLLM's NIXL connector expectations for kv_transfer_params, adds proper error propagation, request ID/auth header forwarding, and streaming response handling.
3. Zero-token memory safety: Fixes an off-by-one bug in uccl_ep.cc:get_dispatch_layout() that cleared num_ranks integers instead of num_rdma_ranks integers for the RDMA rank count buffer when num_tokens == 0.

Additional enhancements include a new CXI/Slingshot transport backend (USE_CXI), low-latency internode path refinements for IPC detection, and expanded test controls for skewed token distributions.

Verdict: Needs changes — two blocking issues identified in streaming resource cleanup and initialization ordering must be addressed before merge.

Research notes

• Reviewed vLLM contribution guidelines and KV transfer design documents. The kv_transfer_params schema used in disagg_proxy.py aligns with the NIXL connector pattern where do_remote_prefill=True signals the decode node to RDMA-read KV cache from prefill.
• The zero-token bug (uccl_ep.cc:546) is confirmed: when num_rdma_ranks < num_ranks (typical multi-node case with 8 GPUs per node), clearing num_ranks integers would write beyond the allocated num_rdma_ranks buffer boundary.
• CXI transport uses libfabric CQ polling with sched_yield() backoff after 1024 spins, which is appropriate for low-latency paths.

Suggested next steps

1. Blocking: Fix streaming session leak in disagg_proxy.py lines 148–159 — wrap the async generator in try/finally to ensure decode_session.close() is called even if iteration fails.
2. Blocking: Add assertion or comment in uccl_ep.cc:543–546 confirming num_rdma_ranks is initialized before the zero-token memset — trace through constructor initialization order.
3. Non-blocking: Document rationale for torch.cuda.current_device() in buffer.py:109 with a brief comment about launcher-owned placement.
4. Non-blocking: Extend test_zero_layout_guard.py to cover the multi-node case explicitly where num_ranks > num_rdma_ranks.

General findings

• The revert chain (commit e527ffa → 56609d8 → c02fb08) shows this device selection issue required multiple iterations; the final approach matches upstream DeepEP's behavior.
• CXI barrier uses dedicated slots in the atomic buffer's last 4KB, isolated from the standard barrier path — no interference expected.
• The skip_network_peer_on_same_ip() helper correctly handles the CXI edge case where host-pinned RDMA buffers prevent CUDA IPC even for same-node peers.
• All changed SMEM/assertion patterns in internode.cu and internode_ll.cu maintain backward compatibility via NUM_MAX_NVL_PEERS <= 8 relaxation.

General findings (auto-demoted from inline due to pre-validation)

• Blocking ep/bench/vllm/disagg_proxy.py:150 — Resource leak in streaming error path.
  • (demoted: code self-check failed at ep/bench/vllm/disagg_proxy.py:150: diff has async def stream_decode():, model claimed async for chunk in decode_resp.content.iter_any():)
• Blocking ep/src/uccl_ep.cc:546 — Potential use of uninitialized num_rdma_ranks in zero-token path.
  • (demoted: path "ep/src/uccl_ep.cc" is not in the PR diff)

[doubleword-code]

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]

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")

[doubleword-code]

Summary

This PR addresses vLLM serving integration with three primary focus areas: (1) device selection fixes for one-container-per-node launchers, (2) disagg proxy error handling and request contract improvements, and (3) a zero-token layout guard regression test with the corresponding runtime fix.

Verdict: Needs changes — The PR contains an important correctness fix (uccl_ep.cc zero-token RDMA buffer clearing), but has a blocking issue in the disagg_proxy.py streaming error path that could cause resource leaks.

Research notes

• PyTorch torch.cuda.current_device(): Returns the index of the currently selected device. Requires torch.cuda.set_device() to be called beforehand; otherwise defaults to device 0. This makes the device selection change correct but introduces an ordering dependency.

• Zero-token bug: The original code cleared num_ranks * sizeof(int) bytes when it should have cleared num_rdma_ranks * sizeof(int) bytes. For a typical 2-node × 4-GPU setup, this would write 8 ints instead of 2, corrupting adjacent memory. The test in test_zero_layout_guard.py correctly validates this with guard values.

• Config parameters: The Config struct has 5 parameters: num_sms, num_max_nvl_chunked_send_tokens, num_max_nvl_chunked_recv_tokens, num_max_rdma_chunked_send_tokens, num_max_rdma_chunked_recv_tokens. The EP8 combine config change (6→8 for rdma send tokens) is valid per constructor assertions.

Suggested next steps

1. Blocking: Fix the streaming session cleanup in disagg_proxy.py by using async with aiohttp.ClientSession() instead of manual session management.

2. Non-blocking: Consider adding a defensive check or documentation in buffer.py and utils.py noting that torch.cuda.set_device() must be called before these functions.

3. Non-blocking: Verify the ctx_.num_local_ranks initialization order in proxy.cpp to ensure the barrier assertion cannot fail spuriously.

General findings

Correctness

• uccl_ep.cc line 546: Critical bug fix — changed from num_ranks * sizeof(int) to num_rdma_ranks * sizeof(int) when clearing the RDMA rank count buffer in the zero-token case. This prevents memory corruption beyond the intended buffer bounds.

• test_zero_layout_guard.py: Well-designed regression test that places guard values around the num_tokens_per_rdma_rank buffer and verifies they weren't clobbered.

Security / Error handling

• disagg_proxy.py streaming path: When decode response returns non-200 status, error handling properly reads body and closes session. However, if the POST request itself throws before returning a response, the session leaks.

Performance

• buffer.py line 774: EP8 combine config change increases num_max_rdma_chunked_send_tokens from 6 to 8. This is a tuning adjustment that should be validated with benchmarks.

Code quality

• proxy.cpp barrier logic: Replaces hard-coded MAX_NUM_GPUS with ctx_.num_local_ranks to support 4-GPU and 8-GPU nodes. The assertion assert(ctx_.num_local_ranks <= 0 || r % ctx_.num_local_ranks == 0) assumes proper initialization.

• uccl_proxy.cpp atomic buffer: Added proper error handling for all allocation paths (cudaHostAlloc, cudaMallocManaged, etc.) with descriptive error messages via std::runtime_error.

Testing

• test_internode.py: Adds --rank-num-tokens for skewed token distribution testing and --smoke-one for quick sanity checks. Good additions for debugging and performance isolation.

• test_zero_layout_guard.py: Should be run as part of multi-node CI to prevent regressions in the zero-token dispatch path.

General findings (auto-demoted from inline due to pre-validation)

• Non-blocking ep/src/uccl_ep.cc:546 — Critical correctness fix — zero-token RDMA buffer clearing.
  • (demoted: path "ep/src/uccl_ep.cc" is not in the PR diff)
• Non-blocking ep/bench/utils.py:49 — Explicitly setting device in init_dist() is correct.
  • (demoted: line 49 (side=RIGHT) is not part of any diff hunk in ep/bench/utils.py)
• Non-blocking ep/bench/test_zero_layout_guard.py:58 — Test design correctly validates the zero-token fix.
  • (demoted: path "ep/bench/test_zero_layout_guard.py" is not in the PR diff)
• Nit ep/bench/buffer.py:774 — EP8 combine config parameter change.
  • (demoted: line 774 (side=RIGHT) is not part of any diff hunk in ep/bench/buffer.py)
• Non-blocking ep/src/proxy.cpp:1552 — Barrier assertion depends on proper initialization.
  • (demoted: path "ep/src/proxy.cpp" is not in the PR diff)

[doubleword-code]

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.

[doubleword-code]

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]

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.

[doubleword-code]

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.

[doubleword-code]

Summary

This PR introduces several important fixes and a major new feature (CXI/libfabric transport support for Grace Hopper systems with Slingshot interconnects). The key changes include:

1. Device selection fix: Replaced LOCAL_RANK environment variable reads with torch.cuda.current_device() for more robust device selection across different launch environments.
2. Zero-token layout guard fix: Corrected a buffer overflow bug where num_ranks * sizeof(int) was used instead of num_rdma_ranks * sizeof(int) when zeroing output buffers in the zero-token case.
3. Disaggregation proxy improvements: Enhanced error propagation, added request ID tracking, and fixed streaming response resource cleanup.
4. CXI transport backend: Added comprehensive libfabric-based RDMA transport support for CXI/Slingshot networks.
5. Configurable NVL peers: Made NUM_MAX_NVL_PEERS configurable via environment variable for 4-GPU node configurations.

The code quality is generally high with proper error handling and thoughtful design. However, there are a few issues that need attention before merging.

Verdict: Needs changes (see Blocking findings below).

Research notes

• Consulted libfabric documentation for fi_getinfo(), fi_mr_regattr(), and FI_MR_ENDPOINT semantics. The implementation correctly handles memory registration with endpoint binding when required.
• Reviewed OWASP guidance on error handling - the disagg_proxy improvements align with best practices for not leaking internal error details while providing sufficient debugging information.
• Cross-referenced CUDA memory allocation patterns - the conditional memset/cudaMemset based on atomic_buffer_is_host_allocated_ is correct.

Suggested next steps

1. Address the Blocking finding in ep/src/uccl_ep.cc - verify the variable name num_rdma_ranks is available in that scope (it appears to be based on the function signature).
2. Consider adding a comment explaining why NUM_MAX_NVL_PEERS <= 8 assertion was relaxed from == 8.
3. Verify the CXI barrier slot initialization is guaranteed before first use in all code paths.

General findings

Positive observations

• The zero-token guard regression test (test_zero_layout_guard.py) is well-designed with a guard value pattern that would catch buffer overruns.
• Error handling in uccl_proxy.cpp for atomic buffer allocation is now comprehensive across all platform variants (CUDA, HIP, EFA, host memory).
• The skip_network_peer_on_same_ip() helper cleanly abstracts the CXI vs non-CXI peer connection logic.
• Device selection using torch.cuda.current_device() is more robust than environment variable parsing, especially under torchrun/srun launchers.

Non-blocking observations

• The CXI barrier uses monotonically increasing counters without reset. While int64_t wraparound would require ~9.2 × 10¹⁸ barriers (effectively never in practice), a comment documenting this assumption would be helpful.
• The strdup() calls in cxi_transport.cpp:Transport::init() are correctly freed by fi_freeinfo(hints) in both success and exception paths - this is correct per libfabric semantics.
• The NUMA pinning logic cfg_.local_rank % 4 assumes 4 GPUs per NUMA node (correct for Grace Hopper), but this assumption isn't documented inline.

General findings (auto-demoted from inline due to pre-validation)

• Non-blocking ep/src/uccl_ep.cc:546 — This fix correctly addresses a buffer overflow bug in the zero-token path.
  • (demoted: path "ep/src/uccl_ep.cc" is not in the PR diff)
• Non-blocking ep/bench/vllm/disagg_proxy.py:154 — Proper session lifecycle management in streaming error path.
  • (demoted: code self-check failed at ep/bench/vllm/disagg_proxy.py:154: diff has finally:, model claimed await decode_session.close())
• Non-blocking ep/src/cxi_transport.cpp:65 — Memory management for strdup'd strings.
  • (demoted: path "ep/src/cxi_transport.cpp" is not in the PR diff)
• Non-blocking ep/include/ep_configs.cuh:4 — Making NUM_MAX_NVL_PEERS configurable is useful for 4-GPU node configurations.
  • (demoted: path "ep/include/ep_configs.cuh" is not in the PR diff)
• Nit ep/src/internode.cu:24 — Assertion relaxation from == 8 to <= 8.
  • (demoted: path "ep/src/internode.cu" is not in the PR diff)
• Non-blocking ep/src/proxy.cpp:91 — Dynamic stride calculation for normal mode peer connections.
  • (demoted: path "ep/src/proxy.cpp" is not in the PR diff)
• Non-blocking ep/src/uccl_proxy.cpp:120 — Correct memory initialization for host-allocated atomic buffers.
  • (demoted: path "ep/src/uccl_proxy.cpp" is not in the PR diff)

[doubleword-code]

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

[fergusfinn]

Closing this rather than rebasing it onto the PR997 line.

The serving-critical device-selection issue is worked around in the current pure-vLLM launcher by unsetting RANK/WORLD_SIZE/LOCAL_RANK before vLLM starts, so PR997 falls back to torch.cuda.current_device(). The remaining pieces here are either harness/test controls or the old disagg proxy path, not part of the current PR997 cleanup. If the LOCAL_RANK issue comes back, we should reopen it as a fresh one-commit PR against pr997-swiss-cxi.