tinymind-native-colab-handoff/bundle/model/omega_multimodal_context.py

"""Omega multimodal/context bridge reference modules.

This is a correctness-first PyTorch graft point for:

- projecting image/audio feature tokens into Omega hidden space,
- compressing long context into block-wise anchors without full historical KV,
- routing text/code/multimodal lanes without zero-work expert collapse.

It is not a claim of native vision/audio quality or faster-than-SOTA kernels.
Those claims remain gated by external encoder checkpoints and benchmark reports.
"""

from __future__ import annotations

from dataclasses import dataclass
import math

import torch
import torch.nn as nn

from .layers import RMSNorm


@dataclass
class OmegaBridgeConfig:
    hidden_dim: int = 5120
    image_feature_dim: int = 768
    audio_feature_dim: int = 768
    local_window: int = 2048
    block_tokens: int = 64
    max_persistent_tokens: int = 2_000_000
    anchor_rank: int = 128
    router_experts: int = 44
    text_experts: int = 12
    code_experts: int = 16
    multimodal_experts: int = 16
    compression_dtype: str = "int4_2:4sp_anchor_reference"

    def __post_init__(self) -> None:
        if self.hidden_dim <= 0:
            raise ValueError("hidden_dim must be positive")
        if self.local_window <= 0 or self.block_tokens <= 0:
            raise ValueError("local_window and block_tokens must be positive")
        if self.text_experts + self.code_experts + self.multimodal_experts != self.router_experts:
            raise ValueError("expert groups must sum to router_experts")
        if self.max_persistent_tokens < self.local_window:
            raise ValueError("max_persistent_tokens must cover local_window")


@dataclass
class BlockWiseKVLedgerState:
    local_exact: torch.Tensor
    anchors: torch.Tensor
    anchor_count: int
    total_tokens_seen: int

    def cached_token_capacity(self) -> int:
        return int(self.local_exact.shape[1] + self.anchors.shape[1])


class CrossModalProjectionGate(nn.Module):
    """Project external encoder features into the language hidden dimension."""

    def __init__(self, cfg: OmegaBridgeConfig):
        super().__init__()
        self.cfg = cfg
        self.image_norm = RMSNorm(cfg.image_feature_dim)
        self.audio_norm = RMSNorm(cfg.audio_feature_dim)
        self.image_proj = nn.Linear(cfg.image_feature_dim, cfg.hidden_dim, bias=False)
        self.audio_proj = nn.Linear(cfg.audio_feature_dim, cfg.hidden_dim, bias=False)
        self.type_embed = nn.Parameter(torch.zeros(3, cfg.hidden_dim))
        nn.init.normal_(self.type_embed, std=0.01)

    def forward(
        self,
        *,
        text_embeds: torch.Tensor | None = None,
        image_features: torch.Tensor | None = None,
        audio_features: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, dict]:
        streams: list[torch.Tensor] = []
        token_counts: dict[str, int] = {"text": 0, "image": 0, "audio": 0}
        if text_embeds is not None:
            if text_embeds.dim() != 3 or text_embeds.shape[-1] != self.cfg.hidden_dim:
                raise ValueError("text_embeds must have shape [batch, seq, hidden_dim]")
            streams.append(text_embeds + self.type_embed[0].view(1, 1, -1))
            token_counts["text"] = int(text_embeds.shape[1])
        if image_features is not None:
            if image_features.dim() != 3 or image_features.shape[-1] != self.cfg.image_feature_dim:
                raise ValueError("image_features must have shape [batch, image_tokens, image_feature_dim]")
            image_tokens = self.image_proj(self.image_norm(image_features)) + self.type_embed[1].view(1, 1, -1)
            streams.append(image_tokens)
            token_counts["image"] = int(image_tokens.shape[1])
        if audio_features is not None:
            if audio_features.dim() != 3 or audio_features.shape[-1] != self.cfg.audio_feature_dim:
                raise ValueError("audio_features must have shape [batch, audio_tokens, audio_feature_dim]")
            audio_tokens = self.audio_proj(self.audio_norm(audio_features)) + self.type_embed[2].view(1, 1, -1)
            streams.append(audio_tokens)
            token_counts["audio"] = int(audio_tokens.shape[1])
        if not streams:
            raise ValueError("at least one modality stream is required")
        batch = streams[0].shape[0]
        if any(stream.shape[0] != batch for stream in streams):
            raise ValueError("all modality streams must have the same batch size")
        return torch.cat(streams, dim=1), {"projected_token_counts": token_counts, "hidden_dim": self.cfg.hidden_dim}


