lightseekorg--tokenspeed/test/runtime/test_flat_host_executor.py

"""FlatMemoryExecutor (M15 Phase D2): flat host tier through the kvstore
transport.

Covers the executor roundtrip against a real (tiny) device pool with
WriteBackDone AND LoadBackDone acks, the ack payload shape riding the
TP-synced commit path, the layer -> mirror-tensor fencing mapping, and the
num_host_pages sizing arithmetic (pure CPU).
"""

from __future__ import annotations

import os
import sys
import unittest
from unittest import mock

# CI Registration (parsed via AST, runtime no-op)
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ci_system.ci_register import register_cuda_ci

register_cuda_ci(est_time=60, suite="runtime-1gpu")

_PKG_FLAT_PROBE = (
    "tokenspeed.runtime.configs.paged_cache_spec.scheduler_ext_flat_kvcache"
)

LAYER_TYPES = ("sliding_attention", "full_attention") * 2

# GDN hybrid: layers 0/2 are state layers (slab pairs 0/1); the KV side
# stays per-layer (linear_attention disables slab pairing) but state layers
# carry no KV tensors under the flat predicate (M18a T4), and the LAST
# layer is an attention layer -- exercising the finish-event pin.
GDN_LAYER_TYPES = ("linear_attention", "full_attention") * 2


class _StubPool:
    """CPU-only device-pool stand-in for sizing arithmetic: 4 layers dedup
    to 2 K + 2 V slabs (paired layers alias the same tensor)."""

    def __init__(self, torch):
        self.page_size = 4
        self.size = 32
        rows = self.size + self.page_size
        k_slabs = [torch.zeros((rows, 1, 8), dtype=torch.bfloat16) for _ in range(2)]
        v_slabs = [torch.zeros((rows, 1, 8), dtype=torch.bfloat16) for _ in range(2)]
        self.k_buffer = [k_slabs[0], k_slabs[0], k_slabs[1], k_slabs[1]]
        self.v_buffer = [v_slabs[0], v_slabs[0], v_slabs[1], v_slabs[1]]


class FlatHostPageSizingTest(unittest.TestCase):
    """num_host_pages budget arithmetic; no CUDA, no scheduler ext."""

    def setUp(self):
        try:
            import torch

            from tokenspeed.runtime.cache.flat_host_mirror import (
                FlatHostMirror,
                flat_bytes_per_host_page,
            )
        except (ImportError, ModuleNotFoundError) as exc:
            self.skipTest(f"needs torch: {exc}")
        self.torch = torch
        self.FlatHostMirror = FlatHostMirror
        self.flat_bytes_per_host_page = flat_bytes_per_host_page

    def test_bytes_per_host_page_matches_mirror(self):
        stub = _StubPool(self.torch)
        # 4 distinct mirrors x page_size 4 x row 1*8 bf16 (16 B) = 256 B.
        self.assertEqual(self.flat_bytes_per_host_page(stub), 256)
        mirror = self.FlatHostMirror(stub, num_host_pages=2)
        self.assertEqual(mirror.bytes_per_host_page(), 256)

    def test_num_host_pages_formula(self):
        from tokenspeed.runtime.cache.executor.flat_memory_executor import (
            flat_num_host_pages,
        )

        # Ratio sizing mirrors the radix token->page align-up:
        # int(size * ratio) // page_size + 1.
        self.assertEqual(
            flat_num_host_pages(
                bytes_per_host_page=256,
                device_pool_size=32,
                page_size=4,
                host_ratio=2.0,
                host_size_gb=0,
            ),
            int(32 * 2.0) // 4 + 1,
        )
        # Explicit GB budget floors to whole mirror pages (never exceeds).
        self.assertEqual(
            flat_num_host_pages(
                bytes_per_host_page=256,
                device_pool_size=32,
                page_size=4,
                host_ratio=2.0,
                host_size_gb=1,
            ),
            int(1e9) // 256,
        )
        # A budget below one page is a configuration error.
        with self.assertRaises(ValueError):
            flat_num_host_pages(
                bytes_per_host_page=int(1e9),
                device_pool_size=32,
                page_size=4,
                host_ratio=2.0,
                host_size_gb=0.5,
            )


