# -*- coding: utf-8 -*- """Ragged, sample-isolated detail-memory compression for OMVT-v2.""" from __future__ import annotations from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F @dataclass(frozen=True) class PackedDetailMemory: """Packed visual memory emitted for a later ragged cross-attention bridge.""" memory: torch.Tensor cu_seqlens: torch.Tensor sample_ids: torch.Tensor token_counts: torch.Tensor source_token_counts: torch.Tensor def _validate_packed_features( features: torch.Tensor, cu_seqlens: torch.Tensor, ) -> torch.Tensor: if features.ndim != 2: raise ValueError("packed features must have shape [N, D]") if cu_seqlens.ndim != 1 or cu_seqlens.numel() < 2: raise ValueError("cu_seqlens must have shape [batch + 1]") if cu_seqlens.dtype not in (torch.int32, torch.int64): raise ValueError("cu_seqlens must use an integer dtype") if cu_seqlens.device.type != "cpu": raise ValueError("cu_seqlens must remain on CPU") if int(cu_seqlens[0].item()) != 0: raise ValueError("cu_seqlens must start at zero") if int(cu_seqlens[-1].item()) != int(features.shape[0]): raise ValueError("cu_seqlens must end at the packed feature count") if bool((cu_seqlens[1:] < cu_seqlens[:-1]).any()): raise ValueError("cu_seqlens must be non-decreasing") return (cu_seqlens[1:] - cu_seqlens[:-1]).to(dtype=torch.long) class _SampleCrossAttention(nn.Module): """One-sample cross attention; callers own ragged segmentation.""" def __init__(self, d_model: int, n_heads: int) -> None: super().__init__() if d_model % n_heads != 0: raise ValueError("d_model must be divisible by n_heads") self.n_heads = int(n_heads) self.head_dim = d_model // n_heads self.norm_query = nn.LayerNorm(d_model) self.norm_context = nn.LayerNorm(d_model) self.query_proj = nn.Linear(d_model, d_model) self.key_value_proj = nn.Linear(d_model, d_model * 2) self.out = nn.Linear(d_model, d_model) def forward(self, query: torch.Tensor, context: torch.Tensor) -> torch.Tensor: if query.ndim != 2 or context.ndim != 2: raise ValueError("query and context must have shape [tokens, D]") if context.shape[0] == 0: raise ValueError("detail compression requires non-empty sample context") n_query, d_model = query.shape projected_query = self.query_proj(self.norm_query(query)) key, value = self.key_value_proj(self.norm_context(context)).chunk(2, dim=-1) projected_query = projected_query.reshape( n_query, self.n_heads, self.head_dim ).transpose(0, 1).unsqueeze(0) key = key.reshape(-1, self.n_heads, self.head_dim).transpose(0, 1).unsqueeze(0) value = value.reshape(-1, self.n_heads, self.head_dim).transpose(0, 1).unsqueeze(0) attended = F.scaled_dot_product_attention(projected_query, key, value) attended = attended.squeeze(0).transpose(0, 1).reshape(n_query, d_model) return self.out(attended) class NativeDetailCompressor(nn.Module): """Compress each sample to a monotonic, hard-capped ragged token budget. The current implementation is deliberately a correctness kernel: it loops over samples and never pads them into a shared dense context. A later optimized kernel must preserve this exact isolation contract. """ def __init__( self, omvt_cfg, *, max_detail_tokens_per_sample: int | None = None, source_tokens_per_detail_token: int = 4, ) -> None: super().__init__() cap = ( int(omvt_cfg.compress_to) if max_detail_tokens_per_sample is None else int(max_detail_tokens_per_sample) ) if cap <= 0: raise ValueError("max_detail_tokens_per_sample must be positive") if source_tokens_per_detail_token <= 0: raise ValueError("source_tokens_per_detail_token must be positive") self.max_detail_tokens_per_sample = cap self.source_tokens_per_detail_token = int(source_tokens_per_detail_token) self.latents = nn.Parameter( torch.randn(cap, omvt_cfg.d_vision) * 0.02 ) self.blocks = nn.ModuleList( [ nn.ModuleDict( { "cross": _SampleCrossAttention( omvt_cfg.d_vision, omvt_cfg.compressor_heads, ), "ffn": nn.Sequential( nn.LayerNorm(omvt_cfg.d_vision), nn.Linear( omvt_cfg.d_vision, omvt_cfg.vision_ffn_hidden, ), nn.GELU(), nn.Linear( omvt_cfg.vision_ffn_hidden, omvt_cfg.d_vision, ), ), } ) for _ in range(omvt_cfg.compressor_layers) ] ) self.final_norm = nn.LayerNorm(omvt_cfg.d_vision) def detail_token_budget(self, source_token_count: int) -> int: count = int(source_token_count) if count < 0: raise ValueError("source_token_count must be non-negative") if count == 0: return 0 requested = ( count + self.source_tokens_per_detail_token - 1 ) // self.source_tokens_per_detail_token return min(self.max_detail_tokens_per_sample, requested) def _one_sample(self, context: torch.Tensor, output_tokens: int) -> torch.Tensor: if output_tokens <= 0 or output_tokens > self.max_detail_tokens_per_sample: raise ValueError("output token count violates the detail-memory cap") latents = self.latents[:output_tokens] for block in self.blocks: latents = latents + block["cross"](latents, context) latents = latents + block["ffn"](latents) return self.final_norm(latents) def forward( self, features: torch.Tensor, cu_seqlens: torch.Tensor, ) -> PackedDetailMemory: source_counts = _validate_packed_features(features, cu_seqlens) output_counts = torch.tensor( [self.detail_token_budget(int(count.item())) for count in source_counts], dtype=torch.long, device="cpu", ) if bool((output_counts > self.max_detail_tokens_per_sample).any()): raise RuntimeError("detail-memory token budget exceeded its hard cap") if bool((source_counts == 0).any()): raise ValueError("native detail compression does not accept empty samples") memories: list[torch.Tensor] = [] sample_id_rows: list[torch.Tensor] = [] output_cumulative = [0] for sample_index in range(int(source_counts.numel())): start = int(cu_seqlens[sample_index].item()) end = int(cu_seqlens[sample_index + 1].item()) output_count = int(output_counts[sample_index].item()) memories.append(self._one_sample(features[start:end], output_count)) sample_id_rows.append( torch.full( (output_count,), sample_index, dtype=torch.long, device=features.device, ) ) output_cumulative.append(output_cumulative[-1] + output_count) memory = torch.cat(memories, dim=0) packed = PackedDetailMemory( memory=memory, cu_seqlens=torch.tensor( output_cumulative, dtype=torch.int32, device="cpu", ), sample_ids=torch.cat(sample_id_rows, dim=0), token_counts=output_counts, source_token_counts=source_counts, ) if int(packed.cu_seqlens[-1].item()) != int(memory.shape[0]): raise RuntimeError("packed detail-memory offsets are inconsistent") return packed __all__ = ["NativeDetailCompressor", "PackedDetailMemory"]