class BlockWiseKVLedger(nn.Module):
    """Bounded exact-local + compressed-anchor long-context state."""

    def __init__(self, cfg: OmegaBridgeConfig):
        super().__init__()
        self.cfg = cfg
        self.norm = RMSNorm(cfg.hidden_dim)
        self.anchor_down = nn.Linear(cfg.hidden_dim, cfg.anchor_rank, bias=False)
        self.anchor_up = nn.Linear(cfg.anchor_rank, cfg.hidden_dim, bias=False)
        self.score = nn.Linear(cfg.hidden_dim, 1, bias=False)

    def empty_state(self, batch: int, device: torch.device, dtype: torch.dtype) -> BlockWiseKVLedgerState:
        return BlockWiseKVLedgerState(
            local_exact=torch.zeros(batch, 0, self.cfg.hidden_dim, device=device, dtype=dtype),
            anchors=torch.zeros(batch, 0, self.cfg.anchor_rank, device=device, dtype=dtype),
            anchor_count=0,
            total_tokens_seen=0,
        )

    def _make_anchors(self, hidden: torch.Tensor) -> torch.Tensor:
        if hidden.shape[1] < self.cfg.block_tokens:
            return torch.zeros(hidden.shape[0], 0, self.cfg.anchor_rank, device=hidden.device, dtype=hidden.dtype)
        usable = hidden[:, : (hidden.shape[1] // self.cfg.block_tokens) * self.cfg.block_tokens]
        blocks = usable.view(hidden.shape[0], -1, self.cfg.block_tokens, self.cfg.hidden_dim)
        weights = torch.softmax(self.score(blocks).squeeze(-1), dim=-1).unsqueeze(-1)
        summaries = (blocks * weights).sum(dim=2)
        return torch.tanh(self.anchor_down(summaries))

    def update(
        self,
        hidden: torch.Tensor,
        state: BlockWiseKVLedgerState | None = None,
    ) -> tuple[BlockWiseKVLedgerState, dict]:
        if hidden.dim() != 3 or hidden.shape[-1] != self.cfg.hidden_dim:
            raise ValueError("hidden must have shape [batch, seq, hidden_dim]")
        hidden = self.norm(hidden)
        if state is None:
            state = self.empty_state(hidden.shape[0], hidden.device, hidden.dtype)
        merged_local = torch.cat([state.local_exact.to(hidden.device, hidden.dtype), hidden.detach()], dim=1)
        next_local = merged_local[:, -self.cfg.local_window :].contiguous()
        historical = merged_local[:, : max(0, merged_local.shape[1] - self.cfg.local_window)]
        new_anchors = self._make_anchors(historical)
        anchors = torch.cat([state.anchors.to(hidden.device, hidden.dtype), new_anchors.detach()], dim=1)
        max_anchor_count = math.ceil(self.cfg.max_persistent_tokens / self.cfg.block_tokens)
        anchors = anchors[:, -max_anchor_count:].contiguous()
        next_state = BlockWiseKVLedgerState(
            local_exact=next_local,
            anchors=anchors,
            anchor_count=int(anchors.shape[1]),
            total_tokens_seen=int(state.total_tokens_seen + hidden.shape[1]),
        )
        metrics = {
            "local_exact_tokens_stored": int(next_local.shape[1]),
            "full_historical_kv_tokens_stored": 0,
            "anchor_count": int(anchors.shape[1]),
            "block_tokens": int(self.cfg.block_tokens),
            "max_persistent_tokens": int(self.cfg.max_persistent_tokens),
            "compression_dtype": self.cfg.compression_dtype,
            "bounded_memory_gate": {
                "passed": bool(next_local.shape[1] <= self.cfg.local_window and anchors.shape[1] <= max_anchor_count),
                "cached_token_capacity": next_state.cached_token_capacity(),
            },
        }
        return next_state, metrics

    def retrieve_anchor_context(self, query: torch.Tensor, state: BlockWiseKVLedgerState, top_k: int = 8) -> torch.Tensor:
        if state.anchors.shape[1] == 0:
            return torch.zeros(query.shape[0], 0, self.cfg.hidden_dim, device=query.device, dtype=query.dtype)
        anchors = self.anchor_up(state.anchors.to(query.device, query.dtype))
        query_summary = query.mean(dim=1)
        scores = torch.einsum("bd,bkd->bk", query_summary, anchors) / math.sqrt(self.cfg.hidden_dim)
        k = min(top_k, anchors.shape[1])
        idx = torch.topk(scores, k=k, dim=-1).indices
        gather_idx = idx.unsqueeze(-1).expand(-1, -1, self.cfg.hidden_dim)
        return torch.gather(anchors, dim=1, index=gather_idx)


class MultiExpertCrossRouter(nn.Module):
    """44-lane router with text/code/multimodal group accounting."""

    def __init__(self, cfg: OmegaBridgeConfig):
        super().__init__()
        self.cfg = cfg
        self.norm = RMSNorm(cfg.hidden_dim)
        self.router = nn.Linear(cfg.hidden_dim, cfg.router_experts, bias=True)
        groups = (
            ["text"] * cfg.text_experts
            + ["code"] * cfg.code_experts
            + ["multimodal"] * cfg.multimodal_experts
        )
        self.groups = groups

    def forward(self, hidden: torch.Tensor, token_type: str = "text") -> tuple[torch.Tensor, dict]:
        if hidden.dim() != 3 or hidden.shape[-1] != self.cfg.hidden_dim:
            raise ValueError("hidden must have shape [batch, seq, hidden_dim]")
        logits = self.router(self.norm(hidden))
        if token_type in {"text", "code", "multimodal"}:
            mask = torch.tensor(
                [0.0 if group == token_type else -1.0 for group in self.groups],
                device=hidden.device,
                dtype=hidden.dtype,
            )
            logits = logits + mask.view(1, 1, -1)
        weights = torch.softmax(logits, dim=-1)
        group_weights: dict[str, torch.Tensor] = {}
        for group in {"text", "code", "multimodal"}:
            idx = [i for i, name in enumerate(self.groups) if name == group]
            group_weights[group] = weights[..., idx].sum(dim=-1).mean()
        metrics = {
            "router_experts": self.cfg.router_experts,
            "group_weights": {key: float(value.detach().cpu()) for key, value in group_weights.items()},
            "zero_work_groups": [key for key, value in group_weights.items() if float(value.detach().cpu()) <= 1e-4],
            "token_type": token_type,
            "no_zero_work_gate": bool(all(float(value.detach().cpu()) > 1e-4 for value in group_weights.values())),
        }
        return weights, metrics


class OmegaMultimodalContextBridge(nn.Module):
    """One integration point for multimodal tokens, bounded long context, and expert routing."""

    def __init__(self, cfg: OmegaBridgeConfig):
        super().__init__()
        self.cfg = cfg
        self.modal_gate = CrossModalProjectionGate(cfg)
        self.ledger = BlockWiseKVLedger(cfg)
        self.router = MultiExpertCrossRouter(cfg)

    def forward(
        self,
        *,
        text_embeds: torch.Tensor | None = None,
        image_features: torch.Tensor | None = None,
        audio_features: torch.Tensor | None = None,
        state: BlockWiseKVLedgerState | None = None,
        token_type: str = "text",
    ) -> tuple[torch.Tensor, BlockWiseKVLedgerState, dict]:
        hidden, modal_metrics = self.modal_gate(
            text_embeds=text_embeds,
            image_features=image_features,
            audio_features=audio_features,
        )
        next_state, ledger_metrics = self.ledger.update(hidden, state)
        _weights, router_metrics = self.router(hidden, token_type=token_type)
        metrics = {
            "modal_projection": modal_metrics,
            "ledger": ledger_metrics,
            "router": router_metrics,
            "claim_gate": {
                "native_multimodal_quality_claim_allowed": False,
                "two_million_context_no_full_kv_claim_allowed": bool(
                    ledger_metrics["bounded_memory_gate"]["passed"]
                    and ledger_metrics["full_historical_kv_tokens_stored"] == 0
                ),
                "beats_flashattention3_claim_allowed": False,
                "reason": "Reference integration only; external encoder and throughput evidence are required for broader claims.",
            },
        }
        return hidden, next_state, metrics