Blame for runtime.py

"""Self-contained mixed global+sliding ModernBERT typed-decision runtime.

Copied as runtime.py into every published Laya V3.1 Hugging Face model.
Does not import TinyCeNN-LM, Laya, or other user project packages.
Depends only on torch, transformers==4.57.6, and safetensors via loader.

Global layers: V2.9 FastWindowMemory (W1024 default; quadratic at 512).
Sliding layers: V3.0 SlidingWindowFusion (W129 default; finite-window SDPA).
"""
from __future__ import annotations

"""Self-contained V2.9 FastWindowMemory typed-decision runtime.

This module is copied into each exported Hugging Face model repository.
It deliberately never imports TinyCeNN-LM, Laya, or local project modules.
Runtime dependencies: torch, transformers==4.57.6 and safetensors.
The Hugging Face Hub client is needed only to download a remote repository.
"""
import copy
import json
import math
from pathlib import Path
from typing import Any, Mapping

import torch
import torch.nn.functional as F
from torch import Tensor, nn
from transformers import AutoConfig, AutoModel, AutoTokenizer
from transformers.models.modernbert.modeling_modernbert import apply_rotary_pos_emb


def _apply_rope_compat(config):
    rope = getattr(config, "rope_parameters", None)
    if not isinstance(rope, dict):
        return
    flat = rope.get("rope_theta")
    for layer_type, attr in (
        ("full_attention", "global_rope_theta"),
        ("sliding_attention", "local_rope_theta"),
    ):
        values = rope.get(layer_type)
        theta = values.get("rope_theta") if isinstance(values, dict) else flat
        if theta is not None and hasattr(config, attr):
            setattr(config, attr, float(theta))


def _valid_tokens(attention_mask: Tensor | None, hidden_states: Tensor) -> Tensor:
    b, t = hidden_states.shape[:2]
    if attention_mask is None:
        return torch.ones((b, t), dtype=torch.bool, device=hidden_states.device)
    if attention_mask.ndim == 2:
        return attention_mask.bool()
    if attention_mask.ndim != 4:
        return torch.ones((b, t), dtype=torch.bool, device=hidden_states.device)
    allowed = (attention_mask if attention_mask.dtype == torch.bool
               else attention_mask > -1e4)
    valid_q = allowed.any(dim=-1).any(dim=1)
    valid_k = allowed.any(dim=-2).any(dim=1)
    valid = valid_q & valid_k
    return valid if valid.shape[-1] == t else torch.ones(
        (b, t), dtype=torch.bool, device=hidden_states.device
    )


