sglang/test/registered/unit/disaggregation/test_pd_state_layer_ids.py

import concurrent.futures
import struct
from types import SimpleNamespace

import numpy as np
import pytest

from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=1, suite="stage-a-test-cpu")

mooncake_conn = pytest.importorskip("sglang.srt.disaggregation.mooncake.conn")
KVArgsRegisterInfo = mooncake_conn.KVArgsRegisterInfo
MooncakeKVManager = mooncake_conn.MooncakeKVManager


def _pack_q(values):
    return b"".join(struct.pack("Q", value) for value in values)


def _pack_i(values):
    return b"".join(struct.pack("i", value) for value in values)


def _pack_u(values):
    return b"".join(struct.pack("I", value) for value in values)


def _register_msg(*, state_layer_ids):
    return [
        b"room-a",
        b"127.0.0.1",
        b"1234",
        b"session-a",
        _pack_q([101, 102]),
        _pack_q([201]),
        _pack_q([301, 302, 303]),
        b"0",
        b"8",
        b"64",
        _pack_u([4096, 4096, 4096]),
        _pack_u([128, 128, 128]),
        _pack_i(state_layer_ids),
    ]


def test_mooncake_register_info_roundtrips_state_layer_ids():
    info = KVArgsRegisterInfo.from_zmq(
        _register_msg(state_layer_ids=[0, 4, 8, -1])
    )

    assert info.dst_state_layer_ids == [0, 4, 8, -1]


def test_mooncake_state_layer_id_mismatch_fails_fast():
    manager = MooncakeKVManager.__new__(MooncakeKVManager)
    manager.kv_args = SimpleNamespace(
        state_type="nsa",
        state_data_ptrs=[101, 102],
        state_item_lens=[64, 64],
        state_layer_ids=[0, 4],
        draft_state_type="none",
        draft_state_buffer_count=0,
    )
    manager.attn_cp_rank = 3
    manager.attn_tp_size = 8
    manager.is_mla_backend = True
    manager._send_kvcache_generic = lambda **_: 0
    req = SimpleNamespace(
        room="room-a",
        mooncake_session_id="session-a",
        dst_state_indices=np.array([11, 12], dtype=np.int32),
    )
    target_info = SimpleNamespace(
        dst_state_layer_ids=[0, 5],
        dst_state_data_ptrs=[201, 202],
        dst_attn_tp_size=8,
        dst_state_item_lens=[64, 64],
        dst_state_dim_per_tensor=[],
        dst_tp_rank=0,
    )

    with pytest.raises(RuntimeError, match=r"\[CP_SHARED_KV_FAIL_FAST\]\[state_layer_ids\].*prefill=\[0, 4\].*decode=\[0, 5\]"):
        MooncakeKVManager.maybe_send_extra(
            manager,
            req,
            prefill_state_indices=[1, 2],
            dst_state_data_ptrs=[201, 202],
            executor=concurrent.futures.ThreadPoolExecutor(max_workers=1),
            target_rank_registration_info=target_info,
        )


def test_mooncake_layer_aware_state_buffer_count_mismatch_fails_fast():
    manager = MooncakeKVManager.__new__(MooncakeKVManager)
    manager.kv_args = SimpleNamespace(
        state_type="nsa",
        state_data_ptrs=[101, 102, 103],
        state_item_lens=[64, 64, 64],
        state_layer_ids=[0, 4, 8],
        draft_state_type="none",
        draft_state_buffer_count=0,
    )
    manager.attn_cp_rank = 3
    manager.attn_tp_size = 8
    manager.is_mla_backend = True
    manager._send_kvcache_generic = lambda **_: 0
    req = SimpleNamespace(
        room="room-b",
        mooncake_session_id="session-b",
        dst_state_indices=np.array([11, 12, 13], dtype=np.int32),
    )
    target_info = SimpleNamespace(
        dst_state_layer_ids=[0, 4, 8],
        dst_state_data_ptrs=[201, 202],
        dst_attn_tp_size=8,
        dst_state_item_lens=[64, 64],
        dst_state_dim_per_tensor=[],
        dst_tp_rank=0,
    )

    with pytest.raises(RuntimeError, match=r"\[CP_SHARED_KV_FAIL_FAST\]\[state_buffer_count\].*src=3.*dst=2"):
        MooncakeKVManager.maybe_send_extra(
            manager,
            req,
            prefill_state_indices=[1, 2, 3],
            dst_state_data_ptrs=[201, 202],
            executor=concurrent.futures.ThreadPoolExecutor(max_workers=1),
            target_rank_registration_info=target_info,
        )


def test_nixl_register_info_roundtrips_state_layer_ids():
    nixl_conn = pytest.importorskip("sglang.srt.disaggregation.nixl.conn")

    msg = [
        b"room-a",
        b"127.0.0.1",
        b"1234",
        b"agent-a",
        b"metadata",
        _pack_q([101, 102]),
        _pack_q([201]),
        _pack_q([301, 302]),
        b"0",
        b"8",
        b"3",
        b"64",
        _pack_i([0, 4, -1]),
    ]

    info = nixl_conn.KVArgsRegisterInfo.from_zmq(msg)

    assert info.dst_state_layer_ids == [0, 4, -1]


def test_nixl_state_layer_id_mismatch_fails_fast():
    nixl_conn = pytest.importorskip("sglang.srt.disaggregation.nixl.conn")
    manager = nixl_conn.NixlKVManager.__new__(nixl_conn.NixlKVManager)
    manager.kv_args = SimpleNamespace(
        state_type="nsa",
        state_data_ptrs=[101, 102],
        state_item_lens=[64, 64],
        state_layer_ids=[0, 4],
    )
    manager.attn_cp_rank = 3
    manager.attn_tp_size = 8
    manager.is_mla_backend = True
    manager._send_kvcache_generic = lambda **_: 0

    with pytest.raises(RuntimeError, match=r"\[CP_SHARED_KV_FAIL_FAST\]\[state_layer_ids\].*prefill=\[0, 4\].*decode=\[0, 5\]"):
        nixl_conn.NixlKVManager.maybe_send_extra(
            manager,
            peer_name="decode-a",
            prefill_state_indices=[1, 2],
            dst_state_data_ptrs=[201, 202],
            dst_state_indices=[11, 12],
            dst_gpu_id=0,
            notif="notif-a",
            decode_tp_size=8,
            dst_state_layer_ids=[0, 5],
        )