test: multimodal tensor-transport + EPD disaggregation coverage#1898
Open
slin1237 wants to merge 23 commits into
Open
test: multimodal tensor-transport + EPD disaggregation coverage#1898slin1237 wants to merge 23 commits into
slin1237 wants to merge 23 commits into
Conversation
…nning GPU-free regression net ahead of generalizing the multimodal tensor transport (crate extraction + adding the RDMA/remote payload path): - transport.rs: resolve_mm_shm_enabled matrix (inline/rdma -> off, shm -> dev-writable, auto -> shm-namespace match), previously untested. - proto_wrapper.rs: vLLM into_proto payload variant — inline vs shm (the shm arm was only exercised with shm_enabled=false before). - epd_encode.rs: extract plan_encode_jobs() (behavior-preserving) so the encode item/assignment count+order checks and bootstrap-info assembly are unit-testable without a live TokenSpeed client; add coverage for item assembly, count/order mismatch, and bootstrap host/port/room. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The multimodal e2e only exercised the default inline transport. Add a
Qwen3-VL/vLLM class that launches the gateway with
`--multimodal-tensor-transport shm` (gateway + worker are co-located in
CI, so they share /dev/shm) and asserts the pixel tensor actually
traveled over shm via `smg_mm_tensors_total{path="shm"}` — not that it
silently fell back to the inline gRPC payload.
Adds a small Prometheus `/metrics` scraper (`Gateway.metric_sum`) to the
e2e infra, since none existed.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Adds a live Encode-Prefill-Decode e2e: the encode worker runs the vision tower, prefill/decode run the LM, and the gateway stitches encode -> prefill -> decode. Covers the EPD path (TokenSpeed-only) that had no e2e. - harness: `WorkerType.ENCODE` + `--disaggregation-mode encode|prefill| decode` (encode/prefill get a mooncake bootstrap port); gateway EPD mode (`--epd-disaggregation --encode/--prefill/--decode`); `_setup_epd` mirrors `_setup_pd` with encode-first GPU offsets; `epd` topology marker. - model: Qwen/Qwen3.6-35B-A3B-FP8 — arch Qwen3_5MoeForConditionalGeneration is in TokenSpeed's multimodal registry, 3B active, FP8 (~35GB, Hopper- native so it runs on h100; NVFP4 would need Blackwell) fits one card at tp=1, so every topology (1e1p1d/1e2p1d/2e1p1d/1e1p2d) fits the 4-GPU lane. - test: the four topologies, one image, assert a correct description. - CI: a tokenspeed leg on e2e-4gpu-chat runs it per-PR. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
slin1237
requested review from
CatherineSue,
XinyueZhang369 and
key4ng
as code owners
github-actions
Bot
added
ci
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds TokenSpeed EPD disaggregation for multimodal workloads, including worker and gateway orchestration, EPD fixtures, encode planning, SHM transport validation, E2E coverage, and CI configuration. ChangesEPD multimodal disaggregation
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
👋 The PR description doesn't fully follow
Please update the PR description so reviewers have the context they need. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
Warning Gemini encountered an error creating the review. You can try again by commenting |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cf5966ef79
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@e2e_test/chat_completions/test_epd_multimodal.py`:
- Around line 39-42: Tighten the `_make_image_content` docstring to match its
actual behavior: it should only describe accepting an already-formed image URL
string, not a raw local path. Update the helper documentation in
`_make_image_content` so callers are directed to convert local files with
`_image_to_base64_url` before passing the result into this function.
In `@model_gateway/src/routers/grpc/epd_encode.rs`:
- Around line 186-196: Add a debug_assert in plan_encode_jobs to enforce the
documented non-empty items precondition before any planning logic runs. Use the
existing plan_encode_jobs symbol and check that items is not empty, so future
callers don’t rely only on build_plan’s guard. Keep the runtime behavior
unchanged; this is just a defensive invariant check that makes the documented
requirement explicit.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: db4de777-d85a-47a6-babe-619443403c58
📒 Files selected for processing (12)
.github/workflows/pr-test-rust.ymle2e_test/chat_completions/test_epd_multimodal.pye2e_test/chat_completions/test_multimodal.pye2e_test/fixtures/hooks.pye2e_test/fixtures/setup_backend.pye2e_test/infra/constants.pye2e_test/infra/gateway.pye2e_test/infra/model_specs.pye2e_test/infra/worker.pymodel_gateway/src/routers/grpc/epd_encode.rsmodel_gateway/src/routers/grpc/multimodal/transport.rsmodel_gateway/src/routers/grpc/proto_wrapper.rs
- setup_backend: pass the EPD topology via request.param, not a class-scoped `epd` marker. The per-param marks weren't visible on the class-scoped fixture, so 1e2p1d/2e1p1d/1e1p2d all silently ran 1/1/1 — the topology matrix wasn't actually testing different topologies. (Codex) - test_epd_multimodal: make it verify *real* EPD, not just that an answer came back (a single-worker fallback would pass that). Mirrors the PD KV-transfer tests: assert the encode worker itself logged accepting the request, and that the image was tokenized into the prompt. Adds the per-request accept log to the TokenSpeed encode servicer. - model_specs + ci_download_model: `skip_tier_download` so the ~35GB EPD model isn't pre-pulled by every tier lane; the tokenspeed 4-GPU lane downloads it by id. (Codex P1) - pr-test-rust: pass a larger pytest `test_timeout` for the EPD lane, so it isn't killed at the 25m default while relaunching worker sets. (Codex P2) - gateway: clearer multi-mode error message. (Claude) - tighten _make_image_content docstring. (CodeRabbit) Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a33e33dc17
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
e2e_test/fixtures/setup_backend.py (1)
329-413: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the shared launch/teardown skeleton with
_setup_pd.
_setup_epdduplicates the "start role workers at computed GPU offset ->_start_gateway-> yield ->finally: shutdown + stop_workers" structure from_setup_pd(Lines 260-321), just with an extra encode stage. A small helper (e.g., one that takes a list of(WorkerType, count, gpu_offset)tuples) would reduce duplication and keep future EPD/PD changes in sync.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@e2e_test/fixtures/setup_backend.py` around lines 329 - 413, _extract the shared launch/teardown flow currently duplicated between _setup_epd and _setup_pd into a small helper that handles starting worker groups from (WorkerType, count, gpu_offset) tuples, then calling _start_gateway, yielding, and performing the finally shutdown/stop_workers cleanup; keep _setup_epd focused on assembling encode/prefill/decode groups and passing them to the helper, using the existing _start_workers_tracked, _start_gateway, and stop_workers symbols to preserve behavior while reducing duplication.e2e_test/infra/gateway.py (1)
90-93: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winNo minimum-count validation for EPD worker lists.
is_epd_modetriggers onencode_workers is not None, but nothing enforces thatencode_workers/prefill_workers/decode_workersare non-empty before building--epd-disaggregationargs. Callingstart(encode_workers=[], prefill_workers=[], decode_workers=[])silently launches an EPD gateway with zero worker args, deferring the failure to an opaque subprocess/CLI error rather than a clear Python-sideValueError(as regular mode does at Lines 216-219 forworker_urls).💡 Proposed validation
elif is_epd_mode: self.pd_mode = False self.epd_mode = True self.igw_mode = False encodes = encode_workers or [] prefills = prefill_workers or [] decodes = decode_workers or [] + if not encodes or not prefills or not decodes: + raise ValueError( + "EPD mode requires non-empty encode_workers, prefill_workers, " + "and decode_workers" + )Also applies to: 166-199
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@e2e_test/infra/gateway.py` around lines 90 - 93, The EPD path in Gateway.start lacks the same minimum-count validation that regular mode applies to worker URLs. In the Gateway.start logic that builds the epd-disaggregation args from encode_workers, prefill_workers, and decode_workers, reject empty lists with a clear ValueError before launching the subprocess, using the existing validation pattern around worker_urls as a guide. Make sure the check runs before is_epd_mode leads into the EPD argument assembly so start(encode_workers=[], prefill_workers=[], decode_workers=[]) fails fast in Python rather than in the CLI.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@e2e_test/fixtures/setup_backend.py`:
- Around line 329-413: _extract the shared launch/teardown flow currently
duplicated between _setup_epd and _setup_pd into a small helper that handles
starting worker groups from (WorkerType, count, gpu_offset) tuples, then calling
_start_gateway, yielding, and performing the finally shutdown/stop_workers
cleanup; keep _setup_epd focused on assembling encode/prefill/decode groups and
passing them to the helper, using the existing _start_workers_tracked,
_start_gateway, and stop_workers symbols to preserve behavior while reducing
duplication.
In `@e2e_test/infra/gateway.py`:
- Around line 90-93: The EPD path in Gateway.start lacks the same minimum-count
validation that regular mode applies to worker URLs. In the Gateway.start logic
that builds the epd-disaggregation args from encode_workers, prefill_workers,
and decode_workers, reject empty lists with a clear ValueError before launching
the subprocess, using the existing validation pattern around worker_urls as a
guide. Make sure the check runs before is_epd_mode leads into the EPD argument
assembly so start(encode_workers=[], prefill_workers=[], decode_workers=[])
fails fast in Python rather than in the CLI.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0c4df285-6b6a-4f55-ac42-f78ed508260d
📒 Files selected for processing (8)
.github/workflows/e2e-gpu-job.yml.github/workflows/pr-test-rust.ymle2e_test/chat_completions/test_epd_multimodal.pye2e_test/fixtures/setup_backend.pye2e_test/infra/gateway.pye2e_test/infra/model_specs.pygrpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.pyscripts/ci_download_model.sh
The EPD e2e's encode worker died at startup with `argument --disaggregation-mode: invalid choice: 'encode'`: the pinned TokenSpeed SHA (5e145af) predated the EPD encode pipeline (tokenspeed #548, b5c762d), whose `--disaggregation-mode encode` the SMG encode servicer already relies on. Bump to current tokenspeed main (69091e1), which includes it. The embedding transfer backend defaults to mooncake, and encode mode requires DP=1 (satisfied by the tp=1 encode worker), so no extra worker args are needed. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 707fca4893
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
The bumped TokenSpeed's kernel failed to compile on CI's CUDA 13 with `error: '__cudaLaunch' was not declared` (trtllm_allreduce_fusion, sm_90a). Root cause matched against TokenSpeed's own test/ci_system/install_deps.sh (which builds cu130 successfully): CUDA 13 relocated the CCCL headers into `include/cccl`, and it sets `TOKENSPEED_KERNEL_BACKEND=cuda` on the kernel build — both of which SMG's install script was missing. Add the `include/cccl` dir to the C/C++ include path (first) and export `TOKENSPEED_KERNEL_BACKEND=cuda`, matching the upstream recipe. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a9d91f528
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
The kernel's `__cudaLaunch was not declared` failure persisted after adding the CCCL include path. Root cause is a torch/CUDA mismatch: the runner installs the CUDA 13.0 toolkit from scratch (nvcc 13), but pip pulled the default PyPI `torch-2.11.0` (a CUDA 12.x build). Its `nvidia-cuda-runtime-cu12` dependency puts a cu12 `crt/host_runtime.h` on the include path, so nvcc 13's cudafe++ generates a host stub that then fails to compile against the cu12 headers. TokenSpeed's own CI avoids this by running on a `cu130-torch-2.11.0` base image; its install_deps.sh only sets `PIP_EXTRA_INDEX_URL=.../cu130` as a fallback. Replicate that on the generic runner: point pip/uv at the cu130 wheel index and install `torch==2.11.0+cu130` explicitly before the `--no-build-isolation` kernel compile. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
scripts/ci_install_tokenspeed.sh (1)
121-144: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winPin torch in the constraints file too, for defense-in-depth.
TOKENSPEED_CONSTRAINTSonly pinsnvidia-cutlass-dsl. The subsequentuv pip install -e tokenspeed-kernel/python/ --no-build-isolation(line 145) and-e "./python"(line 147) resolve their own dependency trees; if either declares a torch spec that isn't satisfied by the already-installed2.11.0+cu130(e.g. an exact/narrower pin inrequirements/cuda.txt), the resolver could silently swap in a different torch build, reintroducing the exact CUDA-13/cu12-header mismatch this script exists to prevent. Adding torch to the same constraints file used fornvidia-cutlass-dslwould guard against that without extra installs.Proposed fix
TOKENSPEED_CONSTRAINTS="$(mktemp)" echo "nvidia-cutlass-dsl==4.5.2" > "$TOKENSPEED_CONSTRAINTS" +echo "torch==2.11.0+cu130" >> "$TOKENSPEED_CONSTRAINTS" export UV_CONSTRAINT="$TOKENSPEED_CONSTRAINTS" export PIP_CONSTRAINT="$TOKENSPEED_CONSTRAINTS"Worth checking TokenSpeed's
requirements/cuda.txtat the pinned ref to see if it declares its own torch version spec that could conflict.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@scripts/ci_install_tokenspeed.sh` around lines 121 - 144, Pin torch==2.11.0+cu130 in TOKENSPEED_CONSTRAINTS alongside nvidia-cutlass-dsl==4.5.2, before exporting UV_CONSTRAINT and PIP_CONSTRAINT. Verify the pinned TokenSpeed requirements do not introduce an incompatible torch constraint, so the subsequent editable installs cannot replace the CUDA 13 torch build.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@scripts/ci_install_tokenspeed.sh`:
- Around line 110-119: Persist the cu130 package indexes for subsequent CI
steps, not only the current shell. Update the environment setup in
scripts/ci_install_tokenspeed.sh, including the corresponding later block, to
append the resolved PIP_EXTRA_INDEX_URL and UV_EXTRA_INDEX_URL values to
GITHUB_ENV alongside the other CUDA variables; verify downstream installs in
e2e-gpu-job.yml inherit them.
---
Outside diff comments:
In `@scripts/ci_install_tokenspeed.sh`:
- Around line 121-144: Pin torch==2.11.0+cu130 in TOKENSPEED_CONSTRAINTS
alongside nvidia-cutlass-dsl==4.5.2, before exporting UV_CONSTRAINT and
PIP_CONSTRAINT. Verify the pinned TokenSpeed requirements do not introduce an
incompatible torch constraint, so the subsequent editable installs cannot
replace the CUDA 13 torch build.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: fbe3ca2a-19f2-4cb3-822f-4e3ef17f3bfe
📒 Files selected for processing (1)
scripts/ci_install_tokenspeed.sh
The `__cudaLaunch was not declared` failure survived both the CCCL include
path and the torch->cu130 switch. The real cause is a header-shadowing
conflict, visible in the compiler note:
site-packages/nvidia/cu13/include/crt/host_runtime.h:130:
note: macro "__cudaLaunch" defined here
torch's +cu130 wheels ship their own crt/host_runtime.h under
nvidia/cu13/include, and torch's CUDAExtension puts that dir on the nvcc
command line, shadowing nvcc's built-in system headers. The pip cu13 patch
level differs from the apt-installed system nvcc, so the stub the system nvcc
generates calls the 2-arg __cudaLaunch while the pip header only defines the
1-arg macro -> "not declared".
Fix, adopting the proven TRT-LLM lane's pattern:
- install the full cuda-toolkit-13-0 (like ci_install_trtllm.sh) so the system
headers are a complete, self-consistent match for the system nvcc, and
- set NVCC_PREPEND_FLAGS=-I$CUDA_HOME/include so the system include wins over
torch's bundled cu13 headers on every nvcc invocation (host stub included),
persisted to GITHUB_ENV for the worker-side JIT builds pytest triggers.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The `__cudaLaunch was not declared` failure was NOT an include-ordering
problem: the full nvcc command shows `-I/usr/local/cuda/include` already
precedes torch's `-I .../nvidia/cu13/include`, so the previous
NVCC_PREPEND_FLAGS attempt was a no-op (identical error). The compiler note
pinpoints the real cause:
.venv/.../nvidia/cu13/include/crt/host_runtime.h:130:
note: macro "__cudaLaunch" defined here (#define __cudaLaunch(fun), 1-arg)
torch's +cu130 wheels bundle a crt/host_runtime.h at a different CUDA patch
level than the apt system nvcc (cuda-13.0.88). nvcc feeds the bundled dir to
the host-stub compile ahead of the system include via its host-compiler flags
(which command-line -I ordering can't override), so the 13.0.88 nvcc emits a
2-arg __cudaLaunch stub the bundled 1-arg macro can't satisfy.
Fix: after installing the +cu130 torch, replace each bundled
site-packages/nvidia/cu*/include/crt with a symlink to the system
$CUDA_HOME/include/crt, so the launch stubs always resolve the crt headers
that match the nvcc generating them, regardless of include order. Drop the
ineffective NVCC_PREPEND_FLAGS and its misleading comment.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
setup.py's _resolve_include_dirs puts $CUDA_HOME/include ahead of torch's bundled nvidia/cu13/include, so for the kernel's host-stub compile to still resolve crt/host_runtime.h to the (mismatched) pip copy, the system header at $CUDA_HOME/include/crt/host_runtime.h must be absent. CUDA_HOME is /usr/local/cuda, and on this runner that symlink is stale/partial: its include/ has cuda_runtime.h (so setup.py adds the dir) but no crt/, so crt falls through to torch's cu13 crt whose 1-arg __cudaLaunch macro can't satisfy the 2-arg stub the 13.0.88 nvcc emits. Point CUDA_HOME at the freshly apt-installed, complete /usr/local/cuda-13.0 (ships cuda-crt-13-0) -- exactly what the working TRT-LLM lane does (ci_install_trtllm.sh). Add a diagnostic that asserts crt/host_runtime.h is present under CUDA_HOME. Drops the earlier crt-symlink block, which never ran: the pip crt headers are pulled by the kernel build itself, after that block. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Diagnostics disproved the include-order theory: with CUDA_HOME=/usr/local/ cuda-13.0 the system crt/host_runtime.h is present AND precedes torch's nvidia/cu13/include in setup.py's -I list, yet the host-stub compile still binds the bundled copy (nvcc resolves it via a path -I ordering can't control). The bundled crt is nvidia-cuda-runtime 13.0.96 -- a newer patch than the apt system nvcc (13.0.88) -- so the 88 nvcc's 2-arg __cudaLaunch stub can't match the 96 header's 1-arg macro. The bundled crt is pulled by the kernel build's own dependency resolution, so external include tweaks can't win and pre-build alignment misses it. Instead: run the kernel build once to materialize the deps (tolerate the expected compile failure), symlink every site-packages/nvidia/cu*/include/crt to the system $CUDA_HOME/include/crt, then build for real -- deps are satisfied so the crt isn't re-pulled, and the stub now compiles against the matching 13.0.88 headers. The site-packages dir is resolved via sysconfig (CWD-independent, since the build runs from the TokenSpeed checkout). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The EPD workers (encode/prefill/decode) loaded the model fine but then hung in mooncake transfer-engine init and never became healthy, timing out the 4-GPU tokenspeed lane. The worker logs show the transfer engine enumerating all 18 RoCE NICs on the H100 runner and logging "Local segment descriptor not found" before dead-ending. The SGLang PD branch already guards against this by passing --disaggregation-ib-device (from detect_ib_device()); the EPD branch detected and stored the device but never put it on the command line. Add the flag, mirroring the PD branch, so mooncake binds a single NIC instead of hanging on the full enumeration. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The EPD ib-device fix was a no-op because detect_ib_device() returned None: it shells out to `ibv_devinfo`, which isn't installed on the tokenspeed GPU runner (the image ships libibverbs but not ibverbs-utils), so the caller never added --disaggregation-ib-device and the mooncake transfer engine kept enumerating all 18 RoCE NICs and hanging. Read /sys/class/infiniband directly instead (works for InfiniBand and RoCE, no CLI dependency), scanning devices in numeric order and returning the first with an ACTIVE port. Keep ibv_devinfo as a fallback and log when nothing is found so a future None is diagnosable. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The EPD workers hang right after mooncake's transfer engine starts listening; the next step is register_memory on a GPU buffer, which needs GPUDirect (nvidia_peermem/dmabuf). Add a timeout-guarded probe to the tokenspeed install verify that dumps the runner's RDMA/GPUDirect state and runs the exact mooncake init + GPU register_memory TokenSpeed does, so the fast 1-GPU tokenspeed lane tells us definitively whether the H100 CI box supports it. Never hangs or fails the lane. Temporary — to be removed once the EPD transport path is settled. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a5e969146d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…RMEM=false) EPD workers loaded the model then hung in mooncake's transfer engine and never became healthy. Root-caused on a GB300 box: mooncake's register_memory fails on GPU buffers with "Bad address" (EFAULT) because it uses the legacy nvidia_peermem GPUDirect path (WITH_NVIDIA_PEERMEM defaults to true), and the NVIDIA Open Kernel driver on the GPU runners has no nvidia_peermem module. TokenSpeed's register() swallows the failure, so the worker deadlocks -> the 600s health timeout. mooncake compiles in the dmabuf path too; it's selected at runtime via WITH_NVIDIA_PEERMEM=false. With that set, GPU-buffer registration succeeds on the open driver (verified on GB300: peermem -> EFAULT, dmabuf -> rc=0, for both cudaMalloc and cuMem allocations). Set it in the encode/prefill/decode worker env so EPD comes up. Host memory and RDMA were always fine; this is purely the GPUDirect registration path, not the hardware. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Extend the temporary tokenspeed probe to register a GPU buffer twice on the H100 runner: once with the default (peermem) path and once with WITH_NVIDIA_PEERMEM=false (dmabuf). Confirms on the actual CI box that dmabuf is the working path (matches the GB300 finding). Also stop relying on an 8.8.8.8 route for the local IP (restricted runners); use hostname -I. Temporary. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Reproduce the TokenSpeed EPD (encode-prefill-decode) multimodal flow on a single local GPU-RDMA node and drive the EPD e2e against it, building the stack from local source so the transport can be iterated in minutes instead of 30-min CI rounds. Captures the non-obvious EPD env: WITH_NVIDIA_PEERMEM=false (mooncake dmabuf GPUDirect on the NVIDIA open kernel driver), LD_PRELOAD libnuma, and MC_INTRANODE_NVLINK for the same-node encode->prefill transfer. Parameterized via env (TORCH_PY, TS_SRC, EPD_MODEL, CUDA_HOME, EPD_TOPOLOGY). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
--system-site-packages inherited the conda env's torch-2.12 C-extensions
(torchcomms), which broke the moment tokenspeed-kernel pinned torch 2.11
("libtorchcomms.so: undefined symbol"). Build in a clean venv and install
torch==2.11.0+cu130 up front instead; verified it runs on GB300 (sm_103).
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
…wnloads) Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a858cd0ed4
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}" | ||
| if [ ! -x "${CUDA_HOME}/bin/nvcc" ]; then | ||
| echo "Installing CUDA toolkit (nvcc not found at ${CUDA_HOME}/bin/nvcc)..." | ||
| if [ ! -x "${CUDA_HOME}/bin/nvcc" ] && [ ! -x "/usr/local/cuda-13.0/bin/nvcc" ]; then |
There was a problem hiding this comment.
When $CUDA_HOME/bin/nvcc already exists but the toolkit is incomplete (the stale /usr/local/cuda case described immediately below), this guard skips cuda-toolkit-13-0; the later host_runtime.h check only warns, and the TokenSpeed kernel build still proceeds into the missing-header failure the comments describe. Include the required header in the install condition or fail before starting the build.
Useful? React with 👍 / 👎.
The EPD prefill worker hangs silently after mooncake's transfer engines come up and never becomes healthy; its worker log dead-ends with no traceback, so neither the job log nor the uploaded artifact reveals where it's stuck. Enable PYTHONFAULTHANDLER on the EPD workers and SIGABRT the process on the health timeout so CPython dumps every thread's Python stack into the worker log before it dies. Temporary — revert once the prefill hang is located and fixed. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The EPD prefill/decode workers loaded the model and stood up mooncake's transfer engines, then hung silently until the 600s health timeout, while the encode worker came up fine. Root cause (server.py:_wait_and_warmup): the smg tokenspeed servicer runs a gRPC warmup that drives stub.Generate to produce a token; the encode role self-skips it (no LM -- which is exactly why encode was healthy), but a disaggregated prefill/decode worker can't complete a lone Generate and blocks forever, so it never reaches SERVING. Mirror tokenspeed's own EPD serve script (serve_qwen35_122b_nvfp4_epd_1e2p1d.sh, tokenspeed#549): set TOKENSPEED_SKIP_GRPC_WARMUP=1 and pass --skip-server-warmup so prefill/decode reach SERVING, plus the runtime knobs that make the actual same-node transfer work -- MC_INTRANODE_NVLINK/MC_INTRA_NVLINK (NVLink-IPC intranode; RoCE loopback fails on one host), NO_PROXY for the local mooncake bootstrap, a 2GB gRPC limit for inline pixels, and the explicit --disaggregation-transfer-backend mooncake / --disaggregation-layerwise-interval. The temporary faulthandler stack-dump is kept for now as a safety net; remove it together with the ci_install_tokenspeed.sh probe once the lane is green. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
With the warmup fix, all EPD workers reach healthy and the request flows, but prefill/decode then OOM mid-generation (GPU 1/2 at 79/79 GiB) -> the scheduler raises, the response comes back empty, and all four topologies fail `assert text is not None`. TokenSpeed's defaults are too aggressive for a 35B FP8 LM on one card: auto gpu-memory-utilization (~0.9), kvstore-ratio 2.0, and a 131K-token KV pool leave no headroom for generation activations + the mooncake/EPD buffers. Size the LM conservatively via the model spec's tokenspeed_args (--gpu-memory-utilization 0.75, --max-model-len 8192, --max-num-seqs 8, --kvstore-ratio 0.5) -- ample for the single-image smoke test -- and set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True on the workers to avoid fragmentation OOM (as the error message recommends). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
Establishes CI ground truth for the multimodal tensor-transport and EPD (encode-prefill-decode) paths before an upcoming refactor that generalizes the RDMA transport (extract it into its own crate and add the remote/RDMA payload for vLLM). Today those paths are thinly covered: the only multimodal end-to-end test is image-only on the default in-message transport, there is no shared-memory or EPD end-to-end test, and the EPD encode planning has no unit tests. Without a net, a refactor bug would ship silently.
What's added (one commit each)
1. Rust unit tests (no GPU) — payload resolution + EPD planning
resolve_mm_shm_enabled) is now covered: in-message and RDMA modes stay off; shared-memory follows/dev/shmwritability; auto requires a matching shared-memory namespace. Previously untested.plan_encode_jobs(behavior-preserving) so the EPD encode item↔worker assignment matching, count/order checks, and bootstrap-info assembly are unit-testable without a live TokenSpeed worker; added coverage for those.2. Shared-memory multimodal e2e (1 GPU, vLLM)
/dev/shmtransport (gateway + worker are co-located in CI, so they share/dev/shm) and asserts the pixel tensor actually traveled over shared memory (via thesmg_mm_tensors_total{path="shm"}metric) — not that it silently fell back to the in-message path./metricsscraper to the e2e infra (none existed).3. EPD multimodal e2e (4 GPUs, TokenSpeed)
Qwen/Qwen3.6-35B-A3B-FP8. Its architecture is in TokenSpeed's multimodal registry, only 3B parameters are active, and FP8 (Hopper-native, so it runs on the h100 runner) fits one card — so every topology fits the existing 4-GPU lane at one worker per card.tokenspeedleg on the 4-GPU chat job runs it per-PR.Verification
transport8,proto_wrappervLLM 3,epd_encode5);cargo fmt/clippyclean for the changed code.ruff check+ruff formatclean;py_compileclean; CI YAML validated.Note
The EPD test needs the
Qwen/Qwen3.6-35B-A3B-FP8checkpoint present in the CI model cache; it fails fast (like other model-dependent e2e tests) if absent.Summary by CodeRabbit