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

"""Flat assembly line for GDN hybrids (M17 C4).

Three contracts: the Qwen3.5 config exposes ``layer_types`` in the
paged-cache label vocabulary; the page-size equalization decision
(``equalized_block_size``) inflates P to cover the GDN state row;
and an MHAConfig carrying state shapes builds a full-coverage pool with
one (conv, ssm) slab pair per state layer, both cache groups published,
and the ctor geometry check enforcing the equalized P. Flat GDN sizing
itself is plan-driven (plan_component_tensors, test_flat_memory_plan).
"""

from __future__ import annotations

import importlib.util
import os
import pathlib
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=15, suite="runtime-1gpu")

_CONFIGS_DIR = (
    pathlib.Path(__file__).resolve().parents[2]
    / "python"
    / "tokenspeed"
    / "runtime"
    / "configs"
)

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


def _load(mod_name: str, file_name: str):
    spec = importlib.util.spec_from_file_location(mod_name, _CONFIGS_DIR / file_name)
    assert spec is not None and spec.loader is not None
    mod = importlib.util.module_from_spec(spec)
    # Register before exec: on py3.9 @dataclass + `from __future__ import
    # annotations` resolves field types via sys.modules[cls.__module__].
    sys.modules[mod_name] = mod
    spec.loader.exec_module(mod)
    return mod


_plan = _load("flat_memory_plan_gdn_assembly_under_test", "flat_memory_plan.py")
equalized_block_size = _plan.equalized_block_size


# Qwen3.5-ish interleaving: 3 linear layers then 1 full, times 12 (48 layers).
QWEN3_5ISH_LAYER_TYPES = (["linear_attention"] * 3 + ["full_attention"]) * 12

# Qwen3.5 defaults at attn TP=1 (configs/qwen3_5_text_base_config.py):
# KV row: 2 (K+V) * 2 kv heads * 256 head_dim * 2 B (bf16) per token-layer.
QWEN3_5ISH_KV_BYTES_PER_SLOT = 2 * 2 * 256 * 2  # 2048
# conv: (2*128*16 + 128*32) x (4 - 1) in bf16; ssm: 32 x 128 x 128 in fp32.
QWEN3_5ISH_CONV_BYTES = (2 * 128 * 16 + 128 * 32) * 3 * 2  # 49152
QWEN3_5ISH_SSM_BYTES = 32 * 128 * 128 * 4  # 2097152


class Qwen3_5LayerTypesTest(unittest.TestCase):
    """The config's layer_types property (interleaving + label vocabulary)."""

    def setUp(self):
        try:
            from tokenspeed.runtime.configs.qwen3_5_text_base_config import (
                Qwen3_5BaseTextConfig,
            )
        except (ImportError, ModuleNotFoundError) as exc:
            self.skipTest(f"needs torch + transformers: {exc}")
        self.config_cls = Qwen3_5BaseTextConfig

    def test_layer_types_interleaving(self):
        cfg = self.config_cls(num_hidden_layers=8, full_attention_interval=4)
        self.assertEqual(
            cfg.layer_types,
            (["linear_attention"] * 3 + ["full_attention"]) * 2,
        )

    def test_layers_block_type_keeps_checkpoint_label(self):
        # models/qwen3_5.py keys layer construction on the checkpoint's
        # "attention" label; layer_types must not change it.
        cfg = self.config_cls(num_hidden_layers=4, full_attention_interval=4)
        self.assertEqual(
            cfg.layers_block_type, ["linear_attention"] * 3 + ["attention"]
        )

    def test_tracks_nextn_interval_override(self):
        # models/qwen3_5_nextn.py rewrites full_attention_interval AFTER
        # construction; layer_types must follow (property, not __init__).
        cfg = self.config_cls(num_hidden_layers=2, full_attention_interval=4)
        cfg.full_attention_interval = 1
        self.assertEqual(cfg.layer_types, ["full_attention"] * 2)


class EqualizedPageSizeTest(unittest.TestCase):
    """Pure equalization decision (no torch)."""

    def _equalized(self, block_size, **kwargs):
        return equalized_block_size(
            layer_types=QWEN3_5ISH_LAYER_TYPES,
            kv_bytes_per_slot=QWEN3_5ISH_KV_BYTES_PER_SLOT,
            state_const_bytes={
                "conv": QWEN3_5ISH_CONV_BYTES,
                "ssm": QWEN3_5ISH_SSM_BYTES,
            },
            block_size=block_size,
            **kwargs,
        )

    def test_inflates_to_cover_state_row(self):
        # ceil((49152 + 2097152) / 2048) = 1048 tokens to cover the state
        # row; default alignment = original block size 64 -> 1088.
        self.assertEqual(self._equalized(64), 1088)

    def test_unchanged_when_kv_row_already_covers(self):
        self.assertEqual(self._equalized(2048), 2048)

    def test_explicit_alignment(self):
        # 1048 rounded up to a multiple of 16 -> 1056.
        self.assertEqual(self._equalized(64, alignment=16), 1056)

    def test_no_state_layers_is_identity(self):
        self.assertEqual(
            equalized_block_size(
                layer_types=["full_attention"] * 4,
                kv_bytes_per_slot=QWEN3_5ISH_KV_BYTES_PER_SLOT,
                state_const_bytes={},
                block_size=64,
            ),
            64,
        )


