source/src/strata/modeling/ph_pat/replacement_block.py

"""Dense, local-only, and physically chart-routed Transformer blocks."""

from __future__ import annotations

import torch
from torch import nn
from torch.nn import functional as F
from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention

from strata.modeling.modules import RMSNorm
from strata.modeling.ph_pat.chart_executor import CompiledPHPATChart
from strata.modeling.ph_pat.chart_read_attention import ChartReadAttention, ChartReadOutput
from strata.modeling.ph_pat.config import PHPATConfig
from strata.modeling.ph_pat.segment_commit import SegmentLayout


_LOCAL_MASK_CACHE: dict[tuple[str, int, int], torch.Tensor] = {}
_FLEX_MASK_CACHE: dict[tuple[str, int, int], BlockMask] = {}
_BLOCK_ISOLATED_MASK_CACHE: dict[tuple[str, int, int], torch.Tensor] = {}
_BLOCK_ISOLATED_FLEX_CACHE: dict[tuple[str, int, int], BlockMask] = {}
_COMPILED_FLEX_ATTENTION = None


class _FeedForward(nn.Module):
    def __init__(self, config: PHPATConfig) -> None:
        super().__init__()
        self.up = nn.Linear(config.d_model, 2 * config.d_ff, bias=False)
        self.down = nn.Linear(config.d_ff, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, hidden: torch.Tensor) -> torch.Tensor:
        gate, value = self.up(hidden).chunk(2, dim=-1)
        return self.down(self.dropout(F.silu(gate) * value))