def _depthwise_sequence_conv(x: Tensor, weight: Tensor) -> Tensor:
    b, h, t, d = x.shape
    kernel = int(weight.shape[-1])
    y = x.permute(0, 1, 3, 2).reshape(b, h * d, t)
    y = F.pad(y, ((kernel - 1) // 2, kernel - 1 - (kernel - 1) // 2))
    y = F.conv1d(y, weight, groups=h * d)
    return y.reshape(b, h, d, t).permute(0, 1, 3, 2).contiguous()


def _orthogonal_maps(heads: int, out_dim: int, in_dim: int) -> Tensor:
    maps = torch.empty(heads, out_dim, in_dim)
    for h in range(heads):
        nn.init.orthogonal_(maps[h])
    return maps


class FastWindowAttention(nn.Module):
    """State-dict-compatible V2.8 full/global attention replacement.

    The local path is quadratic when W >= 2*T-1. W1024 / MAX_LEN512 is
    a full-attention reconstruction control, not an attention-free speedup.
    """

    def __init__(self, original, *, local_window=1024, landmark_stride=16,
                 feature_dim=96, max_landmarks=32):
        super().__init__()
        self.config = original.config
        self.layer_idx = int(getattr(original, "layer_id", -1))
        self.hidden_size = int(self.config.hidden_size)
        self.num_heads = int(self.config.num_attention_heads)
        self.head_dim = int(getattr(original, "head_dim",
                                    self.hidden_size // self.num_heads))
        self.local_window = int(local_window)
        self.landmark_stride = int(landmark_stride)
        self.max_landmarks = int(max_landmarks)
        self.feature_dim = int(feature_dim)
        self.global_input_scale = float(self.head_dim ** -0.25)
        self.Wqkv = copy.deepcopy(original.Wqkv)
        self.Wo = copy.deepcopy(original.Wo)
        self.out_drop = copy.deepcopy(original.out_drop)
        self.rotary_emb = copy.deepcopy(original.rotary_emb)
        h, d = self.num_heads, self.head_dim
        init = _orthogonal_maps(h, self.feature_dim, d)
        self.global_wq = nn.Parameter(init.clone())
        self.global_wk = nn.Parameter(init.clone())
        self.global_log_scale = nn.Parameter(torch.zeros(h))
        # The original constructor initializes identity convolution weights.
        self.local_weight = nn.Parameter(torch.zeros(h * d, 1, 5))
        with torch.no_grad():
            self.local_weight[:, 0, 2] = 1.0
        self.route_logits = nn.Parameter(
            torch.tensor([1.0, 0.2, 0.2, -2.2]).repeat(h, 1)
        )
        self.route_query = nn.Parameter(torch.zeros(h, 4, d))
        self.local_log_temperature = nn.Parameter(torch.zeros(h))
        self.landmark_log_temperature = nn.Parameter(torch.zeros(h))
        self.conv_gain = nn.Parameter(torch.full((h,), -4.0))
        self.Wqkv.requires_grad_(False)
        self.rotary_emb.requires_grad_(False)
        self.last_core_output = None

    def _local_attention(self, q, k, v, valid):
        b, h, t, d = q.shape
        w = self.local_window
        if w < t:
            # Exact blocked centered-local attention inherited by the original
            # PDelta architecture. Never pad a short T up to huge W.
            left, right = (w - 1) // 2, w // 2
            blocks = (t + w - 1) // w
            tail = blocks * w - t
            qp = F.pad(q, (0, 0, 0, tail)).reshape(b, h, blocks, w, d)
            def windows(x):
                return F.pad(x, (0, 0, left, right + tail)).unfold(
                    2, 2 * w - 1, w
                ).transpose(-1, -2)
            kp, vp = windows(k), windows(v)
            key_valid = F.pad(valid, (left, right + tail)).unfold(
                1, 2 * w - 1, w
            )
            qi = torch.arange(w, device=q.device)[:, None]
            ki = torch.arange(2 * w - 1, device=q.device)[None, :]
            band = (ki >= qi) & (ki < qi + w)
            allowed = band[None, None, None] & key_valid[:, None, :, None, :]
            def batch_blocks(x):
                return x.transpose(1, 2).reshape(b * blocks, h, x.shape[-2], d)
            allowed = allowed.transpose(1, 2).reshape(
                b * blocks, 1, w, 2 * w - 1
            )
            z = F.scaled_dot_product_attention(
                batch_blocks(qp), batch_blocks(kp), batch_blocks(vp),
                attn_mask=allowed, dropout_p=0.0
            )
            z = z.reshape(b, blocks, h, w, d).transpose(1, 2)
            return z.reshape(b, h, blocks * w, d)[:, :, :t] * valid[:, None, :, None]
        allowed = valid[:, None, None, :]
        if w < 2 * t - 1:
            qi = torch.arange(t, device=q.device)[:, None]
            ki = torch.arange(t, device=q.device)[None, :]
            band = ((ki >= qi - (w - 1) // 2) &
                    (ki <= qi + w // 2))
            allowed = allowed & band[None, None]
        z = F.scaled_dot_product_attention(q, k, v, attn_mask=allowed, dropout_p=0.0)
        return z * valid[:, None, :, None].to(z.dtype)

    def _landmarks(self, q, k, v, valid):
        b, h, t, d = k.shape
        stride = max(self.landmark_stride, math.ceil(t / self.max_landmarks))
        blocks = math.ceil(t / stride)
        pad = blocks * stride - t
        mask = valid[:, None, :, None].to(k.dtype)
        packed = F.pad(valid, (0, pad), value=False).reshape(b, blocks, stride)
        counts = packed.sum(-1)
        divisor = counts.clamp_min(1).to(k.dtype)[:, None, :, None]
        kp = F.pad(k * mask, (0, 0, 0, pad)).reshape(b, h, blocks, stride, d)
        vp = F.pad(v * mask, (0, 0, 0, pad)).reshape(b, h, blocks, stride, d)
        pooled_k = kp.sum(3) / divisor
        pooled_v = vp.sum(3) / divisor
        allowed = (counts > 0)[:, None, None, :]
        temperature = self.landmark_log_temperature.clamp(-0.4, 0.4).exp()
        qs = q * temperature[None, :, None, None]
        return F.scaled_dot_product_attention(
            qs, pooled_k, pooled_v, attn_mask=allowed, dropout_p=0.0
        ) * mask

    def _global_linear(self, q, k, v, mask):
        scale = self.global_log_scale.to(q.dtype).clamp(-1.5, 1.5).exp()
        qz = torch.einsum("bhtd,hfd->bhtf", q, self.global_wq)
        kz = torch.einsum("bhtd,hfd->bhtf", k, self.global_wk)
        qz = qz * self.global_input_scale * scale[None, :, None, None]
        kz = kz * self.global_input_scale * scale[None, :, None, None]
        qf = torch.cat((torch.softmax(qz, -1), torch.softmax(-qz, -1)), -1).clamp_min(1e-6)
        kf = torch.cat((torch.softmax(kz, -1), torch.softmax(-kz, -1)), -1).clamp_min(1e-6) * mask
        vf = v * mask
        memory = torch.einsum("bhtf,bhtd->bhfd", kf, vf)
        normalizer = kf.sum(2)
        numerator = torch.einsum("bhtf,bhfd->bhtd", qf, memory)
        denominator = torch.einsum("bhtf,bhf->bht", qf, normalizer).unsqueeze(-1)
        return numerator / denominator.clamp_min(1e-6)

    def forward(self, hidden_states, position_embeddings=None,
                attention_mask=None, position_ids=None, sliding_window_mask=None,
                **kwargs):
        b, t, _ = hidden_states.shape
        qkv = self.Wqkv(hidden_states).reshape(
            b, t, 3, self.num_heads, self.head_dim
        )
        q, k, v = qkv.transpose(3, 1).unbind(dim=2)
        if position_embeddings is None:
            cos, sin = self.rotary_emb(qkv, position_ids=position_ids)
        else:
            cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        valid = _valid_tokens(attention_mask, hidden_states)
        mask = valid[:, None, :, None].to(v.dtype)
        q, k, v = q * mask, k * mask, v * mask
        temperature = self.local_log_temperature.clamp(-0.4, 0.4).exp()
        local = self._local_attention(
            q * temperature[None, :, None, None], k, v, valid
        )
        landmarks = self._landmarks(q, k, v, valid)
        global_linear = self._global_linear(q, k, v, mask) * mask
        qnorm = F.normalize(q, dim=-1)
        dynamic = torch.einsum("bhtd,hkd->bhtk", qnorm, self.route_query)
        mixture = torch.softmax(
            self.route_logits.to(q.dtype)[None, :, None, :] + dynamic, dim=-1
        )
        fused = (mixture[..., 0:1] * local
                 + mixture[..., 1:2] * landmarks
                 + mixture[..., 2:3] * global_linear
                 + mixture[..., 3:4] * v)
        local_conv = _depthwise_sequence_conv(v, self.local_weight)
        conv_weight = torch.sigmoid(self.conv_gain).to(v.dtype)[None, :, None, None]
        fused = (fused + conv_weight * (local_conv - v)) * mask
        self.last_core_output = fused
        flat = fused.transpose(1, 2).reshape(b, t, self.hidden_size).to(hidden_states.dtype)
        return self.out_drop(self.Wo(flat)), None


class StandaloneDecisionModel(nn.Module):
    """Original Laya typed-decision head without a project-library dependency."""
    def __init__(self, encoder, head_layers=2, n_act=2, dropout=0.1):
        super().__init__()
        self.encoder = encoder
        d = int(encoder.config.hidden_size)
        nhead = max(1, d // 64)
        if head_layers > 0:
            layer = nn.TransformerEncoderLayer(
                d, nhead, dim_feedforward=4*d, dropout=dropout,
                batch_first=True, norm_first=True
            )
            self.head = nn.TransformerEncoder(
                layer, int(head_layers), enable_nested_tensor=False
            )
        else:
            self.head = None
        self.type_emb = nn.Embedding(3, d)
        self.scorer = nn.Sequential(
            nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1)
        )
        self.act_head = nn.Sequential(
            nn.Linear(d+4, 256), nn.GELU(), nn.Linear(256, int(n_act))
        )
        self.register_buffer("temperature", torch.ones(3))

    def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype):
        h = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
        h = h + self.type_emb(qtype)[:, None, :]
        if self.head is not None:
            h = self.head(h, src_key_padding_mask=~attention_mask.bool())
        index = marker_pos.clamp_min(0)[:, :, None].expand(-1, -1, h.size(-1))
        marker_hidden = torch.gather(h, 1, index)
        logits = self.scorer(marker_hidden).squeeze(-1).float()
        logits = logits.masked_fill(~marker_mask, -1e4)
        p = torch.softmax(logits.detach(), -1)
        k = marker_mask.sum(-1).clamp(min=2).float()
        ent = -(p*torch.log(p.clamp_min(1e-9))).sum(-1) / torch.log(k)
        if p.shape[-1] >= 2:
            top2 = p.topk(2, -1).values
        else:
            top1 = p.topk(1, -1).values
            top2 = torch.cat((top1, torch.zeros_like(top1)), -1)
        features = torch.stack(
            (top2[:, 0], top2[:, 0]-top2[:, 1], ent, k/255.0), -1
        )
        act_logits = self.act_head(torch.cat((h[:, 0].float(), features), -1))
        return logits, act_logits


def _build_global_model(root: str | Path, manifest: dict):
    """Generic exporter protocol: build_model(bundle_directory, manifest)."""
    root = Path(root)
    if manifest["architecture"] != "laya-v29-fastwindow":
        raise ValueError("Expected laya-v29-fastwindow, got " + str(manifest["architecture"]))
    if manifest["runtime_format"] != 1:
        raise ValueError("Unsupported standalone bundle format")
    cfg = AutoConfig.from_pretrained(str(root / "encoder"))
    _apply_rope_compat(cfg)
    cfg.reference_compile = False
    encoder = AutoModel.from_config(cfg, attn_implementation="sdpa")
    model = StandaloneDecisionModel(
        encoder, head_layers=int(manifest["head_layers"]),
        n_act=int(manifest["n_act"]),
    )
    for index in manifest["global_layers"]:
        original = model.encoder.layers[index].attn
        if tuple(original.local_attention) != (-1, -1):
            raise ValueError(f"Layer {index} is not full/global attention")
        model.encoder.layers[index].attn = FastWindowAttention(
            original,
            local_window=int(manifest["window"]),
            landmark_stride=int(manifest["landmark_stride"]),
            feature_dim=int(manifest["feature_dim"]),
            max_landmarks=int(manifest["max_landmarks"]),
        )
    return model


QTYPES = {"choice": 0, "score": 1, "noul": 2}

def normalize_question(question: Mapping[str, Any]):
    qtype = str(question.get("type", question.get("t", "choice"))).lower()
    if qtype not in QTYPES:
        raise ValueError("Unsupported question type: " + qtype)
    criteria = question.get("criteria", question.get("crit"))
    if qtype == "choice":
        if isinstance(criteria, (list, tuple)):
            criteria = {str(x): None for x in criteria}
        if not isinstance(criteria, Mapping) or len(criteria) < 2:
            raise ValueError("Choice requires >=2 options")
        criteria = dict(criteria)
    elif qtype == "score":
        if not isinstance(criteria, (list, tuple)) or len(criteria) < 2:
            raise ValueError("Score requires >=2 levels")
        criteria = list(criteria)
    else:
        if criteria is None:
            criteria = {}
        if not isinstance(criteria, Mapping):
            raise ValueError("Noul criteria must be a mapping")
        criteria = {str(k).lower(): v for k, v in criteria.items()}
    ins = question.get("instructions", question.get("ins", ""))
    if not isinstance(ins, str):
        ins = json.dumps(ins, ensure_ascii=False, default=str)
    result = {"t": qtype, "ins": ins, "crit": criteria}
    if "labels" in question:
        result["labels"] = dict(question["labels"])
    return result


def render_options(q):
    q = normalize_question(q) if "t" not in q else dict(q)
    criteria = q["crit"]
    def render(v):
        return v if isinstance(v, str) else json.dumps(v, ensure_ascii=False, default=str)
    if q["t"] == "choice":
        return [str(k) if v is None or v == "" else f"{k}: {render(v)}"
                for k,v in criteria.items()]
    if q["t"] == "score":
        return [f"level {i}: {render(v)}" for i,v in enumerate(criteria)]
    labels = q.get("labels", {"false": "false", "true": "true"})
    return [
        str(labels.get("false", "false"))+": "+(
            render(criteria.get("false")) if criteria.get("false") not in (None,"")
            else "no, the statement does not hold"),
        str(labels.get("true", "true"))+": "+(
            render(criteria.get("true")) if criteria.get("true") not in (None,"")
            else "yes, the statement holds"),
    ]


def build_sequence(tokenizer, state, question, *, max_len=512, head_max_len=192,
                   truncate_left=False, state_ids=None):
    q = normalize_question(question) if "t" not in question else dict(question)
    options = render_options(q)
    mask = tokenizer.mask_token
    ins = str(q["ins"]).replace(mask, " ")
    head_ids = tokenizer(
        f"{q['t']} question: {ins}", add_special_tokens=False
    )["input_ids"]
    option_ids = []
    for option in options:
        ids = tokenizer(
            " "+option.replace(mask, " "), add_special_tokens=False,
            truncation=True, max_length=48
        )["input_ids"]
        option_ids.append([tokenizer.mask_token_id] + ids)
    budget = head_max_len - sum(len(x) for x in option_ids)
    if budget < 16:
        per = max(4, (head_max_len - 16) // max(1, len(option_ids)))
        option_ids = [x[:per] for x in option_ids]
        budget = head_max_len - sum(len(x) for x in option_ids)
    head_ids = head_ids[:max(8, budget)]
    ids = [tokenizer.cls_token_id] + head_ids + [tokenizer.sep_token_id]
    markers = []
    for opt in option_ids:
        markers.append(len(ids))
        ids.extend(opt)
    ids.append(tokenizer.sep_token_id)
    if state_ids is None:
        payload = state if isinstance(state, str) else json.dumps(
            state, ensure_ascii=False, separators=(",", ":"), default=str
        )
        state_ids = tokenizer(
            payload.replace(mask, " "), add_special_tokens=False
        )["input_ids"]
    room = max(0, max_len-len(ids)-1)
    state_part = list(state_ids[-room:] if truncate_left and room else
                      state_ids[:room])
    ids = (ids+state_part+[tokenizer.sep_token_id])[:max_len]
    markers = [i for i in markers if i < max_len]
    return ids, markers


def predict(model, tokenizer, state, question, *, max_len=512, device=None):
    """Simple task-level inference on a deployed standalone bundle."""
    q = normalize_question(question)
    ids, markers = build_sequence(tokenizer, state, q, max_len=max_len)
    if not markers:
        raise ValueError("No option markers survived truncation")
    parameter = next(model.parameters())
    device = device or parameter.device
    data = {
        "input_ids": torch.tensor([ids], dtype=torch.long, device=device),
        "attention_mask": torch.ones((1, len(ids)), dtype=torch.long, device=device),
        "marker_pos": torch.tensor([markers], dtype=torch.long, device=device),
        "marker_mask": torch.ones((1, len(markers)), dtype=torch.bool, device=device),
        "qtype": torch.tensor([QTYPES[q["t"]]], dtype=torch.long, device=device),
    }
    model.eval()
    with torch.inference_mode():
        scores, actions = model(**data)
    values = torch.softmax(scores.float(), -1)[0].cpu().tolist()
    choice = int(torch.tensor(values).argmax())
    return {
        "choice_index": choice,
        "choice": render_options(q)[choice],
        "probabilities": values,
        "action_logits": actions.float()[0].cpu().tolist(),
    }


# Bundled standalone sliding-fusion implementation
"""ModernBERT SlidingWindowFusion: exact finite-window SDPA plus gated CeNN-style local mixers.

An experiment in approximating the *sliding* ModernBERT layers (not the global
layers). It preserves pretrained QKV, RoPE and output projection. The 64-radius
configuration retains the teacher's 129-position local receptive field, but
remains windowed attention and is NOT an attention-free replacement.
"""

import copy
import math

import torch
from torch import nn, Tensor
import torch.nn.functional as F
from transformers.models.modernbert.modeling_modernbert import apply_rotary_pos_emb


def valid_tokens(mask: Tensor | None, hidden: Tensor) -> Tensor:
    b, t = hidden.shape[:2]
    if mask is None:
        return torch.ones(b, t, dtype=torch.bool, device=hidden.device)
    if mask.ndim == 2:
        return mask.bool()
    if mask.ndim == 4:
        allowed = mask if mask.dtype == torch.bool else mask > -1e4
        return (allowed.any(-1).any(1) & allowed.any(-2).any(1)).bool()
    raise ValueError(f"Unrecognized attention mask shape {tuple(mask.shape)}")


def _identity_depthwise(channels: int, kernel: int) -> Tensor:
    weight = torch.zeros(channels, 1, kernel)
    weight[:, 0, kernel // 2] = 1.0
    return weight


def local_sliding_sdpa(
    q: Tensor, k: Tensor, v: Tensor, valid: Tensor, *,
    radius: int, block_size: int = 64
) -> Tensor:
    """O(T * (block_size + 2*radius)) memory local SDPA, exact ±radius semantics.

    Padding keys are excluded, and invalid query outputs are zero. Use this
    kernel as an alternative to the teacher's full T-by-T additive mask.
    """
    b, h, t, d = q.shape
    w = int(radius)
    bs = int(block_size)
    if w < 0 or bs < 1:
        raise ValueError("radius must be >=0, block_size >=1")
    blocks = math.ceil(t / bs)
    tail = blocks * bs - t
    # [B,H,blocks,BS,D] and [B,H,blocks,BS+2R,D]
    qp = F.pad(q, (0, 0, 0, tail)).reshape(b, h, blocks, bs, d)
    def keys_in_blocks(x):
        return F.pad(x, (0, 0, w, w + tail)).unfold(
            2, bs + 2*w, bs
        ).transpose(-1, -2)
    kp, vp = keys_in_blocks(k), keys_in_blocks(v)
    # SDPA bool mask: True means this key is permitted.
    kv = F.pad(valid, (w, w + tail)).unfold(1, bs + 2*w, bs)
    qi = torch.arange(bs, device=q.device)[:, None]
    ki = torch.arange(bs + 2*w, device=q.device)[None, :]
    band = (ki >= qi) & (ki <= qi + 2*w)
    allowed = (band[None, None, None, :, :]
               & kv[:, None, :, None, :])
    def flatten_blocks(x):
        return x.transpose(1, 2).reshape(
            b * blocks, h, x.shape[-2], d
        )
    allowed = allowed.transpose(1, 2).reshape(
        b * blocks, 1, bs, bs + 2*w
    )
    out = F.scaled_dot_product_attention(
        flatten_blocks(qp), flatten_blocks(kp), flatten_blocks(vp),
        attn_mask=allowed, dropout_p=0.0
    )
    out = out.reshape(b, blocks, h, bs, d).transpose(1, 2)
    return (out.reshape(b, h, blocks*bs, d)[:, :, :t]
            * valid[:, None, :, None].to(out.dtype))


def grouped_value_conv(v: Tensor, weight: Tensor, dilation: int) -> Tensor:
    b, h, t, d = v.shape
    k = weight.shape[-1]
    padding = dilation * (k - 1) // 2
    z = v.permute(0, 1, 3, 2).reshape(b, h*d, t)
    z = F.conv1d(z, weight.to(z.dtype), padding=padding,
                 dilation=dilation, groups=h*d)
    return z.reshape(b, h, d, t).permute(0, 1, 3, 2).contiguous()


class SlidingWindowFusion(nn.Module):
    """Learned per-head, per-token mixture of exact local window and local convs.

    A ±64 window matches ModernBERT's default native local attention radius.
    For radii <64, the context is genuinely truncated; no hidden teacher-mask
    attention path is available. Multiscale convolutions may recover some context.

    Fixed pretrained QKV/RoPE, tiny Wo calibration, FP32-master route weights.
    """

    architecture = "sliding_window_fusion_v30"

    def __init__(self, original: nn.Module, *,
                 radius: int = 64, block_size: int = 64):
        super().__init__()
        if not hasattr(original, "local_attention"):
            raise TypeError("Expected original ModernBERT attention")
        a, z = tuple(original.local_attention)
        if (a, z) == (-1, -1) or a != z:
            raise ValueError("SlidingWindowFusion accepts sliding attention only")
        if not (0 <= radius <= a):
            raise ValueError(f"Candidate radius {radius} exceeds native ±{a}")
        self.config = original.config
        self.layer_idx = int(getattr(original, "layer_id", -1))
        self.num_heads = int(original.num_heads)
        self.head_dim = int(original.head_dim)
        self.hidden_size = self.num_heads * self.head_dim
        self.radius = int(radius)
        self.native_radius = int(a)
        self.block_size = int(block_size)
        self.Wqkv = copy.deepcopy(original.Wqkv)
        self.Wo = copy.deepcopy(original.Wo)
        self.out_drop = copy.deepcopy(original.out_drop)
        self.rotary_emb = copy.deepcopy(original.rotary_emb)
        self.Wqkv.requires_grad_(False)
        self.rotary_emb.requires_grad_(False)
        channels = self.num_heads * self.head_dim
        self.short_weight = nn.Parameter(_identity_depthwise(channels, 7))
        self.dilated_weight = nn.Parameter(_identity_depthwise(channels, 9))
        self.wide_weight = nn.Parameter(_identity_depthwise(channels, 9))
        # 99.7% nominal route mass on the exact local attention path.
        # A small but nonzero residual path can be learned when W is narrowed.
        self.route_logits = nn.Parameter(
            torch.tensor([4.0, -3.0, -3.0, -3.0]).repeat(self.num_heads, 1)
        )
        self.route_query = nn.Parameter(
            torch.zeros(self.num_heads, 4, self.head_dim)
        )
        self.local_log_temperature = nn.Parameter(torch.zeros(self.num_heads))

    def forward(self, hidden_states: Tensor, position_embeddings=None,
                attention_mask=None, sliding_window_mask=None,
                position_ids=None, **kwargs):
        b, t, _ = hidden_states.shape
        qkv = self.Wqkv(hidden_states).reshape(
            b, t, 3, self.num_heads, self.head_dim
        )
        q, k, v = qkv.transpose(3, 1).unbind(2)
        if position_embeddings is None:
            cos, sin = self.rotary_emb(qkv, position_ids=position_ids)
        else:
            cos, sin = position_embeddings
        q, k = apply_rotary_pos_emb(q, k, cos, sin)
        valid = valid_tokens(attention_mask, hidden_states)
        key_mask = valid[:, None, :, None].to(v.dtype)
        # Distinct from native dense mask: no T×T tensor is constructed.
        q, k, v = q * key_mask, k * key_mask, v * key_mask
        temp = self.local_log_temperature.clamp(-0.25, 0.25).exp()
        local = local_sliding_sdpa(
            q * temp[None, :, None, None], k, v, valid,
            radius=self.radius, block_size=self.block_size
        )
        c7 = grouped_value_conv(v, self.short_weight, dilation=1)
        c9 = grouped_value_conv(v, self.dilated_weight, dilation=4)
        cwide = grouped_value_conv(v, self.wide_weight, dilation=8)
        # Dynamic, per-token routing; all branches share the current true
        # input features, so no teacher outputs leak into the student.
        qnorm = F.normalize(q, dim=-1)
        dynamic = torch.einsum("bhtd,hkd->bhtk", qnorm, self.route_query)
        gate = F.softmax(
            self.route_logits.to(q.dtype)[None, :, None, :] + dynamic, -1
        )
        fused = (gate[..., 0:1]*local
                 + gate[..., 1:2]*c7
                 + gate[..., 2:3]*c9
                 + gate[..., 3:4]*cwide) * key_mask
        out = self.Wo(fused.transpose(1, 2).reshape(
            b, t, self.hidden_size
        ).to(hidden_states.dtype))
        return (self.out_drop(out),)

    def config_dict(self):
        return {
            "architecture": self.architecture,
            "layer_idx": self.layer_idx,
            "radius": self.radius,
            "native_radius": self.native_radius,
            "block_size": self.block_size,
            "window_positions": 2*self.radius+1,
            "teacher_window_positions": 2*self.native_radius+1,
            "uses_local_sdpa": True,
            "attention_free": False,
            "full_context_dense_attention": False,
            "receptive_field_conv": [7, 33, 65],
        }



# ============================================================================
# Mixed-runtime reconstruction contract consumed by standalone_hf_loader.py
# ============================================================================
def build_model(root: str | Path, manifest: dict):
    """Build a COMPLETE fully converted model from ordinary config and manifest."""
    if manifest.get("architecture") != "laya-v31-all-attention":
        raise ValueError("Unsupported mixed-runtime architecture")
    if manifest.get("runtime_format") != 1:
        raise ValueError("Unsupported standalone bundle protocol")
    global_layers = [int(x) for x in manifest["global_layers"]]
    sliding_layers = [int(x) for x in manifest["sliding_layers"]]
    if set(global_layers) & set(sliding_layers):
        raise ValueError("Overlap between global and sliding layer indices")
    global_manifest = dict(manifest)
    global_manifest["architecture"] = "laya-v29-fastwindow"
    model = _build_global_model(root, global_manifest)
    if sorted(global_layers + sliding_layers) != list(range(len(model.encoder.layers))):
        raise ValueError("All encoder attention layers must be replaced exactly once")
    for i in sliding_layers:
        original = model.encoder.layers[i].attn
        if tuple(original.local_attention) == (-1, -1):
            raise ValueError(f"Sliding layer {i} is actually global")
        model.encoder.layers[i].attn = SlidingWindowFusion(
            original,
            radius=int(manifest["sliding_radius"]),
            block_size=int(manifest["sliding_block_size"]),
        )
    return model


# --- Jev_CeNN standalone optional Triton kernel ---
"""Jev_CeNN optional CUDA sliding-attention backend.

A single Triton program processes one query/head, evaluates exactly the
centered ±radius QK attention scores, and reduces V without building global
T×T masks, unfolding halo blocks, or allocating attention probabilities.

The kernel is an inference fast path ONLY. Training/autograd uses the
existing differentiable PyTorch block-SDPA implementation. No parameter names,
mathematics, layer type, checkpoints or Hugging Face weights are changed.

The optional import is deliberately lazy/guarded. CPU and environments without
Triton use the passed reference implementation unconditionally.
"""


import math
import torch

try:
    import triton
    import triton.language as tl
except ImportError:  # Portable HF bundles must still work without Triton.
    triton = None
    tl = None


if triton is not None:
    @triton.jit
    def _sliding_forward(
        Q, K, V, MASK, OUT,
        H:tl.constexpr,T:tl.constexpr,D:tl.constexpr,R:tl.constexpr,
        QB:tl.constexpr,QH:tl.constexpr,QT:tl.constexpr,QD:tl.constexpr,
        KB:tl.constexpr,KH:tl.constexpr,KT:tl.constexpr,KD:tl.constexpr,
        VB:tl.constexpr,VH:tl.constexpr,VT:tl.constexpr,VD:tl.constexpr,
        MB:tl.constexpr,MT:tl.constexpr,SCALE:tl.constexpr,
        N:tl.constexpr,DB:tl.constexpr,
    ):
        qi=tl.program_id(0)
        bh=tl.program_id(1)
        b=bh//H
        h=bh%H
        dd=tl.arange(0,DB)
        ks=qi-R+tl.arange(0,N)
        kv=(ks>=0)&(ks<T)&(tl.abs(ks-qi)<=R)
        valid_k=tl.load(MASK+b*MB+ks*MT,mask=kv,other=0)!=0
        q=tl.load(Q+b*QB+h*QH+qi*QT+dd*QD,mask=dd<D,other=0).to(tl.float32)
        keys=tl.load(
            K+b*KB+h*KH+ks[:,None]*KT+dd[None,:]*KD,
            mask=kv[:,None]&(dd[None,:]<D),other=0
        ).to(tl.float32)
        score=tl.sum(keys*q[None,:],axis=1)*SCALE
        score=tl.where(kv&valid_k,score,-1.0e6)
        top=tl.max(score,axis=0)
        weights=tl.exp(score-top)*((kv&valid_k).to(tl.float32))
        denominator=tl.sum(weights,axis=0)
        values=tl.load(
            V+b*VB+h*VH+ks[:,None]*VT+dd[None,:]*VD,
            mask=kv[:,None]&(dd[None,:]<D),other=0
        ).to(tl.float32)
        result=tl.sum(weights[:,None]*values,axis=0)/tl.maximum(denominator,1.0e-20)
        valid_q=tl.load(MASK+b*MB+qi*MT)!=0
        result=tl.where(valid_q,result,0.0)
        tl.store(OUT+((b*H+h)*T+qi)*D+dd,result,mask=dd<D)


def fast_sliding_sdpa(q,k,v,valid,*,radius,block_size=64,reference=None):
    """Same API/shape as original local_sliding_sdpa.

    Accepts all native q/k/v strides (including ModernBERT's interleaved
    QKV views). Never detaches gradients: when needed, delegates to the
    exact reference implementation. Triton errors are deliberately *not*
    swallowed: the Colab preflight must expose failures before deployment.
    """
    if reference is None:
        raise ValueError("reference=local_sliding_sdpa must be supplied")
    supported=(
        triton is not None
        and q.is_cuda and k.is_cuda and v.is_cuda
        and q.dtype in (torch.float16,torch.bfloat16,torch.float32)
        and q.dtype==k.dtype==v.dtype
        and q.ndim==4 and q.shape==k.shape==v.shape
        and valid.ndim==2 and valid.shape==(q.shape[0],q.shape[2])
        and int(radius)>=0 and int(radius)<=256
    )
    # Keep the original, differentiable SDPA route for every training call.
    if not supported or (torch.is_grad_enabled() and
                        (q.requires_grad or k.requires_grad or v.requires_grad)):
        return reference(q,k,v,valid,radius=radius,block_size=block_size)
    b,h,t,d=q.shape
    if d>128:
        return reference(q,k,v,valid,radius=radius,block_size=block_size)
    radius=int(radius)
    output=torch.empty((b,h,t,d),dtype=q.dtype,device=q.device)
    _sliding_forward[(t,b*h)](
        q,k,v,valid,output,
        h,t,d,radius,
        *q.stride(),*k.stride(),*v.stride(),
        *valid.stride(),
        1.0/math.sqrt(d),
        triton.next_power_of_2(2*radius+1),
        triton.next_power_of_2(d),
        num_warps=4,
    )
    return output


def install_fast_sliding_backend(module, *, reference=None):
    """Enable on an existing module namespace, without changing state_dicts.

    Example: import tinycenn_lm.sliding_window_fusion as swf
             install_fast_sliding_backend(swf)
    """
    if reference is None:
        reference=getattr(module,"_jev_original_sliding_sdpa",None)
    if reference is None:
        reference=module.local_sliding_sdpa
    if not hasattr(module,"_jev_original_sliding_sdpa"):
        module._jev_original_sliding_sdpa=reference
    def accelerated(q,k,v,valid,*,radius,block_size=64):
        return fast_sliding_sdpa(
            q,k,v,valid,radius=radius,block_size=block_size,
            reference=module._jev_original_sliding_sdpa
        )
    module.local_sliding_sdpa=accelerated
    return bool(triton is not None and torch.cuda.is_available())


def uninstall_fast_sliding_backend(module):
    if hasattr(module,"_jev_original_sliding_sdpa"):
        module.local_sliding_sdpa=module._jev_original_sliding_sdpa


def triton_available():
    return triton is not None and torch.cuda.is_available()

# Block-SDPA PyTorch was faster on target GPU.

_jev_original_build_model = build_model
def build_model(root,manifest):
    conf=dict(manifest)
    if conf.get("architecture")=="Jev_CeNN":
        conf["architecture"]="laya-v31-all-attention"
    return _jev_original_build_model(root,conf)