dotcache/backends/torch_mps.py

from __future__ import annotations

from collections import OrderedDict
from dataclasses import dataclass
import time
from typing import Any, Literal, Sequence

import numpy as np
from ..tracing import ExecutionTrace
from ..types import EncodedPage, PageHeader
from ..modes.m2_key_sketch import segment_ids_for_token_count
from ..modes.m4_key_project import fixed_project_basis
from ..modes.turbo3 import TURBO3_CENTROIDS, fwht_last_dim
from ..packing import words_per_group

TorchDevice = Literal["mps", "cuda"]
PreparedDevice = Literal["torch_mps", "torch_cuda"]

_UNPACK_METADATA: dict[tuple[TorchDevice, int], tuple[Any, Any]] = {}
_SPILL_UNPACK_METADATA: dict[tuple[TorchDevice, int, int], tuple[Any, Any, Any, Any, Any, Any]] = {}
_TURBO3_CENTROID_TENSORS: dict[TorchDevice, Any] = {}
_FWHT_MATRICES: dict[tuple[TorchDevice, int], Any] = {}
_M4_BASIS_TENSORS: dict[tuple[TorchDevice, int, int, str], Any] = {}
_SEGMENT_ID_TENSORS: dict[tuple[TorchDevice, int, int], Any] = {}
_MAX_PREPARE_PAGES_PER_CHUNK = 128
_MPS_M0_KEY_PREPARE_PAGES_PER_CHUNK = 256
_MAX_PREPARED_CHUNK_CACHE_ENTRIES = 64
_MAX_PREPARED_CHUNK_CACHE_RESIDENT_BYTES = 64 * 1024 * 1024
_PREPARED_CHUNK_CACHE_BUDGET_OVERRIDE_BYTES: int | None = None
_MIN_PREPARED_CHUNK_CACHE_PAGE_COUNT = 4
_PREPARED_CHUNK_CACHE_KINDS = frozenset({"K", "V"})
_PREPARED_CHUNK_CACHE: "OrderedDict[tuple[tuple[int, int], ...], PreparedChunkMPS]" = OrderedDict()
_PREPARED_CHUNK_CACHE_RESIDENT_BYTES = 0
_PREPARED_GROUPED_CHUNK_CACHE: "OrderedDict[tuple[tuple[tuple[int, int], ...], ...], PreparedGroupedChunkMPS]" = OrderedDict()
_PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES = 0
_PREPARED_PAGE_UID = 1
_PREPARED_CHUNK_CACHE_TOUCH_ID = 0
def _load_torch():
    try:
        import torch
    except ImportError as exc:  # pragma: no cover - exercised in environments without torch
        raise RuntimeError("torch is required for the torch accelerator backends") from exc
    return torch

def _backend_name(device_type: TorchDevice) -> PreparedDevice:
    return "torch_cuda" if device_type == "cuda" else "torch_mps"


def torch_device_available(device_type: TorchDevice) -> bool:
    try:
        import torch
    except ImportError:
        return False
    if device_type == "mps":
        return bool(torch.backends.mps.is_available())
    return bool(torch.cuda.is_available())


def mps_available() -> bool:
    return torch_device_available("mps")


def _device_tensor(array: np.ndarray, *, device: TorchDevice):
    torch = _load_torch()
    return torch.from_numpy(np.ascontiguousarray(array)).to(device=device)


def _torch_pack_codes(codes, *, bits: int, layout: str):
    torch = _load_torch()
    if int(codes.ndim) != 4:
        raise ValueError("codes must have shape [page_count, token_count, num_groups, group_size]")
    page_count, token_count, num_groups, group_size = map(int, codes.shape)
    if 32 % bits == 0:
        symbols_per_word = 32 // bits
        if group_size % symbols_per_word != 0:
            raise ValueError("torch-side code packing requires group_size divisible by symbols_per_word")
        word_count = group_size // symbols_per_word
        grouped = codes.to(dtype=torch.int32).reshape(page_count, token_count, num_groups, word_count, symbols_per_word)
        shifts = (torch.arange(symbols_per_word, dtype=torch.int32, device=codes.device) * int(bits)).reshape(1, 1, 1, 1, -1)
        packed = torch.bitwise_left_shift(grouped, shifts).sum(dim=-1).to(dtype=torch.int32)
    else:
        word_count = words_per_group(group_size, bits)
        packed = torch.zeros((page_count, token_count, num_groups, word_count), dtype=torch.int32, device=codes.device)
        codes_i64 = codes.to(dtype=torch.int64)
        for symbol_index in range(group_size):
            bit_offset = symbol_index * int(bits)
            word_index = bit_offset // 32
            bit_index = bit_offset % 32
            value = codes_i64[..., symbol_index]
            packed[..., word_index] = torch.bitwise_or(
                packed[..., word_index],
                torch.bitwise_left_shift(value, bit_index).to(dtype=torch.int32),
            )
            spill = bit_index + int(bits) - 32
            if spill > 0:
                packed[..., word_index + 1] = torch.bitwise_or(
                    packed[..., word_index + 1],
                    torch.bitwise_right_shift(value, int(bits) - spill).to(dtype=torch.int32),
                )
    if layout == "group_major":
        return packed.permute(0, 2, 1, 3).contiguous()
    if layout == "token_major":
        return packed.contiguous()
    raise ValueError("layout must be group_major or token_major")


@dataclass(slots=True)
class PreparedPageTorch:
    device_type: TorchDevice
    source_page: EncodedPage
    header: PageHeader
    payload: Any | None = None
    scales: Any | None = None
    bias: Any | None = None
    codebooks: Any | None = None
    m2_sketch: Any | None = None
    m2_basis: Any | None = None
    m2_mean: Any | None = None
    escape_payload: Any | None = None
    escape_scales: Any | None = None
    unpack_shifts: Any | None = None
    unpack_mask: Any | None = None
    host_to_device_nbytes: int = 0
    resident_nbytes: int = 0
    cache_uid: int = 0

    @property
    def payload_nbytes(self) -> int:
        return self.source_page.payload_nbytes

    @property
    def metadata_nbytes(self) -> int:
        return self.source_page.metadata_nbytes


PreparedPageMPS = PreparedPageTorch


@dataclass(slots=True)
class PreparedChunkMPS:
    header: PageHeader
    payload_groups: tuple[Any, ...]
    codes_groups: tuple[Any, ...] | None
    scales_groups: tuple[Any, ...] | None
    bias_groups: tuple[Any, ...] | None
    escape_payload_batch: Any | None = None
    escape_scales_batch: Any | None = None
    fused_scaled_codes: Any | None = None
    m2_sketch_groups: tuple[Any, ...] | None = None
    m2_basis_groups: tuple[Any, ...] | None = None
    m2_mean_groups: tuple[Any, ...] | None = None
    m2_segment_ids: Any | None = None
    resident_nbytes: int = 0
    touch_id: int = 0


@dataclass(slots=True)
class PreparedGroupedChunkMPS:
    header: PageHeader
    payload_groups: tuple[Any, ...]
    codes_groups: tuple[Any, ...] | None
    scales_groups: tuple[Any, ...] | None
    bias_groups: tuple[Any, ...] | None
    fused_scaled_codes: Any | None = None
    m2_sketch_groups: tuple[Any, ...] | None = None
    m2_basis_groups: tuple[Any, ...] | None = None
    m2_mean_groups: tuple[Any, ...] | None = None
    m2_segment_ids: Any | None = None
    m2_sketch_tensor: Any | None = None
    m2_basis_tensor: Any | None = None
    m2_mean_tensor: Any | None = None
    resident_nbytes: int = 0
    payload_groups_tensor: Any | None = None
    scales_groups_tensor: Any | None = None
    bias_groups_tensor: Any | None = None
    touch_id: int = 0


def _supports_fused_two_group64(header: PageHeader) -> bool:
    return bool(
        header.head_dim == 64
        and header.padded_head_dim == 64
        and header.group_size == 32
        and header.num_groups == 2
    )

def _supports_fused_m0_3bit(header: PageHeader, *, device_type: TorchDevice) -> bool:
    return bool(
        device_type == "mps"
        and header.mode_default == "M0"
        and header.bits == 3
        and header.quant_scheme == "affine"
    )


def _supports_grouped_fused_only_cache(header: PageHeader, *, device_type: TorchDevice) -> bool:
    if _supports_fused_m0_3bit(header, device_type=device_type):
        return True
    return _supports_fused_two_group64(header) and device_type in {"cuda", "mps"}


def _supports_packed_four_group128_cuda(header: PageHeader, *, device_type: TorchDevice) -> bool:
    return bool(
        device_type == "cuda"
        and header.mode_default == "M0"
        and header.quant_scheme == "affine"
        and header.layout == "group_major"
        and header.bits == 4
        and header.group_size == 32
        and header.num_groups == 4
        and header.padded_head_dim == 128
    )


def _fused_two_group64_cache_dtype(*, device_type: TorchDevice):
    torch = _load_torch()
    if device_type == "cuda":
        return torch.float16
    return torch.float32


def _m0_affine_metadata_dtype(*, device_type: TorchDevice):
    return _load_torch().float32


def _escape_scale_dtype(*, device_type: TorchDevice):
    torch = _load_torch()
    if device_type == "mps":
        return torch.float32
    return torch.float16
def prepared_chunk_cache_resident_bytes() -> int:
    return int(_PREPARED_CHUNK_CACHE_RESIDENT_BYTES + _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES)


def _effective_max_prepared_chunk_cache_resident_bytes() -> int:
    if _PREPARED_CHUNK_CACHE_BUDGET_OVERRIDE_BYTES is None:
        return int(_MAX_PREPARED_CHUNK_CACHE_RESIDENT_BYTES)
    return int(min(_MAX_PREPARED_CHUNK_CACHE_RESIDENT_BYTES, _PREPARED_CHUNK_CACHE_BUDGET_OVERRIDE_BYTES))


def _next_prepared_chunk_cache_touch_id() -> int:
    global _PREPARED_CHUNK_CACHE_TOUCH_ID
    _PREPARED_CHUNK_CACHE_TOUCH_ID += 1
    return _PREPARED_CHUNK_CACHE_TOUCH_ID


def _touch_prepared_chunk(chunk: PreparedChunkMPS | PreparedGroupedChunkMPS) -> None:
    chunk.touch_id = _next_prepared_chunk_cache_touch_id()


def _evict_oldest_prepared_chunk_cache_entry() -> bool:
    global _PREPARED_CHUNK_CACHE_RESIDENT_BYTES
    global _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES
    oldest_single = next(iter(_PREPARED_CHUNK_CACHE.items()), None)
    oldest_grouped = next(iter(_PREPARED_GROUPED_CHUNK_CACHE.items()), None)
    if oldest_single is None and oldest_grouped is None:
        return False
    if oldest_grouped is None or (
        oldest_single is not None and oldest_single[1].touch_id <= oldest_grouped[1].touch_id
    ):
        _, evicted_chunk = _PREPARED_CHUNK_CACHE.popitem(last=False)
        _PREPARED_CHUNK_CACHE_RESIDENT_BYTES = max(
            0,
            _PREPARED_CHUNK_CACHE_RESIDENT_BYTES - evicted_chunk.resident_nbytes,
        )
        return True
    _, evicted_chunk = _PREPARED_GROUPED_CHUNK_CACHE.popitem(last=False)
    _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES = max(
        0,
        _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES - evicted_chunk.resident_nbytes,
    )
    return True


def _trim_prepared_chunk_cache() -> None:
    effective_max_resident_bytes = _effective_max_prepared_chunk_cache_resident_bytes()
    while (
        len(_PREPARED_CHUNK_CACHE) + len(_PREPARED_GROUPED_CHUNK_CACHE) > _MAX_PREPARED_CHUNK_CACHE_ENTRIES
        or prepared_chunk_cache_resident_bytes() > effective_max_resident_bytes
    ):
        if not _evict_oldest_prepared_chunk_cache_entry():
            break


def configure_prepared_chunk_cache(
    *,
    max_entries: int | None = None,
    max_resident_bytes: int | None = None,
    min_page_count: int | None = None,
    cached_kinds: Sequence[str] | None = None,
    clear: bool = True,
) -> None:
    global _MAX_PREPARED_CHUNK_CACHE_ENTRIES
    global _MAX_PREPARED_CHUNK_CACHE_RESIDENT_BYTES
    global _MIN_PREPARED_CHUNK_CACHE_PAGE_COUNT
    global _PREPARED_CHUNK_CACHE_KINDS
    if max_entries is not None:
        _MAX_PREPARED_CHUNK_CACHE_ENTRIES = max(0, int(max_entries))
    if max_resident_bytes is not None:
        _MAX_PREPARED_CHUNK_CACHE_RESIDENT_BYTES = max(0, int(max_resident_bytes))
    if min_page_count is not None:
        _MIN_PREPARED_CHUNK_CACHE_PAGE_COUNT = max(1, int(min_page_count))
    if cached_kinds is not None:
        _PREPARED_CHUNK_CACHE_KINDS = frozenset(str(kind) for kind in cached_kinds)
    if clear:
        clear_prepared_chunk_cache()
        return
    _trim_prepared_chunk_cache()


def set_prepared_chunk_cache_budget_override(*, max_resident_bytes: int | None) -> None:
    global _PREPARED_CHUNK_CACHE_BUDGET_OVERRIDE_BYTES
    _PREPARED_CHUNK_CACHE_BUDGET_OVERRIDE_BYTES = None if max_resident_bytes is None else max(0, int(max_resident_bytes))
    _trim_prepared_chunk_cache()


def clear_prepared_chunk_cache() -> None:
    global _PREPARED_CHUNK_CACHE_RESIDENT_BYTES
    global _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES
    _PREPARED_CHUNK_CACHE.clear()
    _PREPARED_CHUNK_CACHE_RESIDENT_BYTES = 0
    _PREPARED_GROUPED_CHUNK_CACHE.clear()
    _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES = 0


def _next_prepared_page_uid() -> int:
    global _PREPARED_PAGE_UID
    page_uid = _PREPARED_PAGE_UID
    _PREPARED_PAGE_UID += 1
    return page_uid


def _prepare_signature(page: EncodedPage | PreparedPageTorch) -> tuple[int | str, ...]:
    source_page = page.source_page if isinstance(page, PreparedPageTorch) else page
    header = source_page.header
    sketch_dim = int(source_page.m2_sketch.shape[-1]) if source_page.m2_sketch is not None else 0
    segment_count = int(source_page.m2_basis.shape[1]) if source_page.m2_basis is not None and source_page.m2_basis.ndim == 4 else 1
    centered = int(source_page.m2_mean is not None)
    return (
        header.kind,
        header.mode_default,
        header.token_count,
        header.head_dim,
        header.padded_head_dim,
        header.group_size,
        header.num_groups,
        header.bits,
        header.words_per_group,
        header.layout,
        header.quant_scheme,
        header.escape_dtype,
        sketch_dim,
        segment_count,
        centered,
    )


def _max_prepare_pages_for_source_page(page: EncodedPage | PreparedPageTorch, *, device_type: TorchDevice) -> int:
    source_page = page.source_page if isinstance(page, PreparedPageTorch) else page
    if device_type == "mps" and source_page.header.mode_default == "M0" and source_page.header.kind == "K":
        return _MPS_M0_KEY_PREPARE_PAGES_PER_CHUNK
    return _MAX_PREPARE_PAGES_PER_CHUNK


def _batched_signature(page: PreparedPageTorch) -> tuple[int | str, ...]:
    header = page.header
    sketch_dim = int(page.m2_sketch.shape[-1]) if page.m2_sketch is not None else 0
    segment_count = int(page.m2_basis.shape[1]) if page.m2_basis is not None and page.m2_basis.dim() == 4 else 1
    centered = int(page.m2_mean is not None)
    return (
        page.device_type,
        header.kind,
        header.mode_default,
        header.token_count,
        header.head_dim,
        header.padded_head_dim,
        header.group_size,
        header.num_groups,
        header.bits,
        header.words_per_group,
        header.layout,
        header.quant_scheme,
        sketch_dim,
        segment_count,
        centered,
    )


def _chunk_compatible_source_pages(
    pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    device_type: TorchDevice,
) -> list[list[EncodedPage | PreparedPageTorch]]:
    chunks: list[list[EncodedPage | PreparedPageTorch]] = []
    current_chunk: list[EncodedPage | PreparedPageTorch] = []
    current_signature: tuple[int | str, ...] | None = None
    current_limit = _MAX_PREPARE_PAGES_PER_CHUNK
    for page in pages:
        signature = _prepare_signature(page)
        if current_chunk and (
            signature != current_signature or len(current_chunk) >= current_limit
        ):
            chunks.append(current_chunk)
            current_chunk = [page]
            current_signature = signature
            current_limit = _max_prepare_pages_for_source_page(page, device_type=device_type)
            continue
        if not current_chunk:
            current_signature = signature
            current_limit = _max_prepare_pages_for_source_page(page, device_type=device_type)
        current_chunk.append(page)
    if current_chunk:
        chunks.append(current_chunk)
    return chunks


def _chunk_compatible_pages(pages: Sequence[PreparedPageTorch]) -> list[list[PreparedPageTorch]]:
    chunks: list[list[PreparedPageTorch]] = []
    current_chunk: list[PreparedPageTorch] = []
    current_signature: tuple[int | str, ...] | None = None
    for page in pages:
        signature = _batched_signature(page)
        if current_chunk and signature != current_signature:
            chunks.append(current_chunk)
            current_chunk = [page]
            current_signature = signature
            continue
        if not current_chunk:
            current_signature = signature
        current_chunk.append(page)
    if current_chunk:
        chunks.append(current_chunk)
    return chunks


def _chunk_lengths_for_pages(pages: Sequence[PreparedPageTorch]) -> tuple[int, ...]:
    return tuple(len(chunk) for chunk in _chunk_compatible_pages(pages))


def _merged_chunk_lengths_for_page_groups(
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
) -> tuple[int, ...]:
    if not pages_by_group:
        return ()
    page_count = len(pages_by_group[0])
    if page_count == 0:
        return ()
    boundaries = {0, page_count}
    for group_pages in pages_by_group:
        if len(group_pages) != page_count:
            raise ValueError("all grouped page streams must have the same page count")
        offset = 0
        for chunk_length in _chunk_lengths_for_pages(group_pages):
            offset += int(chunk_length)
            boundaries.add(offset)
    sorted_boundaries = sorted(boundaries)
    return tuple(
        int(sorted_boundaries[index + 1] - sorted_boundaries[index])
        for index in range(len(sorted_boundaries) - 1)
    )


def _aligned_chunk_lengths_for_page_pairs(
    key_pages: Sequence[PreparedPageTorch],
    value_pages: Sequence[PreparedPageTorch],
) -> tuple[int, ...]:
    if len(key_pages) != len(value_pages):
        raise ValueError("key/value page streams must have matching page counts")
    if not key_pages:
        return ()
    lengths: list[int] = []
    current_length = 0
    current_key_signature: tuple[int | str, ...] | None = None
    current_value_signature: tuple[int | str, ...] | None = None
    for key_page, value_page in zip(key_pages, value_pages, strict=True):
        key_signature = _batched_signature(key_page)
        value_signature = _batched_signature(value_page)
        if (
            current_length > 0
            and (key_signature != current_key_signature or value_signature != current_value_signature)
        ):
            lengths.append(current_length)
            current_length = 0
        if current_length == 0:
            current_key_signature = key_signature
            current_value_signature = value_signature
        current_length += 1
    if current_length > 0:
        lengths.append(current_length)
    return tuple(lengths)


def _signature_buckets_for_page_chunk(
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
) -> tuple[tuple[int, ...], ...]:
    buckets: dict[tuple[tuple[int | str, ...], ...], list[int]] = {}
    for group_index, group_pages in enumerate(pages_by_group):
        signature = tuple(_batched_signature(page) for page in group_pages)
        buckets.setdefault(signature, []).append(int(group_index))
    return tuple(tuple(indices) for indices in buckets.values())


def _prepared_chunk_cache_key(pages: Sequence[PreparedPageTorch]) -> tuple[tuple[int, int], ...] | None:
    if not pages:
        return None
    if len(_chunk_compatible_pages(pages)) != 1:
        return None
    if pages[0].header.mode_default not in ("M0", "M2", "M3", "M4", "T3"):
        return None
    return tuple((int(page.cache_uid), int(page.header.token_count)) for page in pages)


def _segment_ids_tensor(token_count: int, segment_count: int, *, device_type: TorchDevice):
    cache_key = (device_type, int(token_count), int(segment_count))
    cached = _SEGMENT_ID_TENSORS.get(cache_key)
    if cached is not None:
        return cached
    tensor = _load_torch().from_numpy(
        segment_ids_for_token_count(int(token_count), int(segment_count))
    ).to(device=device_type)
    _SEGMENT_ID_TENSORS[cache_key] = tensor
    return tensor


