"""FlatHostMirror (M15 Phase D1): byte-blind pinned-CPU slab mirror. Pins the transport contract only (no engine wiring): one mirror per distinct device KV tensor, whole-page row-range copies both directions, per-tensor load events, and the layer -> tensor-index mapping D2 fences on. """ 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 (pairs 0/1); linear_attention # disables slab pairing, so the KV side stays per-layer -- and under the # flat GDN predicate the state layers' k/v slots are None (M18a T4). GDN_LAYER_TYPES = ("linear_attention", "full_attention") * 2 class FlatHostMirrorTest(unittest.TestCase): """Real (tiny) MHATokenToKVPool on GPU, slab and legacy layouts.""" def setUp(self): try: import torch from tokenspeed.runtime.cache.flat_host_mirror import ( FlatHostMirror, ) from tokenspeed.runtime.layers.attention.kv_cache.mha import ( MHATokenToKVPool, ) except (ImportError, ModuleNotFoundError) as exc: self.skipTest(f"needs torch + tokenspeed_kernel: {exc}") if not torch.cuda.is_available(): self.skipTest("needs a CUDA device") self.torch = torch self.FlatHostMirror = FlatHostMirror self.MHATokenToKVPool = MHATokenToKVPool def _pool(self, *, flat_ext: bool = True): 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=flat_ext): return self.MHATokenToKVPool(**kwargs) def _fill_device_pages(self, mirror, device_pages): # Sentinels distinct per (tensor, page); bf16-exact small ints. 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 _roundtrip_assert(self, mirror, pairs): torch = self.torch p = mirror.page_size device_pages = [d for d, _ in pairs] self._fill_device_pages(mirror, device_pages) before = self._snapshot(mirror, device_pages) stream = torch.cuda.Stream() mirror.store_pages(pairs, stream) stream.synchronize() for dev, _ in mirror.tensor_pairs: for d in device_pages: dev[d * p : (d + 1) * p].zero_() torch.cuda.synchronize() mirror.load_pages(pairs, stream) stream.synchronize() 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", ) def test_slab_roundtrip(self): pool = self._pool(flat_ext=True) mirror = self.FlatHostMirror(pool, num_host_pages=8) # 4 layers dedup to 2 K + 2 V slabs. self.assertEqual(len(mirror.tensor_pairs), 4) self._roundtrip_assert(mirror, [(1, 5), (2, 6), (3, 7)]) # 4 mirrors x page_size 4 x row 1*8 bf16 (16 B) = 256 B per page. self.assertEqual(mirror.bytes_per_host_page(), 4 * 4 * 16) def test_interleaved_groups_roundtrip(self): # Pages owned by different groups: byte-blind copies need no # group awareness (id-exclusivity keeps rows disjoint). pool = self._pool(flat_ext=True) mirror = self.FlatHostMirror(pool, num_host_pages=4) self._roundtrip_assert(mirror, [(2, 0), (3, 1)]) def test_legacy_roundtrip(self): # Legacy layout: all 4+4 per-layer mirrors carry data; copying # rows dead for a page's owner group is harmless (byte-exact). pool = self._pool(flat_ext=False) mirror = self.FlatHostMirror(pool, num_host_pages=8) self.assertEqual(len(mirror.tensor_pairs), 8) self._roundtrip_assert(mirror, [(1, 3), (2, 4)]) def test_events_and_layer_mapping(self): torch = self.torch pool = self._pool(flat_ext=True) mirror = self.FlatHostMirror(pool, num_host_pages=8) self._fill_device_pages(mirror, [1]) stream = torch.cuda.Stream() events = mirror.load_pages_with_events([(1, 5)], stream) self.assertEqual(len(events), len(mirror.tensor_pairs)) stream.synchronize() self.assertTrue(all(event.query() for event in events)) # Slab: paired layers map to the same K-tensor index. self.assertEqual(mirror.num_k_tensors, 2) self.assertEqual( mirror.tensor_index_of_layer(0), mirror.tensor_index_of_layer(1) ) self.assertEqual( mirror.tensor_index_of_layer(2), mirror.tensor_index_of_layer(3) ) self.assertNotEqual( mirror.tensor_index_of_layer(0), mirror.tensor_index_of_layer(2) ) for layer_id in range(4): idx = mirror.tensor_index_of_layer(layer_id) self.assertIs(mirror.tensor_pairs[idx][0], pool.k_buffer[layer_id]) self.assertIs( mirror.tensor_pairs[idx + mirror.num_k_tensors][0], pool.v_buffer[layer_id], ) # Legacy: every layer maps to a distinct index. legacy = self.FlatHostMirror(self._pool(flat_ext=False), num_host_pages=2) self.assertEqual( {legacy.tensor_index_of_layer(i) for i in range(4)}, {0, 1, 2, 3} ) class FlatHostMirrorStateSlabTest(unittest.TestCase): """State slabs join the mirrored set: tensor_pairs order is K*, V*, then (conv, ssm) flattened in slab order; state mirrors use 1-row PAGE spans (state slabs are page-indexed) while KV mirrors span page_size token rows.""" CONV_SHAPE = (2, 4) # 16 B/row bf16 SSM_SHAPE = (2, 8) # 32 B/row bf16 def setUp(self): try: import torch from tokenspeed.runtime.cache.flat_host_mirror import ( FlatHostMirror, flat_bytes_per_host_page, ) from tokenspeed.runtime.layers.attention.kv_cache.mha import ( MHATokenToKVPool, ) except (ImportError, ModuleNotFoundError) as exc: self.skipTest(f"needs torch + tokenspeed_kernel: {exc}") if not torch.cuda.is_available(): self.skipTest("needs a CUDA device") self.torch = torch self.FlatHostMirror = FlatHostMirror self.flat_bytes_per_host_page = flat_bytes_per_host_page self.MHATokenToKVPool = MHATokenToKVPool def _pool(self, *, with_state: bool = True): 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, ) if with_state: kwargs.update( conv_state_shape=self.CONV_SHAPE, temporal_state_shape=self.SSM_SHAPE, ) with mock.patch(_PKG_FLAT_PROBE, return_value=True): return self.MHATokenToKVPool(**kwargs) def _fill_device_pages(self, mirror, device_pages): # Sentinels distinct per (tensor, page); bf16-exact small ints. 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(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_tensors_follow_kv_in_slab_order(self): pool = self._pool() mirror = self.FlatHostMirror(pool, num_host_pages=8) # Flat GDN: state layers carry no KV (k/v slots are None, M18a T4), # so only the 2 attention layers mirror KV (2 K + 2 V), then # conv0, ssm0, conv1, ssm1 -- PINNED order: K*, V*, state tensors # flattened in slab order. self.assertEqual(mirror.num_k_tensors, 2) self.assertEqual(len(mirror.tensor_pairs), 8) self.assertEqual(len(pool.state_slabs), 2) for n, (conv, ssm) in enumerate(pool.state_slabs): self.assertIs(mirror.tensor_pairs[4 + 2 * n][0], conv) self.assertIs(mirror.tensor_pairs[4 + 2 * n + 1][0], ssm) # Per-pair row spans: page_size token rows for KV, 1 page row for # state (state slabs are page-indexed). self.assertEqual(mirror.row_spans, (4,) * 4 + (1,) * 4) for (dev, host), span in zip(mirror.tensor_pairs, mirror.row_spans): if span == 1: self.assertEqual(host.shape, (8, *dev.shape[1:])) else: self.assertEqual(host.shape, (8 * 4, *dev.shape[1:])) def test_bytes_per_host_page_includes_state_rows(self): # Without state shapes the flat GDN predicate is off: all 4 layers # keep KV -> 8 mirrors x page_size 4 x 16 B rows = 512 B. base = self.flat_bytes_per_host_page(self._pool(with_state=False)) self.assertEqual(base, 512) # Flat GDN: state layers carry no KV -> 4 KV mirrors (256 B) plus # 2 state layers x (conv 2*4 + ssm 2*8) bf16 page rows (2 x 48 B). pool = self._pool() with_state = self.flat_bytes_per_host_page(pool) self.assertEqual(with_state, 4 * 4 * 16 + 96) mirror = self.FlatHostMirror(pool, num_host_pages=2) self.assertEqual(mirror.bytes_per_host_page(), with_state) def test_state_roundtrip(self): torch = self.torch pool = self._pool() mirror = self.FlatHostMirror(pool, num_host_pages=8) pairs = [(1, 5), (2, 6), (3, 7)] device_pages = [d for d, _ in pairs] self._fill_device_pages(mirror, device_pages) before = self._snapshot(mirror, device_pages) stream = torch.cuda.Stream() mirror.store_pages(pairs, stream) stream.synchronize() 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() events = mirror.load_pages_with_events(pairs, stream) self.assertEqual(len(events), len(mirror.tensor_pairs)) stream.synchronize() self.assertTrue(all(event.query() for event in events)) 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", ) def test_state_tensor_indices_of_layer(self): pool = self._pool() mirror = self.FlatHostMirror(pool, num_host_pages=2) # State layers 0/2 bind slab pairs 0/1 -> flattened indices after # the 4 KV mirrors; conv immediately precedes its ssm. self.assertEqual(mirror.state_tensor_indices_of_layer(0), (4, 5)) self.assertEqual(mirror.state_tensor_indices_of_layer(2), (6, 7)) self.assertIsNone(mirror.state_tensor_indices_of_layer(1)) self.assertIsNone(mirror.state_tensor_indices_of_layer(3)) # Pools without state slabs expose no state indices for any layer. kv_only = self.FlatHostMirror(self._pool(with_state=False), num_host_pages=2) for layer_id in range(4): self.assertIsNone(kv_only.state_tensor_indices_of_layer(layer_id)) class FlatHostMirrorNoneKVTest(unittest.TestCase): """Flat GDN pools carry None k/v slots on state layers (M18a T4): the mirror's identity-dedup walks must skip them and mirror only the real slabs. CPU stub pool, no CUDA, no scheduler ext -- state mirroring via get_state_buffers is a separate surface and unaffected.""" 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 _stub_pool(self): import types torch = self.torch rows = 8 kv = [torch.zeros((rows, 1, 8), dtype=torch.bfloat16) for _ in range(4)] return types.SimpleNamespace( page_size=4, k_buffer=[None, kv[0], None, kv[1]], v_buffer=[None, kv[2], None, kv[3]], ) def test_mirror_skips_none_kv_entries(self): stub = self._stub_pool() # 4 real mirrors x page_size 4 x 16 B rows = 256 B per host page. self.assertEqual(self.flat_bytes_per_host_page(stub), 256) mirror = self.FlatHostMirror(stub, num_host_pages=2) self.assertEqual(mirror.num_k_tensors, 2) self.assertEqual(len(mirror.tensor_pairs), 4) self.assertIs(mirror.tensor_pairs[0][0], stub.k_buffer[1]) self.assertIs(mirror.tensor_pairs[1][0], stub.k_buffer[3]) # KV layers keep their tensor-index mapping; state layers have no # KV mirror and must fail loud if D2 fencing ever asks for one. self.assertEqual(mirror.tensor_index_of_layer(1), 0) self.assertEqual(mirror.tensor_index_of_layer(3), 1) with self.assertRaisesRegex(ValueError, r"state layer"): mirror.tensor_index_of_layer(0) if __name__ == "__main__": unittest.main()