[PD][MoRI] Drive KV transfers with a sharded synchronous worker pool

[maning00]

Summary

Refactor the MoRI PD KV-transfer path to a bounded, per-room–sharded pool of synchronous workers: send() only enqueues; workers do batch_write → wait_all → notify. This moves both submission and completion off the prefill scheduler thread and bounds in-flight transfers.

Motivation

The old path ran KV submission and status busy-polling on the prefill scheduler thread, so scheduler load slowed transfers, and there was no bound on in-flight requests (RDMA SQ overflow under bursts).

Design

• One worker per shard, FIFO by bootstrap_room (preserves per-request chunk order).
• Outstanding transfers ≤ pool size; poll() becomes O(1).

Testing

Same-machine A/B vs main — DeepSeek-R1 671B FP8, TP8, 2-node PD (1P+1D), mori, --disable-radix-cache.

Single-request latency (output=16), refactor / main:

input	TTFT (ms)	TPOT (ms)	E2E (ms)
1024	118 / 127	10.8 / 10.7	280 / 287
2048	124 / 123	10.8 / 10.7	286 / 284
4096	173 / 187	10.9 / 10.8	337 / 349
8192	286 / 283	11.2 / 11.0	454 / 448

→ within run-to-run noise; no regression.

• GSM8K (2-node PD): 0.97.
• High concurrency (2048 × 8K/1K): refactor ≈ main (2048/2048, ~23.8K input tok/s). Tail timeouts are single-prefill-node capacity (identical on main), not the transfer path.

cc @Duyi-Wang

---

CI States

Latest PR Test (Base): ❌ Run #27112499352
Latest PR Test (Extra): ❌ Run #27112499288

[gemini-code-assist]

Code Review

This pull request updates the documentation index page to feature a new blog post on AMD Instinct MI355X and MoRI, and refactors the MoRI KV transfer mechanism in conn.py to utilize multi-threaded asynchronous queue-based chunk transfers. A critical issue was identified in the _wait_chunk method, where tripping the SLA timeout fails to break the infinite loop, which could cause the worker thread to hang indefinitely if the transfer remains in progress. A code suggestion has been provided to return immediately when the SLA is tripped.

[HaiShaw]

/tag-and-rerun-ci

[maning00]

The failing mi35x-disaggregation-amd CI check is an env mismatch, not a code bug.
This PR uses mori's new wait_all API (mori#341), but the AMD CI image pins an older mori in docker/rocm.Dockerfile (MORI_COMMIT=96ffa169, predates #341), so wait_all is missing.
Fix needs a refreshed AMD image: bump MORI_COMMIT to ≥ #341 and rebuild/publish rocm/sgl-dev.

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #26922

PR: [PD][MoRI] Drive KV transfers with a sharded synchronous worker pool
Changed files: python/sglang/srt/disaggregation/mori/conn.py (+204/-94), python/sglang/srt/environ.py (+26/-0)

AMD: 4 failures (1 likely related) | Others: 5 failures (0 related)

AMD CI Failures

Job	Test File	Test Function	Error	Related?	Explanation	Log
stage-b-test-large-8-gpu-mi35x-disaggregation-amd	test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py	test_generate_smoke	AttributeError: 'IOEngine' object has no attribute 'wait_all'	🔴 Likely	The PR's _wait_chunk calls self.kv_mgr.engine.wait_all(statuses, timeout_ms=...) on the mori IOEngine, but the installed mori library in the CI image does not expose that attribute. Every prefill/decode transfer raises AttributeError → 500 → assertion 500 != 200.	Log
stage-b-test-1-gpu-small-amd (linux-mi325-1gpu-sglang, 5)	test/registered/observability/test_tracing.py	(tracing span test)	AssertionError: ... Expected spans 'dispatch' and 'work', got set()	🟢 Unlikely	Observability/tracing test; PR does not touch tracing.	Log
stage-b-test-2-gpu-large-amd (linux-mi325-2gpu-sglang, 0)	test/registered/quant/test_int4fp8_moe.py	int4fp8 MoE accuracy	AssertionError: 0.5301... not greater than 0.56	🟢 Unlikely	Quantization accuracy threshold failure; PR touches only mori KV transfer code.	Log
stage-b-test-1-gpu-small-amd-nondeterministic (linux-mi325-1gpu-sglang)	test/registered/models/test_vlm_models.py	MiniCPM-V-2_6 mmmu	FileNotFoundError: No JSON result files found in /tmp/test_vlm_mmmu_openbmb_MiniCPM-V-2_6_...	🟢 Unlikely	VLM model test artifact missing; unrelated to KV transfer changes.	Log

