[WIP] EPD routing: pixel lanes (NIXL/SHM), servicer hardening, gateway data-plane perf

[chenht2022]

Summary

Gateway/servicer side of the EPD line, 39 commits:

• EPD request routing (Encode RPC + bootstrap_room injection, mirrors PD routing)
• Pixel delivery lanes: NIXL pixel pull v2 (pre-registered slot arena) + POSIX-SHM handoff; encode pixel ingest moved off the servicer event loop
• Servicer hardening: max_ping_strikes=0 (idle/blocked-loop lanes survived gRPC GOAWAY too_many_pings -> 110s system stalls), landing-ring exhaustion survival under admission bursts
• Gateway data-plane perf: vectorized f32->bf16 pixel down-cast; pooled per-image CHW tensor + patch buffer with cross-thread recycle (exact sizing fixed a 4x patch-buffer underestimate causing ~200MB realloc+copy per request); GPU-free data-plane microbench
• Net effect: gateway minor faults/request halved, effective data-plane parallelism 32 -> 60 cores, microbench 110 -> 192 img/s

Validation (122B, 3-node B200)

• e2e steady-state P8D2E8: 18.03 (pre-pooling) -> 19.01/19.29 req/s; OCRBench acc 0.910 (baseline 0.888-0.894)
• Deterministic GOAWAY reproducer: 0 lane deaths after fix
• Draft: pending TP8 e2e + Kimi-K2.5 processor pooling port

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 0557631f-17bf-4cbb-9997-ae978eefa56b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[gemini-code-assist]

Code Review

This pull request introduces the Encode-Prefill-Decode (EPD) disaggregated mode, adding a vision-only encode stage that processes and ships image embeddings to prefill workers over Mooncake. Key changes include the addition of the TokenSpeedEncoder gRPC service, a high-performance NIXL-based RDMA pixel transport, memory recycling pools for large vision buffers, and EPD pipeline stages in the gateway. The review feedback highlights several improvement opportunities: ensuring image_grid_thw is 2D to prevent a runtime crash, wrapping NIXL transfers in a try...finally block to avoid resource leaks, aligning the default RDMA landing slots between the gateway and encoder, and refactoring concurrent task spawning in the encode stage to use tokio::task::JoinSet instead of join_all and bare tokio::spawn.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If image_grid_thw is a 1D tensor (e.g., when there is only a single image and it was flattened), grid.tolist() will return a flat list of 3 elements instead of a list of lists. Iterating over it will then yield scalar integers, causing a TypeError when attempting to subscript row[0]. Ensuring the tensor is 2D using .view(-1, 3) before calling .tolist() prevents this potential runtime crash.

Suggested change

	grid = model_specific.get("image_grid_thw")
	if grid is None:
	raise ValueError("encode request is missing image_grid_thw")
	# grid is [num_images, 3] = (t, h, w) in patch units, per image.
	grid = model_specific.get("image_grid_thw")
	if grid is None:
	raise ValueError("encode request is missing image_grid_thw")
	grid = grid.view(-1, 3)
	# grid is [num_images, 3] = (t, h, w) in patch units, per image.

[gemini-code-assist]

The NIXL transfer handle h is initialized but not guaranteed to be released if an exception is raised during the transfer or polling loop (e.g., if a timeout occurs or check_xfer_state fails). Wrapping the transfer and polling logic in a try...finally block ensures that agent.release_xfer_handle(h) is always called, preventing potential resource leaks.

            h = agent.initialize_xfer("READ", local, remote, gw_name, str(room).encode())
            try:
                read_deadline = time.monotonic() + float(
                    os.environ.get("SMG_RDMA_READ_TIMEOUT_S", 60)
                )
                spins = 0
                state = agent.transfer(h)
                while state in ("PROC", "IN_PROG"):
                    # Yield between polls: under EPD_INGEST_OFFLOOP up to ~32 worker
                    # threads poll concurrently, and a no-sleep spin starves the GIL
                    # away from the hash/publish work that RETURNS slots -- the ring
                    # runs dry and admission collapses. Spin briefly for small-READ
                    # latency, then back off.
                    spins += 1
                    if spins > 64:
                        time.sleep(0.0005)
                    if time.monotonic() > read_deadline:
                        raise RuntimeError(f"NIXL READ timed out room={room}")
                    state = agent.check_xfer_state(h)
                if state != "DONE":
                    raise RuntimeError(f"NIXL READ state={state} room={room}")
            finally:
                agent.release_xfer_handle(h)

References

1. When managing resources (e.g., workers) that are acquired and need to be released, ensure that all acquired resources are tracked immediately and consistently, especially when exceptions might occur during subsequent operations (e.g., launching new resources). This prevents resource leaks by ensuring that even partially successful acquisitions are accounted for cleanup.

[gemini-code-assist]

There is a default mismatch for the number of RDMA slots between the gateway (which defaults to 64 slots via DEFAULT_POOL_SLOTS) and the encoder servicer (which defaults to 16 slots here). Under high concurrency, this mismatch can lead to unexpected starvation or backpressure on the encoder side. Aligning the default values to 64 ensures consistent capacity across both sides of the transport.

Suggested change

	slot_bytes = int(
	os.environ.get("SMG_RDMA_SLOT_BYTES", 32 * 1024 * 1024)
	)
	n_slots = int(os.environ.get("SMG_RDMA_LANDING_SLOTS", 16))
	slot_bytes = int(
	os.environ.get("SMG_RDMA_SLOT_BYTES", 32 * 1024 * 1024)
	)
	n_slots = int(os.environ.get("SMG_RDMA_LANDING_SLOTS", 64))

[gemini-code-assist]

To adhere to the general rules, prefer using tokio::task::JoinSet over futures::future::join_all when spawning multiple concurrent tasks. This ensures tasks are distributed across the multi-threaded runtime and allows unexpected task failures to be detected immediately.

Suggested change

	let mut sends: Vec<tokio::task::JoinHandle<Result<(), String>>> =
	Vec::with_capacity(splits.len());
	let mut sends = tokio::task::JoinSet::new();

References

1. When spawning multiple concurrent tasks or workers in a Tokio runtime, prefer using tokio::task::JoinSet over futures::future::join_all. This ensures tasks are distributed across the multi-threaded runtime rather than running on a single thread, and allows unexpected task failures to be detected immediately via join_next() instead of being masked. Additionally, JoinSet::spawn can be used when bare tokio::task::spawn is disallowed.

[gemini-code-assist]

Spawn the task directly onto the JoinSet instead of using bare tokio::spawn.

            sends.spawn(async move {

References

1. When spawning multiple concurrent tasks or workers in a Tokio runtime, prefer using tokio::task::JoinSet over futures::future::join_all. This ensures tasks are distributed across the multi-threaded runtime rather than running on a single thread, and allows unexpected task failures to be detected immediately via join_next() instead of being masked. Additionally, JoinSet::spawn can be used when bare tokio::task::spawn is disallowed.

[gemini-code-assist]

Iterate over the JoinSet using join_next() instead of wrapping the handles in join_all.

Suggested change

	tokio::spawn(async move {
	for join_res in join_all(sends).await {
	tokio::spawn(async move {
	while let Some(join_res) = sends.join_next().await {

References

1. When spawning multiple concurrent tasks or workers in a Tokio runtime, prefer using tokio::task::JoinSet over futures::future::join_all. This ensures tasks are distributed across the multi-threaded runtime rather than running on a single thread, and allows unexpected task failures to be detected immediately via join_next() instead of being masked. Additionally, JoinSet::spawn can be used when bare tokio::task::spawn is disallowed.