def _build_prepared_chunk_mps(pages: Sequence[PreparedPageTorch]) -> PreparedChunkMPS:
    torch = _load_torch()
    if not pages:
        raise ValueError("pages must be non-empty")
    header = pages[0].header
    device_type = pages[0].device_type
    if header.mode_default not in ("M0", "M2", "M3", "M4", "T3"):
        raise ValueError("prepared chunk cache currently supports only M0, M2, M3, M4, and T3 pages")
    if header.mode_default == "M3":
        escape_payload_batch = torch.stack(
            [page.escape_payload[: header.token_count, : header.head_dim] for page in pages],
            dim=0,
        ).contiguous()
        escape_scales_batch = None
        resident_nbytes = int(escape_payload_batch.numel() * escape_payload_batch.element_size())
        if header.escape_dtype == "int8":
            escape_scales_batch = torch.stack(
                [page.escape_scales[: header.token_count] for page in pages],
                dim=0,
            ).contiguous()
            escape_scales_batch = escape_scales_batch.to(dtype=_escape_scale_dtype(device_type=device_type))
            resident_nbytes += int(escape_scales_batch.numel() * escape_scales_batch.element_size())
        return PreparedChunkMPS(
            header=header,
            payload_groups=(),
            codes_groups=None,
            scales_groups=None,
            bias_groups=None,
            escape_payload_batch=escape_payload_batch,
            escape_scales_batch=escape_scales_batch,
            fused_scaled_codes=None,
            resident_nbytes=resident_nbytes,
        )
    if header.mode_default == "M2":
        m2_sketch_groups = tuple(
            torch.stack([page.m2_sketch[:, group_index, :] for page in pages], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_basis_groups = tuple(
            torch.stack([page.m2_basis[group_index] for page in pages], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_mean_groups = tuple(
            torch.stack([page.m2_mean[group_index] for page in pages], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_segment_ids = None
        if m2_basis_groups and int(m2_basis_groups[0].dim()) == 4:
            m2_segment_ids = _segment_ids_tensor(
                header.token_count,
                int(m2_basis_groups[0].shape[1]),
                device_type=device_type,
            )
        resident_nbytes = sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_sketch_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_basis_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_mean_groups)
        return PreparedChunkMPS(
            header=header,
            payload_groups=(),
            codes_groups=None,
            scales_groups=None,
            bias_groups=None,
            m2_sketch_groups=m2_sketch_groups,
            m2_basis_groups=m2_basis_groups,
            m2_mean_groups=m2_mean_groups,
            m2_segment_ids=m2_segment_ids,
            fused_scaled_codes=None,
            resident_nbytes=resident_nbytes,
        )
    if header.mode_default == "M4":
        shared_source_basis = pages[0].source_page.m2_basis
        shared_basis = (
            pages[0].header.project_basis == "svd_shared"
            and pages[0].m2_basis is None
            and shared_source_basis is not None
            and all(page.source_page.m2_basis is shared_source_basis for page in pages)
        )
        m2_sketch_groups = tuple(
            torch.stack([page.m2_sketch[:, group_index, :] for page in pages], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_basis_groups = (
            tuple(
                (
                    _device_tensor(np.asarray(pages[0].source_page.m2_basis[group_index]), device=device_type).contiguous()
                    if shared_basis
                    else torch.stack([page.m2_basis[group_index] for page in pages], dim=0).contiguous()
                )
                for group_index in range(header.num_groups)
            )
            if (shared_basis or pages[0].m2_basis is not None)
            else None
        )
        m2_mean_groups = tuple(
            torch.stack([page.m2_mean[group_index] for page in pages], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        resident_nbytes = sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_sketch_groups)
        if m2_basis_groups is not None:
            resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_basis_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_mean_groups)
        return PreparedChunkMPS(
            header=header,
            payload_groups=(),
            codes_groups=None,
            scales_groups=None,
            bias_groups=None,
            m2_sketch_groups=m2_sketch_groups,
            m2_basis_groups=m2_basis_groups,
            m2_mean_groups=m2_mean_groups,
            m2_segment_ids=None,
            fused_scaled_codes=None,
            resident_nbytes=resident_nbytes,
        )
    payload_groups = tuple(torch.stack([page.payload[group_index] for page in pages], dim=0) for group_index in range(header.num_groups))
    codes_groups = tuple(
        _unpack_bits_torch(
            payload_groups[group_index].reshape(-1, header.words_per_group),
            pages[0].unpack_shifts,
            pages[0].unpack_mask,
            header.group_size,
        ).reshape(len(pages), header.token_count, header.group_size)
        for group_index in range(header.num_groups)
    )
    scales_groups = tuple(
        torch.stack([page.scales[:, group_index].to(torch.float32) for page in pages], dim=0)
        for group_index in range(header.num_groups)
    )
    bias_groups = (
        tuple(
            torch.stack([page.bias[:, group_index].to(torch.float32) for page in pages], dim=0)
            for group_index in range(header.num_groups)
        )
        if header.mode_default == "M0"
        else None
    )
    fused_scaled_codes = None
    if header.mode_default == "M0" and (
        _supports_fused_two_group64(header) or _supports_fused_m0_3bit(header, device_type=pages[0].device_type)
    ):
        fused_dtype = _fused_two_group64_cache_dtype(device_type=device_type)
        fused_scaled_codes = torch.cat(
            [
                codes_groups[group_index].to(dtype=fused_dtype) * scales_groups[group_index].to(dtype=fused_dtype)[..., None]
                for group_index in range(header.num_groups)
            ],
            dim=-1,
        ).contiguous()
        bias_groups = tuple(bias.to(dtype=fused_dtype) for bias in bias_groups)
    if fused_scaled_codes is not None:
        # For the fused grouped-64 path, retain only the pre-scaled fused tensor and
        # the affine bias terms. Keeping payload/codes/scales as well doubles memory
        # for data the hot path no longer reads.
        payload_groups = ()
        codes_groups = None
        scales_groups = None
        resident_nbytes = int(fused_scaled_codes.numel() * fused_scaled_codes.element_size())
    else:
        resident_nbytes = sum(int(tensor.numel() * tensor.element_size()) for tensor in payload_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in codes_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in scales_groups)
    if bias_groups is not None:
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in bias_groups)
    return PreparedChunkMPS(
        header=header,
        payload_groups=payload_groups,
        codes_groups=codes_groups,
        scales_groups=scales_groups,
        bias_groups=bias_groups,
        m2_sketch_groups=None,
        m2_basis_groups=None,
        m2_mean_groups=None,
        m2_segment_ids=None,
        fused_scaled_codes=fused_scaled_codes,
        resident_nbytes=resident_nbytes,
    )


def _get_prepared_chunk_mps(pages: Sequence[PreparedPageTorch]) -> PreparedChunkMPS | None:
    global _PREPARED_CHUNK_CACHE_RESIDENT_BYTES
    cache_key = _prepared_chunk_cache_key(pages)
    if cache_key is None:
        return None
    if pages[0].header.kind not in _PREPARED_CHUNK_CACHE_KINDS:
        return None
    if len(pages) < _MIN_PREPARED_CHUNK_CACHE_PAGE_COUNT:
        return None
    cached_chunk = _PREPARED_CHUNK_CACHE.get(cache_key)
    if cached_chunk is not None:
        _touch_prepared_chunk(cached_chunk)
        _PREPARED_CHUNK_CACHE.move_to_end(cache_key)
        return cached_chunk
    prepared_chunk = _build_prepared_chunk_mps(pages)
    effective_max_resident_bytes = _effective_max_prepared_chunk_cache_resident_bytes()
    if (
        _MAX_PREPARED_CHUNK_CACHE_ENTRIES <= 0
        or effective_max_resident_bytes <= 0
        or prepared_chunk.resident_nbytes > effective_max_resident_bytes
    ):
        return prepared_chunk
    _touch_prepared_chunk(prepared_chunk)
    _PREPARED_CHUNK_CACHE[cache_key] = prepared_chunk
    _PREPARED_CHUNK_CACHE_RESIDENT_BYTES += prepared_chunk.resident_nbytes
    _trim_prepared_chunk_cache()
    return prepared_chunk


def _grouped_prepared_chunk_cache_key(
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
) -> tuple[tuple[tuple[int, int], ...], ...] | None:
    if not pages_by_group or not pages_by_group[0]:
        return None
    if pages_by_group[0][0].header.mode_default not in {"M0", "M2", "M4"}:
        return None
    page_count = len(pages_by_group[0])
    cache_key: list[tuple[tuple[int, int], ...]] = []
    for group_pages in pages_by_group:
        if len(group_pages) != page_count:
            return None
        group_key = _prepared_chunk_cache_key(group_pages)
        if group_key is None:
            return None
        cache_key.append(group_key)
    return tuple(cache_key)


def _build_grouped_prepared_chunk_mps(
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
) -> PreparedGroupedChunkMPS | None:
    if not pages_by_group or not pages_by_group[0]:
        return None
    torch = _load_torch()

    def _get_or_build_group_chunk(group_pages: Sequence[PreparedPageTorch]) -> PreparedChunkMPS | None:
        prepared_chunk = _get_prepared_chunk_mps(group_pages)
        if prepared_chunk is not None:
            return prepared_chunk
        cache_key = _prepared_chunk_cache_key(group_pages)
        if cache_key is None:
            return None
        return _build_prepared_chunk_mps(group_pages)
    header = pages_by_group[0][0].header
    device_type = pages_by_group[0][0].device_type
    if header.mode_default == "M2":
        prepared_chunks = [_get_or_build_group_chunk(group_pages) for group_pages in pages_by_group]
        if any(chunk is None for chunk in prepared_chunks):
            return None
        m2_sketch_groups = tuple(
            torch.stack([chunk.m2_sketch_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_basis_groups = tuple(
            torch.stack([chunk.m2_basis_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_mean_groups = tuple(
            torch.stack([chunk.m2_mean_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_sketch_tensor = torch.stack(m2_sketch_groups, dim=3).contiguous()
        m2_basis_tensor = torch.stack(m2_basis_groups, dim=2).contiguous()
        m2_mean_tensor = torch.stack(m2_mean_groups, dim=2).contiguous()
        resident_nbytes = sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_sketch_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_basis_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_mean_groups)
        resident_nbytes += int(m2_sketch_tensor.numel() * m2_sketch_tensor.element_size())
        resident_nbytes += int(m2_basis_tensor.numel() * m2_basis_tensor.element_size())
        resident_nbytes += int(m2_mean_tensor.numel() * m2_mean_tensor.element_size())
        return PreparedGroupedChunkMPS(
            header=header,
            payload_groups=(),
            codes_groups=None,
            scales_groups=None,
            bias_groups=None,
            m2_sketch_groups=m2_sketch_groups,
            m2_basis_groups=m2_basis_groups,
            m2_mean_groups=m2_mean_groups,
            m2_segment_ids=prepared_chunks[0].m2_segment_ids,
            m2_sketch_tensor=m2_sketch_tensor,
            m2_basis_tensor=m2_basis_tensor,
            m2_mean_tensor=m2_mean_tensor,
            payload_groups_tensor=None,
            scales_groups_tensor=None,
            bias_groups_tensor=None,
            fused_scaled_codes=None,
            resident_nbytes=resident_nbytes,
        )
    if header.mode_default == "M4":
        prepared_chunks = [_get_or_build_group_chunk(group_pages) for group_pages in pages_by_group]
        if any(chunk is None for chunk in prepared_chunks):
            return None
        shared_basis = prepared_chunks[0].m2_basis_groups is not None and int(prepared_chunks[0].m2_basis_groups[0].dim()) == 2
        m2_basis_groups = (
            tuple(
                torch.stack([chunk.m2_basis_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
                for group_index in range(header.num_groups)
            )
            if prepared_chunks[0].m2_basis_groups is not None
            else None
        )
        m2_sketch_groups = tuple(
            torch.stack([chunk.m2_sketch_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_mean_groups = tuple(
            torch.stack([chunk.m2_mean_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        m2_sketch_tensor = torch.stack(m2_sketch_groups, dim=3).contiguous()
        m2_basis_tensor = (
            torch.stack(m2_basis_groups, dim=1).contiguous()
            if m2_basis_groups is not None and shared_basis
            else (
                torch.stack(m2_basis_groups, dim=2).contiguous()
                if m2_basis_groups is not None
                else None
            )
        )
        m2_mean_tensor = torch.stack(m2_mean_groups, dim=2).contiguous()
        resident_nbytes = sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_sketch_groups)
        if m2_basis_groups is not None:
            resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_basis_groups)
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in m2_mean_groups)
        resident_nbytes += int(m2_sketch_tensor.numel() * m2_sketch_tensor.element_size())
        if m2_basis_tensor is not None:
            resident_nbytes += int(m2_basis_tensor.numel() * m2_basis_tensor.element_size())
        resident_nbytes += int(m2_mean_tensor.numel() * m2_mean_tensor.element_size())
        return PreparedGroupedChunkMPS(
            header=header,
            payload_groups=(),
            codes_groups=None,
            scales_groups=None,
            bias_groups=None,
            m2_sketch_groups=m2_sketch_groups,
            m2_basis_groups=m2_basis_groups,
            m2_mean_groups=m2_mean_groups,
            m2_segment_ids=None,
            m2_sketch_tensor=m2_sketch_tensor,
            m2_basis_tensor=m2_basis_tensor,
            m2_mean_tensor=m2_mean_tensor,
            resident_nbytes=resident_nbytes,
            payload_groups_tensor=None,
            scales_groups_tensor=None,
            bias_groups_tensor=None,
            fused_scaled_codes=None,
        )
    if _supports_packed_four_group128_cuda(header, device_type=device_type):
        metadata_dtype = _m0_affine_metadata_dtype(device_type=device_type)
        payload_groups_tensor = torch.stack(
            [
                torch.stack(
                    [
                        torch.stack([page.payload[group_index] for page in group_pages], dim=0)
                        for group_pages in pages_by_group
                    ],
                    dim=0,
                )
                for group_index in range(header.num_groups)
            ],
            dim=1,
        ).contiguous()
        scales_groups_tensor = torch.stack(
            [
                torch.stack(
                    [
                        torch.stack([page.scales[:, group_index].to(dtype=metadata_dtype) for page in group_pages], dim=0)
                        for group_pages in pages_by_group
                    ],
                    dim=0,
                )
                for group_index in range(header.num_groups)
            ],
            dim=1,
        ).contiguous()
        bias_groups_tensor = torch.stack(
            [
                torch.stack(
                    [
                        torch.stack([page.bias[:, group_index].to(dtype=metadata_dtype) for page in group_pages], dim=0)
                        for group_pages in pages_by_group
                    ],
                    dim=0,
                )
                for group_index in range(header.num_groups)
            ],
            dim=1,
        ).contiguous()
        payload_groups = tuple(payload_groups_tensor[:, group_index] for group_index in range(header.num_groups))
        scales_groups = tuple(scales_groups_tensor[:, group_index] for group_index in range(header.num_groups))
        bias_groups = tuple(bias_groups_tensor[:, group_index] for group_index in range(header.num_groups))
        resident_nbytes = int(payload_groups_tensor.numel() * payload_groups_tensor.element_size())
        resident_nbytes += int(scales_groups_tensor.numel() * scales_groups_tensor.element_size())
        resident_nbytes += int(bias_groups_tensor.numel() * bias_groups_tensor.element_size())
        return PreparedGroupedChunkMPS(
            header=header,
            payload_groups=payload_groups,
            codes_groups=None,
            scales_groups=scales_groups,
            bias_groups=bias_groups,
            m2_sketch_groups=None,
            m2_basis_groups=None,
            m2_mean_groups=None,
            m2_segment_ids=None,
            m2_sketch_tensor=None,
            m2_basis_tensor=None,
            m2_mean_tensor=None,
            payload_groups_tensor=payload_groups_tensor,
            scales_groups_tensor=scales_groups_tensor,
            bias_groups_tensor=bias_groups_tensor,
            fused_scaled_codes=None,
            resident_nbytes=resident_nbytes,
        )
    prepared_chunks = [_get_or_build_group_chunk(group_pages) for group_pages in pages_by_group]
    if any(chunk is None for chunk in prepared_chunks):
        return None
    if (
        _supports_grouped_fused_only_cache(header, device_type=device_type)
        and all(chunk.fused_scaled_codes is not None for chunk in prepared_chunks)
        and all(chunk.bias_groups is not None for chunk in prepared_chunks)
    ):
        fused_scaled_codes = torch.stack([chunk.fused_scaled_codes for chunk in prepared_chunks], dim=0).contiguous()
        bias_groups = tuple(
            torch.stack([chunk.bias_groups[group_index] for chunk in prepared_chunks], dim=0).contiguous()
            for group_index in range(header.num_groups)
        )
        resident_nbytes = int(fused_scaled_codes.numel() * fused_scaled_codes.element_size())
        resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in bias_groups)
        return PreparedGroupedChunkMPS(
            header=header,
            payload_groups=(),
            codes_groups=None,
            scales_groups=None,
            bias_groups=bias_groups,
            m2_sketch_groups=None,
            m2_basis_groups=None,
            m2_mean_groups=None,
            m2_segment_ids=None,
            m2_sketch_tensor=None,
            m2_basis_tensor=None,
            m2_mean_tensor=None,
            payload_groups_tensor=None,
            scales_groups_tensor=None,
            bias_groups_tensor=None,
            fused_scaled_codes=fused_scaled_codes,
            resident_nbytes=resident_nbytes,
        )
    # Grouped decode uses unpacked codes/scales/bias directly when no fused-only cache
    # is available, so duplicating stacked payload tensors here only burns memory
    # without helping the hot path.
    payload_groups: tuple[Any, ...] = ()
    codes_groups = tuple(
        torch.stack([chunk.codes_groups[group_index] for chunk in prepared_chunks], dim=0)
        for group_index in range(header.num_groups)
    )
    scales_groups = tuple(
        torch.stack([chunk.scales_groups[group_index] for chunk in prepared_chunks], dim=0)
        for group_index in range(header.num_groups)
    )
    bias_groups = tuple(
        torch.stack([chunk.bias_groups[group_index] for chunk in prepared_chunks], dim=0)
        for group_index in range(header.num_groups)
    )
    fused_scaled_codes = None
    if all(chunk.fused_scaled_codes is not None for chunk in prepared_chunks):
        fused_scaled_codes = torch.stack([chunk.fused_scaled_codes for chunk in prepared_chunks], dim=0)
    resident_nbytes = sum(int(tensor.numel() * tensor.element_size()) for tensor in payload_groups)
    resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in codes_groups)
    resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in scales_groups)
    resident_nbytes += sum(int(tensor.numel() * tensor.element_size()) for tensor in bias_groups)
    if fused_scaled_codes is not None:
        resident_nbytes += int(fused_scaled_codes.numel() * fused_scaled_codes.element_size())
    return PreparedGroupedChunkMPS(
        header=header,
        payload_groups=payload_groups,
        codes_groups=codes_groups,
        scales_groups=scales_groups,
        bias_groups=bias_groups,
        m2_sketch_groups=None,
        m2_basis_groups=None,
        m2_mean_groups=None,
        m2_segment_ids=None,
        m2_sketch_tensor=None,
        m2_basis_tensor=None,
        m2_mean_tensor=None,
        payload_groups_tensor=None,
        scales_groups_tensor=None,
        bias_groups_tensor=None,
        fused_scaled_codes=fused_scaled_codes,
        resident_nbytes=resident_nbytes,
    )


def _assemble_grouped_fused_two_group64_components(
    prepared_chunks: Sequence[PreparedChunkMPS],
    *,
    trace: ExecutionTrace | None,
    device_type: TorchDevice,
):
    fused_scaled_codes = _trace_timed_call(
        trace,
        "chunk_assembly",
        device_type=device_type,
        fn=lambda: _load_torch().stack([chunk.fused_scaled_codes for chunk in prepared_chunks], dim=0),
    )
    bias_groups = tuple(
        _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda group_index=group_index: _load_torch().stack([chunk.bias_groups[group_index] for chunk in prepared_chunks], dim=0),
        )
        for group_index in range(2)
    )
    if trace is not None:
        trace.record_temporary(int(fused_scaled_codes.numel() * fused_scaled_codes.element_size()))
        trace.record_temporary(sum(int(tensor.numel() * tensor.element_size()) for tensor in bias_groups))
    return fused_scaled_codes, bias_groups


def _get_grouped_prepared_chunk_mps(
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
) -> PreparedGroupedChunkMPS | None:
    global _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES
    cache_key = _grouped_prepared_chunk_cache_key(pages_by_group)
    if cache_key is None:
        return None
    if pages_by_group[0][0].header.kind not in _PREPARED_CHUNK_CACHE_KINDS:
        return None
    total_page_count = sum(len(group_pages) for group_pages in pages_by_group)
    if total_page_count < _MIN_PREPARED_CHUNK_CACHE_PAGE_COUNT:
        return None
    cached_chunk = _PREPARED_GROUPED_CHUNK_CACHE.get(cache_key)
    if cached_chunk is not None:
        _touch_prepared_chunk(cached_chunk)
        _PREPARED_GROUPED_CHUNK_CACHE.move_to_end(cache_key)
        return cached_chunk
    prepared_chunk = _build_grouped_prepared_chunk_mps(pages_by_group)
    if prepared_chunk is None:
        return None
    effective_max_resident_bytes = _effective_max_prepared_chunk_cache_resident_bytes()
    if (
        _MAX_PREPARED_CHUNK_CACHE_ENTRIES <= 0
        or effective_max_resident_bytes <= 0
        or prepared_chunk.resident_nbytes > effective_max_resident_bytes
    ):
        return prepared_chunk
    _touch_prepared_chunk(prepared_chunk)
    _PREPARED_GROUPED_CHUNK_CACHE[cache_key] = prepared_chunk
    _PREPARED_GROUPED_CHUNK_CACHE_RESIDENT_BYTES += prepared_chunk.resident_nbytes
    _trim_prepared_chunk_cache()
    return prepared_chunk


def page_supported_torch(page: EncodedPage | PreparedPageTorch) -> bool:
    source_page = page.source_page if isinstance(page, PreparedPageTorch) else page
    header = source_page.header
    if header.layout != "group_major":
        return False
    if int(header.group_size) <= 0 or int(header.num_groups) <= 0:
        return False
    if int(header.group_size) * int(header.num_groups) != int(header.padded_head_dim):
        return False
    if header.mode_default == "M3":
        if source_page.escape_payload is None:
            return False
        if header.escape_dtype == "int8":
            return source_page.escape_scales is not None
        return True
    if header.mode_default == "M2":
        return (
            header.kind == "K"
            and header.quant_scheme == "sketch"
            and source_page.m2_sketch is not None
            and source_page.m2_basis is not None
            and source_page.m2_mean is not None
        )
    if header.mode_default == "M4":
        return (
            header.kind == "K"
            and header.quant_scheme == "project"
            and source_page.m2_sketch is not None
            and source_page.m2_mean is not None
        )
    if header.mode_default == "T3":
        return (
            header.quant_scheme == "turbo3"
            and header.bits == 3
            and header.group_size in (32, 64)
            and source_page.payload is not None
            and source_page.scales is not None
            and source_page.codebooks is not None
        )
    return (
        source_page.payload is not None
        and (
            (
                header.mode_default == "M0"
                and header.bits in (2, 3, 4)
                and header.quant_scheme == "affine"
                and source_page.scales is not None
                and source_page.bias is not None
            )
            or (
                header.mode_default == "M1"
                and header.bits in (2, 4)
                and header.quant_scheme == "lut"
                and source_page.codebooks is not None
            )
        )
    )


def page_supported_mps(page: EncodedPage | PreparedPageTorch) -> bool:
    return page_supported_torch(page)


def _unpack_metadata(bits: int, *, device_type: TorchDevice):
    cache_key = (device_type, bits)
    cached = _UNPACK_METADATA.get(cache_key)
    if cached is not None:
        return cached
    torch = _load_torch()
    symbols_per_word = 32 // bits
    shifts = torch.arange(symbols_per_word, dtype=torch.int32, device=device_type) * bits
    mask = torch.tensor((1 << bits) - 1, dtype=torch.int32, device=device_type)
    _UNPACK_METADATA[cache_key] = (shifts, mask)
    return shifts, mask


def _spill_unpack_metadata(bits: int, group_size: int, *, device_type: TorchDevice):
    cache_key = (device_type, bits, group_size)
    cached = _SPILL_UNPACK_METADATA.get(cache_key)
    if cached is not None:
        return cached
    torch = _load_torch()
    bit_offsets = np.arange(group_size, dtype=np.int64) * int(bits)
    word_count = words_per_group(group_size, bits)
    word_indices = torch.as_tensor(bit_offsets // 32, dtype=torch.int64, device=device_type)
    next_word_indices = torch.as_tensor(
        np.minimum(bit_offsets // 32 + 1, word_count - 1),
        dtype=torch.int64,
        device=device_type,
    )
    bit_indices = torch.as_tensor(bit_offsets % 32, dtype=torch.int64, device=device_type)
    spill_width = np.maximum((bit_offsets % 32) + int(bits) - 32, 0).astype(np.int64)
    spill_mask = torch.as_tensor((1 << spill_width) - 1, dtype=torch.int64, device=device_type)
    shift_back = torch.as_tensor(int(bits) - spill_width, dtype=torch.int64, device=device_type)
    spill_flags = torch.as_tensor(spill_width > 0, dtype=torch.bool, device=device_type)
    _SPILL_UNPACK_METADATA[cache_key] = (
        word_indices,
        next_word_indices,
        bit_indices,
        spill_mask,
        shift_back,
        spill_flags,
    )
    return _SPILL_UNPACK_METADATA[cache_key]


def _turbo3_centroids_torch(*, device_type: TorchDevice):
    cached = _TURBO3_CENTROID_TENSORS.get(device_type)
    if cached is not None:
        return cached
    tensor = _device_tensor(TURBO3_CENTROIDS.astype(np.float32, copy=False), device=device_type)
    tensor = tensor.to(dtype=_load_torch().float32)
    _TURBO3_CENTROID_TENSORS[device_type] = tensor
    return tensor


def _fwht_matrix_torch(width: int, *, device_type: TorchDevice):
    cache_key = (device_type, int(width))
    cached = _FWHT_MATRICES.get(cache_key)
    if cached is not None:
        return cached
    if width <= 0 or (width & (width - 1)):
        raise ValueError("FWHT requires the last dimension to be a power of two")
    basis = np.eye(width, dtype=np.float32)
    transformed = fwht_last_dim(basis)
    tensor = _device_tensor(transformed, device=device_type).to(dtype=_load_torch().float32)
    _FWHT_MATRICES[cache_key] = tensor
    return tensor


def _m4_basis_torch(group_size: int, rank: int, *, basis_family: str, device_type: TorchDevice):
    cache_key = (device_type, int(group_size), int(rank), basis_family)
    cached = _M4_BASIS_TENSORS.get(cache_key)
    if cached is not None:
        return cached
    basis = _device_tensor(
        fixed_project_basis(int(group_size), int(rank), basis_family),
        device=device_type,
    ).to(dtype=_load_torch().float32)
    _M4_BASIS_TENSORS[cache_key] = basis
    return basis


def _synchronize_torch_device(device_type: TorchDevice) -> None:
    torch = _load_torch()
    if device_type == "mps" and torch.backends.mps.is_available():
        torch.mps.synchronize()
        return
    if device_type == "cuda" and torch.cuda.is_available():
        torch.cuda.synchronize()


def _trace_timed_call(
    trace: ExecutionTrace | None,
    section: str,
    *,
    device_type: TorchDevice,
    fn,
    synchronize: bool = True,
):
    if trace is None or not trace.capture_timings:
        return fn()
    if synchronize:
        _synchronize_torch_device(device_type)
    start = time.perf_counter()
    result = fn()
    if synchronize:
        _synchronize_torch_device(device_type)
    trace.record_timing(section, (time.perf_counter() - start) * 1000.0)
    return result


def _prepared_page_host_nbytes(page: EncodedPage) -> int:
    total = 0
    if page.payload is not None:
        total += int(page.payload.nbytes)
    if page.scales is not None:
        total += int(page.scales.nbytes)
    if page.bias is not None:
        total += int(page.bias.nbytes)
    if page.codebooks is not None:
        total += int(page.codebooks.nbytes)
    if page.m2_sketch is not None:
        total += int(page.m2_sketch.nbytes)
    if page.m2_basis is not None:
        total += int(page.m2_basis.nbytes)
    if page.m2_mean is not None:
        total += int(page.m2_mean.nbytes)
    if page.escape_payload is not None:
        total += int(page.escape_payload.nbytes)
    if page.escape_scales is not None:
        total += int(page.escape_scales.nbytes)
    return total


def _decode_escape_batch_torch(
    pages: Sequence[PreparedPageTorch],
    *,
    token_count: int,
    head_dim: int,
    promote_float32: bool = True,
):
    torch = _load_torch()
    target_dtype = torch.float32 if promote_float32 else None
    prepared_chunk = _get_prepared_chunk_mps(pages)
    if prepared_chunk is not None and prepared_chunk.escape_payload_batch is not None:
        payload = prepared_chunk.escape_payload_batch[:, :token_count, :head_dim]
        if pages[0].header.escape_dtype == "int8":
            scales = prepared_chunk.escape_scales_batch[:, :token_count]
            result = payload if target_dtype is None else payload.to(dtype=target_dtype)
            scale_values = scales if target_dtype is None else scales.to(dtype=target_dtype)
            return result * scale_values[..., None]
        return payload if target_dtype is None else payload.to(dtype=target_dtype)
    if len(pages) == 1:
        payload = pages[0].escape_payload[:token_count, :head_dim]
        if pages[0].header.escape_dtype == "int8":
            scales = pages[0].escape_scales[:token_count]
            result = payload if target_dtype is None else payload.to(dtype=target_dtype)
            scale_values = scales if target_dtype is None else scales.to(dtype=target_dtype)
            return (result * scale_values[:, None]).unsqueeze(0)
        return (payload if target_dtype is None else payload.to(dtype=target_dtype)).unsqueeze(0)
    payload = torch.stack([page.escape_payload[:token_count, :head_dim] for page in pages], dim=0)
    if pages[0].header.escape_dtype == "int8":
        scales = torch.stack([page.escape_scales[:token_count] for page in pages], dim=0)
        result = payload if target_dtype is None else payload.to(dtype=target_dtype)
        scale_values = scales if target_dtype is None else scales.to(dtype=target_dtype)
        return result * scale_values[..., None]
    return payload if target_dtype is None else payload.to(dtype=target_dtype)


def _m3_native_compute_enabled(pages: Sequence[PreparedPageTorch]) -> bool:
    if not pages:
        return False
    header = pages[0].header
    if pages[0].device_type != "cuda":
        return False
    return header.mode_default == "M3" and header.escape_dtype == "float16"


def _optional_m2_sidecar_batches(
    pages: Sequence[EncodedPage],
    *,
    device_type: TorchDevice,
) -> tuple[Any | None, Any | None, Any | None, int, int]:
    if not pages or not all(page.m2_sketch is not None and page.m2_basis is not None and page.m2_mean is not None for page in pages):
        return None, None, None, 0, 0
    sketch_array = np.stack([np.asarray(page.m2_sketch) for page in pages], axis=0)
    basis_array = np.stack([np.asarray(page.m2_basis) for page in pages], axis=0)
    mean_array = np.stack([np.asarray(page.m2_mean) for page in pages], axis=0)
    sketch_batch = _device_tensor(sketch_array, device=device_type)
    basis_batch = _device_tensor(basis_array, device=device_type)
    mean_batch = _device_tensor(mean_array, device=device_type)
    if device_type == "mps":
        sketch_batch = sketch_batch.to(dtype=_load_torch().float32)
        basis_batch = basis_batch.to(dtype=_load_torch().float32)
        mean_batch = mean_batch.to(dtype=_load_torch().float32)
    return sketch_batch, basis_batch, mean_batch, int(sketch_array.nbytes + basis_array.nbytes + mean_array.nbytes), int(
        sketch_batch.numel() * sketch_batch.element_size()
        + basis_batch.numel() * basis_batch.element_size()
        + mean_batch.numel() * mean_batch.element_size()
    )


def _prepare_page_chunk_torch(
    pages: Sequence[EncodedPage],
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
) -> list[PreparedPageTorch]:
    if not pages:
        return []
    header = pages[0].header
    total_host_to_device_nbytes = 0

    if header.mode_default == "M3":
        escape_array = np.stack([np.asarray(page.escape_payload) for page in pages], axis=0)
        escape_batch = _device_tensor(escape_array, device=device_type)
        total_host_to_device_nbytes += int(escape_batch.numel() * escape_batch.element_size())
        escape_scale_batch = None
        if header.escape_dtype == "int8":
            escape_scale_array = np.stack([np.asarray(page.escape_scales) for page in pages], axis=0)
            escape_scale_batch = _device_tensor(escape_scale_array, device=device_type)
            escape_scale_batch = escape_scale_batch.to(dtype=_escape_scale_dtype(device_type=device_type))
            total_host_to_device_nbytes += int(escape_scale_batch.numel() * escape_scale_batch.element_size())
        prepared_pages = [
            PreparedPageTorch(
                device_type=device_type,
                source_page=page,
                header=page.header,
                escape_payload=escape_batch[index],
                escape_scales=None if escape_scale_batch is None else escape_scale_batch[index],
                host_to_device_nbytes=_prepared_page_host_nbytes(page),
                resident_nbytes=int(escape_batch[index].numel() * escape_batch[index].element_size())
                + (
                    0
                    if escape_scale_batch is None
                    else int(escape_scale_batch[index].numel() * escape_scale_batch[index].element_size())
                ),
                cache_uid=_next_prepared_page_uid(),
            )
            for index, page in enumerate(pages)
        ]
        if trace is not None:
            trace.record_host_to_device(total_host_to_device_nbytes)
        return prepared_pages

    if header.mode_default == "M2":
        sketch_array = np.stack([np.asarray(page.m2_sketch) for page in pages], axis=0)
        basis_array = np.stack([np.asarray(page.m2_basis) for page in pages], axis=0)
        mean_array = np.stack([np.asarray(page.m2_mean) for page in pages], axis=0)
        sketch_batch = _device_tensor(sketch_array, device=device_type)
        basis_batch = _device_tensor(basis_array, device=device_type)
        mean_batch = _device_tensor(mean_array, device=device_type)
        total_host_to_device_nbytes += int(sketch_array.nbytes)
        total_host_to_device_nbytes += int(basis_array.nbytes)
        total_host_to_device_nbytes += int(mean_array.nbytes)
        if device_type == "mps":
            sketch_batch = sketch_batch.to(dtype=_load_torch().float32)
            basis_batch = basis_batch.to(dtype=_load_torch().float32)
            mean_batch = mean_batch.to(dtype=_load_torch().float32)
        prepared_pages = [
            PreparedPageTorch(
                device_type=device_type,
                source_page=page,
                header=page.header,
                m2_sketch=sketch_batch[index],
                m2_basis=basis_batch[index],
                m2_mean=mean_batch[index],
                host_to_device_nbytes=_prepared_page_host_nbytes(page),
                resident_nbytes=(
                    int(sketch_batch[index].numel() * sketch_batch[index].element_size())
                    + int(basis_batch[index].numel() * basis_batch[index].element_size())
                    + int(mean_batch[index].numel() * mean_batch[index].element_size())
                ),
                cache_uid=_next_prepared_page_uid(),
            )
            for index, page in enumerate(pages)
        ]
        if trace is not None:
            trace.record_host_to_device(total_host_to_device_nbytes)
        return prepared_pages

    if header.mode_default == "M4":
        sketch_array = np.stack([np.asarray(page.m2_sketch) for page in pages], axis=0)
        shared_basis = pages[0].header.project_basis == "svd_shared"
        basis_array = None
        if not shared_basis and pages[0].m2_basis is not None:
            basis_array = np.stack([np.asarray(page.m2_basis) for page in pages], axis=0)
        mean_array = np.stack([np.asarray(page.m2_mean) for page in pages], axis=0)
        sketch_batch = _device_tensor(sketch_array, device=device_type)
        basis_batch = None if basis_array is None else _device_tensor(basis_array, device=device_type)
        mean_batch = _device_tensor(mean_array, device=device_type)
        total_host_to_device_nbytes += int(sketch_array.nbytes)
        if basis_array is not None:
            total_host_to_device_nbytes += int(basis_array.nbytes)
        total_host_to_device_nbytes += int(mean_array.nbytes)
        if device_type == "mps":
            sketch_batch = sketch_batch.to(dtype=_load_torch().float32)
            if basis_batch is not None:
                basis_batch = basis_batch.to(dtype=_load_torch().float32)
            mean_batch = mean_batch.to(dtype=_load_torch().float32)
        prepared_pages = [
            PreparedPageTorch(
                device_type=device_type,
                source_page=page,
                header=page.header,
                m2_sketch=sketch_batch[index],
                m2_basis=None if basis_batch is None else basis_batch[index],
                m2_mean=mean_batch[index],
                host_to_device_nbytes=_prepared_page_host_nbytes(page),
                resident_nbytes=(
                    int(sketch_batch[index].numel() * sketch_batch[index].element_size())
                    + (0 if basis_batch is None else int(basis_batch[index].numel() * basis_batch[index].element_size()))
                    + int(mean_batch[index].numel() * mean_batch[index].element_size())
                ),
                cache_uid=_next_prepared_page_uid(),
            )
            for index, page in enumerate(pages)
        ]
        if trace is not None:
            trace.record_host_to_device(total_host_to_device_nbytes)
        return prepared_pages

    if header.mode_default == "M1":
        payload_array = np.stack([np.asarray(page.payload, dtype=np.int32) for page in pages], axis=0)
        codebooks_array = np.stack([np.asarray(page.codebooks) for page in pages], axis=0)
        payload_batch = _device_tensor(payload_array, device=device_type)
        codebooks_batch = _device_tensor(codebooks_array, device=device_type)
        sidecar_sketch_batch, sidecar_basis_batch, sidecar_mean_batch, sidecar_h2d_nbytes, _ = _optional_m2_sidecar_batches(
            pages,
            device_type=device_type,
        )
        total_host_to_device_nbytes += int(payload_array.nbytes)
        total_host_to_device_nbytes += int(codebooks_array.nbytes)
        total_host_to_device_nbytes += sidecar_h2d_nbytes
        if device_type == "mps":
            codebooks_batch = codebooks_batch.to(dtype=_load_torch().float32)
        unpack_shifts, unpack_mask = _unpack_metadata(header.bits, device_type=device_type)
        prepared_pages = [
            PreparedPageTorch(
                device_type=device_type,
                source_page=page,
                header=page.header,
                payload=payload_batch[index],
                codebooks=codebooks_batch[index],
                m2_sketch=None if sidecar_sketch_batch is None else sidecar_sketch_batch[index],
                m2_basis=None if sidecar_basis_batch is None else sidecar_basis_batch[index],
                m2_mean=None if sidecar_mean_batch is None else sidecar_mean_batch[index],
                unpack_shifts=unpack_shifts,
                unpack_mask=unpack_mask,
                host_to_device_nbytes=_prepared_page_host_nbytes(page),
                resident_nbytes=(
                    int(payload_batch[index].numel() * payload_batch[index].element_size())
                    + int(codebooks_batch[index].numel() * codebooks_batch[index].element_size())
                    + (
                        0
                        if sidecar_sketch_batch is None or sidecar_basis_batch is None or sidecar_mean_batch is None
                        else int(sidecar_sketch_batch[index].numel() * sidecar_sketch_batch[index].element_size())
                        + int(sidecar_basis_batch[index].numel() * sidecar_basis_batch[index].element_size())
                        + int(sidecar_mean_batch[index].numel() * sidecar_mean_batch[index].element_size())
                    )
                ),
                cache_uid=_next_prepared_page_uid(),
            )
            for index, page in enumerate(pages)
        ]
        if trace is not None:
            trace.record_host_to_device(total_host_to_device_nbytes)
        return prepared_pages

    if header.mode_default == "T3":
        payload_array = np.stack([np.asarray(page.payload, dtype=np.int32) for page in pages], axis=0)
        scales_array = np.stack([np.asarray(page.scales) for page in pages], axis=0)
        payload_batch = _device_tensor(payload_array, device=device_type)
        scales_batch = _device_tensor(scales_array, device=device_type)
        sidecar_sketch_batch, sidecar_basis_batch, sidecar_mean_batch, sidecar_h2d_nbytes, _ = _optional_m2_sidecar_batches(
            pages,
            device_type=device_type,
        )
        total_host_to_device_nbytes += int(payload_array.nbytes)
        total_host_to_device_nbytes += int(scales_array.nbytes)
        total_host_to_device_nbytes += sidecar_h2d_nbytes
        if device_type == "mps":
            scales_batch = scales_batch.to(dtype=_load_torch().float32)
        codebooks_tensor = _turbo3_centroids_torch(device_type=device_type)
        unpack_shifts, unpack_mask = _unpack_metadata(header.bits, device_type=device_type)
        prepared_pages = [
            PreparedPageTorch(
                device_type=device_type,
                source_page=page,
                header=page.header,
                payload=payload_batch[index],
                scales=scales_batch[index],
                codebooks=codebooks_tensor,
                m2_sketch=None if sidecar_sketch_batch is None else sidecar_sketch_batch[index],
                m2_basis=None if sidecar_basis_batch is None else sidecar_basis_batch[index],
                m2_mean=None if sidecar_mean_batch is None else sidecar_mean_batch[index],
                unpack_shifts=unpack_shifts,
                unpack_mask=unpack_mask,
                host_to_device_nbytes=_prepared_page_host_nbytes(page),
                resident_nbytes=(
                    int(payload_batch[index].numel() * payload_batch[index].element_size())
                    + int(scales_batch[index].numel() * scales_batch[index].element_size())
                    + (
                        0
                        if sidecar_sketch_batch is None or sidecar_basis_batch is None or sidecar_mean_batch is None
                        else int(sidecar_sketch_batch[index].numel() * sidecar_sketch_batch[index].element_size())
                        + int(sidecar_basis_batch[index].numel() * sidecar_basis_batch[index].element_size())
                        + int(sidecar_mean_batch[index].numel() * sidecar_mean_batch[index].element_size())
                    )
                ),
                cache_uid=_next_prepared_page_uid(),
            )
            for index, page in enumerate(pages)
        ]
        if trace is not None:
            trace.record_host_to_device(total_host_to_device_nbytes)
        return prepared_pages

    payload_array = np.stack([np.asarray(page.payload, dtype=np.int32) for page in pages], axis=0)
    scales_array = np.stack([np.asarray(page.scales) for page in pages], axis=0)
    bias_array = np.stack([np.asarray(page.bias) for page in pages], axis=0)
    payload_batch = _device_tensor(payload_array, device=device_type)
    scales_batch = _device_tensor(scales_array, device=device_type)
    bias_batch = _device_tensor(bias_array, device=device_type)
    sidecar_sketch_batch, sidecar_basis_batch, sidecar_mean_batch, sidecar_h2d_nbytes, _ = _optional_m2_sidecar_batches(
        pages,
        device_type=device_type,
    )
    total_host_to_device_nbytes += int(payload_array.nbytes)
    total_host_to_device_nbytes += int(scales_array.nbytes)
    total_host_to_device_nbytes += int(bias_array.nbytes)
    total_host_to_device_nbytes += sidecar_h2d_nbytes
    metadata_dtype = _m0_affine_metadata_dtype(device_type=device_type)
    scales_batch = scales_batch.to(dtype=metadata_dtype)
    bias_batch = bias_batch.to(dtype=metadata_dtype)
    unpack_shifts, unpack_mask = _unpack_metadata(header.bits, device_type=device_type)

    prepared_pages = [
        PreparedPageTorch(
            device_type=device_type,
            source_page=page,
            header=page.header,
            payload=payload_batch[index],
            scales=scales_batch[index],
            bias=bias_batch[index],
            m2_sketch=None if sidecar_sketch_batch is None else sidecar_sketch_batch[index],
            m2_basis=None if sidecar_basis_batch is None else sidecar_basis_batch[index],
            m2_mean=None if sidecar_mean_batch is None else sidecar_mean_batch[index],
            unpack_shifts=unpack_shifts,
            unpack_mask=unpack_mask,
            host_to_device_nbytes=_prepared_page_host_nbytes(page),
            resident_nbytes=(
                int(payload_batch[index].numel() * payload_batch[index].element_size())
                + int(scales_batch[index].numel() * scales_batch[index].element_size())
                + int(bias_batch[index].numel() * bias_batch[index].element_size())
                + (
                    0
                    if sidecar_sketch_batch is None or sidecar_basis_batch is None or sidecar_mean_batch is None
                    else int(sidecar_sketch_batch[index].numel() * sidecar_sketch_batch[index].element_size())
                    + int(sidecar_basis_batch[index].numel() * sidecar_basis_batch[index].element_size())
                    + int(sidecar_mean_batch[index].numel() * sidecar_mean_batch[index].element_size())
                )
            ),
            cache_uid=_next_prepared_page_uid(),
        )
        for index, page in enumerate(pages)
    ]
    if trace is not None:
        trace.record_host_to_device(total_host_to_device_nbytes)
    return prepared_pages


def prepare_pages_torch(
    pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
) -> list[PreparedPageTorch]:
    backend_name = _backend_name(device_type)
    if not torch_device_available(device_type):
        raise RuntimeError(f"{backend_name} is unavailable on this machine")
    prepared_pages: list[PreparedPageTorch] = []
    for page_chunk in _chunk_compatible_source_pages(pages, device_type=device_type):
        if all(isinstance(page, PreparedPageTorch) and page.device_type == device_type for page in page_chunk):
            prepared_pages.extend(page_chunk)  # type: ignore[arg-type]
            continue
        source_pages = []
        for page in page_chunk:
            source_page = page.source_page if isinstance(page, PreparedPageTorch) else page
            if not page_supported_torch(source_page):
                raise ValueError(f"page is unsupported by {backend_name} in this phase")
            source_pages.append(source_page)
        if source_pages:
            prepared_pages.extend(
                _trace_timed_call(
                    trace,
                    "prepare",
                    device_type=device_type,
                    fn=lambda source_pages=source_pages: _prepare_page_chunk_torch(
                        source_pages,
                        device_type=device_type,
                        trace=trace,
                    ),
                )
            )
    return prepared_pages


def prepare_page_torch(
    page: EncodedPage | PreparedPageTorch,
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
) -> PreparedPageTorch:
    if isinstance(page, PreparedPageTorch) and page.device_type == device_type:
        return page
    source_page = page.source_page if isinstance(page, PreparedPageTorch) else page
    return prepare_pages_torch([source_page], device_type=device_type, trace=trace)[0]


def prepare_m0_affine_pages_from_tensor_torch(
    values,
    *,
    config,
    kind: str,
    layer_id: int,
    kv_head_id: int,
    token_start: int,
    device_type: TorchDevice,
    build_runtime_metadata: bool = False,
):
    torch = _load_torch()
    if not torch.is_tensor(values):
        raise TypeError("values must be a torch.Tensor")
    if int(values.ndim) != 3:
        raise ValueError("values must have shape [page_count, token_count, head_dim]")
    if int(values.shape[2]) != int(config.head_dim):
        raise ValueError("values head_dim must match config.head_dim")

    bits = config.bits_k if kind == "K" else config.bits_v
    default_mode = config.default_mode_k if kind == "K" else config.default_mode_v
    quant_scheme = config.quant_scheme_k if kind == "K" else config.quant_scheme_v
    layout = config.payload_layout_k if kind == "K" else config.payload_layout_v
    if default_mode != "M0" or quant_scheme != "affine" or layout != "group_major":
        raise ValueError("direct torch preparation only supports exact M0 affine group_major pages")

    page_count = int(values.shape[0])
    token_count = int(values.shape[1])
    num_groups = int(config.num_groups)
    group_size = int(config.group_size)
    padded_head_dim = int(config.padded_head_dim)
    qmax = float((1 << int(bits)) - 1)
    eps = 1e-8

    values_device = values.to(device=device_type)
    work_dtype = values_device.dtype if values_device.dtype in (torch.float16, torch.bfloat16, torch.float32) else torch.float32
    values_work = values_device.to(dtype=work_dtype)
    if padded_head_dim > int(config.head_dim):
        padded = torch.nn.functional.pad(values_work, (0, padded_head_dim - int(config.head_dim)))
    else:
        padded = values_work
    grouped = padded.reshape(page_count, token_count, num_groups, group_size)
    x_min, x_max = torch.aminmax(grouped, dim=-1)
    scales = torch.clamp(((x_max - x_min).to(dtype=torch.float32) / max(qmax, 1.0)), min=eps).to(dtype=grouped.dtype)
    shifted = (grouped - x_min.unsqueeze(-1)) / scales.unsqueeze(-1)
    codes = torch.clamp(torch.round(shifted), 0.0, qmax).to(dtype=torch.int32)
    payload = _torch_pack_codes(codes, bits=int(bits), layout=layout)
    metadata_dtype = _m0_affine_metadata_dtype(device_type=device_type)
    scales_device = scales.to(dtype=metadata_dtype)
    bias_device = x_min.to(dtype=metadata_dtype)
    unpack_shifts, unpack_mask = _unpack_metadata(int(bits), device_type=device_type)

    prepared_pages: list[PreparedPageTorch] = []
    word_count = int(payload.shape[-1])
    for page_index in range(page_count):
        page_token_start = int(token_start + page_index * token_count)
        header = PageHeader(
            layer_id=layer_id,
            kv_head_id=kv_head_id,
            kind=kind,
            token_start=page_token_start,
            token_count=token_count,
            head_dim=int(config.head_dim),
            padded_head_dim=padded_head_dim,
            group_size=group_size,
            num_groups=num_groups,
            bits=int(bits),
            words_per_group=word_count,
            mode_default="M0",
            layout=layout,
            quant_scheme="affine",
            escape_dtype=config.escape_dtype,
        )
        source_page = EncodedPage(
            header=header,
            payload=np.zeros((num_groups, token_count, word_count), dtype=np.uint32),
            scales=np.zeros((token_count, num_groups), dtype=np.float16),
            bias=np.zeros((token_count, num_groups), dtype=np.float16),
            requested_mode="M0",
        )
        if build_runtime_metadata:
            page_values = values_work[page_index, :, : int(config.head_dim)].detach().cpu().numpy().astype(np.float32, copy=False)
            runtime_page_mean = page_values.mean(axis=0).astype(np.float32, copy=False)
            source_page.runtime_page_mean = runtime_page_mean
            source_page.runtime_page_sketch = runtime_page_mean[None, :]
            source_page.runtime_page_min = page_values.min(axis=0).astype(np.float32, copy=False)
            source_page.runtime_page_max = page_values.max(axis=0).astype(np.float32, copy=False)
        payload_page = payload[page_index]
        scales_page = scales_device[page_index]
        bias_page = bias_device[page_index]
        prepared_pages.append(
            PreparedPageTorch(
                device_type=device_type,
                source_page=source_page,
                header=header,
                payload=payload_page,
                scales=scales_page,
                bias=bias_page,
                unpack_shifts=unpack_shifts,
                unpack_mask=unpack_mask,
                host_to_device_nbytes=0,
                resident_nbytes=(
                    int(payload_page.numel() * payload_page.element_size())
                    + int(scales_page.numel() * scales_page.element_size())
                    + int(bias_page.numel() * bias_page.element_size())
                ),
                cache_uid=_next_prepared_page_uid(),
            )
        )
    return prepared_pages


def _pad_query(query_slice: np.ndarray | Any, padded_head_dim: int, *, device_type: TorchDevice):
    torch = _load_torch()
    if torch.is_tensor(query_slice):
        query = query_slice.to(dtype=torch.float32, device=device_type)
    else:
        query = torch.as_tensor(query_slice, dtype=torch.float32, device=device_type)
    if query.ndim != 1:
        raise ValueError("query_slice must have shape [head_dim]")
    if int(query.shape[0]) > padded_head_dim:
        raise ValueError("query head_dim exceeds padded_head_dim")
    if int(query.shape[0]) == padded_head_dim:
        return query
    padded = torch.zeros(padded_head_dim, dtype=torch.float32, device=device_type)
    padded[: query.shape[0]] = query
    return padded


def _pad_queries(query_slices: np.ndarray | Any, padded_head_dim: int, *, device_type: TorchDevice):
    torch = _load_torch()
    if torch.is_tensor(query_slices):
        queries = query_slices.to(dtype=torch.float32, device=device_type)
    else:
        queries = torch.as_tensor(query_slices, dtype=torch.float32, device=device_type)
    if queries.ndim != 2:
        raise ValueError("query_slices must have shape [query_count, head_dim]")
    if int(queries.shape[1]) > padded_head_dim:
        raise ValueError("query head_dim exceeds padded_head_dim")
    if int(queries.shape[1]) == padded_head_dim:
        return queries
    padded = torch.zeros((queries.shape[0], padded_head_dim), dtype=torch.float32, device=device_type)
    padded[:, : queries.shape[1]] = queries
    return padded


def _coerce_m2_operands(query_groups, group_sketch, group_basis, group_mean):
    torch = _load_torch()
    work_dtype = torch.promote_types(query_groups.dtype, group_sketch.dtype)
    work_dtype = torch.promote_types(work_dtype, group_basis.dtype)
    work_dtype = torch.promote_types(work_dtype, group_mean.dtype)
    if query_groups.dtype != work_dtype:
        query_groups = query_groups.to(dtype=work_dtype)
    if group_sketch.dtype != work_dtype:
        group_sketch = group_sketch.to(dtype=work_dtype)
    if group_basis.dtype != work_dtype:
        group_basis = group_basis.to(dtype=work_dtype)
    if group_mean.dtype != work_dtype:
        group_mean = group_mean.to(dtype=work_dtype)
    return query_groups, group_sketch, group_basis, group_mean


def _coerce_m2_grouped_operands(query_groups, sketch_tensor, basis_tensor, mean_tensor):
    torch = _load_torch()
    work_dtype = torch.promote_types(query_groups.dtype, sketch_tensor.dtype)
    work_dtype = torch.promote_types(work_dtype, basis_tensor.dtype)
    work_dtype = torch.promote_types(work_dtype, mean_tensor.dtype)
    if query_groups.dtype != work_dtype:
        query_groups = query_groups.to(dtype=work_dtype)
    if sketch_tensor.dtype != work_dtype:
        sketch_tensor = sketch_tensor.to(dtype=work_dtype)
    if basis_tensor.dtype != work_dtype:
        basis_tensor = basis_tensor.to(dtype=work_dtype)
    if mean_tensor.dtype != work_dtype:
        mean_tensor = mean_tensor.to(dtype=work_dtype)
    return query_groups, sketch_tensor, basis_tensor, mean_tensor


def _prepare_output_accumulator(out_acc: np.ndarray | None, head_dim: int, padded_head_dim: int, *, device_type: TorchDevice):
    torch = _load_torch()
    output = torch.zeros(padded_head_dim, dtype=torch.float32, device=device_type)
    if out_acc is None:
        return output
    values = torch.as_tensor(out_acc, dtype=torch.float32, device=device_type)
    if values.shape != (head_dim,):
        raise ValueError("out_acc must have shape [head_dim]")
    output[:head_dim] = values
    return output


def _prepare_output_accumulator_tensor(out_acc, head_dim: int, padded_head_dim: int, *, device_type: TorchDevice):
    torch = _load_torch()
    if out_acc is None:
        return torch.zeros(padded_head_dim, dtype=torch.float32, device=device_type)
    if isinstance(out_acc, np.ndarray):
        return _prepare_output_accumulator(out_acc, head_dim, padded_head_dim, device_type=device_type)
    if tuple(out_acc.shape) != (padded_head_dim,):
        raise ValueError("out_acc tensor must have shape [padded_head_dim]")
    return out_acc.to(dtype=torch.float32, device=device_type)


def _prepare_grouped_output_accumulator_tensor(
    out_acc,
    batch_size: int,
    query_count: int,
    head_dim: int,
    padded_head_dim: int,
    *,
    device_type: TorchDevice,
):
    torch = _load_torch()
    expected_shape = (batch_size, query_count, padded_head_dim)
    if out_acc is None:
        return torch.zeros(expected_shape, dtype=torch.float32, device=device_type)
    values = out_acc.to(dtype=torch.float32, device=device_type) if torch.is_tensor(out_acc) else torch.as_tensor(
        out_acc,
        dtype=torch.float32,
        device=device_type,
    )
    if tuple(values.shape) == expected_shape:
        return values
    if tuple(values.shape) != (batch_size, query_count, head_dim):
        raise ValueError("out_acc must have shape [batch_size, query_count, head_dim] or [batch_size, query_count, padded_head_dim]")
    output = torch.zeros(expected_shape, dtype=torch.float32, device=device_type)
    output[:, :, :head_dim] = values
    return output


def _unpack_bits_torch(words, shifts, mask, group_size: int, *, trace: ExecutionTrace | None = None):
    torch = _load_torch()
    if words.ndim != 2:
        raise ValueError("words must have shape [token_count, words_per_group]")
    if shifts is None or mask is None:
        raise ValueError("prepared torch pages require unpack metadata")
    device_type = str(words.device.type)

    def _impl():
        words_u64 = torch.bitwise_and(words.to(dtype=torch.int64), 0xFFFFFFFF)
        mask_i64 = torch.as_tensor(mask, dtype=torch.int64, device=words.device)
        bits = int(mask_i64.item()).bit_length()
        if 32 % bits == 0:
            shifts_i64 = shifts.to(dtype=torch.int64)
            expanded = torch.bitwise_and(torch.bitwise_right_shift(words_u64[..., None], shifts_i64), mask_i64)
            return expanded.reshape(words.shape[0], -1)[:, :group_size].to(torch.float32)
        if words.device.type == "mps":
            word_indices, next_word_indices, bit_indices, spill_mask, shift_back, spill_flags = _spill_unpack_metadata(
                bits,
                group_size,
                device_type=words.device.type,
            )
            bit_indices_2d = bit_indices.unsqueeze(0).expand(words_u64.shape[0], -1)
            spill_mask_2d = spill_mask.unsqueeze(0).expand(words_u64.shape[0], -1)
            shift_back_2d = shift_back.unsqueeze(0).expand(words_u64.shape[0], -1)
            current_words = words_u64[:, word_indices]
            next_words = words_u64[:, next_word_indices]
            values = torch.bitwise_right_shift(current_words, bit_indices_2d)
            spilled = torch.bitwise_left_shift(torch.bitwise_and(next_words, spill_mask_2d), shift_back_2d)
            values = torch.where(spill_flags.unsqueeze(0), torch.bitwise_or(values, spilled), values)
            return torch.bitwise_and(values, mask_i64).to(torch.float32)
        word_indices, next_word_indices, bit_indices, spill_mask, shift_back, spill_flags = _spill_unpack_metadata(
            bits,
            group_size,
            device_type=words.device.type,
        )
        gather_index = word_indices.unsqueeze(0).expand(words_u64.shape[0], -1)
        bit_indices_2d = bit_indices.unsqueeze(0).expand(words_u64.shape[0], -1)
        values = torch.bitwise_right_shift(torch.gather(words_u64, 1, gather_index), bit_indices_2d)
        if bool(spill_flags.any()):
            next_gather_index = next_word_indices.unsqueeze(0).expand(words_u64.shape[0], -1)
            spill_mask_2d = spill_mask.unsqueeze(0).expand(words_u64.shape[0], -1)
            shift_back_2d = shift_back.unsqueeze(0).expand(words_u64.shape[0], -1)
            spilled = torch.bitwise_left_shift(
                torch.bitwise_and(torch.gather(words_u64, 1, next_gather_index), spill_mask_2d),
                shift_back_2d,
            )
            values = torch.where(spill_flags.unsqueeze(0), torch.bitwise_or(values, spilled), values)
        return torch.bitwise_and(values, mask_i64).to(torch.float32)

    return _trace_timed_call(trace, "unpack", device_type=device_type, fn=_impl, synchronize=False)


def _score_m0_logits_flat_torch(codes, queries, scales, bias, query_group_sums):
    torch = _load_torch()
    matmul_dtype = codes.dtype if torch.is_floating_point(codes) else torch.float32
    codes_mm = codes.to(dtype=matmul_dtype)
    queries_mm = queries.to(dtype=matmul_dtype)
    if codes.ndim == 3:
        code_dim = int(codes.shape[-1])
        codes_flat = codes_mm.reshape(-1, code_dim)
        logits = torch.matmul(codes_flat, queries_mm.transpose(0, 1)).transpose(0, 1).to(torch.float32)
        return (
            logits * scales.reshape(1, -1)
            + query_group_sums.reshape(-1, 1) * bias.reshape(1, -1)
        )
    if codes.ndim == 4:
        batch_size = int(codes.shape[0])
        code_dim = int(codes.shape[-1])
        codes_flat = codes_mm.reshape(batch_size, -1, code_dim)
        logits = torch.bmm(codes_flat, queries_mm.transpose(1, 2)).transpose(1, 2).to(torch.float32)
        return (
            logits * scales.reshape(batch_size, 1, -1)
            + query_group_sums.reshape(batch_size, -1, 1) * bias.reshape(batch_size, 1, -1)
        )
    raise ValueError("codes must have shape [page_count, token_count, group_size] or [batch_size, page_count, token_count, group_size]")


def _unpack_packed_word_slice_torch(words, shifts, mask):
    torch = _load_torch()
    words_u64 = torch.bitwise_and(words.to(dtype=torch.int64), 0xFFFFFFFF)
    shifts_i64 = shifts.to(dtype=torch.int64)
    mask_i64 = torch.as_tensor(mask, dtype=torch.int64, device=words.device)
    return torch.bitwise_and(torch.bitwise_right_shift(words_u64[..., None], shifts_i64), mask_i64).to(torch.float32)


def _unpack_packed_group32_torch(payload_words, *, unpack_shifts, unpack_mask, trace=None):
    if payload_words.ndim != 4:
        raise ValueError("payload_words must have shape [batch_size, page_count, token_count, words_per_group]")
    unpacked = _trace_timed_call(
        trace,
        "unpack",
        device_type=str(payload_words.device.type),
        fn=lambda: _unpack_packed_word_slice_torch(payload_words, unpack_shifts, unpack_mask),
        synchronize=False,
    )
    if trace is not None:
        trace.record_temporary(int(unpacked.numel() * unpacked.element_size()))
    return unpacked.reshape(int(payload_words.shape[0]), int(payload_words.shape[1]), int(payload_words.shape[2]), -1)


def _score_m0_logits_packed32_torch(payload_words, queries, scales, bias, query_group_sums, *, unpack_shifts, unpack_mask, trace=None):
    if int(queries.shape[-1]) != int(payload_words.shape[-1]) * int(unpack_shifts.numel()):
        raise ValueError("queries must align with the packed payload layout")
    codes = _unpack_packed_group32_torch(
        payload_words,
        unpack_shifts=unpack_shifts,
        unpack_mask=unpack_mask,
        trace=trace,
    )
    return _score_m0_logits_flat_torch(codes, queries, scales, bias, query_group_sums)


def _score_m0_logits_packed32_grouped_torch(
    payload_words,
    queries,
    scales,
    bias,
    query_group_sums,
    *,
    unpack_shifts,
    unpack_mask,
    trace=None,
):
    torch = _load_torch()
    if payload_words.ndim != 5:
        raise ValueError("payload_words must have shape [batch_size, num_groups, page_count, token_count, words_per_group]")
    if queries.ndim != 4:
        raise ValueError("queries must have shape [batch_size, query_count, num_groups, group_size]")
    batch_size, num_groups, page_count, token_count, words_per_group = map(int, payload_words.shape)
    query_count = int(queries.shape[1])
    group_size = int(queries.shape[-1])
    if int(queries.shape[2]) != num_groups:
        raise ValueError("queries must align with payload_words group count")
    if tuple(scales.shape) != (batch_size, num_groups, page_count, token_count):
        raise ValueError("scales must align with payload_words shape")
    if tuple(bias.shape) != (batch_size, num_groups, page_count, token_count):
        raise ValueError("bias must align with payload_words shape")
    if tuple(query_group_sums.shape) != (batch_size, query_count, num_groups):
        raise ValueError("query_group_sums must align with queries shape")
    if group_size != words_per_group * int(unpack_shifts.numel()):
        raise ValueError("queries must align with the packed payload layout")
    payload_words_flat = payload_words.reshape(batch_size * num_groups, page_count, token_count, words_per_group)
    codes = _unpack_packed_group32_torch(
        payload_words_flat,
        unpack_shifts=unpack_shifts,
        unpack_mask=unpack_mask,
        trace=trace,
    )
    matmul_dtype = codes.dtype if torch.is_floating_point(codes) else torch.float32
    codes_flat = codes.reshape(batch_size * num_groups, -1, group_size).to(dtype=matmul_dtype)
    queries_flat = queries.transpose(1, 2).contiguous().reshape(batch_size * num_groups, query_count, group_size).to(dtype=matmul_dtype)
    logits = torch.bmm(codes_flat, queries_flat.transpose(1, 2)).transpose(1, 2).to(torch.float32)
    scales_flat = scales.reshape(batch_size * num_groups, 1, -1)
    bias_flat = bias.reshape(batch_size * num_groups, 1, -1)
    query_sums_flat = query_group_sums.transpose(1, 2).contiguous().reshape(batch_size * num_groups, query_count, 1)
    logits = logits * scales_flat + query_sums_flat * bias_flat
    return logits.reshape(batch_size, num_groups, query_count, page_count * token_count).sum(dim=1)


def _mix_m0_contribution_torch(weights, codes, scales, bias):
    torch = _load_torch()
    matmul_dtype = codes.dtype if torch.is_floating_point(codes) else torch.float32
    codes_mm = codes.to(dtype=matmul_dtype)
    if codes.ndim == 3:
        code_dim = int(codes.shape[-1])
        codes_flat = codes_mm.reshape(-1, code_dim)
        weights_flat = weights.reshape(weights.shape[0], -1)
        weighted_scales_flat = (weights_flat * scales.reshape(1, -1)).to(dtype=matmul_dtype)
        contribution = torch.matmul(weighted_scales_flat, codes_flat).to(torch.float32)
        bias_term = (weights_flat * bias.reshape(1, -1)).sum(dim=-1, keepdim=True)
        return contribution + bias_term
    if codes.ndim == 4:
        batch_size = int(codes.shape[0])
        code_dim = int(codes.shape[-1])
        codes_flat = codes_mm.reshape(batch_size, -1, code_dim)
        weights_flat = weights.reshape(batch_size, weights.shape[1], -1)
        weighted_scales_flat = (weights_flat * scales.reshape(batch_size, 1, -1)).to(dtype=matmul_dtype)
        contribution = torch.bmm(weighted_scales_flat, codes_flat).to(torch.float32)
        bias_term = (weights_flat * bias.reshape(batch_size, 1, -1)).sum(dim=-1, keepdim=True)
        return contribution + bias_term
    raise ValueError("codes must have shape [page_count, token_count, group_size] or [batch_size, page_count, token_count, group_size]")


def _mix_m0_contribution_packed32_torch(weights, payload_words, scales, bias, *, unpack_shifts, unpack_mask, trace=None):
    codes = _unpack_packed_group32_torch(
        payload_words,
        unpack_shifts=unpack_shifts,
        unpack_mask=unpack_mask,
        trace=trace,
    )
    return _mix_m0_contribution_torch(weights, codes, scales, bias)


def _mix_m0_contribution_packed32_grouped_torch(
    weights,
    payload_words,
    scales,
    bias,
    *,
    unpack_shifts,
    unpack_mask,
    trace=None,
):
    torch = _load_torch()
    if payload_words.ndim != 5:
        raise ValueError("payload_words must have shape [batch_size, num_groups, page_count, token_count, words_per_group]")
    if weights.ndim != 4:
        raise ValueError("weights must have shape [batch_size, query_count, page_count, token_count]")
    batch_size, num_groups, page_count, token_count, words_per_group = map(int, payload_words.shape)
    query_count = int(weights.shape[1])
    if tuple(weights.shape[2:]) != (page_count, token_count):
        raise ValueError("weights must align with payload_words shape")
    if tuple(scales.shape) != (batch_size, num_groups, page_count, token_count):
        raise ValueError("scales must align with payload_words shape")
    if tuple(bias.shape) != (batch_size, num_groups, page_count, token_count):
        raise ValueError("bias must align with payload_words shape")
    payload_words_flat = payload_words.reshape(batch_size * num_groups, page_count, token_count, words_per_group)
    codes = _unpack_packed_group32_torch(
        payload_words_flat,
        unpack_shifts=unpack_shifts,
        unpack_mask=unpack_mask,
        trace=trace,
    )
    group_size = int(codes.shape[-1])
    matmul_dtype = codes.dtype if torch.is_floating_point(codes) else torch.float32
    codes_flat = codes.reshape(batch_size * num_groups, -1, group_size).to(dtype=matmul_dtype)
    weights_flat = weights.reshape(batch_size, query_count, -1)
    weighted_scales = (
        weights_flat[:, :, None, :] * scales.reshape(batch_size, 1, num_groups, -1)
    ).permute(0, 2, 1, 3).contiguous().reshape(batch_size * num_groups, query_count, -1).to(dtype=matmul_dtype)
    contribution = torch.bmm(weighted_scales, codes_flat).to(torch.float32)
    contribution = contribution.reshape(batch_size, num_groups, query_count, group_size).permute(0, 2, 1, 3)
    bias_term = (
        weights_flat[:, :, None, :] * bias.reshape(batch_size, 1, num_groups, -1)
    ).permute(0, 2, 1, 3).sum(dim=-1).permute(0, 2, 1)
    return contribution + bias_term[..., None]


def _score_m0_logits_two_group64_torch(fused_scaled_codes, fused_queries, bias_groups, query_group_sums):
    torch = _load_torch()
    bias0, bias1 = bias_groups
    matmul_dtype = fused_scaled_codes.dtype if torch.is_floating_point(fused_scaled_codes) else torch.float32
    if fused_scaled_codes.ndim == 3:
        fused_queries_mm = fused_queries.to(dtype=matmul_dtype)
        logits = torch.matmul(
            fused_scaled_codes.reshape(-1, int(fused_scaled_codes.shape[-1])),
            fused_queries_mm.transpose(0, 1),
        ).transpose(0, 1).to(torch.float32)
        bias_term = (
            query_group_sums[:, 0:1] * bias0.reshape(1, -1)
            + query_group_sums[:, 1:2] * bias1.reshape(1, -1)
        )
        return logits + bias_term
    if fused_scaled_codes.ndim == 4:
        batch_size = int(fused_scaled_codes.shape[0])
        fused_queries_mm = fused_queries.to(dtype=matmul_dtype)
        logits = torch.bmm(
            fused_scaled_codes.reshape(batch_size, -1, int(fused_scaled_codes.shape[-1])),
            fused_queries_mm.transpose(1, 2),
        ).transpose(1, 2).to(torch.float32)
        bias_term = (
            query_group_sums[:, :, 0:1] * bias0.reshape(batch_size, 1, -1)
            + query_group_sums[:, :, 1:2] * bias1.reshape(batch_size, 1, -1)
        )
        return logits + bias_term
    raise ValueError("fused_scaled_codes must have shape [page_count, token_count, 64] or [batch_size, page_count, token_count, 64]")


def _score_m0_logits_fused_torch(fused_scaled_codes, fused_queries, bias_groups, query_group_sums):
    torch = _load_torch()
    if fused_scaled_codes.ndim == 3:
        logits = torch.matmul(
            fused_scaled_codes.reshape(-1, int(fused_scaled_codes.shape[-1])),
            fused_queries.transpose(0, 1),
        ).transpose(0, 1).to(torch.float32)
        bias_term = torch.zeros_like(logits)
        for group_index, bias_group in enumerate(bias_groups):
            bias_term += query_group_sums[:, group_index : group_index + 1] * bias_group.reshape(1, -1)
        return logits + bias_term
    if fused_scaled_codes.ndim == 4:
        batch_size = int(fused_scaled_codes.shape[0])
        logits = torch.bmm(
            fused_scaled_codes.reshape(batch_size, -1, int(fused_scaled_codes.shape[-1])),
            fused_queries.transpose(1, 2),
        ).transpose(1, 2).to(torch.float32)
        bias_term = torch.zeros_like(logits)
        for group_index, bias_group in enumerate(bias_groups):
            bias_term += query_group_sums[:, :, group_index : group_index + 1] * bias_group.reshape(batch_size, 1, -1)
        return logits + bias_term
    raise ValueError(
        "fused_scaled_codes must have shape [page_count, token_count, padded_head_dim] "
        "or [batch_size, page_count, token_count, padded_head_dim]"
    )


def _mix_m0_contribution_two_group64_torch(weights, fused_scaled_codes, bias_groups):
    torch = _load_torch()
    bias0, bias1 = bias_groups
    matmul_dtype = fused_scaled_codes.dtype if torch.is_floating_point(fused_scaled_codes) else torch.float32
    if fused_scaled_codes.ndim == 3:
        weights_flat = weights.reshape(weights.shape[0], -1)
        weights_mm = weights_flat.to(dtype=matmul_dtype)
        contribution = torch.matmul(
            weights_mm,
            fused_scaled_codes.reshape(-1, int(fused_scaled_codes.shape[-1])),
        ).to(torch.float32)
        bias0_term = (weights_flat * bias0.reshape(1, -1)).sum(dim=-1, keepdim=True)
        bias1_term = (weights_flat * bias1.reshape(1, -1)).sum(dim=-1, keepdim=True)
        return torch.cat([contribution[:, :32] + bias0_term, contribution[:, 32:] + bias1_term], dim=-1)
    if fused_scaled_codes.ndim == 4:
        batch_size = int(fused_scaled_codes.shape[0])
        weights_flat = weights.reshape(batch_size, weights.shape[1], -1)
        weights_mm = weights_flat.to(dtype=matmul_dtype)
        contribution = torch.bmm(
            weights_mm,
            fused_scaled_codes.reshape(batch_size, -1, int(fused_scaled_codes.shape[-1])),
        ).to(torch.float32)
        bias0_term = (weights_flat * bias0.reshape(batch_size, 1, -1)).sum(dim=-1, keepdim=True)
        bias1_term = (weights_flat * bias1.reshape(batch_size, 1, -1)).sum(dim=-1, keepdim=True)
        return torch.cat([contribution[:, :, :32] + bias0_term, contribution[:, :, 32:] + bias1_term], dim=-1)
    raise ValueError("fused_scaled_codes must have shape [page_count, token_count, 64] or [batch_size, page_count, token_count, 64]")


def _mix_m0_contribution_fused_torch(weights, fused_scaled_codes, bias_groups, *, group_size: int):
    torch = _load_torch()
    if fused_scaled_codes.ndim == 3:
        weights_flat = weights.reshape(weights.shape[0], -1)
        output = torch.matmul(
            weights_flat,
            fused_scaled_codes.reshape(-1, int(fused_scaled_codes.shape[-1])),
        ).to(torch.float32)
        for group_index, bias_group in enumerate(bias_groups):
            bias_term = (weights_flat * bias_group.reshape(1, -1)).sum(dim=-1, keepdim=True)
            start = group_index * group_size
            end = start + group_size
            output[:, start:end] += bias_term
        return output
    if fused_scaled_codes.ndim == 4:
        batch_size = int(fused_scaled_codes.shape[0])
        weights_flat = weights.reshape(batch_size, weights.shape[1], -1)
        output = torch.bmm(
            weights_flat,
            fused_scaled_codes.reshape(batch_size, -1, int(fused_scaled_codes.shape[-1])),
        ).to(torch.float32)
        for group_index, bias_group in enumerate(bias_groups):
            bias_term = (weights_flat * bias_group.reshape(batch_size, 1, -1)).sum(dim=-1, keepdim=True)
            start = group_index * group_size
            end = start + group_size
            output[:, :, start:end] += bias_term
        return output
    raise ValueError(
        "fused_scaled_codes must have shape [page_count, token_count, padded_head_dim] "
        "or [batch_size, page_count, token_count, padded_head_dim]"
    )


def _lookup_lut_group_torch(codebooks, codes):
    torch = _load_torch()
    lut = codebooks.to(dtype=torch.float32)
    code_indices = codes.to(dtype=torch.int64)
    if lut.ndim == 1 and code_indices.ndim == 1:
        return lut[code_indices]
    if lut.ndim == 2 and code_indices.ndim == 1 and lut.shape[0] == code_indices.shape[0]:
        page_index = torch.arange(lut.shape[0], device=lut.device)
        return lut[page_index, code_indices]
    if lut.ndim == 2 and code_indices.ndim == 2:
        token_count = int(code_indices.shape[0])
        segment_count = int(lut.shape[0])
        if segment_count == 1:
            return lut[0][code_indices]
        segment_ids = (torch.arange(token_count, device=lut.device, dtype=torch.int64) * segment_count) // max(token_count, 1)
        return lut[segment_ids[:, None], code_indices]
    if lut.ndim == 3 and code_indices.ndim == 3:
        page_count = int(code_indices.shape[0])
        token_count = int(code_indices.shape[1])
        segment_count = int(lut.shape[1])
        if segment_count == 1:
            page_index = torch.arange(page_count, device=lut.device)[:, None, None]
            return lut[page_index, torch.zeros(1, device=lut.device, dtype=torch.int64), code_indices]
        page_index = torch.arange(page_count, device=lut.device)[:, None, None]
        segment_ids = (torch.arange(token_count, device=lut.device, dtype=torch.int64) * segment_count) // max(token_count, 1)
        return lut[page_index, segment_ids[None, :, None], code_indices]
    if lut.ndim == 4 and code_indices.ndim == 4:
        batch_size = int(code_indices.shape[0])
        page_count = int(code_indices.shape[1])
        token_count = int(code_indices.shape[2])
        segment_count = int(lut.shape[2])
        batch_index = torch.arange(batch_size, device=lut.device)[:, None, None, None]
        page_index = torch.arange(page_count, device=lut.device)[None, :, None, None]
        if segment_count == 1:
            return lut[batch_index, page_index, torch.zeros(1, device=lut.device, dtype=torch.int64), code_indices]
        segment_ids = (torch.arange(token_count, device=lut.device, dtype=torch.int64) * segment_count) // max(token_count, 1)
        return lut[batch_index, page_index, segment_ids[None, None, :, None], code_indices]
    raise ValueError("unsupported LUT rank")


def _lookup_turbo_group_torch(codebooks, codes):
    torch = _load_torch()
    lut = codebooks.to(dtype=torch.float32)
    code_indices = codes.to(dtype=torch.int64)
    if lut.ndim == 1 and code_indices.ndim == 2:
        return lut[code_indices]
    if lut.ndim == 1 and code_indices.ndim == 3:
        return lut[code_indices]
    if lut.ndim == 1 and code_indices.ndim == 4:
        return lut[code_indices]
    if lut.ndim == 2 and code_indices.ndim == 3 and lut.shape[0] == code_indices.shape[0]:
        page_index = torch.arange(lut.shape[0], device=lut.device)[:, None, None]
        return lut[page_index, code_indices]
    if lut.ndim == 3 and code_indices.ndim == 4:
        batch_index = torch.arange(lut.shape[0], device=lut.device)[:, None, None, None]
        page_index = torch.arange(lut.shape[1], device=lut.device)[None, :, None, None]
        return lut[batch_index, page_index, code_indices]
    raise ValueError("unsupported turbo3 LUT rank")


def _fwht_last_dim_torch(values, *, trace: ExecutionTrace | None = None):
    torch = _load_torch()
    width = int(values.shape[-1])
    if width <= 0:
        return values
    if width & (width - 1):
        raise ValueError("FWHT requires the last dimension to be a power of two")
    device_type = str(values.device.type)

    def _impl():
        original_shape = tuple(values.shape)
        matrix = _fwht_matrix_torch(width, device_type=device_type)
        transformed = values.to(dtype=torch.float32).reshape(-1, width)
        return torch.matmul(transformed, matrix.T).reshape(original_shape)

    return _trace_timed_call(trace, "fwht", device_type=device_type, fn=_impl, synchronize=False)


def _score_page_chunk_torch(query_slice: np.ndarray | Any, pages: Sequence[PreparedPageTorch], *, trace: ExecutionTrace | None = None):
    torch = _load_torch()
    if not pages:
        raise ValueError("pages must be non-empty")
    header = pages[0].header
    device_type = pages[0].device_type
    if trace is not None:
        trace.record_page_read(
            sum(page.payload_nbytes for page in pages),
            sum(page.metadata_nbytes for page in pages),
        )

    if header.mode_default == "M3":
        use_native_dtype = _m3_native_compute_enabled(pages)
        dense = _decode_escape_batch_torch(
            pages,
            token_count=header.token_count,
            head_dim=header.head_dim,
            promote_float32=not use_native_dtype,
        )
        query = _pad_query(query_slice, header.head_dim, device_type=device_type)
        if use_native_dtype:
            return torch.matmul(dense, query.to(dtype=dense.dtype)).reshape(-1).to(dtype=torch.float32)
        return torch.matmul(dense, query).reshape(-1)

    if header.mode_default == "M2":
        query = _pad_query(query_slice, header.padded_head_dim, device_type=device_type)
        query_groups = query.reshape(header.num_groups, header.group_size)
        page_count = len(pages)
        logits = torch.zeros((page_count, header.token_count), dtype=torch.float32, device=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_sketch = (
                prepared_chunk.m2_sketch_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_sketch_groups is not None
                else torch.stack([page.m2_sketch[:, group_index, :] for page in pages], dim=0)
            )
            group_basis = (
                prepared_chunk.m2_basis_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_basis_groups is not None
                else torch.stack([page.m2_basis[group_index] for page in pages], dim=0)
            )
            group_mean = (
                prepared_chunk.m2_mean_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_mean_groups is not None
                else torch.stack([page.m2_mean[group_index] for page in pages], dim=0)
            )
            qg, group_sketch, group_basis, group_mean = _coerce_m2_operands(
                query_groups[group_index],
                group_sketch,
                group_basis,
                group_mean,
            )
            if group_basis.dim() == 3:
                q_proj = torch.einsum("prg,g->pr", group_basis, qg)
                logits += torch.einsum("ptd,pd->pt", group_sketch, q_proj)
                logits += torch.einsum("pg,g->p", group_mean, qg)[:, None]
                continue
            segment_ids = (
                prepared_chunk.m2_segment_ids
                if prepared_chunk is not None and prepared_chunk.m2_segment_ids is not None
                else _segment_ids_tensor(header.token_count, int(group_basis.shape[1]), device_type=device_type)
            )
            q_proj = torch.einsum("psrg,g->psr", group_basis, qg)
            logits += torch.einsum("ptr,ptr->pt", group_sketch, q_proj[:, segment_ids, :])
            logits += torch.einsum("ptg,g->pt", group_mean[:, segment_ids, :], qg)
        return logits.reshape(-1)

    if header.mode_default == "M4":
        query = _pad_query(query_slice, header.padded_head_dim, device_type=device_type)
        query_groups = query.reshape(header.num_groups, header.group_size)
        page_count = len(pages)
        logits = torch.zeros((page_count, header.token_count), dtype=torch.float32, device=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_sketch = (
                prepared_chunk.m2_sketch_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_sketch_groups is not None
                else torch.stack([page.m2_sketch[:, group_index, :] for page in pages], dim=0)
            )
            group_basis = (
                prepared_chunk.m2_basis_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_basis_groups is not None
                else (
                    torch.stack([page.m2_basis[group_index] for page in pages], dim=0)
                    if pages[0].m2_basis is not None
                    else (
                        _device_tensor(np.asarray(pages[0].source_page.m2_basis[group_index]), device=device_type).contiguous()
                        if (
                            pages[0].header.project_basis == "svd_shared"
                            and pages[0].source_page.m2_basis is not None
                            and all(page.source_page.m2_basis is pages[0].source_page.m2_basis for page in pages)
                        )
                        else None
                    )
                )
            )
            group_mean = (
                prepared_chunk.m2_mean_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_mean_groups is not None
                else torch.stack([page.m2_mean[group_index] for page in pages], dim=0)
            )
            qg = query_groups[group_index]
            if group_basis is not None:
                qg, group_sketch, group_basis, group_mean = _coerce_m2_operands(qg, group_sketch, group_basis, group_mean)
                if int(group_basis.dim()) == 2:
                    q_proj = torch.einsum("rg,g->r", group_basis, qg)
                    logits += torch.einsum("ptr,r->pt", group_sketch, q_proj)
                else:
                    q_proj = torch.einsum("prg,g->pr", group_basis, qg)
                    logits += torch.einsum("ptr,pr->pt", group_sketch, q_proj)
            else:
                basis = _m4_basis_torch(
                    header.group_size,
                    int(pages[0].m2_sketch.shape[-1]),
                    basis_family=header.project_basis,
                    device_type=device_type,
                )
                qg, group_sketch, _, group_mean = _coerce_m2_operands(qg, group_sketch, basis, group_mean)
                q_proj = torch.matmul(qg, basis.to(dtype=qg.dtype).transpose(0, 1))
                logits += torch.einsum("ptr,r->pt", group_sketch, q_proj)
            logits += torch.einsum("pg,g->p", group_mean, qg)[:, None]
        return logits.reshape(-1)

    if header.mode_default == "M1":
        query = _pad_query(query_slice, header.padded_head_dim, device_type=device_type)
        query_groups = query.reshape(header.num_groups, header.group_size)
        page_count = len(pages)
        logits = torch.zeros((page_count, header.token_count), dtype=torch.float32, device=device_type)
        for group_index in range(header.num_groups):
            group_words = torch.stack([page.payload[group_index] for page in pages], dim=0)
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, header.token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            group = _lookup_lut_group_torch(
                torch.stack([page.codebooks[group_index] for page in pages], dim=0),
                codes,
            )
            logits += torch.matmul(group, query_groups[group_index])
        return logits.reshape(-1)

    if header.mode_default == "T3":
        query = _pad_query(query_slice, header.padded_head_dim, device_type=device_type)
        rotated_query_groups = _fwht_last_dim_torch(query.reshape(header.num_groups, header.group_size))
        page_count = len(pages)
        logits = torch.zeros((page_count, header.token_count), dtype=torch.float32, device=device_type)
        codebooks = pages[0].codebooks if pages[0].codebooks is not None else _turbo3_centroids_torch(device_type=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_words = (
                prepared_chunk.payload_groups[group_index]
                if prepared_chunk is not None
                else torch.stack([page.payload[group_index] for page in pages], dim=0)
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, header.token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            scales = (
                prepared_chunk.scales_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.scales_groups is not None
                else torch.stack([page.scales[:, group_index] for page in pages], dim=0)
            )
            corrected = _lookup_turbo_group_torch(codebooks, codes) * scales[..., None]
            logits += torch.matmul(corrected, rotated_query_groups[group_index])
        return logits.reshape(-1)

    query = _pad_query(query_slice, header.padded_head_dim, device_type=device_type)
    query_groups = query.reshape(header.num_groups, header.group_size)
    query_group_sums = query_groups.sum(dim=-1)
    page_count = len(pages)
    logits = torch.zeros((page_count, header.token_count), dtype=torch.float32, device=device_type)

    prepared_chunk = _get_prepared_chunk_mps(pages)
    if (
        prepared_chunk is not None
        and prepared_chunk.fused_scaled_codes is not None
        and prepared_chunk.bias_groups is not None
    ):
        if trace is not None:
            trace.record_per_kv_kernel_variant(section="score", variant="fused_generic")
        fused_query = query.reshape(1, header.padded_head_dim).contiguous()
        fused_query_group_sums = query_group_sums.reshape(1, header.num_groups).contiguous()
        return _score_m0_logits_fused_torch(
            prepared_chunk.fused_scaled_codes,
            fused_query,
            prepared_chunk.bias_groups,
            fused_query_group_sums,
        ).reshape(-1)
    if trace is not None:
        trace.record_per_kv_kernel_variant(section="score", variant="generic")
    for group_index in range(header.num_groups):
        cached_codes = prepared_chunk is not None and prepared_chunk.codes_groups is not None
        if cached_codes:
            codes = prepared_chunk.codes_groups[group_index]
        else:
            codes = _trace_timed_call(
                trace,
                "unpack",
                device_type=device_type,
                fn=lambda group_index=group_index: _unpack_bits_torch(
                    torch.stack([page.payload[group_index] for page in pages], dim=0).reshape(-1, header.words_per_group),
                    pages[0].unpack_shifts,
                    pages[0].unpack_mask,
                    header.group_size,
                ).reshape(page_count, header.token_count, header.group_size),
            )
        if trace is not None and not cached_codes:
            trace.record_temporary(int(codes.numel() * codes.element_size()))
        qg = query_groups[group_index]
        int_dot = torch.matmul(codes, qg)
        scales = (
            prepared_chunk.scales_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.scales_groups is not None
            else torch.stack([page.scales[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        bias = (
            prepared_chunk.bias_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.bias_groups is not None
            else torch.stack([page.bias[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        logits += scales * int_dot + bias * query_group_sums[group_index]

    return logits.reshape(-1)


def _mix_page_chunk_torch(
    attn_weights,
    pages: Sequence[PreparedPageTorch],
    *,
    out_acc=None,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not pages:
        raise ValueError("pages must be non-empty")
    header = pages[0].header
    device_type = pages[0].device_type
    page_count = len(pages)
    token_count = header.token_count
    output = _prepare_output_accumulator_tensor(out_acc, header.head_dim, header.padded_head_dim, device_type=device_type)

    if trace is not None:
        trace.record_page_read(
            sum(page.payload_nbytes for page in pages),
            sum(page.metadata_nbytes for page in pages),
        )

    if not isinstance(attn_weights, torch.Tensor):
        weights = torch.as_tensor(attn_weights, dtype=torch.float32, device=device_type)
    else:
        weights = attn_weights.to(dtype=torch.float32, device=device_type)
    expected_shape = (page_count, token_count)
    if tuple(weights.shape) != expected_shape:
        raise ValueError("attn_weights chunk must have shape [page_count, token_count]")

    if header.mode_default == "M3":
        use_native_dtype = _m3_native_compute_enabled(pages)
        dense = _decode_escape_batch_torch(
            pages,
            token_count=header.token_count,
            head_dim=header.head_dim,
            promote_float32=not use_native_dtype,
        )
        if use_native_dtype:
            output[: header.head_dim] += torch.sum(weights.to(dtype=dense.dtype)[..., None] * dense, dim=(0, 1)).to(
                dtype=torch.float32
            )
        else:
            output[: header.head_dim] += torch.sum(weights[..., None] * dense, dim=(0, 1))
        return output

    if header.mode_default in {"M2", "M4"}:
        raise ValueError(f"{header.mode_default} is only supported for key scoring in this phase")

    if header.mode_default == "M1":
        for group_index in range(header.num_groups):
            group_words = torch.stack([page.payload[group_index] for page in pages], dim=0)
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            group = _lookup_lut_group_torch(
                torch.stack([page.codebooks[group_index] for page in pages], dim=0),
                codes,
            )
            start = group_index * header.group_size
            end = start + header.group_size
            output[start:end] += torch.einsum("pt,ptg->g", weights, group)
        return output

    if header.mode_default == "T3":
        codebooks = pages[0].codebooks if pages[0].codebooks is not None else _turbo3_centroids_torch(device_type=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_words = (
                prepared_chunk.payload_groups[group_index]
                if prepared_chunk is not None
                else torch.stack([page.payload[group_index] for page in pages], dim=0)
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            scales = (
                prepared_chunk.scales_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.scales_groups is not None
                else torch.stack([page.scales[:, group_index] for page in pages], dim=0)
            )
            rotated_group = _lookup_turbo_group_torch(codebooks, codes) * scales[..., None]
            group = _fwht_last_dim_torch(rotated_group)
            start = group_index * header.group_size
            end = start + header.group_size
            output[start:end] += torch.einsum("pt,ptg->g", weights, group)
        return output

    prepared_chunk = _get_prepared_chunk_mps(pages)
    if (
        prepared_chunk is not None
        and prepared_chunk.fused_scaled_codes is not None
        and prepared_chunk.bias_groups is not None
    ):
        if trace is not None:
            trace.record_per_kv_kernel_variant(section="mix", variant="fused_generic")
        output[: header.padded_head_dim] += _mix_m0_contribution_fused_torch(
            weights.reshape(1, page_count, token_count),
            prepared_chunk.fused_scaled_codes,
            prepared_chunk.bias_groups,
            group_size=header.group_size,
        ).reshape(-1)
        return output
    if trace is not None:
        trace.record_per_kv_kernel_variant(section="mix", variant="generic")
    for group_index in range(header.num_groups):
        cached_codes = prepared_chunk is not None and prepared_chunk.codes_groups is not None
        if cached_codes:
            codes = prepared_chunk.codes_groups[group_index]
        else:
            codes = _trace_timed_call(
                trace,
                "unpack",
                device_type=device_type,
                fn=lambda group_index=group_index: _unpack_bits_torch(
                    torch.stack([page.payload[group_index] for page in pages], dim=0).reshape(-1, header.words_per_group),
                    pages[0].unpack_shifts,
                    pages[0].unpack_mask,
                    header.group_size,
                ).reshape(page_count, token_count, header.group_size),
            )
        if trace is not None and not cached_codes:
            trace.record_temporary(int(codes.numel() * codes.element_size()))
        scales = (
            prepared_chunk.scales_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.scales_groups is not None
            else torch.stack([page.scales[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        bias = (
            prepared_chunk.bias_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.bias_groups is not None
            else torch.stack([page.bias[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        weighted_scales = weights * scales
        contribution = torch.sum(weighted_scales[..., None] * codes, dim=(0, 1))
        bias_term = torch.sum(weights * bias)
        start = group_index * header.group_size
        end = start + header.group_size
        output[start:end] += contribution + bias_term

    return output


def score_page_torch(
    query_slice: np.ndarray | Any,
    page: EncodedPage | PreparedPageTorch,
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
) -> np.ndarray:
    prepared = prepare_page_torch(page, device_type=device_type, trace=trace)
    return _score_page_chunk_torch(query_slice, [prepared], trace=trace).detach().cpu().numpy()


def score_pages_torch(
    query_slice: np.ndarray | Any,
    pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
) -> list[np.ndarray]:
    prepared_pages = prepare_pages_torch(pages, device_type=device_type, trace=trace)
    if not prepared_pages:
        return []
    page_logits: list[np.ndarray] = []
    for page_chunk in _chunk_compatible_pages(prepared_pages):
        chunk_logits = _score_page_chunk_torch(query_slice, page_chunk, trace=trace)
        chunk_logits = chunk_logits.reshape(len(page_chunk), page_chunk[0].header.token_count)
        page_logits.extend(chunk_logits[index].detach().cpu().numpy() for index in range(len(page_chunk)))
    return page_logits


def mix_page_torch(
    attn_weights: np.ndarray | Any,
    page: EncodedPage | PreparedPageTorch,
    *,
    device_type: TorchDevice,
    out_acc: np.ndarray | None = None,
    trace: ExecutionTrace | None = None,
) -> np.ndarray:
    prepared = prepare_page_torch(page, device_type=device_type, trace=trace)
    header = prepared.header
    output = _mix_page_chunk_torch(
        np.asarray(attn_weights, dtype=np.float32)[None, :],
        [prepared],
        out_acc=None if out_acc is None else _prepare_output_accumulator(out_acc, header.head_dim, header.padded_head_dim, device_type=device_type),
        trace=trace,
    )
    return output[: header.head_dim].detach().cpu().numpy()


def _page_logits_tensor(page_logits, token_count: int, *, device_type: TorchDevice):
    torch = _load_torch()
    logits = torch.as_tensor(page_logits, dtype=torch.float32, device=device_type)
    if tuple(logits.shape) != (token_count,):
        raise ValueError("precomputed page logits must have shape [token_count]")
    return logits


def decode_step_torch(
    query_slice: np.ndarray | Any,
    key_pages: Sequence[EncodedPage | PreparedPageTorch],
    value_pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    device_type: TorchDevice,
    precomputed_page_logits: Sequence[np.ndarray | Any | None] | None = None,
    trace: ExecutionTrace | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    torch = _load_torch()
    prepared_key_pages = prepare_pages_torch(key_pages, device_type=device_type, trace=trace)
    prepared_value_pages = prepare_pages_torch(value_pages, device_type=device_type, trace=trace)
    if not prepared_key_pages:
        raise ValueError(f"decode_step_{device_type} requires at least one page")
    if precomputed_page_logits is not None and len(precomputed_page_logits) != len(prepared_key_pages):
        raise ValueError("precomputed_page_logits must align with key_pages")

    logits_parts = []
    score_run: list[PreparedPageTorch] = []

    def flush_score_run() -> None:
        nonlocal score_run
        if not score_run:
            return
        chunk_logits = _score_page_chunk_torch(query_slice, score_run, trace=trace)
        chunk_logits = chunk_logits.reshape(len(score_run), score_run[0].header.token_count)
        logits_parts.extend(chunk_logits[index] for index in range(len(score_run)))
        score_run = []

    for index, page in enumerate(prepared_key_pages):
        cached_logits = None if precomputed_page_logits is None else precomputed_page_logits[index]
        if cached_logits is not None:
            flush_score_run()
            logits_parts.append(_page_logits_tensor(cached_logits, page.header.token_count, device_type=device_type))
            continue
        if score_run and _batched_signature(score_run[-1]) != _batched_signature(page):
            flush_score_run()
        score_run.append(page)
    flush_score_run()

    logits = torch.cat(logits_parts, dim=0)
    weights = torch.softmax(logits, dim=0)

    output = torch.zeros(
        prepared_value_pages[0].header.padded_head_dim,
        dtype=torch.float32,
        device=device_type,
    )
    offset = 0
    for page_chunk in _chunk_compatible_pages(prepared_value_pages):
        chunk_token_count = page_chunk[0].header.token_count * len(page_chunk)
        chunk_weights = weights[offset : offset + chunk_token_count].reshape(len(page_chunk), page_chunk[0].header.token_count)
        output = _mix_page_chunk_torch(chunk_weights, page_chunk, out_acc=output, trace=trace)
        offset += chunk_token_count

    head_dim = prepared_value_pages[0].header.head_dim
    return (
        logits.detach().cpu().numpy(),
        weights.detach().cpu().numpy(),
        output[:head_dim].detach().cpu().numpy(),
    )


def _score_page_chunk_multiquery_torch(
    query_slices: np.ndarray | Any,
    pages: Sequence[PreparedPageTorch],
    *,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not pages:
        raise ValueError("pages must be non-empty")
    header = pages[0].header
    device_type = pages[0].device_type
    if torch.is_tensor(query_slices):
        query_count = int(query_slices.shape[0])
    else:
        query_count = int(np.asarray(query_slices).shape[0])
    if trace is not None:
        trace.record_page_read(
            sum(page.payload_nbytes for page in pages),
            sum(page.metadata_nbytes for page in pages),
        )

    if header.mode_default == "M3":
        use_native_dtype = _m3_native_compute_enabled(pages)
        dense = _decode_escape_batch_torch(
            pages,
            token_count=header.token_count,
            head_dim=header.head_dim,
            promote_float32=not use_native_dtype,
        )
        queries = _pad_queries(query_slices, header.head_dim, device_type=device_type)
        if use_native_dtype:
            return torch.einsum("pth,qh->qpt", dense, queries.to(dtype=dense.dtype)).reshape(query_count, -1).to(
                dtype=torch.float32
            )
        return torch.einsum("pth,qh->qpt", dense, queries).reshape(query_count, -1)

    if header.mode_default == "M2":
        queries = _pad_queries(query_slices, header.padded_head_dim, device_type=device_type)
        query_groups = queries.reshape(query_count, header.num_groups, header.group_size)
        page_count = len(pages)
        logits = torch.zeros((query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_sketch = (
                prepared_chunk.m2_sketch_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_sketch_groups is not None
                else torch.stack([page.m2_sketch[:, group_index, :] for page in pages], dim=0)
            )
            group_basis = (
                prepared_chunk.m2_basis_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_basis_groups is not None
                else torch.stack([page.m2_basis[group_index] for page in pages], dim=0)
            )
            group_mean = (
                prepared_chunk.m2_mean_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_mean_groups is not None
                else torch.stack([page.m2_mean[group_index] for page in pages], dim=0)
            )
            qg, group_sketch, group_basis, group_mean = _coerce_m2_operands(
                query_groups[:, group_index, :],
                group_sketch,
                group_basis,
                group_mean,
            )
            if group_basis.dim() == 3:
                q_proj = torch.einsum("prg,qg->qpr", group_basis, qg)
                logits += torch.einsum("ptd,qpd->qpt", group_sketch, q_proj)
                logits += torch.einsum("pg,qg->qp", group_mean, qg)[:, :, None]
                continue
            segment_ids = (
                prepared_chunk.m2_segment_ids
                if prepared_chunk is not None and prepared_chunk.m2_segment_ids is not None
                else _segment_ids_tensor(header.token_count, int(group_basis.shape[1]), device_type=device_type)
            )
            q_proj = torch.einsum("psrg,qg->qpsr", group_basis, qg)
            logits += torch.einsum("ptr,qptr->qpt", group_sketch, q_proj[:, :, segment_ids, :])
            logits += torch.einsum("ptg,qg->qpt", group_mean[:, segment_ids, :], qg)
        return logits.reshape(query_count, -1)

    if header.mode_default == "M4":
        queries = _pad_queries(query_slices, header.padded_head_dim, device_type=device_type)
        query_groups = queries.reshape(query_count, header.num_groups, header.group_size)
        page_count = len(pages)
        logits = torch.zeros((query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_sketch = (
                prepared_chunk.m2_sketch_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_sketch_groups is not None
                else torch.stack([page.m2_sketch[:, group_index, :] for page in pages], dim=0)
            )
            group_basis = (
                prepared_chunk.m2_basis_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_basis_groups is not None
                else (
                    torch.stack([page.m2_basis[group_index] for page in pages], dim=0)
                    if pages[0].m2_basis is not None
                    else (
                        _device_tensor(np.asarray(pages[0].source_page.m2_basis[group_index]), device=device_type).contiguous()
                        if (
                            pages[0].header.project_basis == "svd_shared"
                            and pages[0].source_page.m2_basis is not None
                            and all(page.source_page.m2_basis is pages[0].source_page.m2_basis for page in pages)
                        )
                        else None
                    )
                )
            )
            group_mean = (
                prepared_chunk.m2_mean_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.m2_mean_groups is not None
                else torch.stack([page.m2_mean[group_index] for page in pages], dim=0)
            )
            qg = query_groups[:, group_index, :]
            if group_basis is not None:
                qg, group_sketch, group_basis, group_mean = _coerce_m2_operands(qg, group_sketch, group_basis, group_mean)
                if int(group_basis.dim()) == 2:
                    q_proj = torch.einsum("rg,qg->qr", group_basis, qg)
                    logits += torch.einsum("ptr,qr->qpt", group_sketch, q_proj)
                else:
                    q_proj = torch.einsum("prg,qg->qpr", group_basis, qg)
                    logits += torch.einsum("ptr,qpr->qpt", group_sketch, q_proj)
            else:
                basis = _m4_basis_torch(
                    header.group_size,
                    int(pages[0].m2_sketch.shape[-1]),
                    basis_family=header.project_basis,
                    device_type=device_type,
                )
                qg, group_sketch, _, group_mean = _coerce_m2_operands(qg, group_sketch, basis, group_mean)
                q_proj = torch.matmul(qg, basis.to(dtype=qg.dtype).transpose(0, 1))
                logits += torch.einsum("ptr,qr->qpt", group_sketch, q_proj)
            logits += torch.einsum("pg,qg->qp", group_mean, qg)[:, :, None]
        return logits.reshape(query_count, -1)

    if header.mode_default == "M1":
        queries = _pad_queries(query_slices, header.padded_head_dim, device_type=device_type)
        query_groups = queries.reshape(query_count, header.num_groups, header.group_size)
        page_count = len(pages)
        logits = torch.zeros((query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        for group_index in range(header.num_groups):
            group_words = torch.stack([page.payload[group_index] for page in pages], dim=0)
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, header.token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            group = _lookup_lut_group_torch(
                torch.stack([page.codebooks[group_index] for page in pages], dim=0),
                codes,
            )
            logits += torch.einsum("ptg,qg->qpt", group, query_groups[:, group_index, :])
        return logits.reshape(query_count, -1)

    if header.mode_default == "T3":
        queries = _pad_queries(query_slices, header.padded_head_dim, device_type=device_type)
        rotated_query_groups = _fwht_last_dim_torch(queries.reshape(query_count, header.num_groups, header.group_size))
        page_count = len(pages)
        logits = torch.zeros((query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        codebooks = pages[0].codebooks if pages[0].codebooks is not None else _turbo3_centroids_torch(device_type=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_words = (
                prepared_chunk.payload_groups[group_index]
                if prepared_chunk is not None
                else torch.stack([page.payload[group_index] for page in pages], dim=0)
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, header.token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            scales = (
                prepared_chunk.scales_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.scales_groups is not None
                else torch.stack([page.scales[:, group_index] for page in pages], dim=0)
            )
            corrected = _lookup_turbo_group_torch(codebooks, codes) * scales[..., None]
            logits += torch.einsum("ptg,qg->qpt", corrected, rotated_query_groups[:, group_index, :])
        return logits.reshape(query_count, -1)

    queries = _pad_queries(query_slices, header.padded_head_dim, device_type=device_type)
    query_groups = queries.reshape(query_count, header.num_groups, header.group_size)
    query_group_sums = query_groups.sum(dim=-1)
    page_count = len(pages)
    logits = torch.zeros((query_count, page_count * header.token_count), dtype=torch.float32, device=device_type)

    prepared_chunk = _get_prepared_chunk_mps(pages)
    if (
        prepared_chunk is not None
        and prepared_chunk.fused_scaled_codes is not None
        and prepared_chunk.bias_groups is not None
        and _supports_fused_two_group64(header)
    ):
        if trace is not None:
            trace.record_per_kv_kernel_variant(section="score", variant="fused_two_group64")
        fused_queries = query_groups.reshape(query_count, header.padded_head_dim).contiguous()
        return _score_m0_logits_two_group64_torch(
            prepared_chunk.fused_scaled_codes,
            fused_queries,
            prepared_chunk.bias_groups,
            query_group_sums,
        )
    if trace is not None:
        trace.record_per_kv_kernel_variant(section="score", variant="generic")
    for group_index in range(header.num_groups):
        cached_codes = prepared_chunk is not None and prepared_chunk.codes_groups is not None
        if cached_codes:
            codes = prepared_chunk.codes_groups[group_index]
        else:
            codes = _trace_timed_call(
                trace,
                "unpack",
                device_type=device_type,
                fn=lambda group_index=group_index: _unpack_bits_torch(
                    torch.stack([page.payload[group_index] for page in pages], dim=0).reshape(-1, header.words_per_group),
                    pages[0].unpack_shifts,
                    pages[0].unpack_mask,
                    header.group_size,
                ).reshape(page_count, header.token_count, header.group_size),
            )
        if trace is not None and not cached_codes:
            trace.record_temporary(int(codes.numel() * codes.element_size()))
        qg = query_groups[:, group_index, :]
        scales = (
            prepared_chunk.scales_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.scales_groups is not None
            else torch.stack([page.scales[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        bias = (
            prepared_chunk.bias_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.bias_groups is not None
            else torch.stack([page.bias[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        logits += _score_m0_logits_flat_torch(
            codes,
            qg,
            scales,
            bias,
            query_group_sums[:, group_index],
        )

    return logits


def _mix_page_chunk_multiquery_torch(
    attn_weights,
    pages: Sequence[PreparedPageTorch],
    *,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not pages:
        raise ValueError("pages must be non-empty")
    header = pages[0].header
    device_type = pages[0].device_type
    page_count = len(pages)
    token_count = header.token_count

    if trace is not None:
        trace.record_page_read(
            sum(page.payload_nbytes for page in pages),
            sum(page.metadata_nbytes for page in pages),
        )

    if not isinstance(attn_weights, torch.Tensor):
        weights = torch.as_tensor(attn_weights, dtype=torch.float32, device=device_type)
    else:
        weights = attn_weights.to(dtype=torch.float32, device=device_type)
    if weights.ndim != 3 or tuple(weights.shape[1:]) != (page_count, token_count):
        raise ValueError("attn_weights chunk must have shape [query_count, page_count, token_count]")

    query_count = int(weights.shape[0])
    output = torch.zeros((query_count, header.padded_head_dim), dtype=torch.float32, device=device_type)

    if header.mode_default == "M3":
        use_native_dtype = _m3_native_compute_enabled(pages)
        dense = _decode_escape_batch_torch(
            pages,
            token_count=header.token_count,
            head_dim=header.head_dim,
            promote_float32=not use_native_dtype,
        )
        if use_native_dtype:
            output[:, : header.head_dim] += torch.einsum(
                "qpt,pth->qh",
                weights.to(dtype=dense.dtype),
                dense,
            ).to(dtype=torch.float32)
        else:
            output[:, : header.head_dim] += torch.einsum("qpt,pth->qh", weights, dense)
        return output

    if header.mode_default in {"M2", "M4"}:
        raise ValueError(f"{header.mode_default} is only supported for key scoring in this phase")

    if header.mode_default == "M1":
        for group_index in range(header.num_groups):
            group_words = torch.stack([page.payload[group_index] for page in pages], dim=0)
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            group = _lookup_lut_group_torch(
                torch.stack([page.codebooks[group_index] for page in pages], dim=0),
                codes,
            )
            start = group_index * header.group_size
            end = start + header.group_size
            output[:, start:end] += torch.einsum("qpt,ptg->qg", weights, group)
        return output

    if header.mode_default == "T3":
        codebooks = pages[0].codebooks if pages[0].codebooks is not None else _turbo3_centroids_torch(device_type=device_type)
        prepared_chunk = _get_prepared_chunk_mps(pages)
        for group_index in range(header.num_groups):
            group_words = (
                prepared_chunk.payload_groups[group_index]
                if prepared_chunk is not None
                else torch.stack([page.payload[group_index] for page in pages], dim=0)
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages[0].unpack_shifts,
                pages[0].unpack_mask,
                header.group_size,
            ).reshape(page_count, token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            scales = (
                prepared_chunk.scales_groups[group_index]
                if prepared_chunk is not None and prepared_chunk.scales_groups is not None
                else torch.stack([page.scales[:, group_index] for page in pages], dim=0)
            )
            rotated_group = _lookup_turbo_group_torch(codebooks, codes) * scales[..., None]
            group = _fwht_last_dim_torch(rotated_group)
            start = group_index * header.group_size
            end = start + header.group_size
            output[:, start:end] += torch.einsum("qpt,ptg->qg", weights, group)
        return output

    prepared_chunk = _get_prepared_chunk_mps(pages)
    if (
        prepared_chunk is not None
        and prepared_chunk.fused_scaled_codes is not None
        and prepared_chunk.bias_groups is not None
        and _supports_fused_two_group64(header)
    ):
        if trace is not None:
            trace.record_per_kv_kernel_variant(section="mix", variant="fused_two_group64")
        output[:, : header.padded_head_dim] += _mix_m0_contribution_two_group64_torch(
            weights,
            prepared_chunk.fused_scaled_codes,
            prepared_chunk.bias_groups,
        )
        return output
    if trace is not None:
        trace.record_per_kv_kernel_variant(section="mix", variant="generic")
    for group_index in range(header.num_groups):
        cached_codes = prepared_chunk is not None and prepared_chunk.codes_groups is not None
        if cached_codes:
            codes = prepared_chunk.codes_groups[group_index]
        else:
            codes = _trace_timed_call(
                trace,
                "unpack",
                device_type=device_type,
                fn=lambda group_index=group_index: _unpack_bits_torch(
                    torch.stack([page.payload[group_index] for page in pages], dim=0).reshape(-1, header.words_per_group),
                    pages[0].unpack_shifts,
                    pages[0].unpack_mask,
                    header.group_size,
                ).reshape(page_count, token_count, header.group_size),
            )
        if trace is not None and not cached_codes:
            trace.record_temporary(int(codes.numel() * codes.element_size()))
        scales = (
            prepared_chunk.scales_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.scales_groups is not None
            else torch.stack([page.scales[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        bias = (
            prepared_chunk.bias_groups[group_index]
            if prepared_chunk is not None and prepared_chunk.bias_groups is not None
            else torch.stack([page.bias[:, group_index].to(torch.float32) for page in pages], dim=0)
        )
        start = group_index * header.group_size
        end = start + header.group_size
        output[:, start:end] += _mix_m0_contribution_torch(weights, codes, scales, bias)

    return output


def _score_page_chunk_grouped_multiquery_torch(
    query_groups,
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    *,
    prepared_query_groups_tensor=None,
    query_group_sums=None,
    compact_grouped_chunk: bool = False,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not pages_by_group or not pages_by_group[0]:
        raise ValueError("pages_by_group must be non-empty")
    batch_size = len(pages_by_group)
    page_count = len(pages_by_group[0])
    header = pages_by_group[0][0].header
    device_type = pages_by_group[0][0].device_type

    for group_pages in pages_by_group:
        if len(group_pages) != page_count:
            raise ValueError("all page groups must have the same page count")

    if torch.is_tensor(query_groups):
        queries = query_groups.to(dtype=torch.float32, device=device_type)
    else:
        queries = torch.stack(
            [
                group.to(dtype=torch.float32, device=device_type)
                if torch.is_tensor(group)
                else torch.as_tensor(group, dtype=torch.float32, device=device_type)
                for group in query_groups
            ],
            dim=0,
        )
    if queries.ndim != 3:
        raise ValueError("query_groups must have shape [batch_size, query_count, head_dim]")
    query_count = int(queries.shape[1])
    if int(queries.shape[0]) != batch_size:
        raise ValueError("query_groups batch size must align with pages_by_group")

    signature_buckets = _signature_buckets_for_page_chunk(pages_by_group)
    if len(signature_buckets) > 1:
        bucket_logits = torch.zeros(
            (batch_size, query_count, page_count * header.token_count),
            dtype=torch.float32,
            device=device_type,
        )
        for group_indices in signature_buckets:
            group_index_list = list(group_indices)
            sub_prepared_queries = (
                None
                if prepared_query_groups_tensor is None
                else prepared_query_groups_tensor[group_index_list]
            )
            sub_query_sums = None if query_group_sums is None else query_group_sums[group_index_list]
            bucket_logits[group_index_list] = _score_page_chunk_grouped_multiquery_torch(
                queries[group_index_list],
                [pages_by_group[group_index] for group_index in group_index_list],
                prepared_query_groups_tensor=sub_prepared_queries,
                query_group_sums=sub_query_sums,
                compact_grouped_chunk=compact_grouped_chunk,
                trace=trace,
            )
        return bucket_logits

    if trace is not None:
        trace.record_page_read(
            sum(page.payload_nbytes for group_pages in pages_by_group for page in group_pages),
            sum(page.metadata_nbytes for group_pages in pages_by_group for page in group_pages),
        )

    if header.mode_default == "M3":
        use_native_dtype = _m3_native_compute_enabled(pages_by_group[0])
        dense = torch.stack(
            [
                _decode_escape_batch_torch(
                    group_pages,
                    token_count=header.token_count,
                    head_dim=header.head_dim,
                    promote_float32=not use_native_dtype,
                )
                for group_pages in pages_by_group
            ],
            dim=0,
        )
        if use_native_dtype:
            return torch.einsum("bpth,bqh->bqpt", dense, queries.to(dtype=dense.dtype)).reshape(
                batch_size, query_count, -1
            ).to(dtype=torch.float32)
        return torch.einsum("bpth,bqh->bqpt", dense, queries).reshape(batch_size, query_count, -1)

    if header.mode_default == "M2":
        query_groups_tensor = prepared_query_groups_tensor
        if query_groups_tensor is None:
            padded_queries = _pad_queries(
                queries.reshape(batch_size * query_count, header.head_dim),
                header.padded_head_dim,
                device_type=device_type,
            ).reshape(batch_size, query_count, header.padded_head_dim)
            query_groups_tensor = padded_queries.reshape(batch_size, query_count, header.num_groups, header.group_size)
        logits = torch.zeros((batch_size, query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        grouped_prepared_chunk = (
            _build_grouped_prepared_chunk_mps(pages_by_group)
            if compact_grouped_chunk
            else _get_grouped_prepared_chunk_mps(pages_by_group)
        )
        if (
            grouped_prepared_chunk is not None
            and grouped_prepared_chunk.m2_sketch_tensor is not None
            and grouped_prepared_chunk.m2_basis_tensor is not None
            and grouped_prepared_chunk.m2_mean_tensor is not None
        ):
            qg, sketch_tensor, basis_tensor, mean_tensor = _coerce_m2_grouped_operands(
                query_groups_tensor,
                grouped_prepared_chunk.m2_sketch_tensor,
                grouped_prepared_chunk.m2_basis_tensor,
                grouped_prepared_chunk.m2_mean_tensor,
            )
            if int(basis_tensor.shape[3]) == 1:
                squeezed_basis = basis_tensor.squeeze(3)
                squeezed_mean = mean_tensor.squeeze(3)
                q_proj = torch.einsum("bpgrc,bqgc->bqpgr", squeezed_basis, qg)
                logits += torch.einsum("bptgr,bqpgr->bqpt", sketch_tensor, q_proj)
                logits += torch.einsum("bpgc,bqgc->bqp", squeezed_mean, qg)[:, :, :, None]
                return logits.reshape(batch_size, query_count, -1)
            segment_ids = (
                grouped_prepared_chunk.m2_segment_ids
                if grouped_prepared_chunk.m2_segment_ids is not None
                else _segment_ids_tensor(header.token_count, int(basis_tensor.shape[3]), device_type=device_type)
            )
            q_proj = torch.einsum("bpgsrc,bqgc->bqpgsr", basis_tensor, qg)
            logits += torch.einsum("bptgr,bqpgtr->bqpt", sketch_tensor, q_proj[:, :, :, :, segment_ids, :])
            logits += torch.einsum("bpgtc,bqgc->bqpt", mean_tensor[:, :, :, segment_ids, :], qg)
            return logits.reshape(batch_size, query_count, -1)
        prepared_chunks = None if grouped_prepared_chunk is not None else [
            (
                _build_prepared_chunk_mps(group_pages)
                if compact_grouped_chunk
                else _get_prepared_chunk_mps(group_pages) or (
                    _build_prepared_chunk_mps(group_pages) if _prepared_chunk_cache_key(group_pages) is not None else None
                )
            )
            for group_pages in pages_by_group
        ]
        for group_index in range(header.num_groups):
            if grouped_prepared_chunk is not None and grouped_prepared_chunk.m2_sketch_groups is not None:
                group_sketch = grouped_prepared_chunk.m2_sketch_groups[group_index]
                group_basis = grouped_prepared_chunk.m2_basis_groups[group_index]
                group_mean = grouped_prepared_chunk.m2_mean_groups[group_index]
            elif prepared_chunks is not None and all(
                chunk is not None and chunk.m2_sketch_groups is not None for chunk in prepared_chunks
            ):
                group_sketch = torch.stack([chunk.m2_sketch_groups[group_index] for chunk in prepared_chunks], dim=0)
                group_basis = torch.stack([chunk.m2_basis_groups[group_index] for chunk in prepared_chunks], dim=0)
                group_mean = torch.stack([chunk.m2_mean_groups[group_index] for chunk in prepared_chunks], dim=0)
            else:
                group_sketch = torch.stack(
                    [torch.stack([page.m2_sketch[:, group_index, :] for page in group_pages], dim=0) for group_pages in pages_by_group],
                    dim=0,
                )
                group_basis = torch.stack(
                    [torch.stack([page.m2_basis[group_index] for page in group_pages], dim=0) for group_pages in pages_by_group],
                    dim=0,
                )
                group_mean = torch.stack(
                    [torch.stack([page.m2_mean[group_index] for page in group_pages], dim=0) for group_pages in pages_by_group],
                    dim=0,
                )
            qg, group_sketch, group_basis, group_mean = _coerce_m2_operands(
                query_groups_tensor[:, :, group_index, :],
                group_sketch,
                group_basis,
                group_mean,
            )
            if group_basis.dim() == 4:
                q_proj = torch.einsum("bprg,bqg->bqpr", group_basis, qg)
                logits += torch.einsum("bptd,bqpd->bqpt", group_sketch, q_proj)
                logits += torch.einsum("bpg,bqg->bqp", group_mean, qg)[:, :, :, None]
                continue
            segment_ids = (
                grouped_prepared_chunk.m2_segment_ids
                if grouped_prepared_chunk is not None and grouped_prepared_chunk.m2_segment_ids is not None
                else (
                    prepared_chunks[0].m2_segment_ids
                    if prepared_chunks is not None and prepared_chunks[0] is not None and prepared_chunks[0].m2_segment_ids is not None
                    else _segment_ids_tensor(header.token_count, int(group_basis.shape[2]), device_type=device_type)
                )
            )
            q_proj = torch.einsum("bpsrg,bqg->bqpsr", group_basis, qg)
            logits += torch.einsum("bptr,bqptr->bqpt", group_sketch, q_proj[:, :, :, segment_ids, :])
            logits += torch.einsum("bptg,bqg->bqpt", group_mean[:, :, segment_ids, :], qg)
        return logits.reshape(batch_size, query_count, -1)

    if header.mode_default == "M4":
        query_groups_tensor = prepared_query_groups_tensor
        if query_groups_tensor is None:
            padded_queries = _pad_queries(
                queries.reshape(batch_size * query_count, header.head_dim),
                header.padded_head_dim,
                device_type=device_type,
            ).reshape(batch_size, query_count, header.padded_head_dim)
            query_groups_tensor = padded_queries.reshape(batch_size, query_count, header.num_groups, header.group_size)
        logits = torch.zeros((batch_size, query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        grouped_prepared_chunk = (
            _build_grouped_prepared_chunk_mps(pages_by_group)
            if compact_grouped_chunk
            else _get_grouped_prepared_chunk_mps(pages_by_group)
        )
        if (
            grouped_prepared_chunk is not None
            and grouped_prepared_chunk.m2_sketch_tensor is not None
            and grouped_prepared_chunk.m2_mean_tensor is not None
        ):
            if grouped_prepared_chunk.m2_basis_tensor is not None:
                qg, sketch_tensor, basis_tensor, mean_tensor = _coerce_m2_grouped_operands(
                    query_groups_tensor,
                    grouped_prepared_chunk.m2_sketch_tensor,
                    grouped_prepared_chunk.m2_basis_tensor,
                    grouped_prepared_chunk.m2_mean_tensor,
                )
                if int(basis_tensor.dim()) == 4:
                    q_proj = torch.einsum("bgrc,bqgc->bqgr", basis_tensor, qg)
                    logits += torch.einsum("bptgr,bqgr->bqpt", sketch_tensor, q_proj)
                else:
                    q_proj = torch.einsum("bpgrc,bqgc->bqpgr", basis_tensor, qg)
                    logits += torch.einsum("bptgr,bqpgr->bqpt", sketch_tensor, q_proj)
            else:
                basis = _m4_basis_torch(
                    header.group_size,
                    int(pages_by_group[0][0].m2_sketch.shape[-1]),
                    basis_family=header.project_basis,
                    device_type=device_type,
                )
                qg, sketch_tensor, _, mean_tensor = _coerce_m2_grouped_operands(
                    query_groups_tensor,
                    grouped_prepared_chunk.m2_sketch_tensor,
                    basis.view(1, 1, 1, *basis.shape),
                    grouped_prepared_chunk.m2_mean_tensor,
                )
                basis_work = basis.to(dtype=qg.dtype)
                q_proj = torch.einsum("rc,bqgc->bqgr", basis_work, qg)
                logits += torch.einsum("bptgr,bqgr->bqpt", sketch_tensor, q_proj)
            logits += torch.einsum("bpgc,bqgc->bqp", mean_tensor, qg)[:, :, :, None]
            return logits.reshape(batch_size, query_count, -1)
        prepared_chunks = None if grouped_prepared_chunk is not None else [
            (
                _build_prepared_chunk_mps(group_pages)
                if compact_grouped_chunk
                else _get_prepared_chunk_mps(group_pages) or (
                    _build_prepared_chunk_mps(group_pages) if _prepared_chunk_cache_key(group_pages) is not None else None
                )
            )
            for group_pages in pages_by_group
        ]
        for group_index in range(header.num_groups):
            if grouped_prepared_chunk is not None and grouped_prepared_chunk.m2_sketch_groups is not None:
                group_sketch = grouped_prepared_chunk.m2_sketch_groups[group_index]
                group_basis = grouped_prepared_chunk.m2_basis_groups[group_index]
                group_mean = grouped_prepared_chunk.m2_mean_groups[group_index]
            elif prepared_chunks is not None and all(
                chunk is not None and chunk.m2_sketch_groups is not None for chunk in prepared_chunks
            ):
                group_sketch = torch.stack([chunk.m2_sketch_groups[group_index] for chunk in prepared_chunks], dim=0)
                group_basis = (
                    torch.stack([chunk.m2_basis_groups[group_index] for chunk in prepared_chunks], dim=0)
                    if prepared_chunks[0].m2_basis_groups is not None
                    else None
                )
                group_mean = torch.stack([chunk.m2_mean_groups[group_index] for chunk in prepared_chunks], dim=0)
            else:
                group_sketch = torch.stack(
                    [torch.stack([page.m2_sketch[:, group_index, :] for page in group_pages], dim=0) for group_pages in pages_by_group],
                    dim=0,
                )
                group_basis = (
                    torch.stack(
                        [torch.stack([page.m2_basis[group_index] for page in group_pages], dim=0) for group_pages in pages_by_group],
                        dim=0,
                    )
                    if pages_by_group[0][0].m2_basis is not None
                    else None
                )
                group_mean = torch.stack(
                    [torch.stack([page.m2_mean[group_index] for page in group_pages], dim=0) for group_pages in pages_by_group],
                    dim=0,
                )
            qg = query_groups_tensor[:, :, group_index, :]
            if group_basis is not None:
                qg, group_sketch, group_basis, group_mean = _coerce_m2_operands(qg, group_sketch, group_basis, group_mean)
                if int(group_basis.dim()) == 2:
                    q_proj = torch.einsum("rg,bqg->bqr", group_basis, qg)
                    logits += torch.einsum("bptr,bqr->bqpt", group_sketch, q_proj)
                elif int(group_basis.dim()) == 3:
                    q_proj = torch.einsum("brg,bqg->bqr", group_basis, qg)
                    logits += torch.einsum("bptr,bqr->bqpt", group_sketch, q_proj)
                else:
                    q_proj = torch.einsum("bprg,bqg->bqpr", group_basis, qg)
                    logits += torch.einsum("bptr,bqpr->bqpt", group_sketch, q_proj)
            else:
                basis = _m4_basis_torch(
                    header.group_size,
                    int(pages_by_group[0][0].m2_sketch.shape[-1]),
                    basis_family=header.project_basis,
                    device_type=device_type,
                )
                qg, group_sketch, _, group_mean = _coerce_m2_operands(qg, group_sketch, basis, group_mean)
                basis_work = basis.to(dtype=qg.dtype)
                q_proj = torch.einsum("rc,bqc->bqr", basis_work, qg)
                logits += torch.einsum("bptr,bqr->bqpt", group_sketch, q_proj)
            logits += torch.einsum("bpg,bqg->bqp", group_mean, qg)[:, :, :, None]
        return logits.reshape(batch_size, query_count, -1)

    if header.mode_default == "M1":
        query_groups_tensor = prepared_query_groups_tensor
        if query_groups_tensor is None:
            padded_queries = _pad_queries(
                queries.reshape(batch_size * query_count, header.head_dim),
                header.padded_head_dim,
                device_type=device_type,
            ).reshape(batch_size, query_count, header.padded_head_dim)
            query_groups_tensor = padded_queries.reshape(batch_size, query_count, header.num_groups, header.group_size)
        logits = torch.zeros((batch_size, query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        for group_index in range(header.num_groups):
            group_words = torch.stack(
                [torch.stack([page.payload[group_index] for page in group_pages], dim=0) for group_pages in pages_by_group],
                dim=0,
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages_by_group[0][0].unpack_shifts,
                pages_by_group[0][0].unpack_mask,
                header.group_size,
                trace=trace,
            ).reshape(batch_size, page_count, header.token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            group = _lookup_lut_group_torch(
                torch.stack(
                    [
                        torch.stack([page.codebooks[group_index] for page in group_pages], dim=0)
                        for group_pages in pages_by_group
                    ],
                    dim=0,
                ),
                codes,
            )
            logits += torch.einsum("bptg,bqg->bqpt", group, query_groups_tensor[:, :, group_index, :])
        return logits.reshape(batch_size, query_count, -1)

    if header.mode_default == "T3":
        query_groups_tensor = prepared_query_groups_tensor
        if query_groups_tensor is None:
            padded_queries = _pad_queries(
                queries.reshape(batch_size * query_count, header.head_dim),
                header.padded_head_dim,
                device_type=device_type,
            ).reshape(batch_size, query_count, header.padded_head_dim)
            query_groups_tensor = padded_queries.reshape(batch_size, query_count, header.num_groups, header.group_size)
        rotated_query_groups = _fwht_last_dim_torch(
            query_groups_tensor,
            trace=trace,
        )
        logits = torch.zeros((batch_size, query_count, page_count, header.token_count), dtype=torch.float32, device=device_type)
        codebooks = pages_by_group[0][0].codebooks if pages_by_group[0][0].codebooks is not None else _turbo3_centroids_torch(device_type=device_type)
        prepared_chunks = [_get_prepared_chunk_mps(group_pages) for group_pages in pages_by_group]
        for group_index in range(header.num_groups):
            group_words = torch.stack(
                [
                    prepared_chunks[group_id].payload_groups[group_index]
                    if prepared_chunks[group_id] is not None
                    else torch.stack([page.payload[group_index] for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages_by_group[0][0].unpack_shifts,
                pages_by_group[0][0].unpack_mask,
                header.group_size,
                trace=trace,
            ).reshape(batch_size, page_count, header.token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            scales = torch.stack(
                [
                    prepared_chunks[group_id].scales_groups[group_index]
                    if prepared_chunks[group_id] is not None and prepared_chunks[group_id].scales_groups is not None
                    else torch.stack([page.scales[:, group_index] for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
            corrected = _lookup_turbo_group_torch(codebooks, codes) * scales[..., None]
            logits += torch.einsum("bptg,bqg->bqpt", corrected, rotated_query_groups[:, :, group_index, :])
        return logits.reshape(batch_size, query_count, -1)

    query_groups_tensor = prepared_query_groups_tensor
    if query_groups_tensor is None:
        padded_queries = _pad_queries(
            queries.reshape(batch_size * query_count, header.head_dim),
            header.padded_head_dim,
            device_type=device_type,
        ).reshape(batch_size, query_count, header.padded_head_dim)
        query_groups_tensor = padded_queries.reshape(batch_size, query_count, header.num_groups, header.group_size)
    query_group_sums_tensor = query_group_sums if query_group_sums is not None else query_groups_tensor.sum(dim=-1)
    logits = torch.zeros((batch_size, query_count, page_count * header.token_count), dtype=torch.float32, device=device_type)
    grouped_prepared_chunk = (
        _build_grouped_prepared_chunk_mps(pages_by_group)
        if compact_grouped_chunk
        else _get_grouped_prepared_chunk_mps(pages_by_group)
    )
    prepared_chunks = None if grouped_prepared_chunk is not None else [
        (
            _build_prepared_chunk_mps(group_pages)
            if compact_grouped_chunk
            else _get_prepared_chunk_mps(group_pages) or (
                _build_prepared_chunk_mps(group_pages) if _prepared_chunk_cache_key(group_pages) is not None else None
            )
        )
        for group_pages in pages_by_group
    ]
    if (
        grouped_prepared_chunk is not None
        and grouped_prepared_chunk.payload_groups
        and grouped_prepared_chunk.scales_groups is not None
        and grouped_prepared_chunk.bias_groups is not None
        and grouped_prepared_chunk.payload_groups_tensor is not None
        and grouped_prepared_chunk.scales_groups_tensor is not None
        and grouped_prepared_chunk.bias_groups_tensor is not None
        and _supports_packed_four_group128_cuda(header, device_type=device_type)
    ):
        unpack_shifts = pages_by_group[0][0].unpack_shifts
        unpack_mask = pages_by_group[0][0].unpack_mask
        if unpack_shifts is None or unpack_mask is None:
            raise ValueError("packed grouped CUDA path requires unpack metadata")
        if trace is not None:
            trace.record_grouped_kernel_variant(section="score", variant="packed_cuda")
        for group_start in range(0, header.num_groups, 2):
            group_end = min(header.num_groups, group_start + 2)
            logits += _score_m0_logits_packed32_grouped_torch(
                grouped_prepared_chunk.payload_groups_tensor[:, group_start:group_end],
                query_groups_tensor[:, :, group_start:group_end, :],
                grouped_prepared_chunk.scales_groups_tensor[:, group_start:group_end],
                grouped_prepared_chunk.bias_groups_tensor[:, group_start:group_end],
                query_group_sums_tensor[:, :, group_start:group_end],
                unpack_shifts=unpack_shifts,
                unpack_mask=unpack_mask,
                trace=trace,
            )
        return logits
    if (
        grouped_prepared_chunk is not None
        and grouped_prepared_chunk.fused_scaled_codes is not None
        and grouped_prepared_chunk.bias_groups is not None
        and _supports_fused_two_group64(header)
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="score", variant="fused_two_group64")
        fused_queries = query_groups_tensor.reshape(batch_size, query_count, header.padded_head_dim).contiguous()
        return _score_m0_logits_two_group64_torch(
            grouped_prepared_chunk.fused_scaled_codes,
            fused_queries,
            grouped_prepared_chunk.bias_groups,
            query_group_sums_tensor,
        )
    if (
        grouped_prepared_chunk is not None
        and grouped_prepared_chunk.fused_scaled_codes is not None
        and grouped_prepared_chunk.bias_groups is not None
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="score", variant="fused_generic")
        fused_queries = query_groups_tensor.reshape(batch_size, query_count, header.padded_head_dim).contiguous()
        return _score_m0_logits_fused_torch(
            grouped_prepared_chunk.fused_scaled_codes,
            fused_queries,
            grouped_prepared_chunk.bias_groups,
            query_group_sums_tensor,
        )
    if (
        grouped_prepared_chunk is None
        and prepared_chunks is not None
        and _supports_fused_two_group64(header)
        and all(chunk is not None and chunk.fused_scaled_codes is not None and chunk.bias_groups is not None for chunk in prepared_chunks)
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="score", variant="fused_two_group64")
        fused_scaled_codes, bias_groups = _assemble_grouped_fused_two_group64_components(
            prepared_chunks,
            trace=trace,
            device_type=device_type,
        )
        fused_queries = query_groups_tensor.reshape(batch_size, query_count, header.padded_head_dim).contiguous()
        return _score_m0_logits_two_group64_torch(
            fused_scaled_codes,
            fused_queries,
            bias_groups,
            query_group_sums_tensor,
        )
    if (
        grouped_prepared_chunk is None
        and prepared_chunks is not None
        and all(chunk is not None and chunk.fused_scaled_codes is not None and chunk.bias_groups is not None for chunk in prepared_chunks)
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="score", variant="fused_generic")
        fused_scaled_codes = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda: _load_torch().stack([chunk.fused_scaled_codes for chunk in prepared_chunks], dim=0),
        )
        bias_groups = tuple(
            _trace_timed_call(
                trace,
                "chunk_assembly",
                device_type=device_type,
                fn=lambda group_index=group_index: _load_torch().stack([chunk.bias_groups[group_index] for chunk in prepared_chunks], dim=0),
            )
            for group_index in range(header.num_groups)
        )
        if trace is not None:
            trace.record_temporary(int(fused_scaled_codes.numel() * fused_scaled_codes.element_size()))
            trace.record_temporary(sum(int(tensor.numel() * tensor.element_size()) for tensor in bias_groups))
        fused_queries = query_groups_tensor.reshape(batch_size, query_count, header.padded_head_dim).contiguous()
        return _score_m0_logits_fused_torch(
            fused_scaled_codes,
            fused_queries,
            bias_groups,
            query_group_sums_tensor,
        )
    if trace is not None:
        trace.record_grouped_kernel_variant(section="score", variant="generic")
    for group_index in range(header.num_groups):
        cached_codes = grouped_prepared_chunk is not None and grouped_prepared_chunk.codes_groups is not None
        if cached_codes:
            codes = grouped_prepared_chunk.codes_groups[group_index]
        else:
            def _build_codes(group_index: int = group_index):
                return torch.stack(
                    [
                        prepared_chunks[group_id].codes_groups[group_index]
                        if prepared_chunks is not None and prepared_chunks[group_id] is not None and prepared_chunks[group_id].codes_groups is not None
                        else _trace_timed_call(
                            trace,
                            "unpack",
                            device_type=device_type,
                            fn=lambda group_pages=group_pages, group_index=group_index: _unpack_bits_torch(
                                torch.stack([page.payload[group_index] for page in group_pages], dim=0).reshape(-1, header.words_per_group),
                                pages_by_group[0][0].unpack_shifts,
                                pages_by_group[0][0].unpack_mask,
                                header.group_size,
                            ).reshape(page_count, header.token_count, header.group_size),
                        )
                        for group_id, group_pages in enumerate(pages_by_group)
                    ],
                    dim=0,
                )

            codes = _trace_timed_call(
                trace,
                "chunk_assembly",
                device_type=device_type,
                fn=_build_codes,
            )
        if trace is not None and not cached_codes:
            trace.record_temporary(int(codes.numel() * codes.element_size()))
        qg = query_groups_tensor[:, :, group_index, :]
        scales = (
            grouped_prepared_chunk.scales_groups[group_index]
            if grouped_prepared_chunk is not None and grouped_prepared_chunk.scales_groups is not None
            else torch.stack(
                [
                    prepared_chunks[group_id].scales_groups[group_index]
                    if prepared_chunks is not None and prepared_chunks[group_id] is not None and prepared_chunks[group_id].scales_groups is not None
                    else torch.stack([page.scales[:, group_index].to(torch.float32) for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
        )
        bias = (
            grouped_prepared_chunk.bias_groups[group_index]
            if grouped_prepared_chunk is not None and grouped_prepared_chunk.bias_groups is not None
            else torch.stack(
                [
                    prepared_chunks[group_id].bias_groups[group_index]
                    if prepared_chunks is not None and prepared_chunks[group_id] is not None and prepared_chunks[group_id].bias_groups is not None
                    else torch.stack([page.bias[:, group_index].to(torch.float32) for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
        )
        logits += _score_m0_logits_flat_torch(
            codes,
            qg,
            scales,
            bias,
            query_group_sums_tensor[:, :, group_index],
        )

    return logits


def _mix_page_chunk_grouped_multiquery_torch(
    attn_weights,
    pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    *,
    out_acc=None,
    compact_grouped_chunk: bool = False,
    disable_packed_grouped_cuda: bool = False,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not pages_by_group or not pages_by_group[0]:
        raise ValueError("pages_by_group must be non-empty")
    batch_size = len(pages_by_group)
    page_count = len(pages_by_group[0])
    header = pages_by_group[0][0].header
    device_type = pages_by_group[0][0].device_type
    token_count = header.token_count

    weights = attn_weights if isinstance(attn_weights, torch.Tensor) else torch.as_tensor(attn_weights, dtype=torch.float32, device=device_type)
    weights = weights.to(dtype=torch.float32, device=device_type)
    if weights.ndim != 4 or tuple(weights.shape[2:]) != (page_count, token_count):
        raise ValueError("grouped attn_weights chunk must have shape [batch_size, query_count, page_count, token_count]")

    query_count = int(weights.shape[1])
    output = _prepare_grouped_output_accumulator_tensor(
        out_acc,
        batch_size,
        query_count,
        header.head_dim,
        header.padded_head_dim,
        device_type=device_type,
    )

    signature_buckets = _signature_buckets_for_page_chunk(pages_by_group)
    if len(signature_buckets) > 1:
        for group_indices in signature_buckets:
            group_index_list = list(group_indices)
            output[group_index_list] = _mix_page_chunk_grouped_multiquery_torch(
                weights[group_index_list],
                [pages_by_group[group_index] for group_index in group_index_list],
                out_acc=output[group_index_list],
                compact_grouped_chunk=compact_grouped_chunk,
                disable_packed_grouped_cuda=disable_packed_grouped_cuda,
                trace=trace,
            )
        return output

    if trace is not None:
        trace.record_page_read(
            sum(page.payload_nbytes for group_pages in pages_by_group for page in group_pages),
            sum(page.metadata_nbytes for group_pages in pages_by_group for page in group_pages),
        )

    if header.mode_default == "M3":
        use_native_dtype = _m3_native_compute_enabled(pages_by_group[0])
        dense = torch.stack(
            [
                _decode_escape_batch_torch(
                    group_pages,
                    token_count=header.token_count,
                    head_dim=header.head_dim,
                    promote_float32=not use_native_dtype,
                )
                for group_pages in pages_by_group
            ],
            dim=0,
        )
        if use_native_dtype:
            output[:, :, : header.head_dim] += torch.einsum(
                "bqpt,bpth->bqh",
                weights.to(dtype=dense.dtype),
                dense,
            ).to(dtype=torch.float32)
        else:
            output[:, :, : header.head_dim] += torch.einsum("bqpt,bpth->bqh", weights, dense)
        return output

    if header.mode_default in {"M2", "M4"}:
        raise ValueError(f"{header.mode_default} is only supported for key scoring in this phase")

    if header.mode_default == "M1":
        for group_index in range(header.num_groups):
            group_words = torch.stack(
                [torch.stack([page.payload[group_index] for page in group_pages], dim=0) for group_pages in pages_by_group],
                dim=0,
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages_by_group[0][0].unpack_shifts,
                pages_by_group[0][0].unpack_mask,
                header.group_size,
                trace=trace,
            ).reshape(batch_size, page_count, token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            group = _lookup_lut_group_torch(
                torch.stack(
                    [
                        torch.stack([page.codebooks[group_index] for page in group_pages], dim=0)
                        for group_pages in pages_by_group
                    ],
                    dim=0,
                ),
                codes,
            )
            start = group_index * header.group_size
            end = start + header.group_size
            output[:, :, start:end] += torch.einsum("bqpt,bptg->bqg", weights, group)
        return output

    if header.mode_default == "T3":
        codebooks = pages_by_group[0][0].codebooks if pages_by_group[0][0].codebooks is not None else _turbo3_centroids_torch(device_type=device_type)
        prepared_chunks = [_get_prepared_chunk_mps(group_pages) for group_pages in pages_by_group]
        for group_index in range(header.num_groups):
            group_words = torch.stack(
                [
                    prepared_chunks[group_id].payload_groups[group_index]
                    if prepared_chunks[group_id] is not None
                    else torch.stack([page.payload[group_index] for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
            codes = _unpack_bits_torch(
                group_words.reshape(-1, header.words_per_group),
                pages_by_group[0][0].unpack_shifts,
                pages_by_group[0][0].unpack_mask,
                header.group_size,
                trace=trace,
            ).reshape(batch_size, page_count, token_count, header.group_size)
            if trace is not None:
                trace.record_temporary(int(codes.numel() * codes.element_size()))
            scales = torch.stack(
                [
                    prepared_chunks[group_id].scales_groups[group_index]
                    if prepared_chunks[group_id] is not None and prepared_chunks[group_id].scales_groups is not None
                    else torch.stack([page.scales[:, group_index] for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
            rotated_group = _lookup_turbo_group_torch(codebooks, codes) * scales[..., None]
            group = _fwht_last_dim_torch(rotated_group, trace=trace)
            start = group_index * header.group_size
            end = start + header.group_size
            output[:, :, start:end] += torch.einsum("bqpt,bptg->bqg", weights, group)
        return output

    grouped_prepared_chunk = (
        _build_grouped_prepared_chunk_mps(pages_by_group)
        if compact_grouped_chunk
        else _get_grouped_prepared_chunk_mps(pages_by_group)
    )
    prepared_chunks = None if grouped_prepared_chunk is not None else [
        (
            _build_prepared_chunk_mps(group_pages)
            if compact_grouped_chunk
            else _get_prepared_chunk_mps(group_pages) or (
                _build_prepared_chunk_mps(group_pages) if _prepared_chunk_cache_key(group_pages) is not None else None
            )
        )
        for group_pages in pages_by_group
    ]
    if (
        grouped_prepared_chunk is not None
        and grouped_prepared_chunk.payload_groups
        and grouped_prepared_chunk.scales_groups is not None
        and grouped_prepared_chunk.bias_groups is not None
        and grouped_prepared_chunk.payload_groups_tensor is not None
        and grouped_prepared_chunk.scales_groups_tensor is not None
        and grouped_prepared_chunk.bias_groups_tensor is not None
        and _supports_packed_four_group128_cuda(header, device_type=device_type)
        and not disable_packed_grouped_cuda
    ):
        unpack_shifts = pages_by_group[0][0].unpack_shifts
        unpack_mask = pages_by_group[0][0].unpack_mask
        if unpack_shifts is None or unpack_mask is None:
            raise ValueError("packed grouped CUDA path requires unpack metadata")
        if trace is not None:
            trace.record_grouped_kernel_variant(section="mix", variant="packed_cuda")
        for group_start in range(0, header.num_groups, 2):
            group_end = min(header.num_groups, group_start + 2)
            start = group_start * header.group_size
            end = group_end * header.group_size
            output[:, :, start:end] += _mix_m0_contribution_packed32_grouped_torch(
                weights,
                grouped_prepared_chunk.payload_groups_tensor[:, group_start:group_end],
                grouped_prepared_chunk.scales_groups_tensor[:, group_start:group_end],
                grouped_prepared_chunk.bias_groups_tensor[:, group_start:group_end],
                unpack_shifts=unpack_shifts,
                unpack_mask=unpack_mask,
                trace=trace,
            ).reshape(batch_size, query_count, end - start)
        return output
    if (
        grouped_prepared_chunk is not None
        and grouped_prepared_chunk.fused_scaled_codes is not None
        and grouped_prepared_chunk.bias_groups is not None
        and _supports_fused_two_group64(header)
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="mix", variant="fused_two_group64")
        output[:, :, : header.padded_head_dim] += _mix_m0_contribution_two_group64_torch(
            weights,
            grouped_prepared_chunk.fused_scaled_codes,
            grouped_prepared_chunk.bias_groups,
        )
        return output
    if (
        grouped_prepared_chunk is not None
        and grouped_prepared_chunk.fused_scaled_codes is not None
        and grouped_prepared_chunk.bias_groups is not None
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="mix", variant="fused_generic")
        output[:, :, : header.padded_head_dim] += _mix_m0_contribution_fused_torch(
            weights,
            grouped_prepared_chunk.fused_scaled_codes,
            grouped_prepared_chunk.bias_groups,
            group_size=header.group_size,
        )
        return output
    if (
        grouped_prepared_chunk is None
        and prepared_chunks is not None
        and _supports_fused_two_group64(header)
        and all(chunk is not None and chunk.fused_scaled_codes is not None and chunk.bias_groups is not None for chunk in prepared_chunks)
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="mix", variant="fused_two_group64")
        fused_scaled_codes, bias_groups = _assemble_grouped_fused_two_group64_components(
            prepared_chunks,
            trace=trace,
            device_type=device_type,
        )
        output[:, :, : header.padded_head_dim] += _mix_m0_contribution_two_group64_torch(
            weights,
            fused_scaled_codes,
            bias_groups,
        )
        return output
    if (
        grouped_prepared_chunk is None
        and prepared_chunks is not None
        and all(chunk is not None and chunk.fused_scaled_codes is not None and chunk.bias_groups is not None for chunk in prepared_chunks)
    ):
        if trace is not None:
            trace.record_grouped_kernel_variant(section="mix", variant="fused_generic")
        fused_scaled_codes = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda: _load_torch().stack([chunk.fused_scaled_codes for chunk in prepared_chunks], dim=0),
        )
        bias_groups = tuple(
            _trace_timed_call(
                trace,
                "chunk_assembly",
                device_type=device_type,
                fn=lambda group_index=group_index: _load_torch().stack([chunk.bias_groups[group_index] for chunk in prepared_chunks], dim=0),
            )
            for group_index in range(header.num_groups)
        )
        if trace is not None:
            trace.record_temporary(int(fused_scaled_codes.numel() * fused_scaled_codes.element_size()))
            trace.record_temporary(sum(int(tensor.numel() * tensor.element_size()) for tensor in bias_groups))
        output[:, :, : header.padded_head_dim] += _mix_m0_contribution_fused_torch(
            weights,
            fused_scaled_codes,
            bias_groups,
            group_size=header.group_size,
        )
        return output
    if trace is not None:
        trace.record_grouped_kernel_variant(section="mix", variant="generic")
    for group_index in range(header.num_groups):
        cached_codes = grouped_prepared_chunk is not None and grouped_prepared_chunk.codes_groups is not None
        if cached_codes:
            codes = grouped_prepared_chunk.codes_groups[group_index]
        else:
            def _build_codes(group_index: int = group_index):
                return torch.stack(
                    [
                        prepared_chunks[group_id].codes_groups[group_index]
                        if prepared_chunks is not None and prepared_chunks[group_id] is not None and prepared_chunks[group_id].codes_groups is not None
                        else _trace_timed_call(
                            trace,
                            "unpack",
                            device_type=device_type,
                            fn=lambda group_pages=group_pages, group_index=group_index: _unpack_bits_torch(
                                torch.stack([page.payload[group_index] for page in group_pages], dim=0).reshape(-1, header.words_per_group),
                                pages_by_group[0][0].unpack_shifts,
                                pages_by_group[0][0].unpack_mask,
                                header.group_size,
                            ).reshape(page_count, token_count, header.group_size),
                        )
                        for group_id, group_pages in enumerate(pages_by_group)
                    ],
                    dim=0,
                )

            codes = _trace_timed_call(
                trace,
                "chunk_assembly",
                device_type=device_type,
                fn=_build_codes,
            )
        if trace is not None and not cached_codes:
            trace.record_temporary(int(codes.numel() * codes.element_size()))
        scales = (
            grouped_prepared_chunk.scales_groups[group_index]
            if grouped_prepared_chunk is not None and grouped_prepared_chunk.scales_groups is not None
            else torch.stack(
                [
                    prepared_chunks[group_id].scales_groups[group_index]
                    if prepared_chunks is not None and prepared_chunks[group_id] is not None and prepared_chunks[group_id].scales_groups is not None
                    else torch.stack([page.scales[:, group_index].to(torch.float32) for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
        )
        bias = (
            grouped_prepared_chunk.bias_groups[group_index]
            if grouped_prepared_chunk is not None and grouped_prepared_chunk.bias_groups is not None
            else torch.stack(
                [
                    prepared_chunks[group_id].bias_groups[group_index]
                    if prepared_chunks is not None and prepared_chunks[group_id] is not None and prepared_chunks[group_id].bias_groups is not None
                    else torch.stack([page.bias[:, group_index].to(torch.float32) for page in group_pages], dim=0)
                    for group_id, group_pages in enumerate(pages_by_group)
                ],
                dim=0,
            )
        )
        start = group_index * header.group_size
        end = start + header.group_size
        output[:, :, start:end] += _mix_m0_contribution_torch(weights, codes, scales, bias)

    return output


def decode_multi_query_step_torch_tensor(
    query_slices: np.ndarray | Any,
    key_pages: Sequence[EncodedPage | PreparedPageTorch],
    value_pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    prepared_key_pages = prepare_pages_torch(key_pages, device_type=device_type, trace=trace)
    prepared_value_pages = prepare_pages_torch(value_pages, device_type=device_type, trace=trace)
    if not prepared_key_pages:
        raise ValueError(f"decode_multi_query_step_{device_type} requires at least one page")
    if trace is not None:
        trace.record_per_kv_decode_call()

    logits_parts = []
    if torch.is_tensor(query_slices):
        query_count = int(query_slices.shape[0])
    else:
        query_count = int(np.asarray(query_slices).shape[0])
    for page_chunk in _chunk_compatible_pages(prepared_key_pages):
        if trace is not None:
            trace.record_per_kv_score_chunk(
                query_count=query_count,
                page_count=len(page_chunk),
                token_count=page_chunk[0].header.token_count,
            )
        logits_parts.append(
            _trace_timed_call(
                trace,
                "score",
                device_type=device_type,
                fn=lambda page_chunk=page_chunk: _score_page_chunk_multiquery_torch(query_slices, page_chunk, trace=trace),
            )
        )
    logits = torch.cat(logits_parts, dim=1)
    weights = _trace_timed_call(
        trace,
        "softmax",
        device_type=device_type,
        fn=lambda: torch.softmax(logits, dim=1),
    )
    output = torch.zeros(
        (query_count, prepared_value_pages[0].header.padded_head_dim),
        dtype=torch.float32,
        device=device_type,
    )
    offset = 0
    for page_chunk in _chunk_compatible_pages(prepared_value_pages):
        chunk_token_count = page_chunk[0].header.token_count * len(page_chunk)
        chunk_weights = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda page_chunk=page_chunk, offset=offset, chunk_token_count=chunk_token_count: weights[
                :,
                offset : offset + chunk_token_count,
            ].reshape(
                weights.shape[0],
                len(page_chunk),
                page_chunk[0].header.token_count,
            ),
        )
        if trace is not None:
            trace.record_per_kv_mix_chunk(
                query_count=query_count,
                page_count=len(page_chunk),
                token_count=page_chunk[0].header.token_count,
                head_dim=prepared_value_pages[0].header.padded_head_dim,
            )
        output += _trace_timed_call(
            trace,
            "mix",
            device_type=device_type,
            fn=lambda chunk_weights=chunk_weights, page_chunk=page_chunk: _mix_page_chunk_multiquery_torch(
                chunk_weights,
                page_chunk,
                trace=trace,
            ),
        )
        offset += chunk_token_count

    head_dim = prepared_value_pages[0].header.head_dim
    return logits, weights, output[:, :head_dim]


def decode_grouped_multiquery_step_torch_tensor(
    query_groups,
    key_pages_by_group: Sequence[Sequence[EncodedPage | PreparedPageTorch]],
    value_pages_by_group: Sequence[Sequence[EncodedPage | PreparedPageTorch]],
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
):
    if not key_pages_by_group or not value_pages_by_group:
        raise ValueError("grouped decode requires non-empty key/value page groups")
    if len(key_pages_by_group) != len(value_pages_by_group):
        raise ValueError("key_pages_by_group and value_pages_by_group must have the same group count")
    group_count = len(key_pages_by_group)
    if len(query_groups) != group_count:
        raise ValueError("query_groups must align with key/value group count")

    prepared_key_groups = [prepare_pages_torch(group_pages, device_type=device_type, trace=trace) for group_pages in key_pages_by_group]
    prepared_value_groups = [prepare_pages_torch(group_pages, device_type=device_type, trace=trace) for group_pages in value_pages_by_group]
    if not prepared_key_groups[0]:
        raise ValueError("grouped decode requires at least one key page per group")

    return decode_grouped_multiquery_step_prepared_torch_tensor(
        query_groups,
        prepared_key_groups,
        prepared_value_groups,
        trace=trace,
    )


def decode_grouped_multiquery_step_prepared_torch_tensor(
    query_groups,
    key_pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    value_pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    *,
    key_chunk_lengths: Sequence[int] | None = None,
    value_chunk_lengths: Sequence[int] | None = None,
    compact_grouped_chunk: bool = False,
    compact_grouped_mix_chunk: bool = False,
    disable_packed_grouped_cuda_mix: bool = False,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not key_pages_by_group or not value_pages_by_group:
        raise ValueError("grouped decode requires non-empty key/value page groups")
    if len(key_pages_by_group) != len(value_pages_by_group):
        raise ValueError("key_pages_by_group and value_pages_by_group must have the same group count")
    group_count = len(key_pages_by_group)
    if len(query_groups) != group_count:
        raise ValueError("query_groups must align with key/value group count")
    if not key_pages_by_group[0]:
        raise ValueError("grouped decode requires at least one key page per group")
    device_type = key_pages_by_group[0][0].device_type

    query_tensors = [
        group.to(dtype=torch.float32, device=device_type)
        if torch.is_tensor(group)
        else torch.as_tensor(group, dtype=torch.float32, device=device_type)
        for group in query_groups
    ]
    query_count = int(query_tensors[0].shape[0])
    for group_query in query_tensors[1:]:
        if int(group_query.shape[0]) != query_count:
            raise ValueError("all query groups must have the same query count for batched grouped decode")
    stacked_queries = torch.stack(query_tensors, dim=0)
    if trace is not None:
        trace.record_grouped_decode_call(output_only=False)

    first_key_group = key_pages_by_group[0]
    first_value_group = value_pages_by_group[0]
    first_header = first_key_group[0].header
    prepared_query_groups_tensor = _trace_timed_call(
        trace,
        "chunk_assembly",
        device_type=device_type,
        fn=lambda: _pad_queries(
            stacked_queries.reshape(group_count * query_count, first_header.head_dim),
            first_header.padded_head_dim,
            device_type=device_type,
        ).reshape(group_count, query_count, first_header.num_groups, first_header.group_size),
    )
    query_group_sums = prepared_query_groups_tensor.sum(dim=-1)
    if key_chunk_lengths is None:
        key_chunk_lengths = _merged_chunk_lengths_for_page_groups(key_pages_by_group)
    else:
        key_chunk_lengths = tuple(int(length) for length in key_chunk_lengths)
    if value_chunk_lengths is None:
        value_chunk_lengths = _merged_chunk_lengths_for_page_groups(value_pages_by_group)
    else:
        value_chunk_lengths = tuple(int(length) for length in value_chunk_lengths)
    if sum(key_chunk_lengths) != len(first_key_group) or sum(value_chunk_lengths) != len(first_value_group):
        raise ValueError("grouped decode chunk lengths must cover all key/value pages exactly")
    # Score and mix can use different chunk schedules as long as both cover the
    # same flattened token stream in order.

    logits_parts = []
    key_offset = 0
    for chunk_length in key_chunk_lengths:
        chunk_pages = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda key_offset=key_offset, chunk_length=chunk_length: [
                group_pages[key_offset : key_offset + chunk_length] for group_pages in key_pages_by_group
            ],
        )
        if trace is not None:
            trace.record_grouped_score_chunk(
                batch_size=group_count,
                query_count=query_count,
                page_count=chunk_length,
                token_count=first_key_group[0].header.token_count,
            )
        logits_parts.append(
            _trace_timed_call(
                trace,
                "score",
                device_type=device_type,
                fn=lambda chunk_pages=chunk_pages: _score_page_chunk_grouped_multiquery_torch(
                    stacked_queries,
                    chunk_pages,
                    prepared_query_groups_tensor=prepared_query_groups_tensor,
                    query_group_sums=query_group_sums,
                    compact_grouped_chunk=compact_grouped_chunk,
                    trace=trace,
                ),
            )
        )
        key_offset += chunk_length
    logits = torch.cat(logits_parts, dim=2)
    weights = _trace_timed_call(
        trace,
        "softmax",
        device_type=device_type,
        fn=lambda: torch.softmax(logits, dim=2),
    )

    head_dim = first_value_group[0].header.head_dim
    padded_head_dim = first_value_group[0].header.padded_head_dim
    output = torch.zeros((group_count, query_count, padded_head_dim), dtype=torch.float32, device=device_type)
    offset = 0
    value_offset = 0
    for chunk_index, chunk_length in enumerate(value_chunk_lengths):
        chunk_token_count = first_value_group[value_offset].header.token_count * chunk_length
        chunk_weights = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda offset=offset, chunk_token_count=chunk_token_count, chunk_length=chunk_length, token_count=first_value_group[value_offset].header.token_count: weights[
                :,
                :,
                offset : offset + chunk_token_count,
            ].reshape(
                group_count,
                query_count,
                chunk_length,
                token_count,
            ),
        )
        chunk_pages = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda value_offset=value_offset, chunk_length=chunk_length: [
                group_pages[value_offset : value_offset + chunk_length] for group_pages in value_pages_by_group
            ],
        )
        if trace is not None:
            trace.record_grouped_mix_chunk(
                batch_size=group_count,
                query_count=query_count,
                page_count=chunk_length,
                token_count=first_value_group[value_offset].header.token_count,
                head_dim=padded_head_dim,
            )
        output = _trace_timed_call(
            trace,
            "mix",
            device_type=device_type,
            fn=lambda chunk_weights=chunk_weights, chunk_pages=chunk_pages, output=output: _mix_page_chunk_grouped_multiquery_torch(
                chunk_weights,
                chunk_pages,
                out_acc=output,
                compact_grouped_chunk=(compact_grouped_chunk or compact_grouped_mix_chunk),
                disable_packed_grouped_cuda=disable_packed_grouped_cuda_mix,
                trace=trace,
            ),
        )
        offset += chunk_token_count
        value_offset += chunk_length

    return logits, weights, output[:, :, :head_dim]


def decode_grouped_multiquery_step_prepared_torch_tensor_output_only(
    query_groups,
    key_pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    value_pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    *,
    compact_grouped_chunk: bool = False,
    compact_grouped_mix_chunk: bool = False,
    disable_packed_grouped_cuda_mix: bool = False,
    trace: ExecutionTrace | None = None,
):
    torch = _load_torch()
    if not key_pages_by_group or not value_pages_by_group:
        raise ValueError("grouped decode requires non-empty key/value page groups")
    if len(key_pages_by_group) != len(value_pages_by_group):
        raise ValueError("key_pages_by_group and value_pages_by_group must have the same group count")
    group_count = len(key_pages_by_group)
    if len(query_groups) != group_count:
        raise ValueError("query_groups must align with key/value group count")
    if not key_pages_by_group[0]:
        raise ValueError("grouped decode requires at least one key page per group")
    device_type = key_pages_by_group[0][0].device_type

    query_tensors = [
        group.to(dtype=torch.float32, device=device_type)
        if torch.is_tensor(group)
        else torch.as_tensor(group, dtype=torch.float32, device=device_type)
        for group in query_groups
    ]
    query_count = int(query_tensors[0].shape[0])
    for group_query in query_tensors[1:]:
        if int(group_query.shape[0]) != query_count:
            raise ValueError("all query groups must have the same query count for batched grouped decode")
    stacked_queries = torch.stack(query_tensors, dim=0)
    if trace is not None:
        trace.record_grouped_decode_call(output_only=True)

    first_key_group = key_pages_by_group[0]
    first_value_group = value_pages_by_group[0]
    first_header = first_key_group[0].header
    prepared_query_groups_tensor = _trace_timed_call(
        trace,
        "chunk_assembly",
        device_type=device_type,
        fn=lambda: _pad_queries(
            stacked_queries.reshape(group_count * query_count, first_header.head_dim),
            first_header.padded_head_dim,
            device_type=device_type,
        ).reshape(group_count, query_count, first_header.num_groups, first_header.group_size),
    )
    query_group_sums = prepared_query_groups_tensor.sum(dim=-1)
    shared_chunk_lengths = _aligned_chunk_lengths_for_page_pairs(first_key_group, first_value_group)

    head_dim = first_value_group[0].header.head_dim
    padded_head_dim = first_value_group[0].header.padded_head_dim
    output = torch.zeros((group_count, query_count, padded_head_dim), dtype=torch.float32, device=device_type)
    running_max = torch.full((group_count, query_count), float("-inf"), dtype=torch.float32, device=device_type)
    running_denom = torch.zeros((group_count, query_count), dtype=torch.float32, device=device_type)

    key_offset = 0
    value_offset = 0
    for chunk_index, chunk_length in enumerate(shared_chunk_lengths):
        key_chunk_pages = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda key_offset=key_offset, chunk_length=chunk_length: [
                group_pages[key_offset : key_offset + chunk_length] for group_pages in key_pages_by_group
            ],
        )
        logits_chunk = _trace_timed_call(
            trace,
            "score",
            device_type=device_type,
            fn=lambda key_chunk_pages=key_chunk_pages: _score_page_chunk_grouped_multiquery_torch(
                stacked_queries,
                key_chunk_pages,
                prepared_query_groups_tensor=prepared_query_groups_tensor,
                query_group_sums=query_group_sums,
                compact_grouped_chunk=compact_grouped_chunk,
                trace=trace,
            ),
        )
        if trace is not None:
            trace.record_grouped_score_chunk(
                batch_size=group_count,
                query_count=query_count,
                page_count=chunk_length,
                token_count=first_key_group[0].header.token_count,
            )
        value_chunk_pages = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda value_offset=value_offset, chunk_length=chunk_length: [
                group_pages[value_offset : value_offset + chunk_length] for group_pages in value_pages_by_group
            ],
        )
        chunk_template = value_chunk_pages[0]
        chunk_token_count = chunk_template[0].header.token_count * chunk_length

        prev_max = running_max
        chunk_max = torch.amax(logits_chunk, dim=2)

        def _normalize_chunk():
            new_max = torch.maximum(prev_max, chunk_max)
            prev_scale = torch.where(torch.isfinite(prev_max), torch.exp(prev_max - new_max), torch.zeros_like(new_max))
            chunk_exp = torch.exp(logits_chunk - new_max[:, :, None])
            return new_max, prev_scale, chunk_exp

        new_max, prev_scale, chunk_exp = _trace_timed_call(
            trace,
            "softmax",
            device_type=device_type,
            fn=_normalize_chunk,
        )
        output *= prev_scale[:, :, None]
        running_denom = (running_denom * prev_scale) + torch.sum(chunk_exp, dim=2)
        chunk_weights = _trace_timed_call(
            trace,
            "chunk_assembly",
            device_type=device_type,
            fn=lambda chunk_exp=chunk_exp, chunk_length=chunk_length, chunk_token_count=chunk_template[0].header.token_count: chunk_exp.reshape(
                group_count,
                query_count,
                chunk_length,
                chunk_token_count,
            ),
        )
        output = _trace_timed_call(
            trace,
            "mix",
            device_type=device_type,
            fn=lambda chunk_weights=chunk_weights, value_chunk_pages=value_chunk_pages, output=output: _mix_page_chunk_grouped_multiquery_torch(
                chunk_weights,
                value_chunk_pages,
                out_acc=output,
                compact_grouped_chunk=(compact_grouped_chunk or compact_grouped_mix_chunk),
                disable_packed_grouped_cuda=disable_packed_grouped_cuda_mix,
                trace=trace,
            ),
        )
        if trace is not None:
            trace.record_grouped_mix_chunk(
                batch_size=group_count,
                query_count=query_count,
                page_count=chunk_length,
                token_count=chunk_template[0].header.token_count,
                head_dim=padded_head_dim,
            )
        running_max = new_max
        key_offset += chunk_length
        value_offset += chunk_length

    output = _trace_timed_call(
        trace,
        "softmax",
        device_type=device_type,
        fn=lambda: output / torch.clamp(running_denom[:, :, None], min=1e-12),
    )
    return output[:, :, :head_dim]


def decode_multi_query_step_torch(
    query_slices: np.ndarray,
    key_pages: Sequence[EncodedPage | PreparedPageTorch],
    value_pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    device_type: TorchDevice,
    trace: ExecutionTrace | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    logits, weights, output = decode_multi_query_step_torch_tensor(
        query_slices,
        key_pages,
        value_pages,
        device_type=device_type,
        trace=trace,
    )
    return (
        logits.detach().cpu().numpy(),
        weights.detach().cpu().numpy(),
        output.detach().cpu().numpy(),
    )


def prepare_pages_mps(
    pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    trace: ExecutionTrace | None = None,
) -> list[PreparedPageTorch]:
    return prepare_pages_torch(pages, device_type="mps", trace=trace)


def prepare_page_mps(page: EncodedPage | PreparedPageTorch, *, trace: ExecutionTrace | None = None) -> PreparedPageTorch:
    return prepare_page_torch(page, device_type="mps", trace=trace)


def score_page_mps(
    query_slice: np.ndarray | Any,
    page: EncodedPage | PreparedPageTorch,
    *,
    trace: ExecutionTrace | None = None,
) -> np.ndarray:
    return score_page_torch(query_slice, page, device_type="mps", trace=trace)


def score_pages_mps(
    query_slice: np.ndarray | Any,
    pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    trace: ExecutionTrace | None = None,
) -> list[np.ndarray]:
    return score_pages_torch(query_slice, pages, device_type="mps", trace=trace)


def mix_page_mps(
    attn_weights: np.ndarray | Any,
    page: EncodedPage | PreparedPageTorch,
    *,
    out_acc: np.ndarray | None = None,
    trace: ExecutionTrace | None = None,
) -> np.ndarray:
    return mix_page_torch(attn_weights, page, device_type="mps", out_acc=out_acc, trace=trace)


def decode_step_mps(
    query_slice: np.ndarray | Any,
    key_pages: Sequence[EncodedPage | PreparedPageTorch],
    value_pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    precomputed_page_logits: Sequence[np.ndarray | Any | None] | None = None,
    trace: ExecutionTrace | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    return decode_step_torch(
        query_slice,
        key_pages,
        value_pages,
        device_type="mps",
        precomputed_page_logits=precomputed_page_logits,
        trace=trace,
    )


def decode_multi_query_step_mps_tensor(
    query_slices: np.ndarray | Any,
    key_pages: Sequence[EncodedPage | PreparedPageTorch],
    value_pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    trace: ExecutionTrace | None = None,
):
    return decode_multi_query_step_torch_tensor(
        query_slices,
        key_pages,
        value_pages,
        device_type="mps",
        trace=trace,
    )


def decode_grouped_multiquery_step_mps_tensor(
    query_groups,
    key_pages_by_group: Sequence[Sequence[EncodedPage | PreparedPageTorch]],
    value_pages_by_group: Sequence[Sequence[EncodedPage | PreparedPageTorch]],
    *,
    trace: ExecutionTrace | None = None,
):
    return decode_grouped_multiquery_step_torch_tensor(
        query_groups,
        key_pages_by_group,
        value_pages_by_group,
        device_type="mps",
        trace=trace,
    )


def decode_grouped_multiquery_step_prepared_mps_tensor(
    query_groups,
    key_pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    value_pages_by_group: Sequence[Sequence[PreparedPageTorch]],
    *,
    trace: ExecutionTrace | None = None,
):
    return decode_grouped_multiquery_step_prepared_torch_tensor(
        query_groups,
        key_pages_by_group,
        value_pages_by_group,
        trace=trace,
    )


def decode_multi_query_step_mps(
    query_slices: np.ndarray,
    key_pages: Sequence[EncodedPage | PreparedPageTorch],
    value_pages: Sequence[EncodedPage | PreparedPageTorch],
    *,
    trace: ExecutionTrace | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    return decode_multi_query_step_torch(query_slices, key_pages, value_pages, device_type="mps", trace=trace)