Other CI Failures

Job	Test File	Test Function	Error	Related?	Explanation	Log
base-c-test-8-gpu-h20 (1)	test/registered/disaggregation/test_disaggregation_different_tp.py	test_gsm8k	KVTransferError: Failed to get kvcache from prefill instance	🟢 Unlikely	Uses disaggregation_transfer_backend='mooncake', not MoRI. PR only touches the mori backend.	Log
multimodal-gen-test-8-npu-a3	sglang/multimodal_gen/test/server/ascend/test_server_8_npu.py	test_diffusion_generation[wan2_2_t2v_14b_w8a8_8npu]	Diffusion testcase failed checks	🟢 Unlikely	NPU diffusion test, unrelated to KV transfer.	Log
stage-b-test-1-npu-a2 (0)	test/registered/ascend/basic_function/quant/test_npu_w8a8_quantization.py	NPU w8a8 perf	AssertionError: 528.29... not greater than or equal to 700	🟢 Unlikely	NPU w8a8 perf threshold; unrelated.	Log
multimodal-gen-test-2-npu-a3	sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py	test_diffusion_generation[qwen_image_t2i_2npu]	Diffusion testcase failed checks	🟢 Unlikely	NPU diffusion test, unrelated.	Log
multimodal-gen-test-1-npu-a3	sglang/multimodal_gen/test/server/ascend/test_server_1_npu.py	test_diffusion_generation[wan2_1_t2v_1.3b_1_npu]	Diffusion testcase failed checks	🟢 Unlikely	NPU diffusion test, unrelated.	Log

Details

🔴 stage-b-test-large-8-gpu-mi35x-disaggregation-amd — this is a real, PR-caused regression and must be fixed before merge.

• The PR adds _wait_chunk in python/sglang/srt/disaggregation/mori/conn.py:

  rc = self.kv_mgr.engine.wait_all(statuses, timeout_ms=self.kv_mgr._wait_poll_ms ...)
  

• At runtime on the MI35x image the mori library raises AttributeError: 'IOEngine' object has no attribute 'wait_all' on every transfer. Both prefill and decode sides return HTTP 500 immediately, and test_generate_smoke asserts 500 != 200.
• Likely fixes (pick one): (a) bump the mori dependency to the version that ships IOEngine.wait_all, (b) call the actually-available wait API on the current mori build, or (c) gate the new path on a feature check / version probe before calling wait_all.

All other failures are infrastructure or unrelated test issues:

• The H20 disaggregation test runs the mooncake backend (disaggregation_transfer_backend='mooncake'), so KVTransferError there cannot be caused by mori-only changes.
• The libavutil/torio errors in the mi35x job log are non-fatal (the actual fatal failure is the mori wait_all AttributeError above).
• The AMD 1gpu/2gpu/nondeterministic and NPU failures touch tracing, int4fp8 MoE, VLM, NPU diffusion, and NPU w8a8 quantization — none of which exercise the modified MoRI KV transfer path.

---

Generated by amd-bot using Claude Code CLI

[HaiShaw]

@yctseng0211 @bingxche Would you please help to check the PR image used here?
Ideally we should use a new staged image to test the PR with Dockerfile changes.

[HaiShaw]

@amd-bot ci-status

[amd-bot]

@HaiShaw

CI Status for PR #26922

Merge verdict: ❌ Do not merge. This PR has a real, self-inflicted bug: the AMD MoRI disaggregation e2e test (test_mori_transfer_engine_e2e.py) fails because the PR's new _wait_chunk code calls engine.wait_all(...), a method that does not exist on the IOEngine from the pinned mori commit. The other 6 test failures are all pre-existing/unrelated to this PR. The good news: this PR's changed code is exercised by PR CI, and CI caught the bug.

Note

The PR's core changed code path (mori/conn.py sharded worker pool) is covered — the stage-b-test-large-8-gpu-mi35x-disaggregation-amd suite runs test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py, which exercises the new transfer code end-to-end and red-flagged it. The 7 new SGLANG_MORI_* env vars are default-valued, but the e2e test drives the default path, so green-after-fix would be meaningful here.

