| """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)) |
| |
| 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: |
| |
| |
| 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(), |
| } |
|
|
|
|
| |
| """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 |
| |
| 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) |
| |
| 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)) |
| |
| |
| 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) |
| |
| 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) |
| |
| |
| 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], |
| } |
|
|
|
|
|
|
| |
| |
| |
| 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 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: |
| 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 |
| ) |
| |
| 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() |
|
|
| |
|
|
| _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) |
|
|