class LoadBackDonePayloadTest(unittest.TestCase):
    """LoadBackDone rides the same TP-synced commit path as WriteBackDone;
    needs the scheduler ext, not a GPU."""

    def setUp(self):
        try:
            from tokenspeed_scheduler import Cache

            from tokenspeed.runtime.engine.scheduler_utils import (
                cache_event_from_payload,
                cache_event_to_payload,
                pop_common_cache_event_payloads,
            )
        except (ImportError, ModuleNotFoundError, AttributeError) as exc:
            self.skipTest(f"needs tokenspeed_scheduler ext: {exc}")
        if not hasattr(Cache, "LoadBackDoneEvent"):
            self.skipTest("scheduler ext predates the LoadBackDoneEvent binding")
        self.Cache = Cache
        self.to_payload = cache_event_to_payload
        self.from_payload = cache_event_from_payload
        self.pop_common = pop_common_cache_event_payloads

    def _load_done(self, op_id: int):
        evt = self.Cache.LoadBackDoneEvent()
        evt.op_id = op_id
        evt.success = True
        return evt

    def test_payload_shape_and_roundtrip(self):
        payload = self.to_payload(self._load_done(3))
        self.assertEqual(
            payload,
            {
                "kind": "LoadBackDoneEvent",
                "op_id": 3,
                "success": True,
                "request_id": "",
            },
        )
        ready = self.pop_common([[payload]])  # world_size=1 gather
        self.assertEqual(ready, [payload])
        evt = self.from_payload(ready[0])
        self.assertIsInstance(evt, self.Cache.LoadBackDoneEvent)
        self.assertEqual(int(evt.op_id), 3)
        self.assertTrue(evt.success)

    def test_pop_common_requires_all_ranks(self):
        payload = self.to_payload(self._load_done(4))
        self.assertEqual(self.pop_common([[payload], []]), [])
        both = self.pop_common([[payload], [dict(payload)]])
        self.assertEqual(both, [payload])