Changed files: python/sglang/srt/disaggregation/mori/conn.py (+207/-95), python/sglang/srt/environ.py (+26/-0), docker/rocm.Dockerfile (+1/-1, bumps MORI_COMMIT → d87651c)

AMD: 3 failures (1 related) · Others: 4 failures (0 related) · (cascade/aggregator jobs pr-test-amd-finish, pr-test-npu-finish, pr-test-extra-finish, call-gate / pr-gate collapsed into their root causes below)

AMD CI Failures

Job	Test File	Test Function	Error	Related?	Why
stage-b-...-mi35x-disaggregation-amd	test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py	TestMoriTransferEngineE2E.test_generate_smoke	AttributeError: 'IOEngine' object has no attribute 'wait_all' → 500	🔴	PR's new _wait_chunk calls self.kv_mgr.engine.wait_all(statuses, timeout_ms=...); the IOEngine from pinned mori d87651c has no such method
stage-c-test-large-8-gpu-amd (3)	test/registered/ops/test_aiter_allreduce_fusion_amd.py	test_fused_ar_rms_residual_accuracy	fused residual output diverges from unfused path	🟢	aiter allreduce-fusion numerical kernel; PR touches none of this path
stage-c-test-large-8-gpu-amd-mi35x (0)	test/registered/amd/test_deepseek_r1_mxfp4_8gpu.py	test_a_gsm8k	ConnectionError: port 11000 refused (server didn't come up)	🟢	DeepSeek MXFP4 8-GPU server startup; unrelated to MoRI/disaggregation

Other CI Failures

Job	Test File	Test Function	Error	Related?	Why
build-test (xeon-gnr, base-b-test-cpu)	test/registered/scheduler/test_routing_key_scheduling.py	TestRoutingKeyScheduling.setUpClass	RuntimeError: decode: expect req_lens to be int64, got Int → server exit -9	🟢	CPU-backend scheduler/decode dtype issue; PR has no CPU or scheduler changes
stage-b-test-1-npu-a2 (0)	test/registered/ascend/basic_function/quant/test_npu_w8a8_quantization.py	(suite exit 1)	NPU W8A8 quant failure	🟢	NPU quantization path; PR is AMD/MoRI only
multimodal-gen-test-1-npu-a3	multimodal_gen/test/server/ascend/test_server_1_npu.py	test_diffusion_generation[flux_image_t2i_npu]	assert 46092.4 <= 25838.4 (perf threshold)	🟢	NPU diffusion latency perf gate; unrelated
multimodal-gen-test-2-npu-a3	multimodal_gen/test/server/ascend/test_server_2_npu.py	test_diffusion_generation[flux_2_image_t2i_2npu]	assert 63363.3 <= 58197.1 (perf threshold)	🟢	NPU diffusion latency perf gate; unrelated

Details / what to do before merge

• 🔴 Fix the wait_all bug (blocking). In conn.py, _wait_chunk calls self.kv_mgr.engine.wait_all(statuses, timeout_ms=self.kv_mgr._wait_poll_ms), but the IOEngine exposed by the pinned mori commit d87651c does not implement wait_all. Verify whether (a) the intended mori commit that adds wait_all was actually pinned in rocm.Dockerfile, or (b) the method name/signature on the mori side differs. Re-run test_mori_transfer_engine_e2e.py in the disaggregation suite after fixing. Note: the prefill/decode KVTransferError 500s seen during the run all trace back to this same missing attribute — it's a single root cause, not multiple.
• 🟢 The other 6 failures are pre-existing/unrelated (CPU decode-dtype, aiter allreduce numerics, DeepSeek-MXFP4 server startup, NPU quant, two NPU diffusion perf gates). They are not introduced by this PR and should not gate it — but they also mean a fully-green run isn't achievable on this branch regardless; judge merge-readiness on the disaggregation suite specifically.

---

Generated by amd-bot using Claude Code CLI

[HaiShaw]

@bingxche @yctseng0211 left review comment above - this PR changed MoRI version, according to @Duyi-Wang if we use image with the change included, the CI failure in test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py should not happen. Please check.

[bingxche]

@bingxche @yctseng0211 left review comment above - this PR changed MoRI version, according to @Duyi-Wang if we use image with the change included, the CI failure in test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py should not happen. Please check.

Hi @HaiShaw Test passed with new MORI commit.
https://github.com/sgl-project/sglang/actions/runs/27115280593/job/80020964373#step:9:27354

cc @yctseng0211