Deterministic CPU reproducer: SGLang PD releases KV pages while a Mooncake transfer is still writing to them (sgl-project/sglang)

sglang_pd_page_reuse_repro.py

#!/usr/bin/env python3
"""Deterministic reproducer: SGLang PD releases KV pages while a Mooncake
transfer is still writing to them.

Run against an UNPATCHED sglang checkout. CPU only -- no GPU, no network, no
second process, ~1 second.

    cd /path/to/sglang && PYTHONPATH=python python sglang_pd_page_reuse_repro.py

Invariant under test
--------------------
While a transfer worker can still write to a request's KV pages, no rank may be
told the request is terminal. Reporting a terminal state is exactly what lets
the scheduler hand those pages back to the allocator, so a terminal state here
means a request allocated the same pages next has its KV silently overwritten.

Rather than trying to observe corruption probabilistically, this asserts the
ordering directly, so it fails identically every run.

Exit 1 = bug reproduced.  Exit 0 = pages are held until the transfer returns.
"""
import sys
import threading
import time
from collections import OrderedDict
from types import SimpleNamespace
from unittest.mock import Mock, patch

import numpy as np

from sglang.srt.disaggregation.base.conn import KVPoll
from sglang.srt.disaggregation.common.utils import FastQueue, TransferKVChunk
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager, MooncakeKVSender
from sglang.srt.disaggregation.utils import poll_and_all_reduce

ROOM = 7
DST_ADDR = 0xDEAD0000  # the decode KV page this transfer targets

entered_write = threading.Event()
release_write = threading.Event()


def build_prefill_manager():
    """A real MooncakeKVManager with just the prefill transfer path wired up."""
    m = MooncakeKVManager.__new__(MooncakeKVManager)
    m.request_status = {}
    m.check_status = lambda room: m.request_status[room]

    def update_status(room, status):
        cur = m.request_status.get(room)
        if status == KVPoll.Failed or cur is None:
            m.request_status[room] = status
        else:
            m.request_status[room] = max(cur, status)

    m.update_status = update_status
    m.record_failure = Mock()
    m.failure_records, m.failure_lock = {}, threading.Lock()
    m.transfer_infos = {
        ROOM: {
            "session:1": SimpleNamespace(
                room=ROOM,
                endpoint="127.0.0.1",
                dst_port=1,
                mooncake_session_id="session:1",
                dst_kv_indices=np.array([0], dtype=np.int32),
                is_dummy=False,
            )
        }
    }
    m.decode_kv_args_table = {
        "session:1": SimpleNamespace(
            dst_attn_tp_size=1,
            dst_kv_ptrs=[DST_ADDR],
            dst_kv_layer_ids=[0],
            dst_state_layer_ids=[],
        )
    }
    m.kv_args = SimpleNamespace(
        kv_data_ptrs=[0x1000], kv_layer_ids=[0], state_layer_ids=[], state_data_ptrs=[]
    )
    m._get_dsa_cache_transfer_skip_flags = Mock(return_value=(False, False))
    m.req_to_decode_prefix_len = {}
    m.session_lock, m.failed_sessions = threading.Lock(), set()
    m.is_mla_backend, m.is_hybrid_mla_backend = True, False
    m.attn_tp_size = m.attn_cp_size = m.pp_size = 1
    m.attn_tp_rank = m.attn_cp_rank = m.pp_rank = m.attn_dp_rank = 0
    m.enable_trace = m.enable_staging = False
    m.bootstrap_port = 0
    m.transfer_queues = [FastQueue()]

    def blocking_send_kvcache(*_a, **_k):
        """Stands in for batch_transfer_sync: an RDMA write that is in flight."""
        entered_write.set()
        release_write.wait(5)
        return 0

    m.send_kvcache = blocking_send_kvcache

    # Attributes only present once the ownership barrier exists; harmless before.
    optional = [
        ("_room_lifetimes", OrderedDict()),
        ("_room_lifetimes_lock", threading.Lock()),
        ("quiesce_timeout", 60.0),
        ("max_unquiesced", 256),
        ("_unquiesced_rooms", set()),
        ("_unquiesced_lock", threading.Lock()),
        ("_room_sweep_ttl", 300.0),
    ]
    try:
        from sglang.srt.environ import TransferBarrierLevel

        optional.append(("transfer_barrier", TransferBarrierLevel.WARN))
    except ImportError:
        pass
    for name, val in optional:
        if not hasattr(m, name):
            setattr(m, name, val)
    return m


def build_sender(mgr):
    s = MooncakeKVSender.__new__(MooncakeKVSender)
    s.kv_mgr, s.bootstrap_room, s.conclude_state = mgr, ROOM, None
    s.init_time = time.time()
    s.trace_ctx = SimpleNamespace(
        abort=Mock(), trace_req_finish=Mock(), copy_for_thread=Mock(return_value=Mock())
    )
    for name, val in [
        ("_quiescing", False),
        ("_quiesce_deadline", float("inf")),
        ("_quiesce_timed_out", False),
    ]:
        if not hasattr(s, name):
            setattr(s, name, val)
    if hasattr(mgr, "open_room_transfers"):
        mgr.open_room_transfers(ROOM)
    mgr.request_status[ROOM] = KVPoll.WaitingForInput
    return s


def poll(sender):
    """What poll_and_all_reduce would report to the scheduler (single rank)."""
    with patch(
        "sglang.srt.disaggregation.utils.dist.all_reduce",
        side_effect=lambda tensor, **kw: None,
    ):
        return int(poll_and_all_reduce([sender], object())[0])


NAMES = {0: "Failed", 1: "Bootstrapping", 2: "WaitingForInput", 3: "Transferring", 4: "Success"}


def main():
    mgr = build_prefill_manager()
    sender = build_sender(mgr)
    threading.Thread(
        target=mgr.transfer_worker, args=(mgr.transfer_queues[0], Mock()), daemon=True
    ).start()
    mgr.transfer_queues[0].put(
        TransferKVChunk(
            room=ROOM,
            prefill_kv_indices=np.array([0], dtype=np.int32),
            index_slice=slice(0, 1),
            is_last_chunk=False,
            prefill_aux_index=None,
            state_indices=None,
            trace_ctx=Mock(),
        )
    )

    if not entered_write.wait(5):
        print("INCONCLUSIVE: the worker never entered the transfer")
        return 2
    print(f"[T1] transfer worker is inside a write to decode page {hex(DST_ADDR)}")

    sender.abort()
    print("[T2] request aborted while that write is still in flight")

    state = poll(sender)
    print(f"[T3] scheduler was told: {NAMES[state]}")

    if state == int(KVPoll.Failed):
        print(
            f"\nRESULT: REPRODUCED -- the scheduler may now release {hex(DST_ADDR)}\n"
            "        while the transfer is still writing to it. A request that is\n"
            "        allocated that page next has its KV silently overwritten."
        )
        release_write.set()
        return 1

    print("[T4] terminal state withheld; the pages stay owned")
    release_write.set()
    for _ in range(500):
        if poll(sender) == int(KVPoll.Failed):
            print("[T5] transfer returned; request terminal, pages now releasable")
            print("\nRESULT: NOT REPRODUCED -- release strictly follows completion.")
            return 0
        time.sleep(0.01)
    print("\nRESULT: INCONCLUSIVE -- never became terminal after completion (hang).")
    return 2


if __name__ == "__main__":
    sys.exit(main())