class _AttentionProjection(nn.Module):
    def __init__(self, config: PHPATConfig) -> None:
        super().__init__()
        self.config = config
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out = nn.Linear(config.d_model, config.d_model, bias=False)

    def split(self, hidden: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        batch, seq, _ = hidden.shape
        q, k, v = self.qkv(hidden).chunk(3, dim=-1)
        shape = (batch, seq, self.config.num_heads, self.config.head_dim)
        return tuple(x.view(shape).transpose(1, 2) for x in (q, k, v))  # type: ignore[return-value]

    def merge(self, hidden: torch.Tensor) -> torch.Tensor:
        batch, heads, seq, dim = hidden.shape
        return self.out(hidden.transpose(1, 2).contiguous().view(batch, seq, heads * dim))


class DenseGlobalCausalAttention(_AttentionProjection):
    """Global attention exists only in this concrete class."""

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        q, k, v = self.split(hidden)
        if bool(attention_mask.all()):
            value = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        else:
            seq = hidden.shape[1]
            causal = torch.ones(seq, seq, device=hidden.device, dtype=torch.bool).tril()
            mask = causal.view(1, 1, seq, seq) & attention_mask.to(torch.bool).view(hidden.shape[0], 1, 1, seq)
            value = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.merge(value)


class ChunkedLocalCausalAttention(_AttentionProjection):
    """Fused local causal SDPA with a shared immutable band mask."""

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        q, k, v = self.split(hidden)
        seq = hidden.shape[1]
        if hidden.device.type == "cuda" and bool(attention_mask.all()) and seq >= 256:
            block_mask = _flex_local_mask(seq, self.config.local_window, hidden.device)
            value = _compiled_flex_attention()(q, k, v, block_mask=block_mask)
        else:
            local = _local_causal_mask(seq, self.config.local_window, hidden.device)
            valid_keys = attention_mask.to(torch.bool).view(hidden.shape[0], 1, 1, seq)
            mask = local.view(1, 1, seq, seq) & valid_keys
            value = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.merge(value) * attention_mask.to(hidden.dtype).unsqueeze(-1)


class BlockIsolatedLocalCausalAttention(_AttentionProjection):
    """Causal fixed blocks with no hidden-state relay across local windows."""

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        q, k, v = self.split(hidden)
        seq = hidden.shape[1]
        if hidden.device.type == "cuda" and bool(attention_mask.all()) and seq >= 256:
            block_mask = _block_isolated_flex_mask(seq, self.config.local_window, hidden.device)
            value = _compiled_flex_attention()(q, k, v, block_mask=block_mask)
        else:
            local = _block_isolated_causal_mask(seq, self.config.local_window, hidden.device)
            valid_keys = attention_mask.to(torch.bool).view(hidden.shape[0], 1, 1, seq)
            mask = local.view(1, 1, seq, seq) & valid_keys
            value = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        return self.merge(value) * attention_mask.to(hidden.dtype).unsqueeze(-1)


def _local_causal_mask(seq_len: int, window: int, device: torch.device) -> torch.Tensor:
    key = (str(device), seq_len, window)
    cached = _LOCAL_MASK_CACHE.get(key)
    if cached is None:
        positions = torch.arange(seq_len, device=device)
        delta = positions.view(-1, 1) - positions.view(1, -1)
        cached = (delta >= 0) & (delta < window)
        _LOCAL_MASK_CACHE[key] = cached
    return cached


def _flex_local_mask(seq_len: int, window: int, device: torch.device) -> BlockMask:
    key = (str(device), seq_len, window)
    cached = _FLEX_MASK_CACHE.get(key)
    if cached is None:
        def mask_mod(_batch, _head, query_index, key_index):
            distance = query_index - key_index
            return (distance >= 0) & (distance < window)

        cached = create_block_mask(
            mask_mod,
            B=None,
            H=None,
            Q_LEN=seq_len,
            KV_LEN=seq_len,
            device=str(device),
            BLOCK_SIZE=128,
            _compile=True,
        )
        _FLEX_MASK_CACHE[key] = cached
    return cached


def _block_isolated_causal_mask(seq_len: int, window: int, device: torch.device) -> torch.Tensor:
    key = (str(device), seq_len, window)
    cached = _BLOCK_ISOLATED_MASK_CACHE.get(key)
    if cached is None:
        positions = torch.arange(seq_len, device=device)
        query = positions.view(-1, 1)
        key_position = positions.view(1, -1)
        cached = (key_position <= query) & ((query // window) == (key_position // window))
        _BLOCK_ISOLATED_MASK_CACHE[key] = cached
    return cached


def _block_isolated_flex_mask(seq_len: int, window: int, device: torch.device) -> BlockMask:
    key = (str(device), seq_len, window)
    cached = _BLOCK_ISOLATED_FLEX_CACHE.get(key)
    if cached is None:
        def mask_mod(_batch, _head, query_index, key_index):
            return (key_index <= query_index) & ((query_index // window) == (key_index // window))

        cached = create_block_mask(
            mask_mod,
            B=None,
            H=None,
            Q_LEN=seq_len,
            KV_LEN=seq_len,
            device=str(device),
            BLOCK_SIZE=128,
            _compile=True,
        )
        _BLOCK_ISOLATED_FLEX_CACHE[key] = cached
    return cached


def _compiled_flex_attention():
    global _COMPILED_FLEX_ATTENTION
    if _COMPILED_FLEX_ATTENTION is None:
        _COMPILED_FLEX_ATTENTION = torch.compile(flex_attention, dynamic=False)
    return _COMPILED_FLEX_ATTENTION


class _BaseBlock(nn.Module):
    def __init__(self, config: PHPATConfig, attention: nn.Module) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.d_model)
        self.attention = attention
        self.ffn_norm = RMSNorm(config.d_model)
        self.ffn = _FeedForward(config)
        self.dropout = nn.Dropout(config.dropout)

    def token_path(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        hidden = hidden + self.dropout(self.attention(self.attn_norm(hidden), attention_mask))
        return hidden + self.dropout(self.ffn(self.ffn_norm(hidden)))


class DenseGlobalBlock(_BaseBlock):
    def __init__(self, config: PHPATConfig) -> None:
        super().__init__(config, DenseGlobalCausalAttention(config))

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> tuple[torch.Tensor, None]:
        return self.token_path(hidden, attention_mask), None


class LocalOnlyBlock(_BaseBlock):
    def __init__(self, config: PHPATConfig) -> None:
        super().__init__(config, ChunkedLocalCausalAttention(config))

    def forward(self, hidden: torch.Tensor, attention_mask: torch.Tensor) -> tuple[torch.Tensor, None]:
        return self.token_path(hidden, attention_mask), None


class LocalChartReplacementBlock(_BaseBlock):
    """No dense attention module or fallback path is instantiated here."""

    def __init__(self, config: PHPATConfig) -> None:
        super().__init__(config, ChunkedLocalCausalAttention(config))
        self.chart_norm = RMSNorm(config.d_model)
        self.read_gate = nn.Linear(config.d_model, 2)
        self.event_scale = nn.Parameter(torch.tensor(float(config.register_injection_scale)))
        self.primitive_scale = nn.Parameter(torch.tensor(float(config.register_injection_scale)))

    def forward(
        self,
        hidden: torch.Tensor,
        attention_mask: torch.Tensor,
        read: ChartReadOutput,
        *,
        detach_memory: bool,
    ) -> tuple[torch.Tensor, ChartReadOutput]:
        hidden = self.token_path(hidden, attention_mask)
        gates = torch.sigmoid(self.read_gate(self.chart_norm(hidden)))
        event = read.event_role.detach() if detach_memory else read.event_role
        primitive = read.primitive.detach() if detach_memory else read.primitive
        update = torch.tanh(self.event_scale) * gates[..., :1] * event
        update = update + torch.tanh(self.primitive_scale) * gates[..., 1:] * primitive
        return hidden + self.dropout(update), read


__all__ = [
    "BlockIsolatedLocalCausalAttention",
    "ChunkedLocalCausalAttention",
    "DenseGlobalBlock",
    "DenseGlobalCausalAttention",
    "LocalChartReplacementBlock",
    "LocalOnlyBlock",
]