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

"""Unit tests for A1: the async per-layer transfer primitive
(MooncakeKVManager._transfer_layers_async).

It must submit each layer's transfer non-blocking (one batch_transfer_async_submit
per non-empty layer), then wait for all once; skip empty layers; and on a submit
failure, drain what was submitted and return -1. Engine is mocked (no RDMA/GPU).
"""

import unittest

from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager


class _FakeEngine:
    def __init__(self, submit_rets=None, wait_ret=0):
        self.submits = []
        self.waits = []
        self._submit_rets = submit_rets
        self._wait_ret = wait_ret
        self._n = 0

    def batch_transfer_async_submit(self, sid, src, dst, lens):
        self.submits.append((sid, list(src), list(dst), list(lens)))
        if self._submit_rets is not None:
            r = self._submit_rets[self._n]
            self._n += 1
            return r
        return 100 + len(self.submits)

    def wait_batch_transfers(self, ids):
        self.waits.append(list(ids))
        return self._wait_ret


def _mgr(engine):
    m = MooncakeKVManager.__new__(MooncakeKVManager)  # skip heavy init
    m.engine = engine
    return m


def _one_block_per_layer(src, dst, item_len):
    return [(src, dst, item_len)]


class TestPerLayerAsyncTransfer(unittest.TestCase):
    def test_one_submit_per_layer_then_single_wait(self):
        eng = _FakeEngine()
        status = _mgr(eng)._transfer_layers_async(
            "sess", [(1000, 2000, 64), (1064, 2064, 64), (1128, 2128, 64)],
            _one_block_per_layer,
        )
        self.assertEqual(status, 0)
        self.assertEqual(len(eng.submits), 3)          # one async submit per layer
        self.assertEqual(len(eng.waits), 1)            # single wait-for-all
        self.assertEqual(eng.waits[0], [101, 102, 103])
        self.assertEqual(eng.submits[0], ("sess", [1000], [2000], [64]))

    def test_skips_empty_layers(self):
        eng = _FakeEngine()

        def blocks(src, dst, il):
            return [] if src == 1064 else [(src, dst, il)]

        _mgr(eng)._transfer_layers_async(
            "s", [(1000, 2000, 64), (1064, 2064, 64), (1128, 2128, 64)], blocks
        )
        self.assertEqual(len(eng.submits), 2)          # middle layer skipped
        self.assertEqual(eng.waits[0], [101, 102])

    def test_submit_failure_drains_and_returns_minus_one(self):
        eng = _FakeEngine(submit_rets=[101, -1])       # 2nd submit fails
        status = _mgr(eng)._transfer_layers_async(
            "s", [(1, 2, 3), (4, 5, 6), (7, 8, 9)], _one_block_per_layer
        )
        self.assertEqual(status, -1)
        self.assertEqual(len(eng.submits), 2)          # stopped at the failure
        self.assertEqual(eng.waits, [[101]])           # drained the first submit

    def test_wait_status_propagates(self):
        eng = _FakeEngine(wait_ret=-7)
        status = _mgr(eng)._transfer_layers_async("s", [(1, 2, 3)], _one_block_per_layer)
        self.assertEqual(status, -7)

    def test_no_layers_is_zero(self):
        eng = _FakeEngine()
        self.assertEqual(_mgr(eng)._transfer_layers_async("s", [], _one_block_per_layer), 0)
        self.assertEqual(eng.waits, [[]])


class _Reg:
    pass


class _KVArgs:
    pass


def _mgr_for_ctx(src_ptrs, dst_ptrs, item_lens, is_mla=True, registered=True):
    m = MooncakeKVManager.__new__(MooncakeKVManager)
    m.is_mla_backend = is_mla
    m.engine = object()
    kv = _KVArgs()
    kv.kv_data_ptrs = src_ptrs
    kv.kv_item_lens = item_lens
    kv.prefill_start_layer = 0
    m.kv_args = kv
    if registered:
        reg = _Reg()
        reg.dst_kv_ptrs = dst_ptrs
        m.decode_kv_args_table = {"sess": reg}
    else:
        m.decode_kv_args_table = {}
    return m


class TestBuildPerLayerContext(unittest.TestCase):
    def test_get_blocks_matches_per_layer_math(self):
        import numpy as np

        m = _mgr_for_ctx([1000, 9000], [5000, 8000], [64, 64])
        ctx = m.build_per_layer_context(
            "sess",
            np.array([0, 1, 2], dtype=np.int32),  # one contiguous prefill run
            np.array([10, 11, 12], dtype=np.int32),  # -> one contiguous dst run
        )
        self.assertIsNotNone(ctx)
        # layer 0: src=1000+0*64, dst=5000+10*64, len=64*3 (3-page run)
        self.assertEqual(ctx.get_blocks(0), ([1000], [5000 + 10 * 64], [64 * 3]))
        self.assertEqual(ctx.get_blocks(1), ([9000], [8000 + 10 * 64], [64 * 3]))
        self.assertIsNone(ctx.get_blocks(2))  # only 2 layers

    def test_none_for_mha(self):
        import numpy as np

        m = _mgr_for_ctx([1], [2], [3], is_mla=False)
        self.assertIsNone(
            m.build_per_layer_context("sess", np.array([0], np.int32), np.array([0], np.int32))
        )

    def test_none_when_decode_not_registered(self):
        import numpy as np

        m = _mgr_for_ctx([1], [2], [3], registered=False)
        self.assertIsNone(
            m.build_per_layer_context("sess", np.array([0], np.int32), np.array([0], np.int32))
        )

    def test_none_for_empty_indices(self):
        import numpy as np

        m = _mgr_for_ctx([1], [2], [3])
        self.assertIsNone(
            m.build_per_layer_context(
                "sess", np.array([], np.int32), np.array([], np.int32)
            )
        )


if __name__ == "__main__":
    unittest.main()