class FlatMemoryExecutorTest(unittest.TestCase):
    """Real (tiny) MHATokenToKVPool on GPU driving the flat executor."""

    def setUp(self):
        try:
            import torch
            from tokenspeed_scheduler import Cache

            from tokenspeed.runtime.cache.executor.flat_memory_executor import (
                FlatMemoryExecutor,
            )
            from tokenspeed.runtime.cache.transfer.types import CacheKind
            from tokenspeed.runtime.layers.attention.kv_cache.mha import (
                MHATokenToKVPool,
            )
        except (ImportError, ModuleNotFoundError) as exc:
            self.skipTest(f"needs torch + tokenspeed_kernel + scheduler ext: {exc}")
        if not hasattr(Cache, "LoadBackDoneEvent"):
            self.skipTest("scheduler ext predates the LoadBackDoneEvent binding")
        if not torch.cuda.is_available():
            self.skipTest("needs a CUDA device")
        self.torch = torch
        self.Cache = Cache
        self.CacheKind = CacheKind
        self.FlatMemoryExecutor = FlatMemoryExecutor
        self.MHATokenToKVPool = MHATokenToKVPool

    def _pool(self):
        kwargs = dict(
            size=32,
            dtype=self.torch.bfloat16,
            head_num=1,
            head_dim=8,
            layer_num=4,
            device="cuda",
            enable_memory_saver=False,
            max_batch_size=2,
            max_context_len=64,
            page_size=4,
            rank=0,
            layer_types=LAYER_TYPES,
            sliding_window_tokens=128,
            enable_alt_stream=False,
        )
        with mock.patch(_PKG_FLAT_PROBE, return_value=True):
            return self.MHATokenToKVPool(**kwargs)

    def _executor(self, pool):
        return self.FlatMemoryExecutor(device_pool=pool, host_ratio=2.0, host_size_gb=0)

    def _fill_device_pages(self, mirror, device_pages):
        p = mirror.page_size
        for tensor_idx, (dev, _) in enumerate(mirror.tensor_pairs):
            for d in device_pages:
                dev[d * p : (d + 1) * p].fill_(tensor_idx * 16 + d + 1)
        self.torch.cuda.synchronize()

    def _snapshot(self, mirror, device_pages):
        p = mirror.page_size
        return [
            {d: dev[d * p : (d + 1) * p].cpu().clone() for d in device_pages}
            for dev, _ in mirror.tensor_pairs
        ]

    def _drain(self, executor, expect: int) -> list:
        results = []
        for _ in range(1000):
            results.extend(executor.poll_results())
            if len(results) >= expect:
                return results
            self.torch.cuda.synchronize()
        self.fail(f"expected {expect} acks, drained {len(results)}")

    def test_roundtrip_with_acks_and_fencing(self):
        torch = self.torch
        pool = self._pool()
        executor = self._executor(pool)
        mirror = executor.mirror
        # Ratio sizing: int(32 * 2.0) // 4 + 1.
        self.assertEqual(executor.num_host_pages, 17)
        self.assertTrue(executor.emits_loadback_acks)
        # The fencing counter is registered where the radix pool would be.
        self.assertIs(pool.layer_transfer_counter, executor._counter)

        device_pages = [1, 2]
        self._fill_device_pages(mirror, device_pages)
        before = self._snapshot(mirror, device_pages)

        # WriteBack: device pages [1, 2] -> host pages [5, 6].
        executor.submit_writeback([7], [[1, 2]], [[5, 6]])
        executor.flush()
        results = self._drain(executor, 1)
        self.assertEqual(len(results), 1)
        self.assertIsInstance(results[0], self.Cache.WriteBackDoneEvent)
        self.assertEqual(int(results[0].op_id), 7)
        self.assertTrue(results[0].success)

        for dev, _ in mirror.tensor_pairs:
            p = mirror.page_size
            for d in device_pages:
                dev[d * p : (d + 1) * p].zero_()
        torch.cuda.synchronize()

        # LoadBack: host pages [5, 6] -> device pages [1, 2] (wire order:
        # src=host, dst=device, as C++ FlatLoadBackOperation emits).
        executor.submit_loadback([9], [[5, 6]], [[1, 2]])
        executor.flush()

        # Layerwise fencing: producer registered under the op, consumer waits
        # gate per layer through the pool's registered counter.
        producer_idx = executor.get_producer_index(self.CacheKind.KV, 9)
        self.assertIsNotNone(producer_idx)
        executor.set_consumer(self.CacheKind.KV, [producer_idx])
        for layer_id in range(4):
            pool.layer_transfer_counter.wait_until(layer_id)
        torch.cuda.synchronize()

        results = self._drain(executor, 1)
        self.assertEqual(len(results), 1)
        self.assertIsInstance(results[0], self.Cache.LoadBackDoneEvent)
        self.assertEqual(int(results[0].op_id), 9)
        self.assertTrue(results[0].success)

        after = self._snapshot(mirror, device_pages)
        for tensor_idx in range(len(mirror.tensor_pairs)):
            for d in device_pages:
                self.assertTrue(
                    torch.equal(
                        before[tensor_idx][d].view(torch.uint8),
                        after[tensor_idx][d].view(torch.uint8),
                    ),
                    f"tensor {tensor_idx} device page {d} not byte-exact",
                )

        # Producer index is popped exactly once (event_loop consumes it in
        # _setup_layerwise_loadback right after submit).
        self.assertIsNone(executor.get_producer_index(self.CacheKind.KV, 9))

    def test_layer_event_mapping(self):
        pool = self._pool()
        executor = self._executor(pool)
        mirror = executor.mirror
        self._fill_device_pages(mirror, [3])
        executor.submit_writeback([1], [[3]], [[0]])
        executor.flush()
        self._drain(executor, 1)

        executor.submit_loadback([2], [[0]], [[3]])
        executor.flush()
        producer_idx = executor.get_producer_index(self.CacheKind.KV, 2)
        producer_event = executor._counter.events[producer_idx]
        # Paired slab layers share their V slab's event; distinct groups
        # get distinct events; finish_event covers every copy.
        self.assertIs(producer_event.load_events[0], producer_event.load_events[1])
        self.assertIs(producer_event.load_events[2], producer_event.load_events[3])
        self.assertIsNot(producer_event.load_events[0], producer_event.load_events[2])
        self.torch.cuda.synchronize()
        self.assertTrue(producer_event.finish_event.query())
        self._drain(executor, 1)

    def _state_pool(self):
        kwargs = dict(
            size=32,
            dtype=self.torch.bfloat16,
            head_num=1,
            head_dim=8,
            layer_num=4,
            device="cuda",
            enable_memory_saver=False,
            max_batch_size=2,
            max_context_len=64,
            page_size=4,
            rank=0,
            layer_types=GDN_LAYER_TYPES,
            sliding_window_tokens=None,
            enable_alt_stream=False,
            conv_state_shape=(2, 4),
            temporal_state_shape=(2, 8),
        )
        with mock.patch(_PKG_FLAT_PROBE, return_value=True):
            return self.MHATokenToKVPool(**kwargs)

    def _fill_spans(self, mirror, device_pages):
        for tensor_idx, ((dev, _), span) in enumerate(
            zip(mirror.tensor_pairs, mirror.row_spans)
        ):
            for d in device_pages:
                dev[d * span : (d + 1) * span].fill_(tensor_idx * 16 + d + 1)
        self.torch.cuda.synchronize()

    def _snapshot_spans(self, mirror, device_pages):
        return [
            {d: dev[d * span : (d + 1) * span].cpu().clone() for d in device_pages}
            for (dev, _), span in zip(mirror.tensor_pairs, mirror.row_spans)
        ]

    def test_state_layer_event_mapping(self):
        pool = self._state_pool()
        executor = self._executor(pool)
        mirror = executor.mirror
        # Flat GDN KV (2 K + 2 V; state layers carry no KV, M18a T4) +
        # conv0, ssm0, conv1, ssm1.
        self.assertEqual(len(mirror.tensor_pairs), 8)
        self._fill_spans(mirror, [3])
        executor.submit_writeback([1], [[3]], [[0]])
        executor.flush()
        self._drain(executor, 1)

        # Spy on the per-tensor events to pin the layer -> event mapping.
        captured = {}
        orig = mirror.load_pages_with_events

        def spy(pairs, stream):
            events = orig(pairs, stream)
            captured["events"] = events
            return events

        mirror.load_pages_with_events = spy
        executor.submit_loadback([2], [[0]], [[3]])
        executor.flush()
        events = captured["events"]
        self.assertEqual(len(events), 8)
        producer_idx = executor.get_producer_index(self.CacheKind.KV, 2)
        producer_event = executor._counter.events[producer_idx]
        # State layers 0/2 ack on their ssm event (conv precedes ssm on the
        # serial stream, so it covers the pair); attention layer 1 keeps its
        # V-tensor event (num_k=2, k-index 0 -> events[2]); the LAST layer
        # pins events[-1] so finish_event (producer-slot reuse fence) covers
        # the trailing state copies.
        self.assertIs(producer_event.load_events[0], events[5])
        self.assertIs(producer_event.load_events[1], events[2])
        self.assertIs(producer_event.load_events[2], events[7])
        self.assertIs(producer_event.load_events[3], events[-1])
        self.torch.cuda.synchronize()
        self.assertTrue(producer_event.finish_event.query())
        self._drain(executor, 1)

    def test_state_pool_roundtrip_with_fencing(self):
        torch = self.torch
        pool = self._state_pool()
        executor = self._executor(pool)
        mirror = executor.mirror

        device_pages = [1, 2]
        self._fill_spans(mirror, device_pages)
        before = self._snapshot_spans(mirror, device_pages)

        executor.submit_writeback([21], [[1, 2]], [[5, 6]])
        executor.flush()
        self._drain(executor, 1)

        for (dev, _), span in zip(mirror.tensor_pairs, mirror.row_spans):
            for d in device_pages:
                dev[d * span : (d + 1) * span].zero_()
        torch.cuda.synchronize()

        executor.submit_loadback([23], [[5, 6]], [[1, 2]])
        executor.flush()
        producer_idx = executor.get_producer_index(self.CacheKind.KV, 23)
        self.assertIsNotNone(producer_idx)
        executor.set_consumer(self.CacheKind.KV, [producer_idx])
        for layer_id in range(4):
            pool.layer_transfer_counter.wait_until(layer_id)
        torch.cuda.synchronize()
        self._drain(executor, 1)

        after = self._snapshot_spans(mirror, device_pages)
        for tensor_idx in range(len(mirror.tensor_pairs)):
            for d in device_pages:
                self.assertTrue(
                    torch.equal(
                        before[tensor_idx][d].view(torch.uint8),
                        after[tensor_idx][d].view(torch.uint8),
                    ),
                    f"tensor {tensor_idx} device page {d} not byte-exact",
                )

    def test_empty_op_acks_immediately(self):
        pool = self._pool()
        executor = self._executor(pool)
        # C++ dedups transfers across ops of one batched operation, so an op
        # can arrive with empty page lists; it still owes exactly one ack.
        executor.submit_writeback([11], [[]], [[]])
        executor.submit_loadback([12], [[]], [[]])
        executor.flush()
        results = self._drain(executor, 2)
        kinds = {type(r).__name__: int(r.op_id) for r in results}
        self.assertEqual(kinds, {"WriteBackDoneEvent": 11, "LoadBackDoneEvent": 12})


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