fix(disagg): mori dropped decode-side abort notifications

[ShangmingCai]

Motivation

Mori's prefill side has been silently dropping every decode-side abort notification.

All backends inherit CommonKVReceiver.abort(), which calls _send_abort_notification() and sends an unguarded 4-frame message:

sock.send_multipart([b"ABORT", room, decode_ip, decode_port])   # no guard frame

The three backends read that socket differently:

Backend	Guard	ABORT handling
nixl	GUARD = b"NixlMsgGuard"	checks WATERMARK / STAGING_RSP / _handle_abort_notification before its assert msg[0] == GUARD — handled
mooncake	none	decodes frame 0, compares room == "ABORT" — handled
mori	MORI_GUARD = b"MoriMsgGuard"	_validate_message checks the guard first, so ABORT is rejected — dropped

Mori's own sends put MORI_GUARD in frame 0 (lines 657, 1686, 1745), but the inherited abort path never does. So on mori every abort produced

Received malformed bootstrap message

and was discarded. Two consequences:

1. The prefill side never learned the request was aborted, so the room stayed in request_status until the bootstrap/waiting timeout fired instead of being released promptly.
2. The log blamed malformed or foreign traffic for sglang's own message, which is actively misleading when triaging.

Nixl's ordering shows this was a known hazard — its abort check is deliberately placed above its guard assert. Mori never got the same treatment.

Modifications

1 file, +47 / -0. Adds MoriKVManager._handle_abort_notification and calls it in _start_bootstrap_thread ahead of _validate_message.

Why the receive side rather than adding MORI_GUARD to the sender

Adding the guard to _send_abort_notification would break rolling upgrades in both directions: an old prefill would reject a newly guarded abort, and a new prefill would reject an unguarded one from an old decode. Intercepting on the receive side keeps the wire format frozen and matches what nixl already does.

Correctness details

• check_status() indexes request_status directly and raises KeyError for an unknown room, so the active check short-circuits on membership first — same as nixl.
• Mori guards request_status with transfer_lock at four other sites, so the read-and-flip takes that lock. record_failure is called outside it, so transfer_lock and failure_lock never nest.
• Unknown rooms cannot be resurrected: update_status deliberately ignores a Failed transition for an absent room, so a late abort cannot pollute a future request that reuses the same bootstrap_room.

Deliberately not implemented: the deferred-KV-release ack

Nixl's handler also registers a deferred ack target and may send ABORT_ACK. That path depends on _staging_outstanding and _maybe_ack_drained_abort; mori has neither, and no staging at all. Acking before a mori transfer has drained could free decode pages while a write is still in flight — exactly the hazard nixl's own comment warns about.

With this change, decode falls back to its release timeout, which is the documented behavior when no ack arrives. That is strictly better than today (the abort is dropped entirely), but a mori owner with hardware should add the ack path as a follow-up.

Accuracy Tests

Not applicable — no model output, kernel, or forward path is touched. This is a control-plane message handler.

Speed Tests and Profiling

Not applicable. One byte-compare per bootstrap message on an already-ZMQ-bound thread.

Verification performed

The handler's logic was exercised against a stand-in manager reproducing mori's transfer_lock and CommonKVManager.update_status's exact semantics — 16 assertions, all passing:

Case	Expected	Result
guarded mori message	not consumed, status untouched	pass
empty message	not consumed	pass
active room	consumed, marked Failed, reason recorded	pass
already-Success room	consumed, stays Success, no failure recorded	pass
unknown room	consumed, request_status stays empty (no resurrection)	pass
malformed / truncated room id	consumed, no raise	pass
repeated calls	no deadlock, lock released	pass

Pinned ruff 0.15.1 --select=F401,F821,UP037 (the pre-commit config) passes, plus pinned black 26.1.0 and isort 7.0.0.

What CI will and will not cover here

test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py launches real prefill+decode servers on 8xMI35x. It has no abort coverage, so it will not validate the fix — but it is meaningful regression cover for the risky part of this change: every transfer-info message now passes through _handle_abort_notification before _validate_message, so a mistake in dispatch ordering or return value would break both test_generate_smoke and test_generate_smoke_tp_mismatch.

No mori hardware was available to me, so the end-to-end abort behavior is unverified. Please run the AMD disaggregation stage before merging.

Why there is no unit test in this PR

I wanted to add test/registered/unit/disaggregation/test_mori_abort_notification.py modelled on the existing test_nixl_deferred_kv_release.py, which unit-tests _handle_abort_notification on CPU CI via NixlKVManager.__new__(cls).

That is not currently possible for mori. nixl/conn.py imports nixl._api lazily inside a method, so the module imports fine without NIXL installed. mori/conn.py imports eagerly at module level:

from mori.cpp import TransferStatus
from mori.io import (BackendType, EngineDesc, IOEngine, IOEngineConfig,
                     MemoryDesc, MemoryLocationType, PollCqMode,
                     RdmaBackendConfig, StatusCode)

so import sglang.srt.disaggregation.mori.conn fails without the mori package. That is very likely why mori has 0 unit test files today while nixl has 3 — and why this bug survived: the shared abort() reaches all three backends, nixl proved its handler with unit tests, mooncake happens to parse it, and mori was never checked.

A CPU unit test would need patch.dict stubs for all ten names. Making mori's imports lazy (as nixl does) would be the cleaner fix and would unlock unit testing for the whole backend. Happy to do either as a follow-up — say which you prefer.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. — blocked, see "Why there is no unit test in this PR" above. Logic was verified out-of-tree with the 16 assertions listed.
• Update documentation according to Write documentations. — N/A: internal control-plane handler.
• Provide accuracy and speed benchmark results. — N/A, see above.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #32561230455
Latest PR Test (Extra): ❌ Run #32561230376
Latest PR Test (AMD ROCm 7.2): ❌ Run #32561230429

[ShangmingCai]

/rerun-test test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py

[github-actions]

Results for /rerun-test test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py:

⛔ test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py: No register_cuda_ci(runner_config=...) or register_cpu_ci() found in test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py. This file may not be a registered CI test.

[billishyahao]

Hi @ShangmingCai Shangming, Thanks for good catch! Actually there is another similar fix which is in review queue inside amd #29133 . cc previous reviewer and author @HaiShaw @Duyi-Wang @maning00 for their input

[ShangmingCai]

Included in #29133