class GdnFlatPoolAssemblyTest(unittest.TestCase):
    """MHAConfig with state shapes -> create_pool: full-coverage pool with
    state slabs, both published groups, and the equalized-P geometry gate."""

    # 3 linear + 1 full; kv row = 2 * 1 head * 8 dim * 2 B = 32 B/slot;
    # state row = conv (4*3 bf16 = 24 B) + ssm (2*4*4 fp32 = 128 B) = 152 B;
    # ceil(152 / 32) = 5 > block 4 -> equalized P = 8 (multiple of 4).
    LAYER_TYPES = ("linear_attention",) * 3 + ("full_attention",)
    CONV_SHAPE = (4, 3)
    TEMPORAL_SHAPE = (2, 4, 4)

    def setUp(self):
        try:
            import torch

            from tokenspeed.runtime.layers.attention.configs.mha import (
                MHAConfig,
            )
        except (ImportError, ModuleNotFoundError) as exc:
            self.skipTest(f"needs torch + tokenspeed_kernel: {exc}")
        self.torch = torch
        self.MHAConfig = MHAConfig

    def _config(self, page_size: int):
        torch = self.torch
        return self.MHAConfig(
            device="cpu",
            backend_name=None,
            num_attention_heads=2,
            num_kv_heads=1,
            head_dim=8,
            attn_tp_size=1,
            dtype=torch.bfloat16,
            kv_cache_dtype=torch.bfloat16,
            page_size=page_size,
            context_len=64,
            max_bs=2,
            max_graph_bs=2,
            kv_cache_quant_method=None,
            layer_types=self.LAYER_TYPES,
            max_scheduled_tokens=16,
            conv_state_shape=self.CONV_SHAPE,
            temporal_state_shape=self.TEMPORAL_SHAPE,
            conv_dtype=torch.bfloat16,
            ssm_dtype=torch.float32,
        )

    def _pool(self, page_size: int):
        with mock.patch(_PKG_FLAT_PROBE, return_value=True):
            return self._config(page_size).create_pool(
                len(self.LAYER_TYPES), 32, 0, False
            )

    def test_equalization_decision_matches_pure_helper(self):
        self.assertEqual(
            equalized_block_size(
                layer_types=list(self.LAYER_TYPES),
                kv_bytes_per_slot=32,
                state_const_bytes={"conv": 24, "ssm": 128},
                block_size=4,
            ),
            8,
        )

    def test_assembly_at_equalized_page_size(self):
        pool = self._pool(page_size=8)
        # One (conv, ssm) slab pair per state layer, rows over the shared
        # page-id space (size // P + 1 null row).
        self.assertEqual(len(pool.state_slabs), 3)
        conv, ssm = pool.state_slabs[0]
        self.assertEqual(tuple(conv.shape), (5, *self.CONV_SHAPE))
        self.assertEqual(conv.dtype, self.torch.bfloat16)
        self.assertEqual(tuple(ssm.shape), (5, *self.TEMPORAL_SHAPE))
        self.assertEqual(ssm.dtype, self.torch.float32)
        # Both groups published (upstream signal for flat state paging).
        self.assertEqual(
            sorted(spec.group_id for spec in pool.paged_cache_group_specs),
            ["full_attention", "linear_attention"],
        )
        # Plan-sized coverage (M18a T4): the k/v lists stay layer-indexed,
        # but state layers carry no KV tensors (None slots) -- only the
        # full-attention layer allocates.
        self.assertEqual(len(pool.k_buffer), len(self.LAYER_TYPES))
        for layer_id, label in enumerate(self.LAYER_TYPES):
            if label == "linear_attention":
                self.assertIsNone(pool.k_buffer[layer_id])
                self.assertIsNone(pool.v_buffer[layer_id])
            else:
                self.assertIsNotNone(pool.k_buffer[layer_id])
                self.assertIsNotNone(pool.v_buffer[layer_id])

    def test_geometry_raises_at_original_page_size(self):
        with self.assertRaisesRegex(ValueError, "pre-equalized"):
            self._pool(page_size=4)


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