Fix race on shared ZMQ sockets in mooncake KV manager senders

[yhyang201]

Motivation

The mooncake KV manager caches one ZMQ PUSH socket per destination endpoint (CommonKVManager._connect), and this socket is shared by all sender threads on the prefill side (up to SGLANG_DISAGGREGATION_THREAD_POOL_SIZE transfer workers plus the bootstrap thread). send_multipart is called on the shared socket without any lock. pyzmq sockets are not thread-safe, so concurrent sends can interleave the parts of two multipart messages and corrupt the stream.

Depending on how the corrupted message lands on the decode side, this causes:

• decode_thread dies with struct.error: unpack requires a buffer of 4 bytes in _handle_aux_data, after which all in-flight transfers on that rank hit KVTransferError ... timed out after 300.0s in KVPoll.WaitingForInput;
• corrupted aux data (output token ids) is written into the metadata buffers, leading to CUDA error: unspecified launch failure on one rank while the remaining EP ranks spin forever at 100% GPU in the MoE A2A combine kernel — the whole decode worker appears deadlocked until the watchdog kills it.

The probability scales with request rate and sender-thread count. On a 1P1D setup (Kimi-K2.5-NVFP4, decode TP8/EP8/DP8 across 2 GB300 nodes, streaming, concurrency 8192) it reproduced within minutes on every run; cuda-gdb on the hung ranks showed 7 ranks spinning in moeA2ACombineKernel waiting for the one rank whose context had faulted.

Note the decode→prefill direction already handles this: _connect_to_bootstrap_server returns (sock, lock) and all its callers send under the lock. This PR applies the same discipline to the prefill→decode direction.

Modifications

• common/conn.py: maintain a per-endpoint send lock next to the socket cache and add _send_multipart_locked.
• mooncake/conn.py: route the four unlocked sender paths (transfer status sync, aux data, CHUNK_READY, ABORT_ACK) through it.

Verification

On the reproducer above (previously failed 3/3 within 4–8 min), the patched build completed 40960/40960 requests over 12 min with zero struct.error / launch failures / KV timeouts, at unchanged throughput (52.2k output tok/s, TPOT 27.2 ms).

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #28627692160
Latest PR Test (Extra): ❌ Run #28627692097

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[AmazeQiu]

We hit the same root cause in production, but with a different crash signature — sharing our case as additional evidence.

Environment: GLM-based model, PD disaggregation with Mooncake backend, 2-node deployment (prefill TP16 + decode TP16), sgl-model-gateway as router.

Symptom: The whole decode instance became permanently unavailable (100% reproducible until restart). Every request deterministically failed on TP rank 7 with:

KVTransferError(bootstrap_room=...): Request ... timed out after 300.0s in KVPoll.WaitingForInput

All other ranks then failed with Failed due to an unknown reason from another rank. The decode server kept returning 500, which tripped the router's circuit breaker, and since this was our only decode worker, the router returned 503 No available decode workers (all circuits open or unhealthy) for all traffic.

Root cause investigation:

1. py-spy dump on the TP7 scheduler process showed the decode_thread was gone, while healthy ranks (e.g. TP6) still had Thread-2 (decode_thread) blocked in recv_multipart.
2. Digging into the service log, we found the thread had died hours earlier during a burst of client-side aborts (many concurrent AbortReq → ABORT / ABORT_ACK exchanges):

Exception in thread Thread-2 (decode_thread):
Traceback (most recent call last):
  File "/usr/lib/python3.12/threading.py", line 1073, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.12/threading.py", line 1010, in run
    self._target(*self._args, **self._kwargs)
  File "/sgl-workspace/sglang/python/sglang/srt/disaggregation/mooncake/conn.py", line 1586, in decode_thread
    bootstrap_room, status, prefill_rank = msg
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: not enough values to unpack (expected 3, got 1)

3. A 1-frame message cannot be produced by any legitimate sender in the protocol (status sync = 3 frames, ABORT_ACK = 2, AUX_DATA = 6, CHUNK_READY = 7). The only way to get one is frame interleaving from concurrent unlocked send_multipart calls on the shared cached PUSH socket — exactly what this PR fixes. In our case the bootstrap_thread (sending 2-frame ABORT_ACK during the abort storm) raced with a transfer worker (sending the 3-frame status sync) on the same endpoint socket, and the receiver ended up with a stray single-frame message.
4. Once decode_thread died, all subsequent KV-completion notifications for that rank piled up unconsumed → every request hit the 300s WaitingForInput timeout → permanent instance-wide outage. The "always rank 7" pattern was just the deterministic residue of the dead thread, not a hardware issue.

So this is the same bug as the struct.error in _handle_aux_data described here — just a different way the interleaved frames happened to be split.

Two suggestions:

1. Hope this PR can be prioritized for review/merge — the bug takes down the entire decode instance permanently and we can confirm it happens in production, not just stress tests.
2. Beyond the send-side locking (which fixes the corruption itself), it would be worth wrapping the decode_thread / bootstrap_thread loop bodies in try/except as a defense-in-depth measure. These threads are single consumers of the per-rank ZMQ socket with no supervision or restart; any single malformed message (from any future cause) kills the thread silently and turns into a 300s-timeout-per-request full outage that is very hard to diagnose. Logging and dropping the bad message would keep the instance alive.