Clean up: Delete 40 training checkpoint directories

checkpoint-1000/chat_template.jinja

@@ -1,4 +0,0 @@
-{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
-' + message['content'] + '<|im_end|>
-' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
-' }}{% endif %}

checkpoint-1000/config.json

@@ -1,62 +0,0 @@
-{
-  "_use_flash_xsa": true,
-  "architectures": [
-    "SurjoForCausalLM"
-  ],
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_surjo.SurjoConfig",
-    "AutoModel": "modeling_surjo.SurjoForCausalLM",
-    "AutoModelForCausalLM": "modeling_surjo.SurjoForCausalLM"
-  },
-  "coda_layers": 1,
-  "dtype": "float32",
-  "gdn_allow_neg_eigval": false,
-  "gdn_conv_kernel_size": 4,
-  "gdn_key_dim": 64,
-  "gdn_num_v_heads": 8,
-  "gdn_per_xsa": 3,
-  "gdn_value_dim": 64,
-  "gradient_checkpointing": false,
-  "head_dim": 64,
-  "hidden_act": "silu",
-  "hidden_size": 512,
-  "initializer_range": 0.02,
-  "intermediate_size": 1536,
-  "layer_types": [
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "full_attention"
-  ],
-  "loop_noise_std": 0.02,
-  "max_position_embeddings": 2048,
-  "model_type": "surjo",
-  "num_attention_heads": 8,
-  "num_groups": 2,
-  "num_hidden_layers": 10,
-  "num_key_value_heads": 4,
-  "partial_rotary_factor": 1.0,
-  "prelude_layers": 1,
-  "recurrent_layers": 8,
-  "recurrent_passes": 2,
-  "rms_norm_eps": 1e-05,
-  "rope_parameters": {
-    "partial_rotary_factor": 1.0,
-    "rope_theta": 10000.0,
-    "rope_type": "default"
-  },
-  "rope_theta": 10000.0,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.16.1",
-  "use_cache": false,
-  "vocab_size": 32768,
-  "xsa_projection": true
-}

checkpoint-1000/configuration_surjo.py

@@ -1,148 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""Surjo gen-3 model configuration."""
-
-from typing import List, Optional
-from transformers import PretrainedConfig
-
-
-class SurjoConfig(PretrainedConfig):
-    r"""Configuration for the Surjo gen-3 hybrid architecture."""
-
-    model_type = "surjo"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size: int = 50257,
-        hidden_size: int = 384,
-        intermediate_size: int = 1024,
-        num_attention_heads: int = 6,
-        num_key_value_heads: int = 2,
-        head_dim: int = 64,
-        hidden_act: str = "silu",
-        max_position_embeddings: int = 2048,
-        initializer_range: float = 0.02,
-        rms_norm_eps: float = 1e-6,
-        tie_word_embeddings: bool = True,
-        rope_theta: float = 10000.0,
-        attention_dropout: float = 0.0,
-        attention_bias: bool = False,
-        prelude_layers: int = 1,
-        recurrent_layers: int = 8,
-        coda_layers: int = 1,
-        recurrent_passes: int = 2,
-        loop_noise_std: float = 0.02,
-        gdn_per_xsa: int = 3,
-        gdn_key_dim: Optional[int] = None,
-        gdn_value_dim: Optional[int] = None,
-        gdn_num_v_heads: Optional[int] = None,
-        gdn_conv_kernel_size: int = 4,
-        gdn_allow_neg_eigval: bool = False,
-        xsa_projection: bool = True,
-        rope_parameters: Optional[dict] = None,
-        rope_scaling: Optional[dict] = None,
-        partial_rotary_factor: float = 1.0,
-        gradient_checkpointing: bool = False,
-        use_cache: bool = True,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.tie_word_embeddings = tie_word_embeddings
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.attention_bias = attention_bias
-
-        self.prelude_layers = prelude_layers
-        self.recurrent_layers = recurrent_layers
-        self.coda_layers = coda_layers
-        self.recurrent_passes = recurrent_passes
-        self.loop_noise_std = loop_noise_std
-        self.gdn_per_xsa = gdn_per_xsa
-
-        self.gdn_key_dim = gdn_key_dim if gdn_key_dim is not None else self.head_dim
-        self.gdn_value_dim = gdn_value_dim if gdn_value_dim is not None else self.head_dim
-        self.gdn_num_v_heads = (
-            gdn_num_v_heads if gdn_num_v_heads is not None else num_attention_heads
-        )
-        self.gdn_conv_kernel_size = gdn_conv_kernel_size
-        self.gdn_allow_neg_eigval = gdn_allow_neg_eigval
-
-        self.xsa_projection = xsa_projection
-
-        if rope_parameters is not None:
-            self.rope_parameters = dict(rope_parameters)
-        elif rope_scaling is not None:
-            self.rope_parameters = dict(rope_scaling)
-        else:
-            self.rope_parameters = {"rope_type": "default"}
-        self.rope_parameters.setdefault("rope_type", "default")
-        self.rope_parameters.setdefault("rope_theta", rope_theta)
-        if self.rope_parameters.get("rope_type") != "default":
-            self.rope_parameters.setdefault("factor", 1.0)
-        self.rope_scaling = dict(self.rope_parameters)
-        self.partial_rotary_factor = partial_rotary_factor
-
-        self.gradient_checkpointing = gradient_checkpointing
-        self.use_cache = use_cache
-
-        group = gdn_per_xsa + 1
-        if recurrent_layers % group != 0:
-            raise ValueError(
-                f"recurrent_layers ({recurrent_layers}) must be divisible by "
-                f"gdn_per_xsa + 1 = {group}"
-            )
-        if num_attention_heads % num_key_value_heads != 0:
-            raise ValueError(
-                f"num_attention_heads ({num_attention_heads}) must be divisible by "
-                f"num_key_value_heads ({num_key_value_heads})"
-            )
-        if (
-            self.gdn_num_v_heads < num_attention_heads
-            or self.gdn_num_v_heads % num_attention_heads != 0
-        ):
-            raise ValueError(
-                f"gdn_num_v_heads ({self.gdn_num_v_heads}) must be >= "
-                f"num_attention_heads ({num_attention_heads}) and divisible by it"
-            )
-        valid_rope_types = {"default", "linear", "dynamic", "yarn", "longrope", "llama3"}
-        rt = self.rope_parameters.get("rope_type", "default")
-        if rt not in valid_rope_types:
-            raise ValueError(f"rope_type must be one of {valid_rope_types}, got {rt}")
-
-        self.num_groups = recurrent_layers // group
-        self.num_hidden_layers = prelude_layers + recurrent_layers + coda_layers
-        self.layer_types = self._build_layer_types()
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    def _build_layer_types(self) -> List[str]:
-        types = ["full_attention"] * self.prelude_layers
-        for _ in range(self.num_groups):
-            types.extend(["linear_attention"] * self.gdn_per_xsa)
-            types.append("full_attention")
-        types.extend(["full_attention"] * self.coda_layers)
-        assert len(types) == self.num_hidden_layers
-        return types
-
-    @property
-    def surjo_layer_types(self) -> List[str]:
-        return ["xsa" if t == "full_attention" else "gdn" for t in self.layer_types]
-
-    @property
-    def effective_num_layers(self) -> int:
-        return (
-            self.prelude_layers
-            + self.recurrent_layers * self.recurrent_passes
-            + self.coda_layers
-        )

checkpoint-1000/generation_config.json

@@ -1,7 +0,0 @@
-{
-  "_from_model_config": true,
-  "output_attentions": false,
-  "output_hidden_states": false,
-  "transformers_version": "5.16.1",
-  "use_cache": false
-}

checkpoint-1000/model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a5abd20d992673c35432854a3c7b3c0aa29b82981b300637d2e587e1330a57f4
-size 215196728

checkpoint-1000/modeling_surjo.py

@@ -1,926 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""SurjoLabs Surjo Architecture"""
-
-import math
-from typing import Dict, List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-
-from transformers import PreTrainedModel
-from transformers.activations import ACT2FN
-
-try:
-    from transformers import GenerationMixin
-except ImportError:
-    from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-
-try:
-    from .configuration_surjo import SurjoConfig
-except (ImportError, ValueError):
-    from configuration_surjo import SurjoConfig
-
-try:
-    from fla.modules.fused_linear_cross_entropy import (
-        fused_linear_cross_entropy_loss as _fla_flce,
-    )
-except ImportError:
-    try:
-        from fla.modules import fused_linear_cross_entropy_loss as _fla_flce
-    except ImportError:
-        _fla_flce = None
-
-try:
-    from fla.modules.norm import rms_norm as _fla_rms_norm
-except ImportError:
-    try:
-        from fla.modules import rms_norm as _fla_rms_norm
-    except ImportError:
-        _fla_rms_norm = None
-
-try:
-    from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
-    FLA_AVAILABLE = True
-except ImportError:
-    chunk_gdn2 = None
-    fused_recurrent_gdn2 = None
-    FLA_AVAILABLE = False
-
-_USE_FLA_FUSED_CE = True
-_USE_FLA_RMS_NORM = False
-
-if FLA_AVAILABLE:
-
-    @torch.compiler.disable
-    def _fla_chunk(q, k, v, g, b, w, initial_state=None, output_final_state=False,
-                   cu_seqlens=None):
-        return chunk_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
-        )
-
-    @torch.compiler.disable
-    def _fla_recurrent(q, k, v, g, b, w, initial_state=None, output_final_state=False):
-        return fused_recurrent_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-        )
-
-try:
-    from flash_attn import flash_attn_varlen_func
-    FLASH_ATTN_AVAILABLE = True
-except ImportError:
-    flash_attn_varlen_func = None
-    FLASH_ATTN_AVAILABLE = False
-
-
-def _mask_padding_states(hidden_states, padding_mask):
-    if padding_mask is not None:
-        return hidden_states * padding_mask.unsqueeze(-1).to(hidden_states.dtype)
-    return hidden_states
-
-
-def naive_gdn2_reference(q, k, v, g, b, w, initial_state=None):
-    """fp32 reference for GDN-2."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    S = (
-        torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
-        if initial_state is None
-        else initial_state.float()
-    )
-    scale = K ** -0.5
-    outs = []
-    for t in range(T):
-        q_t = F.normalize(q[:, t], dim=-1)
-        k_t = F.normalize(k[:, t], dim=-1)
-        v_t, g_t, b_t, w_t = v[:, t], g[:, t], b[:, t], w[:, t]
-        S = S * torch.exp(g_t).unsqueeze(-1)
-        e = b_t * k_t
-        r = torch.einsum("bhkv,bhk->bhv", S, e)
-        z = w_t * v_t
-        S = S + torch.einsum("bhk,bhv->bhkv", k_t, z - r)
-        outs.append(scale * torch.einsum("bhkv,bhk->bhv", S, q_t))
-    return torch.stack(outs, dim=1), S
-
-
-class SurjoRMSNorm(nn.Module):
-    """FLA Triton RMSNorm when available; else fused ATen F.rms_norm."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, x):
-        if x.dtype != self.weight.dtype:
-            x = x.to(self.weight.dtype)
-        if _USE_FLA_RMS_NORM and _fla_rms_norm is not None and x.is_cuda:
-            return _fla_rms_norm(x, self.weight, None, eps=self.eps)
-        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
-
-
-class ClampedMLP(nn.Module):
-    def __init__(self, config):
-        hidden_size = config.hidden_size
-        intermediate_size = config.intermediate_size
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        return self.down_proj(
-            self.act_fn(self.gate_proj(x).clamp(-15.0, 15.0)) * self.up_proj(x)
-        )
-
-
-class SurjoRotaryEmbedding(nn.Module):
-    """RoPE with frequencies computed in forward() (no stored buffer)."""
-
-    def __init__(self, config, device=None):
-        super().__init__()
-        self.config = config
-        self.rope_type = config.rope_parameters.get("rope_type", "default")
-        self.max_position_embeddings = config.max_position_embeddings
-        self.base = config.rope_parameters.get("rope_theta", config.rope_theta)
-        self.factor = config.rope_parameters.get("factor", 1.0)
-        self.dim = int(config.head_dim * config.partial_rotary_factor)
-
-    def _inv_freq(self, device):
-        inv = 1.0 / (
-            self.base
-            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
-        )
-        if self.rope_type == "linear":
-            inv = inv / self.factor
-        return inv
-
-    @torch.no_grad()
-    def forward(self, position_ids):
-        inv = self._inv_freq(position_ids.device)
-        freqs = torch.einsum("bt,f->btf", position_ids.float(), inv)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        return emb.cos(), emb.sin()
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rope(q, k, cos, sin):
-    cos = cos[:, None, :, :].to(q.dtype)
-    sin = sin[:, None, :, :].to(q.dtype)
-    q = q * cos + rotate_half(q) * sin
-    k = k * cos + rotate_half(k) * sin
-    return q, k
-
-
-def apply_rope_flat(q, k, cos, sin):
-    """RoPE on flat varlen tensors. q, k: [total, H, D]; cos, sin: [B, T, D]."""
-    total = q.shape[0]
-    cos_f = cos.reshape(total, -1)[:, None, :].to(q.dtype)
-    sin_f = sin.reshape(total, -1)[:, None, :].to(q.dtype)
-    q = q * cos_f + rotate_half(q) * sin_f
-    k = k * cos_f + rotate_half(k) * sin_f
-    return q, k
-
-
-@torch.compiler.disable
-def _flash_varlen(q, k, v, cu_seqlens, max_seqlen):
-    """FA2 varlen. cu_seqlens device int32; max_seqlen a constant (no sync)."""
-    return flash_attn_varlen_func(
-        q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
-        dropout_p=0.0, causal=True,
-    )
-
-
-class ShortConvolution(nn.Module):
-    def __init__(self, dim: int, kernel_size: int = 4):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.conv = nn.Conv1d(dim, dim, kernel_size, groups=dim, bias=False)
-
-    def forward(self, x, cache=None, output_state=False):
-        B, T, C = x.shape
-        xt = x.transpose(1, 2)
-        if T == 1 and cache is not None:
-            comb = torch.cat([cache, xt], dim=2)
-            w = self.conv.weight.squeeze(1)
-            out = F.silu(torch.einsum("bck,ck->bc", comb, w).unsqueeze(-1))
-            new = comb[:, :, 1:] if output_state else None
-            return out.transpose(1, 2), new
-        if cache is not None:
-            xt_in = torch.cat([cache, xt], dim=2)
-        else:
-            xt_in = F.pad(xt, (self.kernel_size - 1, 0))
-        out = F.silu(self.conv(xt_in))
-        new = xt_in[:, :, -(self.kernel_size - 1):] if output_state else None
-        return out.transpose(1, 2), new
-
-
-class SurjoCache:
-    def __init__(self):
-        self.kv = {}
-        self.gdn_state = {}
-        self.gdn_conv = {}
-
-    def get_seq_length(self):
-        if not self.kv:
-            return 0
-        return next(iter(self.kv.values()))[0].shape[2]
-
-    @staticmethod
-    def num_xsa_slots(config):
-        return (
-            config.prelude_layers
-            + config.num_groups * config.recurrent_passes
-            + config.coda_layers
-        )
-
-
-class XSAAttention(nn.Module):
-    """Full attention with XSA value-projection subtraction."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.num_kv_groups = self.num_heads // self.num_kv_heads
-        self.head_dim = config.head_dim
-        self.xsa_projection = config.xsa_projection
-        self.dropout = config.attention_dropout
-        self.hidden_size = config.hidden_size
-        self.use_flash = getattr(config, "_use_flash_xsa", False)
-        self.q_proj = nn.Linear(
-            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = SurjoRMSNorm(self.head_dim)
-        self.k_norm = SurjoRMSNorm(self.head_dim)
-
-    def _xsa_projection_step(self, attn, v):
-        dv = (attn.float() * v.float()).sum(dim=-1, keepdim=True)
-        vv = (v.float() * v.float()).sum(dim=-1, keepdim=True).clamp_min(1e-4)
-        scale = (dv / vv).to(attn.dtype)
-        return attn - scale * v
-
-    def forward(self, hidden_states, attention_mask, cache, slot, use_cache,
-                cache_position, cos, sin, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        H, Hkv, D = self.num_heads, self.num_kv_heads, self.head_dim
-
-        q = self.q_proj(hidden_states).view(B, T, H, D).transpose(1, 2)
-        k = self.k_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        use_varlen = (
-            cu_seqlens is not None
-            and cache is None
-            and self.use_flash
-            and FLASH_ATTN_AVAILABLE
-            and hidden_states.is_cuda
-        )
-        if use_varlen:
-            total = B * T
-            q_f = q.transpose(1, 2).reshape(total, H, D).to(torch.bfloat16)
-            k_f = k.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            v_f = v.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            if Hkv != H:
-                k_f = k_f.repeat_interleave(self.num_kv_groups, dim=1)
-                v_f = v_f.repeat_interleave(self.num_kv_groups, dim=1)
-            q_f, k_f = apply_rope_flat(q_f, k_f, cos, sin)
-            cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-            o_f = _flash_varlen(q_f, k_f, v_f, cu_dev, T)
-            attn_out = o_f.view(B, T, H, D).to(hidden_states.dtype)
-            v_xsa = v_f.view(B, T, H, D).to(hidden_states.dtype)
-            if self.xsa_projection:
-                attn_out = self._xsa_projection_step(attn_out, v_xsa)
-            out = attn_out.reshape(B, T, H * D)
-            return self.o_proj(out)
-
-        q, k = apply_rope(q, k, cos, sin)
-        v_cur = v
-        if cache is not None and use_cache:
-            if slot in cache.kv:
-                pk, pv = cache.kv[slot]
-                k = torch.cat([pk, k.to(pk.dtype)], dim=2)
-                v = torch.cat([pv, v.to(pv.dtype)], dim=2)
-            cache.kv[slot] = (k, v)
-        S = k.shape[2]
-        if Hkv != H:
-            k_e = k.repeat_interleave(self.num_kv_groups, dim=1)
-            v_e = v.repeat_interleave(self.num_kv_groups, dim=1)
-            v_cur_e = v_cur.repeat_interleave(self.num_kv_groups, dim=1)
-        else:
-            k_e, v_e, v_cur_e = k, v, v_cur
-
-        is_causal = False
-        mask4 = None
-        if attention_mask is not None:
-            L = attention_mask.shape[-1]
-            if L < S:
-                attention_mask = F.pad(attention_mask, (S - L, 0), value=1)
-            elif L > S:
-                attention_mask = attention_mask[:, -S:]
-            neg = torch.finfo(q.dtype).min
-            m = attention_mask[:, None, None, :].to(q.dtype)
-            if T == S:
-                if T > 1:
-                    causal = torch.ones(T, S, device=q.device, dtype=q.dtype).tril()
-                    allow = causal[None, None] * m
-                    mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                        torch.full_like(allow, neg))
-                else:
-                    mask4 = torch.where(m > 0, torch.zeros_like(m),
-                                        torch.full_like(m, neg))
-            else:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                causal = kv_pos[None, :] <= cache_position[:, None].to(q.device)
-                allow = causal[None, None].to(q.dtype) * m
-                mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                    torch.full_like(allow, neg))
-        elif cu_seqlens is not None and T > 1:
-            neg = torch.finfo(q.dtype).min
-            cu_dev = cu_seqlens.to(q.device, torch.int32)
-            gpos = torch.arange(B * T, device=q.device)
-            doc_id = torch.searchsorted(cu_dev, gpos, right=True) - 1
-            doc_id = doc_id.clamp(min=0).view(B, T)
-            same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
-            idx = torch.arange(T, device=q.device)
-            causal = idx.unsqueeze(0) <= idx.unsqueeze(1)
-            allow = (same & causal).unsqueeze(1)
-            mask4 = torch.where(
-                allow,
-                torch.zeros(B, 1, T, T, dtype=q.dtype, device=q.device),
-                torch.full((B, 1, T, T), neg, dtype=q.dtype, device=q.device),
-            )
-        else:
-            if T == S and T > 1:
-                is_causal = True
-            elif T > 1:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                neg = torch.finfo(q.dtype).min
-                causal = torch.zeros(T, S, device=q.device, dtype=q.dtype)
-                causal = causal.masked_fill(
-                    kv_pos[None, :] > cache_position[:, None], neg
-                )
-                mask4 = causal[None, None, :, :]
-
-        attn = F.scaled_dot_product_attention(
-            q, k_e, v_e, attn_mask=mask4, is_causal=is_causal,
-            dropout_p=self.dropout if self.training else 0.0,
-        )
-
-        if self.xsa_projection:
-            attn = self._xsa_projection_step(attn, v_cur_e)
-
-        out = attn.transpose(1, 2).reshape(B, T, H * D)
-        return self.o_proj(out)
-
-
-class GDN2Attention(nn.Module):
-    """Gated DeltaNet-2 linear attention."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        H = config.num_attention_heads
-        Hv = config.gdn_num_v_heads
-        self.num_heads = H
-        self.num_v_heads = Hv
-        self.gva = Hv // H
-        self.head_k_dim = config.gdn_key_dim
-        self.head_v_dim = config.gdn_value_dim
-        self.key_dim = H * self.head_k_dim
-        self.value_dim = Hv * self.head_v_dim
-        self.conv_kernel = config.gdn_conv_kernel_size
-        self.allow_neg_eigval = config.gdn_allow_neg_eigval
-        hidden = config.hidden_size
-        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.q_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.k_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.v_conv = ShortConvolution(self.value_dim, self.conv_kernel)
-        self.f_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.key_dim, bias=False),
-        )
-        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.A_log = nn.Parameter(
-            torch.log(torch.empty(H, dtype=torch.float32).uniform_(1, 16))
-        )
-        self.A_log._no_weight_decay = True
-        dt = torch.exp(
-            torch.rand(self.key_dim, dtype=torch.float32)
-            * (math.log(0.1) - math.log(0.001))
-            + math.log(0.001)
-        ).clamp(min=1e-4)
-        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
-        self.dt_bias._no_weight_decay = True
-        self.g_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
-        )
-        self.o_norm = SurjoRMSNorm(self.head_v_dim, eps=config.rms_norm_eps)
-        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)
-
-    def forward(self, hidden_states, attention_mask, gdn_state, conv_state,
-                output_state, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != T:
-                attention_mask = (
-                    attention_mask[:, -T:] if attention_mask.shape[-1] > T else None
-                )
-            if attention_mask is not None:
-                hidden_states = _mask_padding_states(hidden_states, attention_mask)
-        cq = conv_state[0] if conv_state is not None else None
-        ck = conv_state[1] if conv_state is not None else None
-        cvs = conv_state[2] if conv_state is not None else None
-        q, nq = self.q_conv(self.q_proj(hidden_states), cq, output_state)
-        k, nk = self.k_conv(self.k_proj(hidden_states), ck, output_state)
-        v, nv = self.v_conv(self.v_proj(hidden_states), cvs, output_state)
-        g = F.softplus(self.f_proj(hidden_states).float() + self.dt_bias)
-        b = self.b_proj(hidden_states).sigmoid()
-        w = self.w_proj(hidden_states).sigmoid()
-        q = q.reshape(B, T, self.num_heads, self.head_k_dim)
-        k = k.reshape(B, T, self.num_heads, self.head_k_dim)
-        v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
-        g = g.view(B, T, self.num_heads, self.head_k_dim)
-        b = b.view(B, T, self.num_heads, self.head_k_dim)
-        w = w.view(B, T, self.num_v_heads, self.head_v_dim)
-        g = -self.A_log.float().exp().view(1, 1, self.num_heads, 1) * g
-        if self.allow_neg_eigval:
-            b = b * 2.0
-        if self.gva > 1:
-            q = q.repeat_interleave(self.gva, dim=2)
-            k = k.repeat_interleave(self.gva, dim=2)
-            g = g.repeat_interleave(self.gva, dim=2)
-            b = b.repeat_interleave(self.gva, dim=2)
-        init = gdn_state.float() if gdn_state is not None else None
-
-        if FLA_AVAILABLE and hidden_states.is_cuda:
-            if T == 1 and init is not None:
-                o, fs = _fla_recurrent(q, k, v, g, b, w, init, output_state)
-            elif cu_seqlens is not None:
-                total = B * T
-                cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-                q_f = q.reshape(1, total, self.num_heads, self.head_k_dim)
-                k_f = k.reshape(1, total, self.num_heads, self.head_k_dim)
-                v_f = v.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                g_f = g.reshape(1, total, self.num_heads, self.head_k_dim)
-                b_f = b.reshape(1, total, self.num_heads, self.head_k_dim)
-                w_f = w.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                o, fs = _fla_chunk(q_f, k_f, v_f, g_f, b_f, w_f, init,
-                                   output_state, cu_dev)
-                o = o.reshape(B, T, self.num_v_heads, self.head_v_dim)
-            else:
-                o, fs = _fla_chunk(q, k, v, g, b, w, init, output_state, None)
-        else:
-            o, fs = naive_gdn2_reference(q.float(), k.float(), v.float(), g,
-                                         b.float(), w.float(), init)
-
-        # Cast back to activation precision
-        o = o.to(hidden_states.dtype)
-
-        fs = fs if output_state else None
-        gate = self.g_proj(hidden_states).view(B, T, self.num_v_heads, self.head_v_dim)
-        o = self.o_norm(o)
-        o = o * torch.sigmoid(gate)
-        o = o.reshape(B, T, self.value_dim)
-        conv_out = (nq, nk, nv) if output_state else None
-        return self.o_proj(o), fs, conv_out
-
-
-def _is_xsa(config, li: int) -> bool:
-    if li < config.prelude_layers:
-        return True
-    if li >= config.prelude_layers + config.recurrent_layers:
-        return True
-    return (li - config.prelude_layers) % (config.gdn_per_xsa + 1) == config.gdn_per_xsa
-
-
-class SurjoDecoderLayer(nn.Module):
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_type = "xsa" if _is_xsa(config, layer_idx) else "gdn"
-        if self.layer_type == "xsa":
-            self.self_attn = XSAAttention(config, layer_idx)
-        else:
-            self.linear_attn = GDN2Attention(config, layer_idx)
-        self.mlp = ClampedMLP(config)
-        self.input_layernorm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SurjoRMSNorm(config.hidden_size,
-                                                     eps=config.rms_norm_eps)
-
-    def forward(self, hidden_states, attention_mask=None, cos=None, sin=None,
-                cu_seqlens=None):
-        h = self.input_layernorm(hidden_states)
-        if self.layer_type == "xsa":
-            out = self.self_attn(h, attention_mask, None, -1, False, None,
-                                 cos, sin, cu_seqlens)
-        else:
-            out, _, _ = self.linear_attn(h, attention_mask, None, None, False,
-                                         cu_seqlens)
-        hidden_states = hidden_states + out
-        h2 = self.post_attention_layernorm(hidden_states)
-        return hidden_states + self.mlp(h2)
-
-
-class SurjoPreTrainedModel(PreTrainedModel):
-    config_class = SurjoConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["SurjoDecoderLayer"]
-    _supports_flash_attn = False
-    _supports_flash_attn_2 = False
-    _supports_sdpa = True
-    _is_stateful = True
-    _can_compile_fullgraph = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, GDN2Attention):
-            module.A_log.data.copy_(
-                torch.log(
-                    torch.empty(module.num_heads, dtype=torch.float32,
-                                device=module.A_log.device).uniform_(1, 16)
-                )
-            )
-            dt = torch.exp(
-                torch.rand(module.key_dim, dtype=torch.float32,
-                           device=module.dt_bias.device)
-                * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
-            ).clamp(min=1e-4)
-            module.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt)))
-
-
-class SurjoModel(SurjoPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList(
-            [SurjoDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
-        )
-        self.norm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary = SurjoRotaryEmbedding(config)
-        self.grad_checkpointing = False
-
-        slot = 0
-        self.xsa_slot_map = {}
-        self.exec_plan = []
-        for li in range(config.prelude_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        for p in range(config.recurrent_passes):
-            for g in range(config.num_groups):
-                base = config.prelude_layers + g * (config.gdn_per_xsa + 1)
-                for j in range(config.gdn_per_xsa):
-                    self.exec_plan.append((self.layers[base + j], base + j, p, -1))
-                li = base + config.gdn_per_xsa
-                self.xsa_slot_map[(li, p)] = slot
-                self.exec_plan.append((self.layers[li], li, p, slot))
-                slot += 1
-        for li in range(config.prelude_layers + config.recurrent_layers,
-                        config.num_hidden_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        self.num_xsa_slots = slot
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, use_cache=None,
-                cache_position=None, return_dict=True,
-                output_hidden_states=False, output_attentions=False,
-                cu_seqlens=None, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        if use_cache is None:
-            use_cache = past_key_values is not None
-        if use_cache and past_key_values is None:
-            past_key_values = SurjoCache()
-        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
-
-        if past_len > 0:
-            if input_ids is not None:
-                n_ids = input_ids.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_ids
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    input_ids = input_ids[:, -n_new:] if n_new <= n_ids else input_ids
-                elif n_ids > past_len:
-                    input_ids = input_ids[:, past_len:]
-                else:
-                    input_ids = input_ids[:, -1:]
-            elif inputs_embeds is not None:
-                n_emb = inputs_embeds.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_emb
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    inputs_embeds = (inputs_embeds[:, -n_new:] if n_new <= n_emb
-                                     else inputs_embeds)
-                elif n_emb > past_len:
-                    inputs_embeds = inputs_embeds[:, past_len:]
-                else:
-                    inputs_embeds = inputs_embeds[:, -1:]
-
-        if inputs_embeds is None:
-            if input_ids is None:
-                raise ValueError("Specify input_ids or inputs_embeds")
-            inputs_embeds = self.embed_tokens(input_ids)
-        B, T = inputs_embeds.shape[0], inputs_embeds.shape[1]
-
-        cache_position = torch.arange(past_len, past_len + T,
-                                      device=inputs_embeds.device)
-        if position_ids is not None and position_ids.shape[-1] != T:
-            position_ids = None
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0).expand(B, -1)
-        cos, sin = self.rotary(position_ids)
-
-        gdn_states = dict(past_key_values.gdn_state) if past_key_values is not None else {}
-        conv_states = dict(past_key_values.gdn_conv) if past_key_values is not None else {}
-
-        noise_std = self.config.loop_noise_std
-        hidden = inputs_embeds
-        prev_p = 0
-        use_ckpt = self.training and self.grad_checkpointing
-        all_hidden_states = () if output_hidden_states else None
-
-        for (layer, li, p, slot) in self.exec_plan:
-            if output_hidden_states:
-                all_hidden_states += (hidden,)
-            if p != prev_p and p > 0 and self.training and noise_std > 0.0:
-                hidden = hidden + torch.randn_like(hidden) * noise_std
-            prev_p = p
-
-            if use_ckpt:
-                hidden = torch.utils.checkpoint.checkpoint(
-                    self._train_layer, layer, hidden, attention_mask, cos, sin,
-                    cu_seqlens, use_reentrant=False,
-                )
-            elif use_cache:
-                h = layer.input_layernorm(hidden)
-                if layer.layer_type == "xsa":
-                    out = layer.self_attn(h, attention_mask, past_key_values,
-                                          slot, use_cache, cache_position,
-                                          cos, sin, cu_seqlens)
-                else:
-                    gs = gdn_states.get((li, p))
-                    cs = conv_states.get((li, p))
-                    out, ngs, ncs = layer.linear_attn(h, attention_mask, gs, cs,
-                                                      use_cache, cu_seqlens)
-                    if ngs is not None:
-                        gdn_states[(li, p)] = ngs
-                    if ncs is not None:
-                        conv_states[(li, p)] = ncs
-                hidden = hidden + out
-                h2 = layer.post_attention_layernorm(hidden)
-                hidden = hidden + layer.mlp(h2)
-            else:
-                hidden = layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-        hidden = self.norm(hidden)
-        if output_hidden_states:
-            all_hidden_states += (hidden,)
-
-        if use_cache and past_key_values is not None:
-            past_key_values.gdn_state.update(gdn_states)
-            past_key_values.gdn_conv.update(conv_states)
-
-        if not return_dict:
-            out = tuple(v for v in [hidden, past_key_values if use_cache else None,
-                                    all_hidden_states] if v is not None)
-            return out
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=None,
-        )
-
-    def _train_layer(self, layer, hidden, attention_mask, cos, sin, cu_seqlens):
-        return layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-
-class SurjoForCausalLM(SurjoPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = SurjoModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.post_init()
-
-    def tie_weights(self, **kwargs):
-        try:
-            super().tie_weights(**kwargs)
-        except TypeError:
-            pass
-        except Exception:
-            pass
-        if (getattr(self.config, "tie_word_embeddings", False)
-                and getattr(self, "model", None) is not None
-                and getattr(self, "lm_head", None) is not None):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-    def save_pretrained(self, *args, **kwargs):
-        was_tied = (getattr(self, "lm_head", None) is not None
-                    and getattr(self, "model", None) is not None
-                    and self.lm_head.weight is self.model.embed_tokens.weight)
-        if was_tied and self.config.tie_word_embeddings:
-            self.lm_head.weight = nn.Parameter(
-                self.model.embed_tokens.weight.detach().clone(),
-                requires_grad=self.model.embed_tokens.weight.requires_grad,
-            )
-        try:
-            return super().save_pretrained(*args, **kwargs)
-        finally:
-            if was_tied and self.config.tie_word_embeddings:
-                self.tie_weights()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        self.lm_head = value
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
-        self.model.grad_checkpointing = True
-
-    def gradient_checkpointing_disable(self):
-        self.model.grad_checkpointing = False
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, labels=None,
-                use_cache=None, output_hidden_states=None, output_attentions=None,
-                return_dict=None, logits_to_keep=0, cache_position=None,
-                cu_seqlens=None, **kwargs):
-        if logits_to_keep in (None, 0) and kwargs.get("num_logits_to_keep"):
-            logits_to_keep = kwargs.pop("num_logits_to_keep")
-        return_dict = True if return_dict is None else return_dict
-        if use_cache is None:
-            use_cache = not (self.training or labels is not None)
-
-        outputs = self.model(
-            input_ids=input_ids, attention_mask=attention_mask,
-            position_ids=position_ids, past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds, use_cache=use_cache,
-            cache_position=cache_position,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            cu_seqlens=cu_seqlens,
-            **kwargs,
-        )
-        hidden_states = outputs[0]
-
-        loss, logits = None, None
-        if labels is not None:
-            sh = hidden_states[..., :-1, :].contiguous()
-            sl = labels[..., 1:].contiguous()
-            if _USE_FLA_FUSED_CE and _fla_flce is not None:
-                loss = _fla_flce(
-                    sh.view(-1, sh.shape[-1]),
-                    sl.view(-1),
-                    self.lm_head.weight,
-                    num_chunks=8,
-                    reduction="mean",
-                )
-            else:
-                n_chunks = max(8, sh.shape[0] // 8)
-                tot = hidden_states.new_zeros((), dtype=torch.float32)
-                n_tok = 0
-                for h_c, l_c in zip(sh.chunk(n_chunks, dim=0), sl.chunk(n_chunks, dim=0)):
-                    if l_c.numel() == 0:
-                        continue
-                    lg = self.lm_head(h_c)
-                    tot = tot + F.cross_entropy(
-                        lg.float().view(-1, lg.size(-1)), l_c.view(-1), reduction="sum"
-                    )
-                    n_tok += l_c.numel()
-                loss = tot / max(n_tok, 1)
-        else:
-            if isinstance(logits_to_keep, int) and logits_to_keep > 0:
-                sl_hidden = hidden_states[:, -logits_to_keep:, :]
-            elif torch.is_tensor(logits_to_keep):
-                sl_hidden = hidden_states[:, logits_to_keep, :]
-            else:
-                sl_hidden = hidden_states
-            logits = self.lm_head(sl_hidden)
-
-        if not return_dict:
-            out = (logits,) + outputs[1:]
-            return (loss,) + out if loss is not None else out
-        return CausalLMOutputWithPast(
-            loss=loss, logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
-                                      attention_mask=None, inputs_embeds=None,
-                                      cache_position=None, use_cache=True, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        has_state = past_key_values is not None and len(past_key_values.kv) > 0
-        if has_state:
-            input_ids = input_ids[:, -1:]
-            if kwargs.get("logits_to_keep") is None:
-                kwargs["logits_to_keep"] = 1
-        past_len = past_key_values.get_seq_length() if has_state else 0
-        cache_position = torch.arange(past_len, past_len + input_ids.shape[1],
-                                      device=input_ids.device)
-        model_inputs = {
-            "input_ids": input_ids.contiguous(),
-            "past_key_values": past_key_values,
-            "use_cache": use_cache,
-            "cache_position": cache_position,
-            "attention_mask": attention_mask,
-        }
-        model_inputs.update(kwargs)
-        return model_inputs
-
-    def _reorder_cache(self, past_key_values: SurjoCache, beam_idx):
-        for slot, (k, v) in past_key_values.kv.items():
-            past_key_values.kv[slot] = (k.index_select(0, beam_idx),
-                                        v.index_select(0, beam_idx))
-        for key in list(past_key_values.gdn_state.keys()):
-            s = past_key_values.gdn_state[key]
-            if s is not None:
-                past_key_values.gdn_state[key] = s.index_select(0, beam_idx)
-        for key in list(past_key_values.gdn_conv.keys()):
-            cs = past_key_values.gdn_conv[key]
-            if cs is not None:
-                past_key_values.gdn_conv[key] = tuple(
-                    c.index_select(0, beam_idx) if c is not None else None for c in cs
-                )
-        return past_key_values
-
-
-SurjoConfig.register_for_auto_class()
-SurjoForCausalLM.register_for_auto_class()

checkpoint-1000/optimizer.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:967151e939f37797c229ea14fa72f92a68435418b62251eb097af6488e532fad
-size 148034507

checkpoint-1000/optimizer_1.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:157010ade80de4660d8fe9c541a7e1262cee4c9b3435c396a76abed8ef9c5cde
-size 134405851

checkpoint-1000/random_states_0.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:d40bd742577a21d618dea12387d03f8b61a77aca9c823ed92049267177110e08
-size 14821

checkpoint-1000/tokenizer_config.json

@@ -1,85 +0,0 @@
-{
-  "backend": "tokenizers",
-  "bos_token": "<|bos|>",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|eos|>",
-  "extra_special_tokens": [
-    "<|unk|>",
-    "<|pad|>",
-    "<|bos|>",
-    "<|eos|>",
-    "<|mask|>",
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|system|>",
-    "<|user|>",
-    "<|assistant|>",
-    "<think>",
-    "</think>",
-    "<|begin_of_thought|>",
-    "<|end_of_thought|>",
-    "<answer>",
-    "</answer>",
-    "<|step|>",
-    "<|/step|>",
-    "<context>",
-    "</context>",
-    "<|doc_start|>",
-    "<|doc_end|>",
-    "<|search|>",
-    "<|search_results|>",
-    "<|tool_list_start|>",
-    "<|tool_list_end|>",
-    "<tools>",
-    "</tools>",
-    "<|tool_call_start|>",
-    "<|tool_call_end|>",
-    "<|tool_call|>",
-    "<|/tool_call|>",
-    "<|tool_response_start|>",
-    "<|tool_response_end|>",
-    "<|tool_response|>",
-    "<|/tool_response|>",
-    "<|image|>",
-    "<|image_pad|>",
-    "<|image_placeholder|>",
-    "<|audio|>",
-    "<|audio_pad|>",
-    "<|audio_placeholder|>",
-    "<|video|>",
-    "<|video_pad|>",
-    "<|fim_prefix|>",
-    "<|fim_suffix|>",
-    "<|fim_middle|>",
-    "<|repo_name|>",
-    "<|file_separator|>",
-    "<|reward|>",
-    "<|reserved_0|>",
-    "<|reserved_1|>",
-    "<|reserved_2|>",
-    "<|reserved_3|>",
-    "<|reserved_4|>",
-    "<|reserved_5|>",
-    "<|reserved_6|>",
-    "<|reserved_7|>",
-    "<|reserved_8|>",
-    "<|reserved_9|>",
-    "<|reserved_10|>",
-    "<|reserved_11|>",
-    "<|reserved_12|>",
-    "<|reserved_13|>",
-    "<|reserved_14|>",
-    "<|reserved_15|>",
-    "<|reserved_16|>",
-    "<|reserved_17|>",
-    "<|reserved_18|>",
-    "<|reserved_19|>"
-  ],
-  "is_local": false,
-  "local_files_only": false,
-  "mask_token": "<|mask|>",
-  "model_max_length": 10000000,
-  "pad_token": "<|pad|>",
-  "tokenizer_class": "TokenizersBackend",
-  "unk_token": "<|unk|>"
-}

checkpoint-10000/chat_template.jinja

@@ -1,4 +0,0 @@
-{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
-' + message['content'] + '<|im_end|>
-' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
-' }}{% endif %}

checkpoint-10000/config.json

@@ -1,62 +0,0 @@
-{
-  "_use_flash_xsa": true,
-  "architectures": [
-    "SurjoForCausalLM"
-  ],
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_surjo.SurjoConfig",
-    "AutoModel": "modeling_surjo.SurjoForCausalLM",
-    "AutoModelForCausalLM": "modeling_surjo.SurjoForCausalLM"
-  },
-  "coda_layers": 1,
-  "dtype": "float32",
-  "gdn_allow_neg_eigval": false,
-  "gdn_conv_kernel_size": 4,
-  "gdn_key_dim": 64,
-  "gdn_num_v_heads": 8,
-  "gdn_per_xsa": 3,
-  "gdn_value_dim": 64,
-  "gradient_checkpointing": false,
-  "head_dim": 64,
-  "hidden_act": "silu",
-  "hidden_size": 512,
-  "initializer_range": 0.02,
-  "intermediate_size": 1536,
-  "layer_types": [
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "full_attention"
-  ],
-  "loop_noise_std": 0.02,
-  "max_position_embeddings": 2048,
-  "model_type": "surjo",
-  "num_attention_heads": 8,
-  "num_groups": 2,
-  "num_hidden_layers": 10,
-  "num_key_value_heads": 4,
-  "partial_rotary_factor": 1.0,
-  "prelude_layers": 1,
-  "recurrent_layers": 8,
-  "recurrent_passes": 2,
-  "rms_norm_eps": 1e-05,
-  "rope_parameters": {
-    "partial_rotary_factor": 1.0,
-    "rope_theta": 10000.0,
-    "rope_type": "default"
-  },
-  "rope_theta": 10000.0,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.16.1",
-  "use_cache": false,
-  "vocab_size": 32768,
-  "xsa_projection": true
-}

checkpoint-10000/configuration_surjo.py

@@ -1,148 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""Surjo gen-3 model configuration."""
-
-from typing import List, Optional
-from transformers import PretrainedConfig
-
-
-class SurjoConfig(PretrainedConfig):
-    r"""Configuration for the Surjo gen-3 hybrid architecture."""
-
-    model_type = "surjo"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size: int = 50257,
-        hidden_size: int = 384,
-        intermediate_size: int = 1024,
-        num_attention_heads: int = 6,
-        num_key_value_heads: int = 2,
-        head_dim: int = 64,
-        hidden_act: str = "silu",
-        max_position_embeddings: int = 2048,
-        initializer_range: float = 0.02,
-        rms_norm_eps: float = 1e-6,
-        tie_word_embeddings: bool = True,
-        rope_theta: float = 10000.0,
-        attention_dropout: float = 0.0,
-        attention_bias: bool = False,
-        prelude_layers: int = 1,
-        recurrent_layers: int = 8,
-        coda_layers: int = 1,
-        recurrent_passes: int = 2,
-        loop_noise_std: float = 0.02,
-        gdn_per_xsa: int = 3,
-        gdn_key_dim: Optional[int] = None,
-        gdn_value_dim: Optional[int] = None,
-        gdn_num_v_heads: Optional[int] = None,
-        gdn_conv_kernel_size: int = 4,
-        gdn_allow_neg_eigval: bool = False,
-        xsa_projection: bool = True,
-        rope_parameters: Optional[dict] = None,
-        rope_scaling: Optional[dict] = None,
-        partial_rotary_factor: float = 1.0,
-        gradient_checkpointing: bool = False,
-        use_cache: bool = True,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.tie_word_embeddings = tie_word_embeddings
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.attention_bias = attention_bias
-
-        self.prelude_layers = prelude_layers
-        self.recurrent_layers = recurrent_layers
-        self.coda_layers = coda_layers
-        self.recurrent_passes = recurrent_passes
-        self.loop_noise_std = loop_noise_std
-        self.gdn_per_xsa = gdn_per_xsa
-
-        self.gdn_key_dim = gdn_key_dim if gdn_key_dim is not None else self.head_dim
-        self.gdn_value_dim = gdn_value_dim if gdn_value_dim is not None else self.head_dim
-        self.gdn_num_v_heads = (
-            gdn_num_v_heads if gdn_num_v_heads is not None else num_attention_heads
-        )
-        self.gdn_conv_kernel_size = gdn_conv_kernel_size
-        self.gdn_allow_neg_eigval = gdn_allow_neg_eigval
-
-        self.xsa_projection = xsa_projection
-
-        if rope_parameters is not None:
-            self.rope_parameters = dict(rope_parameters)
-        elif rope_scaling is not None:
-            self.rope_parameters = dict(rope_scaling)
-        else:
-            self.rope_parameters = {"rope_type": "default"}
-        self.rope_parameters.setdefault("rope_type", "default")
-        self.rope_parameters.setdefault("rope_theta", rope_theta)
-        if self.rope_parameters.get("rope_type") != "default":
-            self.rope_parameters.setdefault("factor", 1.0)
-        self.rope_scaling = dict(self.rope_parameters)
-        self.partial_rotary_factor = partial_rotary_factor
-
-        self.gradient_checkpointing = gradient_checkpointing
-        self.use_cache = use_cache
-
-        group = gdn_per_xsa + 1
-        if recurrent_layers % group != 0:
-            raise ValueError(
-                f"recurrent_layers ({recurrent_layers}) must be divisible by "
-                f"gdn_per_xsa + 1 = {group}"
-            )
-        if num_attention_heads % num_key_value_heads != 0:
-            raise ValueError(
-                f"num_attention_heads ({num_attention_heads}) must be divisible by "
-                f"num_key_value_heads ({num_key_value_heads})"
-            )
-        if (
-            self.gdn_num_v_heads < num_attention_heads
-            or self.gdn_num_v_heads % num_attention_heads != 0
-        ):
-            raise ValueError(
-                f"gdn_num_v_heads ({self.gdn_num_v_heads}) must be >= "
-                f"num_attention_heads ({num_attention_heads}) and divisible by it"
-            )
-        valid_rope_types = {"default", "linear", "dynamic", "yarn", "longrope", "llama3"}
-        rt = self.rope_parameters.get("rope_type", "default")
-        if rt not in valid_rope_types:
-            raise ValueError(f"rope_type must be one of {valid_rope_types}, got {rt}")
-
-        self.num_groups = recurrent_layers // group
-        self.num_hidden_layers = prelude_layers + recurrent_layers + coda_layers
-        self.layer_types = self._build_layer_types()
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    def _build_layer_types(self) -> List[str]:
-        types = ["full_attention"] * self.prelude_layers
-        for _ in range(self.num_groups):
-            types.extend(["linear_attention"] * self.gdn_per_xsa)
-            types.append("full_attention")
-        types.extend(["full_attention"] * self.coda_layers)
-        assert len(types) == self.num_hidden_layers
-        return types
-
-    @property
-    def surjo_layer_types(self) -> List[str]:
-        return ["xsa" if t == "full_attention" else "gdn" for t in self.layer_types]
-
-    @property
-    def effective_num_layers(self) -> int:
-        return (
-            self.prelude_layers
-            + self.recurrent_layers * self.recurrent_passes
-            + self.coda_layers
-        )

checkpoint-10000/generation_config.json

@@ -1,7 +0,0 @@
-{
-  "_from_model_config": true,
-  "output_attentions": false,
-  "output_hidden_states": false,
-  "transformers_version": "5.16.1",
-  "use_cache": false
-}

checkpoint-10000/model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7446de985a1bc372cc3b96e80543b79a5082dc2d24c3677f34fd994783fb6a95
-size 215196728

checkpoint-10000/modeling_surjo.py

@@ -1,926 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""SurjoLabs Surjo Architecture"""
-
-import math
-from typing import Dict, List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-
-from transformers import PreTrainedModel
-from transformers.activations import ACT2FN
-
-try:
-    from transformers import GenerationMixin
-except ImportError:
-    from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-
-try:
-    from .configuration_surjo import SurjoConfig
-except (ImportError, ValueError):
-    from configuration_surjo import SurjoConfig
-
-try:
-    from fla.modules.fused_linear_cross_entropy import (
-        fused_linear_cross_entropy_loss as _fla_flce,
-    )
-except ImportError:
-    try:
-        from fla.modules import fused_linear_cross_entropy_loss as _fla_flce
-    except ImportError:
-        _fla_flce = None
-
-try:
-    from fla.modules.norm import rms_norm as _fla_rms_norm
-except ImportError:
-    try:
-        from fla.modules import rms_norm as _fla_rms_norm
-    except ImportError:
-        _fla_rms_norm = None
-
-try:
-    from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
-    FLA_AVAILABLE = True
-except ImportError:
-    chunk_gdn2 = None
-    fused_recurrent_gdn2 = None
-    FLA_AVAILABLE = False
-
-_USE_FLA_FUSED_CE = True
-_USE_FLA_RMS_NORM = False
-
-if FLA_AVAILABLE:
-
-    @torch.compiler.disable
-    def _fla_chunk(q, k, v, g, b, w, initial_state=None, output_final_state=False,
-                   cu_seqlens=None):
-        return chunk_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
-        )
-
-    @torch.compiler.disable
-    def _fla_recurrent(q, k, v, g, b, w, initial_state=None, output_final_state=False):
-        return fused_recurrent_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-        )
-
-try:
-    from flash_attn import flash_attn_varlen_func
-    FLASH_ATTN_AVAILABLE = True
-except ImportError:
-    flash_attn_varlen_func = None
-    FLASH_ATTN_AVAILABLE = False
-
-
-def _mask_padding_states(hidden_states, padding_mask):
-    if padding_mask is not None:
-        return hidden_states * padding_mask.unsqueeze(-1).to(hidden_states.dtype)
-    return hidden_states
-
-
-def naive_gdn2_reference(q, k, v, g, b, w, initial_state=None):
-    """fp32 reference for GDN-2."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    S = (
-        torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
-        if initial_state is None
-        else initial_state.float()
-    )
-    scale = K ** -0.5
-    outs = []
-    for t in range(T):
-        q_t = F.normalize(q[:, t], dim=-1)
-        k_t = F.normalize(k[:, t], dim=-1)
-        v_t, g_t, b_t, w_t = v[:, t], g[:, t], b[:, t], w[:, t]
-        S = S * torch.exp(g_t).unsqueeze(-1)
-        e = b_t * k_t
-        r = torch.einsum("bhkv,bhk->bhv", S, e)
-        z = w_t * v_t
-        S = S + torch.einsum("bhk,bhv->bhkv", k_t, z - r)
-        outs.append(scale * torch.einsum("bhkv,bhk->bhv", S, q_t))
-    return torch.stack(outs, dim=1), S
-
-
-class SurjoRMSNorm(nn.Module):
-    """FLA Triton RMSNorm when available; else fused ATen F.rms_norm."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, x):
-        if x.dtype != self.weight.dtype:
-            x = x.to(self.weight.dtype)
-        if _USE_FLA_RMS_NORM and _fla_rms_norm is not None and x.is_cuda:
-            return _fla_rms_norm(x, self.weight, None, eps=self.eps)
-        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
-
-
-class ClampedMLP(nn.Module):
-    def __init__(self, config):
-        hidden_size = config.hidden_size
-        intermediate_size = config.intermediate_size
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        return self.down_proj(
-            self.act_fn(self.gate_proj(x).clamp(-15.0, 15.0)) * self.up_proj(x)
-        )
-
-
-class SurjoRotaryEmbedding(nn.Module):
-    """RoPE with frequencies computed in forward() (no stored buffer)."""
-
-    def __init__(self, config, device=None):
-        super().__init__()
-        self.config = config
-        self.rope_type = config.rope_parameters.get("rope_type", "default")
-        self.max_position_embeddings = config.max_position_embeddings
-        self.base = config.rope_parameters.get("rope_theta", config.rope_theta)
-        self.factor = config.rope_parameters.get("factor", 1.0)
-        self.dim = int(config.head_dim * config.partial_rotary_factor)
-
-    def _inv_freq(self, device):
-        inv = 1.0 / (
-            self.base
-            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
-        )
-        if self.rope_type == "linear":
-            inv = inv / self.factor
-        return inv
-
-    @torch.no_grad()
-    def forward(self, position_ids):
-        inv = self._inv_freq(position_ids.device)
-        freqs = torch.einsum("bt,f->btf", position_ids.float(), inv)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        return emb.cos(), emb.sin()
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rope(q, k, cos, sin):
-    cos = cos[:, None, :, :].to(q.dtype)
-    sin = sin[:, None, :, :].to(q.dtype)
-    q = q * cos + rotate_half(q) * sin
-    k = k * cos + rotate_half(k) * sin
-    return q, k
-
-
-def apply_rope_flat(q, k, cos, sin):
-    """RoPE on flat varlen tensors. q, k: [total, H, D]; cos, sin: [B, T, D]."""
-    total = q.shape[0]
-    cos_f = cos.reshape(total, -1)[:, None, :].to(q.dtype)
-    sin_f = sin.reshape(total, -1)[:, None, :].to(q.dtype)
-    q = q * cos_f + rotate_half(q) * sin_f
-    k = k * cos_f + rotate_half(k) * sin_f
-    return q, k
-
-
-@torch.compiler.disable
-def _flash_varlen(q, k, v, cu_seqlens, max_seqlen):
-    """FA2 varlen. cu_seqlens device int32; max_seqlen a constant (no sync)."""
-    return flash_attn_varlen_func(
-        q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
-        dropout_p=0.0, causal=True,
-    )
-
-
-class ShortConvolution(nn.Module):
-    def __init__(self, dim: int, kernel_size: int = 4):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.conv = nn.Conv1d(dim, dim, kernel_size, groups=dim, bias=False)
-
-    def forward(self, x, cache=None, output_state=False):
-        B, T, C = x.shape
-        xt = x.transpose(1, 2)
-        if T == 1 and cache is not None:
-            comb = torch.cat([cache, xt], dim=2)
-            w = self.conv.weight.squeeze(1)
-            out = F.silu(torch.einsum("bck,ck->bc", comb, w).unsqueeze(-1))
-            new = comb[:, :, 1:] if output_state else None
-            return out.transpose(1, 2), new
-        if cache is not None:
-            xt_in = torch.cat([cache, xt], dim=2)
-        else:
-            xt_in = F.pad(xt, (self.kernel_size - 1, 0))
-        out = F.silu(self.conv(xt_in))
-        new = xt_in[:, :, -(self.kernel_size - 1):] if output_state else None
-        return out.transpose(1, 2), new
-
-
-class SurjoCache:
-    def __init__(self):
-        self.kv = {}
-        self.gdn_state = {}
-        self.gdn_conv = {}
-
-    def get_seq_length(self):
-        if not self.kv:
-            return 0
-        return next(iter(self.kv.values()))[0].shape[2]
-
-    @staticmethod
-    def num_xsa_slots(config):
-        return (
-            config.prelude_layers
-            + config.num_groups * config.recurrent_passes
-            + config.coda_layers
-        )
-
-
-class XSAAttention(nn.Module):
-    """Full attention with XSA value-projection subtraction."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.num_kv_groups = self.num_heads // self.num_kv_heads
-        self.head_dim = config.head_dim
-        self.xsa_projection = config.xsa_projection
-        self.dropout = config.attention_dropout
-        self.hidden_size = config.hidden_size
-        self.use_flash = getattr(config, "_use_flash_xsa", False)
-        self.q_proj = nn.Linear(
-            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = SurjoRMSNorm(self.head_dim)
-        self.k_norm = SurjoRMSNorm(self.head_dim)
-
-    def _xsa_projection_step(self, attn, v):
-        dv = (attn.float() * v.float()).sum(dim=-1, keepdim=True)
-        vv = (v.float() * v.float()).sum(dim=-1, keepdim=True).clamp_min(1e-4)
-        scale = (dv / vv).to(attn.dtype)
-        return attn - scale * v
-
-    def forward(self, hidden_states, attention_mask, cache, slot, use_cache,
-                cache_position, cos, sin, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        H, Hkv, D = self.num_heads, self.num_kv_heads, self.head_dim
-
-        q = self.q_proj(hidden_states).view(B, T, H, D).transpose(1, 2)
-        k = self.k_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        use_varlen = (
-            cu_seqlens is not None
-            and cache is None
-            and self.use_flash
-            and FLASH_ATTN_AVAILABLE
-            and hidden_states.is_cuda
-        )
-        if use_varlen:
-            total = B * T
-            q_f = q.transpose(1, 2).reshape(total, H, D).to(torch.bfloat16)
-            k_f = k.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            v_f = v.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            if Hkv != H:
-                k_f = k_f.repeat_interleave(self.num_kv_groups, dim=1)
-                v_f = v_f.repeat_interleave(self.num_kv_groups, dim=1)
-            q_f, k_f = apply_rope_flat(q_f, k_f, cos, sin)
-            cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-            o_f = _flash_varlen(q_f, k_f, v_f, cu_dev, T)
-            attn_out = o_f.view(B, T, H, D).to(hidden_states.dtype)
-            v_xsa = v_f.view(B, T, H, D).to(hidden_states.dtype)
-            if self.xsa_projection:
-                attn_out = self._xsa_projection_step(attn_out, v_xsa)
-            out = attn_out.reshape(B, T, H * D)
-            return self.o_proj(out)
-
-        q, k = apply_rope(q, k, cos, sin)
-        v_cur = v
-        if cache is not None and use_cache:
-            if slot in cache.kv:
-                pk, pv = cache.kv[slot]
-                k = torch.cat([pk, k.to(pk.dtype)], dim=2)
-                v = torch.cat([pv, v.to(pv.dtype)], dim=2)
-            cache.kv[slot] = (k, v)
-        S = k.shape[2]
-        if Hkv != H:
-            k_e = k.repeat_interleave(self.num_kv_groups, dim=1)
-            v_e = v.repeat_interleave(self.num_kv_groups, dim=1)
-            v_cur_e = v_cur.repeat_interleave(self.num_kv_groups, dim=1)
-        else:
-            k_e, v_e, v_cur_e = k, v, v_cur
-
-        is_causal = False
-        mask4 = None
-        if attention_mask is not None:
-            L = attention_mask.shape[-1]
-            if L < S:
-                attention_mask = F.pad(attention_mask, (S - L, 0), value=1)
-            elif L > S:
-                attention_mask = attention_mask[:, -S:]
-            neg = torch.finfo(q.dtype).min
-            m = attention_mask[:, None, None, :].to(q.dtype)
-            if T == S:
-                if T > 1:
-                    causal = torch.ones(T, S, device=q.device, dtype=q.dtype).tril()
-                    allow = causal[None, None] * m
-                    mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                        torch.full_like(allow, neg))
-                else:
-                    mask4 = torch.where(m > 0, torch.zeros_like(m),
-                                        torch.full_like(m, neg))
-            else:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                causal = kv_pos[None, :] <= cache_position[:, None].to(q.device)
-                allow = causal[None, None].to(q.dtype) * m
-                mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                    torch.full_like(allow, neg))
-        elif cu_seqlens is not None and T > 1:
-            neg = torch.finfo(q.dtype).min
-            cu_dev = cu_seqlens.to(q.device, torch.int32)
-            gpos = torch.arange(B * T, device=q.device)
-            doc_id = torch.searchsorted(cu_dev, gpos, right=True) - 1
-            doc_id = doc_id.clamp(min=0).view(B, T)
-            same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
-            idx = torch.arange(T, device=q.device)
-            causal = idx.unsqueeze(0) <= idx.unsqueeze(1)
-            allow = (same & causal).unsqueeze(1)
-            mask4 = torch.where(
-                allow,
-                torch.zeros(B, 1, T, T, dtype=q.dtype, device=q.device),
-                torch.full((B, 1, T, T), neg, dtype=q.dtype, device=q.device),
-            )
-        else:
-            if T == S and T > 1:
-                is_causal = True
-            elif T > 1:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                neg = torch.finfo(q.dtype).min
-                causal = torch.zeros(T, S, device=q.device, dtype=q.dtype)
-                causal = causal.masked_fill(
-                    kv_pos[None, :] > cache_position[:, None], neg
-                )
-                mask4 = causal[None, None, :, :]
-
-        attn = F.scaled_dot_product_attention(
-            q, k_e, v_e, attn_mask=mask4, is_causal=is_causal,
-            dropout_p=self.dropout if self.training else 0.0,
-        )
-
-        if self.xsa_projection:
-            attn = self._xsa_projection_step(attn, v_cur_e)
-
-        out = attn.transpose(1, 2).reshape(B, T, H * D)
-        return self.o_proj(out)
-
-
-class GDN2Attention(nn.Module):
-    """Gated DeltaNet-2 linear attention."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        H = config.num_attention_heads
-        Hv = config.gdn_num_v_heads
-        self.num_heads = H
-        self.num_v_heads = Hv
-        self.gva = Hv // H
-        self.head_k_dim = config.gdn_key_dim
-        self.head_v_dim = config.gdn_value_dim
-        self.key_dim = H * self.head_k_dim
-        self.value_dim = Hv * self.head_v_dim
-        self.conv_kernel = config.gdn_conv_kernel_size
-        self.allow_neg_eigval = config.gdn_allow_neg_eigval
-        hidden = config.hidden_size
-        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.q_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.k_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.v_conv = ShortConvolution(self.value_dim, self.conv_kernel)
-        self.f_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.key_dim, bias=False),
-        )
-        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.A_log = nn.Parameter(
-            torch.log(torch.empty(H, dtype=torch.float32).uniform_(1, 16))
-        )
-        self.A_log._no_weight_decay = True
-        dt = torch.exp(
-            torch.rand(self.key_dim, dtype=torch.float32)
-            * (math.log(0.1) - math.log(0.001))
-            + math.log(0.001)
-        ).clamp(min=1e-4)
-        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
-        self.dt_bias._no_weight_decay = True
-        self.g_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
-        )
-        self.o_norm = SurjoRMSNorm(self.head_v_dim, eps=config.rms_norm_eps)
-        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)
-
-    def forward(self, hidden_states, attention_mask, gdn_state, conv_state,
-                output_state, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != T:
-                attention_mask = (
-                    attention_mask[:, -T:] if attention_mask.shape[-1] > T else None
-                )
-            if attention_mask is not None:
-                hidden_states = _mask_padding_states(hidden_states, attention_mask)
-        cq = conv_state[0] if conv_state is not None else None
-        ck = conv_state[1] if conv_state is not None else None
-        cvs = conv_state[2] if conv_state is not None else None
-        q, nq = self.q_conv(self.q_proj(hidden_states), cq, output_state)
-        k, nk = self.k_conv(self.k_proj(hidden_states), ck, output_state)
-        v, nv = self.v_conv(self.v_proj(hidden_states), cvs, output_state)
-        g = F.softplus(self.f_proj(hidden_states).float() + self.dt_bias)
-        b = self.b_proj(hidden_states).sigmoid()
-        w = self.w_proj(hidden_states).sigmoid()
-        q = q.reshape(B, T, self.num_heads, self.head_k_dim)
-        k = k.reshape(B, T, self.num_heads, self.head_k_dim)
-        v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
-        g = g.view(B, T, self.num_heads, self.head_k_dim)
-        b = b.view(B, T, self.num_heads, self.head_k_dim)
-        w = w.view(B, T, self.num_v_heads, self.head_v_dim)
-        g = -self.A_log.float().exp().view(1, 1, self.num_heads, 1) * g
-        if self.allow_neg_eigval:
-            b = b * 2.0
-        if self.gva > 1:
-            q = q.repeat_interleave(self.gva, dim=2)
-            k = k.repeat_interleave(self.gva, dim=2)
-            g = g.repeat_interleave(self.gva, dim=2)
-            b = b.repeat_interleave(self.gva, dim=2)
-        init = gdn_state.float() if gdn_state is not None else None
-
-        if FLA_AVAILABLE and hidden_states.is_cuda:
-            if T == 1 and init is not None:
-                o, fs = _fla_recurrent(q, k, v, g, b, w, init, output_state)
-            elif cu_seqlens is not None:
-                total = B * T
-                cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-                q_f = q.reshape(1, total, self.num_heads, self.head_k_dim)
-                k_f = k.reshape(1, total, self.num_heads, self.head_k_dim)
-                v_f = v.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                g_f = g.reshape(1, total, self.num_heads, self.head_k_dim)
-                b_f = b.reshape(1, total, self.num_heads, self.head_k_dim)
-                w_f = w.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                o, fs = _fla_chunk(q_f, k_f, v_f, g_f, b_f, w_f, init,
-                                   output_state, cu_dev)
-                o = o.reshape(B, T, self.num_v_heads, self.head_v_dim)
-            else:
-                o, fs = _fla_chunk(q, k, v, g, b, w, init, output_state, None)
-        else:
-            o, fs = naive_gdn2_reference(q.float(), k.float(), v.float(), g,
-                                         b.float(), w.float(), init)
-
-        # Cast back to activation precision
-        o = o.to(hidden_states.dtype)
-
-        fs = fs if output_state else None
-        gate = self.g_proj(hidden_states).view(B, T, self.num_v_heads, self.head_v_dim)
-        o = self.o_norm(o)
-        o = o * torch.sigmoid(gate)
-        o = o.reshape(B, T, self.value_dim)
-        conv_out = (nq, nk, nv) if output_state else None
-        return self.o_proj(o), fs, conv_out
-
-
-def _is_xsa(config, li: int) -> bool:
-    if li < config.prelude_layers:
-        return True
-    if li >= config.prelude_layers + config.recurrent_layers:
-        return True
-    return (li - config.prelude_layers) % (config.gdn_per_xsa + 1) == config.gdn_per_xsa
-
-
-class SurjoDecoderLayer(nn.Module):
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_type = "xsa" if _is_xsa(config, layer_idx) else "gdn"
-        if self.layer_type == "xsa":
-            self.self_attn = XSAAttention(config, layer_idx)
-        else:
-            self.linear_attn = GDN2Attention(config, layer_idx)
-        self.mlp = ClampedMLP(config)
-        self.input_layernorm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SurjoRMSNorm(config.hidden_size,
-                                                     eps=config.rms_norm_eps)
-
-    def forward(self, hidden_states, attention_mask=None, cos=None, sin=None,
-                cu_seqlens=None):
-        h = self.input_layernorm(hidden_states)
-        if self.layer_type == "xsa":
-            out = self.self_attn(h, attention_mask, None, -1, False, None,
-                                 cos, sin, cu_seqlens)
-        else:
-            out, _, _ = self.linear_attn(h, attention_mask, None, None, False,
-                                         cu_seqlens)
-        hidden_states = hidden_states + out
-        h2 = self.post_attention_layernorm(hidden_states)
-        return hidden_states + self.mlp(h2)
-
-
-class SurjoPreTrainedModel(PreTrainedModel):
-    config_class = SurjoConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["SurjoDecoderLayer"]
-    _supports_flash_attn = False
-    _supports_flash_attn_2 = False
-    _supports_sdpa = True
-    _is_stateful = True
-    _can_compile_fullgraph = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, GDN2Attention):
-            module.A_log.data.copy_(
-                torch.log(
-                    torch.empty(module.num_heads, dtype=torch.float32,
-                                device=module.A_log.device).uniform_(1, 16)
-                )
-            )
-            dt = torch.exp(
-                torch.rand(module.key_dim, dtype=torch.float32,
-                           device=module.dt_bias.device)
-                * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
-            ).clamp(min=1e-4)
-            module.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt)))
-
-
-class SurjoModel(SurjoPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList(
-            [SurjoDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
-        )
-        self.norm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary = SurjoRotaryEmbedding(config)
-        self.grad_checkpointing = False
-
-        slot = 0
-        self.xsa_slot_map = {}
-        self.exec_plan = []
-        for li in range(config.prelude_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        for p in range(config.recurrent_passes):
-            for g in range(config.num_groups):
-                base = config.prelude_layers + g * (config.gdn_per_xsa + 1)
-                for j in range(config.gdn_per_xsa):
-                    self.exec_plan.append((self.layers[base + j], base + j, p, -1))
-                li = base + config.gdn_per_xsa
-                self.xsa_slot_map[(li, p)] = slot
-                self.exec_plan.append((self.layers[li], li, p, slot))
-                slot += 1
-        for li in range(config.prelude_layers + config.recurrent_layers,
-                        config.num_hidden_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        self.num_xsa_slots = slot
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, use_cache=None,
-                cache_position=None, return_dict=True,
-                output_hidden_states=False, output_attentions=False,
-                cu_seqlens=None, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        if use_cache is None:
-            use_cache = past_key_values is not None
-        if use_cache and past_key_values is None:
-            past_key_values = SurjoCache()
-        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
-
-        if past_len > 0:
-            if input_ids is not None:
-                n_ids = input_ids.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_ids
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    input_ids = input_ids[:, -n_new:] if n_new <= n_ids else input_ids
-                elif n_ids > past_len:
-                    input_ids = input_ids[:, past_len:]
-                else:
-                    input_ids = input_ids[:, -1:]
-            elif inputs_embeds is not None:
-                n_emb = inputs_embeds.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_emb
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    inputs_embeds = (inputs_embeds[:, -n_new:] if n_new <= n_emb
-                                     else inputs_embeds)
-                elif n_emb > past_len:
-                    inputs_embeds = inputs_embeds[:, past_len:]
-                else:
-                    inputs_embeds = inputs_embeds[:, -1:]
-
-        if inputs_embeds is None:
-            if input_ids is None:
-                raise ValueError("Specify input_ids or inputs_embeds")
-            inputs_embeds = self.embed_tokens(input_ids)
-        B, T = inputs_embeds.shape[0], inputs_embeds.shape[1]
-
-        cache_position = torch.arange(past_len, past_len + T,
-                                      device=inputs_embeds.device)
-        if position_ids is not None and position_ids.shape[-1] != T:
-            position_ids = None
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0).expand(B, -1)
-        cos, sin = self.rotary(position_ids)
-
-        gdn_states = dict(past_key_values.gdn_state) if past_key_values is not None else {}
-        conv_states = dict(past_key_values.gdn_conv) if past_key_values is not None else {}
-
-        noise_std = self.config.loop_noise_std
-        hidden = inputs_embeds
-        prev_p = 0
-        use_ckpt = self.training and self.grad_checkpointing
-        all_hidden_states = () if output_hidden_states else None
-
-        for (layer, li, p, slot) in self.exec_plan:
-            if output_hidden_states:
-                all_hidden_states += (hidden,)
-            if p != prev_p and p > 0 and self.training and noise_std > 0.0:
-                hidden = hidden + torch.randn_like(hidden) * noise_std
-            prev_p = p
-
-            if use_ckpt:
-                hidden = torch.utils.checkpoint.checkpoint(
-                    self._train_layer, layer, hidden, attention_mask, cos, sin,
-                    cu_seqlens, use_reentrant=False,
-                )
-            elif use_cache:
-                h = layer.input_layernorm(hidden)
-                if layer.layer_type == "xsa":
-                    out = layer.self_attn(h, attention_mask, past_key_values,
-                                          slot, use_cache, cache_position,
-                                          cos, sin, cu_seqlens)
-                else:
-                    gs = gdn_states.get((li, p))
-                    cs = conv_states.get((li, p))
-                    out, ngs, ncs = layer.linear_attn(h, attention_mask, gs, cs,
-                                                      use_cache, cu_seqlens)
-                    if ngs is not None:
-                        gdn_states[(li, p)] = ngs
-                    if ncs is not None:
-                        conv_states[(li, p)] = ncs
-                hidden = hidden + out
-                h2 = layer.post_attention_layernorm(hidden)
-                hidden = hidden + layer.mlp(h2)
-            else:
-                hidden = layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-        hidden = self.norm(hidden)
-        if output_hidden_states:
-            all_hidden_states += (hidden,)
-
-        if use_cache and past_key_values is not None:
-            past_key_values.gdn_state.update(gdn_states)
-            past_key_values.gdn_conv.update(conv_states)
-
-        if not return_dict:
-            out = tuple(v for v in [hidden, past_key_values if use_cache else None,
-                                    all_hidden_states] if v is not None)
-            return out
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=None,
-        )
-
-    def _train_layer(self, layer, hidden, attention_mask, cos, sin, cu_seqlens):
-        return layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-
-class SurjoForCausalLM(SurjoPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = SurjoModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.post_init()
-
-    def tie_weights(self, **kwargs):
-        try:
-            super().tie_weights(**kwargs)
-        except TypeError:
-            pass
-        except Exception:
-            pass
-        if (getattr(self.config, "tie_word_embeddings", False)
-                and getattr(self, "model", None) is not None
-                and getattr(self, "lm_head", None) is not None):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-    def save_pretrained(self, *args, **kwargs):
-        was_tied = (getattr(self, "lm_head", None) is not None
-                    and getattr(self, "model", None) is not None
-                    and self.lm_head.weight is self.model.embed_tokens.weight)
-        if was_tied and self.config.tie_word_embeddings:
-            self.lm_head.weight = nn.Parameter(
-                self.model.embed_tokens.weight.detach().clone(),
-                requires_grad=self.model.embed_tokens.weight.requires_grad,
-            )
-        try:
-            return super().save_pretrained(*args, **kwargs)
-        finally:
-            if was_tied and self.config.tie_word_embeddings:
-                self.tie_weights()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        self.lm_head = value
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
-        self.model.grad_checkpointing = True
-
-    def gradient_checkpointing_disable(self):
-        self.model.grad_checkpointing = False
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, labels=None,
-                use_cache=None, output_hidden_states=None, output_attentions=None,
-                return_dict=None, logits_to_keep=0, cache_position=None,
-                cu_seqlens=None, **kwargs):
-        if logits_to_keep in (None, 0) and kwargs.get("num_logits_to_keep"):
-            logits_to_keep = kwargs.pop("num_logits_to_keep")
-        return_dict = True if return_dict is None else return_dict
-        if use_cache is None:
-            use_cache = not (self.training or labels is not None)
-
-        outputs = self.model(
-            input_ids=input_ids, attention_mask=attention_mask,
-            position_ids=position_ids, past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds, use_cache=use_cache,
-            cache_position=cache_position,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            cu_seqlens=cu_seqlens,
-            **kwargs,
-        )
-        hidden_states = outputs[0]
-
-        loss, logits = None, None
-        if labels is not None:
-            sh = hidden_states[..., :-1, :].contiguous()
-            sl = labels[..., 1:].contiguous()
-            if _USE_FLA_FUSED_CE and _fla_flce is not None:
-                loss = _fla_flce(
-                    sh.view(-1, sh.shape[-1]),
-                    sl.view(-1),
-                    self.lm_head.weight,
-                    num_chunks=8,
-                    reduction="mean",
-                )
-            else:
-                n_chunks = max(8, sh.shape[0] // 8)
-                tot = hidden_states.new_zeros((), dtype=torch.float32)
-                n_tok = 0
-                for h_c, l_c in zip(sh.chunk(n_chunks, dim=0), sl.chunk(n_chunks, dim=0)):
-                    if l_c.numel() == 0:
-                        continue
-                    lg = self.lm_head(h_c)
-                    tot = tot + F.cross_entropy(
-                        lg.float().view(-1, lg.size(-1)), l_c.view(-1), reduction="sum"
-                    )
-                    n_tok += l_c.numel()
-                loss = tot / max(n_tok, 1)
-        else:
-            if isinstance(logits_to_keep, int) and logits_to_keep > 0:
-                sl_hidden = hidden_states[:, -logits_to_keep:, :]
-            elif torch.is_tensor(logits_to_keep):
-                sl_hidden = hidden_states[:, logits_to_keep, :]
-            else:
-                sl_hidden = hidden_states
-            logits = self.lm_head(sl_hidden)
-
-        if not return_dict:
-            out = (logits,) + outputs[1:]
-            return (loss,) + out if loss is not None else out
-        return CausalLMOutputWithPast(
-            loss=loss, logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
-                                      attention_mask=None, inputs_embeds=None,
-                                      cache_position=None, use_cache=True, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        has_state = past_key_values is not None and len(past_key_values.kv) > 0
-        if has_state:
-            input_ids = input_ids[:, -1:]
-            if kwargs.get("logits_to_keep") is None:
-                kwargs["logits_to_keep"] = 1
-        past_len = past_key_values.get_seq_length() if has_state else 0
-        cache_position = torch.arange(past_len, past_len + input_ids.shape[1],
-                                      device=input_ids.device)
-        model_inputs = {
-            "input_ids": input_ids.contiguous(),
-            "past_key_values": past_key_values,
-            "use_cache": use_cache,
-            "cache_position": cache_position,
-            "attention_mask": attention_mask,
-        }
-        model_inputs.update(kwargs)
-        return model_inputs
-
-    def _reorder_cache(self, past_key_values: SurjoCache, beam_idx):
-        for slot, (k, v) in past_key_values.kv.items():
-            past_key_values.kv[slot] = (k.index_select(0, beam_idx),
-                                        v.index_select(0, beam_idx))
-        for key in list(past_key_values.gdn_state.keys()):
-            s = past_key_values.gdn_state[key]
-            if s is not None:
-                past_key_values.gdn_state[key] = s.index_select(0, beam_idx)
-        for key in list(past_key_values.gdn_conv.keys()):
-            cs = past_key_values.gdn_conv[key]
-            if cs is not None:
-                past_key_values.gdn_conv[key] = tuple(
-                    c.index_select(0, beam_idx) if c is not None else None for c in cs
-                )
-        return past_key_values
-
-
-SurjoConfig.register_for_auto_class()
-SurjoForCausalLM.register_for_auto_class()

checkpoint-10000/optimizer.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:329160acb14ad89d30b66d262c4f163a6a38360ff4be17f0602534fa887e01d0
-size 148034507

checkpoint-10000/optimizer_1.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f4e8fe1c9ad35ca5d4cb90c625ac6e1351e65e92d848090f51eadf7d99f6e4b4
-size 134405851

checkpoint-10000/random_states_0.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a56e4c6f6a6f8275a982682f65388f684d03e1e8168844e5f2c802364cc8e1f3
-size 14821

checkpoint-10000/tokenizer_config.json

@@ -1,85 +0,0 @@
-{
-  "backend": "tokenizers",
-  "bos_token": "<|bos|>",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|eos|>",
-  "extra_special_tokens": [
-    "<|unk|>",
-    "<|pad|>",
-    "<|bos|>",
-    "<|eos|>",
-    "<|mask|>",
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|system|>",
-    "<|user|>",
-    "<|assistant|>",
-    "<think>",
-    "</think>",
-    "<|begin_of_thought|>",
-    "<|end_of_thought|>",
-    "<answer>",
-    "</answer>",
-    "<|step|>",
-    "<|/step|>",
-    "<context>",
-    "</context>",
-    "<|doc_start|>",
-    "<|doc_end|>",
-    "<|search|>",
-    "<|search_results|>",
-    "<|tool_list_start|>",
-    "<|tool_list_end|>",
-    "<tools>",
-    "</tools>",
-    "<|tool_call_start|>",
-    "<|tool_call_end|>",
-    "<|tool_call|>",
-    "<|/tool_call|>",
-    "<|tool_response_start|>",
-    "<|tool_response_end|>",
-    "<|tool_response|>",
-    "<|/tool_response|>",
-    "<|image|>",
-    "<|image_pad|>",
-    "<|image_placeholder|>",
-    "<|audio|>",
-    "<|audio_pad|>",
-    "<|audio_placeholder|>",
-    "<|video|>",
-    "<|video_pad|>",
-    "<|fim_prefix|>",
-    "<|fim_suffix|>",
-    "<|fim_middle|>",
-    "<|repo_name|>",
-    "<|file_separator|>",
-    "<|reward|>",
-    "<|reserved_0|>",
-    "<|reserved_1|>",
-    "<|reserved_2|>",
-    "<|reserved_3|>",
-    "<|reserved_4|>",
-    "<|reserved_5|>",
-    "<|reserved_6|>",
-    "<|reserved_7|>",
-    "<|reserved_8|>",
-    "<|reserved_9|>",
-    "<|reserved_10|>",
-    "<|reserved_11|>",
-    "<|reserved_12|>",
-    "<|reserved_13|>",
-    "<|reserved_14|>",
-    "<|reserved_15|>",
-    "<|reserved_16|>",
-    "<|reserved_17|>",
-    "<|reserved_18|>",
-    "<|reserved_19|>"
-  ],
-  "is_local": false,
-  "local_files_only": false,
-  "mask_token": "<|mask|>",
-  "model_max_length": 10000000,
-  "pad_token": "<|pad|>",
-  "tokenizer_class": "TokenizersBackend",
-  "unk_token": "<|unk|>"
-}

checkpoint-10500/chat_template.jinja

@@ -1,4 +0,0 @@
-{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
-' + message['content'] + '<|im_end|>
-' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
-' }}{% endif %}

checkpoint-10500/config.json

@@ -1,62 +0,0 @@
-{
-  "_use_flash_xsa": true,
-  "architectures": [
-    "SurjoForCausalLM"
-  ],
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_surjo.SurjoConfig",
-    "AutoModel": "modeling_surjo.SurjoForCausalLM",
-    "AutoModelForCausalLM": "modeling_surjo.SurjoForCausalLM"
-  },
-  "coda_layers": 1,
-  "dtype": "float32",
-  "gdn_allow_neg_eigval": false,
-  "gdn_conv_kernel_size": 4,
-  "gdn_key_dim": 64,
-  "gdn_num_v_heads": 8,
-  "gdn_per_xsa": 3,
-  "gdn_value_dim": 64,
-  "gradient_checkpointing": false,
-  "head_dim": 64,
-  "hidden_act": "silu",
-  "hidden_size": 512,
-  "initializer_range": 0.02,
-  "intermediate_size": 1536,
-  "layer_types": [
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "full_attention"
-  ],
-  "loop_noise_std": 0.02,
-  "max_position_embeddings": 2048,
-  "model_type": "surjo",
-  "num_attention_heads": 8,
-  "num_groups": 2,
-  "num_hidden_layers": 10,
-  "num_key_value_heads": 4,
-  "partial_rotary_factor": 1.0,
-  "prelude_layers": 1,
-  "recurrent_layers": 8,
-  "recurrent_passes": 2,
-  "rms_norm_eps": 1e-05,
-  "rope_parameters": {
-    "partial_rotary_factor": 1.0,
-    "rope_theta": 10000.0,
-    "rope_type": "default"
-  },
-  "rope_theta": 10000.0,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.16.1",
-  "use_cache": false,
-  "vocab_size": 32768,
-  "xsa_projection": true
-}

checkpoint-10500/configuration_surjo.py

@@ -1,148 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""Surjo gen-3 model configuration."""
-
-from typing import List, Optional
-from transformers import PretrainedConfig
-
-
-class SurjoConfig(PretrainedConfig):
-    r"""Configuration for the Surjo gen-3 hybrid architecture."""
-
-    model_type = "surjo"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size: int = 50257,
-        hidden_size: int = 384,
-        intermediate_size: int = 1024,
-        num_attention_heads: int = 6,
-        num_key_value_heads: int = 2,
-        head_dim: int = 64,
-        hidden_act: str = "silu",
-        max_position_embeddings: int = 2048,
-        initializer_range: float = 0.02,
-        rms_norm_eps: float = 1e-6,
-        tie_word_embeddings: bool = True,
-        rope_theta: float = 10000.0,
-        attention_dropout: float = 0.0,
-        attention_bias: bool = False,
-        prelude_layers: int = 1,
-        recurrent_layers: int = 8,
-        coda_layers: int = 1,
-        recurrent_passes: int = 2,
-        loop_noise_std: float = 0.02,
-        gdn_per_xsa: int = 3,
-        gdn_key_dim: Optional[int] = None,
-        gdn_value_dim: Optional[int] = None,
-        gdn_num_v_heads: Optional[int] = None,
-        gdn_conv_kernel_size: int = 4,
-        gdn_allow_neg_eigval: bool = False,
-        xsa_projection: bool = True,
-        rope_parameters: Optional[dict] = None,
-        rope_scaling: Optional[dict] = None,
-        partial_rotary_factor: float = 1.0,
-        gradient_checkpointing: bool = False,
-        use_cache: bool = True,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.tie_word_embeddings = tie_word_embeddings
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.attention_bias = attention_bias
-
-        self.prelude_layers = prelude_layers
-        self.recurrent_layers = recurrent_layers
-        self.coda_layers = coda_layers
-        self.recurrent_passes = recurrent_passes
-        self.loop_noise_std = loop_noise_std
-        self.gdn_per_xsa = gdn_per_xsa
-
-        self.gdn_key_dim = gdn_key_dim if gdn_key_dim is not None else self.head_dim
-        self.gdn_value_dim = gdn_value_dim if gdn_value_dim is not None else self.head_dim
-        self.gdn_num_v_heads = (
-            gdn_num_v_heads if gdn_num_v_heads is not None else num_attention_heads
-        )
-        self.gdn_conv_kernel_size = gdn_conv_kernel_size
-        self.gdn_allow_neg_eigval = gdn_allow_neg_eigval
-
-        self.xsa_projection = xsa_projection
-
-        if rope_parameters is not None:
-            self.rope_parameters = dict(rope_parameters)
-        elif rope_scaling is not None:
-            self.rope_parameters = dict(rope_scaling)
-        else:
-            self.rope_parameters = {"rope_type": "default"}
-        self.rope_parameters.setdefault("rope_type", "default")
-        self.rope_parameters.setdefault("rope_theta", rope_theta)
-        if self.rope_parameters.get("rope_type") != "default":
-            self.rope_parameters.setdefault("factor", 1.0)
-        self.rope_scaling = dict(self.rope_parameters)
-        self.partial_rotary_factor = partial_rotary_factor
-
-        self.gradient_checkpointing = gradient_checkpointing
-        self.use_cache = use_cache
-
-        group = gdn_per_xsa + 1
-        if recurrent_layers % group != 0:
-            raise ValueError(
-                f"recurrent_layers ({recurrent_layers}) must be divisible by "
-                f"gdn_per_xsa + 1 = {group}"
-            )
-        if num_attention_heads % num_key_value_heads != 0:
-            raise ValueError(
-                f"num_attention_heads ({num_attention_heads}) must be divisible by "
-                f"num_key_value_heads ({num_key_value_heads})"
-            )
-        if (
-            self.gdn_num_v_heads < num_attention_heads
-            or self.gdn_num_v_heads % num_attention_heads != 0
-        ):
-            raise ValueError(
-                f"gdn_num_v_heads ({self.gdn_num_v_heads}) must be >= "
-                f"num_attention_heads ({num_attention_heads}) and divisible by it"
-            )
-        valid_rope_types = {"default", "linear", "dynamic", "yarn", "longrope", "llama3"}
-        rt = self.rope_parameters.get("rope_type", "default")
-        if rt not in valid_rope_types:
-            raise ValueError(f"rope_type must be one of {valid_rope_types}, got {rt}")
-
-        self.num_groups = recurrent_layers // group
-        self.num_hidden_layers = prelude_layers + recurrent_layers + coda_layers
-        self.layer_types = self._build_layer_types()
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    def _build_layer_types(self) -> List[str]:
-        types = ["full_attention"] * self.prelude_layers
-        for _ in range(self.num_groups):
-            types.extend(["linear_attention"] * self.gdn_per_xsa)
-            types.append("full_attention")
-        types.extend(["full_attention"] * self.coda_layers)
-        assert len(types) == self.num_hidden_layers
-        return types
-
-    @property
-    def surjo_layer_types(self) -> List[str]:
-        return ["xsa" if t == "full_attention" else "gdn" for t in self.layer_types]
-
-    @property
-    def effective_num_layers(self) -> int:
-        return (
-            self.prelude_layers
-            + self.recurrent_layers * self.recurrent_passes
-            + self.coda_layers
-        )

checkpoint-10500/generation_config.json

@@ -1,7 +0,0 @@
-{
-  "_from_model_config": true,
-  "output_attentions": false,
-  "output_hidden_states": false,
-  "transformers_version": "5.16.1",
-  "use_cache": false
-}

checkpoint-10500/model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:f9aee4cb4e131c9a6a77f9a56b3db728622e4450c2b318043c9804031008949d
-size 215196728

checkpoint-10500/modeling_surjo.py

@@ -1,926 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""SurjoLabs Surjo Architecture"""
-
-import math
-from typing import Dict, List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-
-from transformers import PreTrainedModel
-from transformers.activations import ACT2FN
-
-try:
-    from transformers import GenerationMixin
-except ImportError:
-    from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-
-try:
-    from .configuration_surjo import SurjoConfig
-except (ImportError, ValueError):
-    from configuration_surjo import SurjoConfig
-
-try:
-    from fla.modules.fused_linear_cross_entropy import (
-        fused_linear_cross_entropy_loss as _fla_flce,
-    )
-except ImportError:
-    try:
-        from fla.modules import fused_linear_cross_entropy_loss as _fla_flce
-    except ImportError:
-        _fla_flce = None
-
-try:
-    from fla.modules.norm import rms_norm as _fla_rms_norm
-except ImportError:
-    try:
-        from fla.modules import rms_norm as _fla_rms_norm
-    except ImportError:
-        _fla_rms_norm = None
-
-try:
-    from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
-    FLA_AVAILABLE = True
-except ImportError:
-    chunk_gdn2 = None
-    fused_recurrent_gdn2 = None
-    FLA_AVAILABLE = False
-
-_USE_FLA_FUSED_CE = True
-_USE_FLA_RMS_NORM = False
-
-if FLA_AVAILABLE:
-
-    @torch.compiler.disable
-    def _fla_chunk(q, k, v, g, b, w, initial_state=None, output_final_state=False,
-                   cu_seqlens=None):
-        return chunk_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
-        )
-
-    @torch.compiler.disable
-    def _fla_recurrent(q, k, v, g, b, w, initial_state=None, output_final_state=False):
-        return fused_recurrent_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-        )
-
-try:
-    from flash_attn import flash_attn_varlen_func
-    FLASH_ATTN_AVAILABLE = True
-except ImportError:
-    flash_attn_varlen_func = None
-    FLASH_ATTN_AVAILABLE = False
-
-
-def _mask_padding_states(hidden_states, padding_mask):
-    if padding_mask is not None:
-        return hidden_states * padding_mask.unsqueeze(-1).to(hidden_states.dtype)
-    return hidden_states
-
-
-def naive_gdn2_reference(q, k, v, g, b, w, initial_state=None):
-    """fp32 reference for GDN-2."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    S = (
-        torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
-        if initial_state is None
-        else initial_state.float()
-    )
-    scale = K ** -0.5
-    outs = []
-    for t in range(T):
-        q_t = F.normalize(q[:, t], dim=-1)
-        k_t = F.normalize(k[:, t], dim=-1)
-        v_t, g_t, b_t, w_t = v[:, t], g[:, t], b[:, t], w[:, t]
-        S = S * torch.exp(g_t).unsqueeze(-1)
-        e = b_t * k_t
-        r = torch.einsum("bhkv,bhk->bhv", S, e)
-        z = w_t * v_t
-        S = S + torch.einsum("bhk,bhv->bhkv", k_t, z - r)
-        outs.append(scale * torch.einsum("bhkv,bhk->bhv", S, q_t))
-    return torch.stack(outs, dim=1), S
-
-
-class SurjoRMSNorm(nn.Module):
-    """FLA Triton RMSNorm when available; else fused ATen F.rms_norm."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, x):
-        if x.dtype != self.weight.dtype:
-            x = x.to(self.weight.dtype)
-        if _USE_FLA_RMS_NORM and _fla_rms_norm is not None and x.is_cuda:
-            return _fla_rms_norm(x, self.weight, None, eps=self.eps)
-        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
-
-
-class ClampedMLP(nn.Module):
-    def __init__(self, config):
-        hidden_size = config.hidden_size
-        intermediate_size = config.intermediate_size
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        return self.down_proj(
-            self.act_fn(self.gate_proj(x).clamp(-15.0, 15.0)) * self.up_proj(x)
-        )
-
-
-class SurjoRotaryEmbedding(nn.Module):
-    """RoPE with frequencies computed in forward() (no stored buffer)."""
-
-    def __init__(self, config, device=None):
-        super().__init__()
-        self.config = config
-        self.rope_type = config.rope_parameters.get("rope_type", "default")
-        self.max_position_embeddings = config.max_position_embeddings
-        self.base = config.rope_parameters.get("rope_theta", config.rope_theta)
-        self.factor = config.rope_parameters.get("factor", 1.0)
-        self.dim = int(config.head_dim * config.partial_rotary_factor)
-
-    def _inv_freq(self, device):
-        inv = 1.0 / (
-            self.base
-            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
-        )
-        if self.rope_type == "linear":
-            inv = inv / self.factor
-        return inv
-
-    @torch.no_grad()
-    def forward(self, position_ids):
-        inv = self._inv_freq(position_ids.device)
-        freqs = torch.einsum("bt,f->btf", position_ids.float(), inv)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        return emb.cos(), emb.sin()
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rope(q, k, cos, sin):
-    cos = cos[:, None, :, :].to(q.dtype)
-    sin = sin[:, None, :, :].to(q.dtype)
-    q = q * cos + rotate_half(q) * sin
-    k = k * cos + rotate_half(k) * sin
-    return q, k
-
-
-def apply_rope_flat(q, k, cos, sin):
-    """RoPE on flat varlen tensors. q, k: [total, H, D]; cos, sin: [B, T, D]."""
-    total = q.shape[0]
-    cos_f = cos.reshape(total, -1)[:, None, :].to(q.dtype)
-    sin_f = sin.reshape(total, -1)[:, None, :].to(q.dtype)
-    q = q * cos_f + rotate_half(q) * sin_f
-    k = k * cos_f + rotate_half(k) * sin_f
-    return q, k
-
-
-@torch.compiler.disable
-def _flash_varlen(q, k, v, cu_seqlens, max_seqlen):
-    """FA2 varlen. cu_seqlens device int32; max_seqlen a constant (no sync)."""
-    return flash_attn_varlen_func(
-        q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
-        dropout_p=0.0, causal=True,
-    )
-
-
-class ShortConvolution(nn.Module):
-    def __init__(self, dim: int, kernel_size: int = 4):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.conv = nn.Conv1d(dim, dim, kernel_size, groups=dim, bias=False)
-
-    def forward(self, x, cache=None, output_state=False):
-        B, T, C = x.shape
-        xt = x.transpose(1, 2)
-        if T == 1 and cache is not None:
-            comb = torch.cat([cache, xt], dim=2)
-            w = self.conv.weight.squeeze(1)
-            out = F.silu(torch.einsum("bck,ck->bc", comb, w).unsqueeze(-1))
-            new = comb[:, :, 1:] if output_state else None
-            return out.transpose(1, 2), new
-        if cache is not None:
-            xt_in = torch.cat([cache, xt], dim=2)
-        else:
-            xt_in = F.pad(xt, (self.kernel_size - 1, 0))
-        out = F.silu(self.conv(xt_in))
-        new = xt_in[:, :, -(self.kernel_size - 1):] if output_state else None
-        return out.transpose(1, 2), new
-
-
-class SurjoCache:
-    def __init__(self):
-        self.kv = {}
-        self.gdn_state = {}
-        self.gdn_conv = {}
-
-    def get_seq_length(self):
-        if not self.kv:
-            return 0
-        return next(iter(self.kv.values()))[0].shape[2]
-
-    @staticmethod
-    def num_xsa_slots(config):
-        return (
-            config.prelude_layers
-            + config.num_groups * config.recurrent_passes
-            + config.coda_layers
-        )
-
-
-class XSAAttention(nn.Module):
-    """Full attention with XSA value-projection subtraction."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.num_kv_groups = self.num_heads // self.num_kv_heads
-        self.head_dim = config.head_dim
-        self.xsa_projection = config.xsa_projection
-        self.dropout = config.attention_dropout
-        self.hidden_size = config.hidden_size
-        self.use_flash = getattr(config, "_use_flash_xsa", False)
-        self.q_proj = nn.Linear(
-            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = SurjoRMSNorm(self.head_dim)
-        self.k_norm = SurjoRMSNorm(self.head_dim)
-
-    def _xsa_projection_step(self, attn, v):
-        dv = (attn.float() * v.float()).sum(dim=-1, keepdim=True)
-        vv = (v.float() * v.float()).sum(dim=-1, keepdim=True).clamp_min(1e-4)
-        scale = (dv / vv).to(attn.dtype)
-        return attn - scale * v
-
-    def forward(self, hidden_states, attention_mask, cache, slot, use_cache,
-                cache_position, cos, sin, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        H, Hkv, D = self.num_heads, self.num_kv_heads, self.head_dim
-
-        q = self.q_proj(hidden_states).view(B, T, H, D).transpose(1, 2)
-        k = self.k_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        use_varlen = (
-            cu_seqlens is not None
-            and cache is None
-            and self.use_flash
-            and FLASH_ATTN_AVAILABLE
-            and hidden_states.is_cuda
-        )
-        if use_varlen:
-            total = B * T
-            q_f = q.transpose(1, 2).reshape(total, H, D).to(torch.bfloat16)
-            k_f = k.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            v_f = v.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            if Hkv != H:
-                k_f = k_f.repeat_interleave(self.num_kv_groups, dim=1)
-                v_f = v_f.repeat_interleave(self.num_kv_groups, dim=1)
-            q_f, k_f = apply_rope_flat(q_f, k_f, cos, sin)
-            cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-            o_f = _flash_varlen(q_f, k_f, v_f, cu_dev, T)
-            attn_out = o_f.view(B, T, H, D).to(hidden_states.dtype)
-            v_xsa = v_f.view(B, T, H, D).to(hidden_states.dtype)
-            if self.xsa_projection:
-                attn_out = self._xsa_projection_step(attn_out, v_xsa)
-            out = attn_out.reshape(B, T, H * D)
-            return self.o_proj(out)
-
-        q, k = apply_rope(q, k, cos, sin)
-        v_cur = v
-        if cache is not None and use_cache:
-            if slot in cache.kv:
-                pk, pv = cache.kv[slot]
-                k = torch.cat([pk, k.to(pk.dtype)], dim=2)
-                v = torch.cat([pv, v.to(pv.dtype)], dim=2)
-            cache.kv[slot] = (k, v)
-        S = k.shape[2]
-        if Hkv != H:
-            k_e = k.repeat_interleave(self.num_kv_groups, dim=1)
-            v_e = v.repeat_interleave(self.num_kv_groups, dim=1)
-            v_cur_e = v_cur.repeat_interleave(self.num_kv_groups, dim=1)
-        else:
-            k_e, v_e, v_cur_e = k, v, v_cur
-
-        is_causal = False
-        mask4 = None
-        if attention_mask is not None:
-            L = attention_mask.shape[-1]
-            if L < S:
-                attention_mask = F.pad(attention_mask, (S - L, 0), value=1)
-            elif L > S:
-                attention_mask = attention_mask[:, -S:]
-            neg = torch.finfo(q.dtype).min
-            m = attention_mask[:, None, None, :].to(q.dtype)
-            if T == S:
-                if T > 1:
-                    causal = torch.ones(T, S, device=q.device, dtype=q.dtype).tril()
-                    allow = causal[None, None] * m
-                    mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                        torch.full_like(allow, neg))
-                else:
-                    mask4 = torch.where(m > 0, torch.zeros_like(m),
-                                        torch.full_like(m, neg))
-            else:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                causal = kv_pos[None, :] <= cache_position[:, None].to(q.device)
-                allow = causal[None, None].to(q.dtype) * m
-                mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                    torch.full_like(allow, neg))
-        elif cu_seqlens is not None and T > 1:
-            neg = torch.finfo(q.dtype).min
-            cu_dev = cu_seqlens.to(q.device, torch.int32)
-            gpos = torch.arange(B * T, device=q.device)
-            doc_id = torch.searchsorted(cu_dev, gpos, right=True) - 1
-            doc_id = doc_id.clamp(min=0).view(B, T)
-            same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
-            idx = torch.arange(T, device=q.device)
-            causal = idx.unsqueeze(0) <= idx.unsqueeze(1)
-            allow = (same & causal).unsqueeze(1)
-            mask4 = torch.where(
-                allow,
-                torch.zeros(B, 1, T, T, dtype=q.dtype, device=q.device),
-                torch.full((B, 1, T, T), neg, dtype=q.dtype, device=q.device),
-            )
-        else:
-            if T == S and T > 1:
-                is_causal = True
-            elif T > 1:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                neg = torch.finfo(q.dtype).min
-                causal = torch.zeros(T, S, device=q.device, dtype=q.dtype)
-                causal = causal.masked_fill(
-                    kv_pos[None, :] > cache_position[:, None], neg
-                )
-                mask4 = causal[None, None, :, :]
-
-        attn = F.scaled_dot_product_attention(
-            q, k_e, v_e, attn_mask=mask4, is_causal=is_causal,
-            dropout_p=self.dropout if self.training else 0.0,
-        )
-
-        if self.xsa_projection:
-            attn = self._xsa_projection_step(attn, v_cur_e)
-
-        out = attn.transpose(1, 2).reshape(B, T, H * D)
-        return self.o_proj(out)
-
-
-class GDN2Attention(nn.Module):
-    """Gated DeltaNet-2 linear attention."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        H = config.num_attention_heads
-        Hv = config.gdn_num_v_heads
-        self.num_heads = H
-        self.num_v_heads = Hv
-        self.gva = Hv // H
-        self.head_k_dim = config.gdn_key_dim
-        self.head_v_dim = config.gdn_value_dim
-        self.key_dim = H * self.head_k_dim
-        self.value_dim = Hv * self.head_v_dim
-        self.conv_kernel = config.gdn_conv_kernel_size
-        self.allow_neg_eigval = config.gdn_allow_neg_eigval
-        hidden = config.hidden_size
-        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.q_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.k_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.v_conv = ShortConvolution(self.value_dim, self.conv_kernel)
-        self.f_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.key_dim, bias=False),
-        )
-        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.A_log = nn.Parameter(
-            torch.log(torch.empty(H, dtype=torch.float32).uniform_(1, 16))
-        )
-        self.A_log._no_weight_decay = True
-        dt = torch.exp(
-            torch.rand(self.key_dim, dtype=torch.float32)
-            * (math.log(0.1) - math.log(0.001))
-            + math.log(0.001)
-        ).clamp(min=1e-4)
-        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
-        self.dt_bias._no_weight_decay = True
-        self.g_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
-        )
-        self.o_norm = SurjoRMSNorm(self.head_v_dim, eps=config.rms_norm_eps)
-        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)
-
-    def forward(self, hidden_states, attention_mask, gdn_state, conv_state,
-                output_state, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != T:
-                attention_mask = (
-                    attention_mask[:, -T:] if attention_mask.shape[-1] > T else None
-                )
-            if attention_mask is not None:
-                hidden_states = _mask_padding_states(hidden_states, attention_mask)
-        cq = conv_state[0] if conv_state is not None else None
-        ck = conv_state[1] if conv_state is not None else None
-        cvs = conv_state[2] if conv_state is not None else None
-        q, nq = self.q_conv(self.q_proj(hidden_states), cq, output_state)
-        k, nk = self.k_conv(self.k_proj(hidden_states), ck, output_state)
-        v, nv = self.v_conv(self.v_proj(hidden_states), cvs, output_state)
-        g = F.softplus(self.f_proj(hidden_states).float() + self.dt_bias)
-        b = self.b_proj(hidden_states).sigmoid()
-        w = self.w_proj(hidden_states).sigmoid()
-        q = q.reshape(B, T, self.num_heads, self.head_k_dim)
-        k = k.reshape(B, T, self.num_heads, self.head_k_dim)
-        v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
-        g = g.view(B, T, self.num_heads, self.head_k_dim)
-        b = b.view(B, T, self.num_heads, self.head_k_dim)
-        w = w.view(B, T, self.num_v_heads, self.head_v_dim)
-        g = -self.A_log.float().exp().view(1, 1, self.num_heads, 1) * g
-        if self.allow_neg_eigval:
-            b = b * 2.0
-        if self.gva > 1:
-            q = q.repeat_interleave(self.gva, dim=2)
-            k = k.repeat_interleave(self.gva, dim=2)
-            g = g.repeat_interleave(self.gva, dim=2)
-            b = b.repeat_interleave(self.gva, dim=2)
-        init = gdn_state.float() if gdn_state is not None else None
-
-        if FLA_AVAILABLE and hidden_states.is_cuda:
-            if T == 1 and init is not None:
-                o, fs = _fla_recurrent(q, k, v, g, b, w, init, output_state)
-            elif cu_seqlens is not None:
-                total = B * T
-                cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-                q_f = q.reshape(1, total, self.num_heads, self.head_k_dim)
-                k_f = k.reshape(1, total, self.num_heads, self.head_k_dim)
-                v_f = v.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                g_f = g.reshape(1, total, self.num_heads, self.head_k_dim)
-                b_f = b.reshape(1, total, self.num_heads, self.head_k_dim)
-                w_f = w.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                o, fs = _fla_chunk(q_f, k_f, v_f, g_f, b_f, w_f, init,
-                                   output_state, cu_dev)
-                o = o.reshape(B, T, self.num_v_heads, self.head_v_dim)
-            else:
-                o, fs = _fla_chunk(q, k, v, g, b, w, init, output_state, None)
-        else:
-            o, fs = naive_gdn2_reference(q.float(), k.float(), v.float(), g,
-                                         b.float(), w.float(), init)
-
-        # Cast back to activation precision
-        o = o.to(hidden_states.dtype)
-
-        fs = fs if output_state else None
-        gate = self.g_proj(hidden_states).view(B, T, self.num_v_heads, self.head_v_dim)
-        o = self.o_norm(o)
-        o = o * torch.sigmoid(gate)
-        o = o.reshape(B, T, self.value_dim)
-        conv_out = (nq, nk, nv) if output_state else None
-        return self.o_proj(o), fs, conv_out
-
-
-def _is_xsa(config, li: int) -> bool:
-    if li < config.prelude_layers:
-        return True
-    if li >= config.prelude_layers + config.recurrent_layers:
-        return True
-    return (li - config.prelude_layers) % (config.gdn_per_xsa + 1) == config.gdn_per_xsa
-
-
-class SurjoDecoderLayer(nn.Module):
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_type = "xsa" if _is_xsa(config, layer_idx) else "gdn"
-        if self.layer_type == "xsa":
-            self.self_attn = XSAAttention(config, layer_idx)
-        else:
-            self.linear_attn = GDN2Attention(config, layer_idx)
-        self.mlp = ClampedMLP(config)
-        self.input_layernorm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SurjoRMSNorm(config.hidden_size,
-                                                     eps=config.rms_norm_eps)
-
-    def forward(self, hidden_states, attention_mask=None, cos=None, sin=None,
-                cu_seqlens=None):
-        h = self.input_layernorm(hidden_states)
-        if self.layer_type == "xsa":
-            out = self.self_attn(h, attention_mask, None, -1, False, None,
-                                 cos, sin, cu_seqlens)
-        else:
-            out, _, _ = self.linear_attn(h, attention_mask, None, None, False,
-                                         cu_seqlens)
-        hidden_states = hidden_states + out
-        h2 = self.post_attention_layernorm(hidden_states)
-        return hidden_states + self.mlp(h2)
-
-
-class SurjoPreTrainedModel(PreTrainedModel):
-    config_class = SurjoConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["SurjoDecoderLayer"]
-    _supports_flash_attn = False
-    _supports_flash_attn_2 = False
-    _supports_sdpa = True
-    _is_stateful = True
-    _can_compile_fullgraph = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, GDN2Attention):
-            module.A_log.data.copy_(
-                torch.log(
-                    torch.empty(module.num_heads, dtype=torch.float32,
-                                device=module.A_log.device).uniform_(1, 16)
-                )
-            )
-            dt = torch.exp(
-                torch.rand(module.key_dim, dtype=torch.float32,
-                           device=module.dt_bias.device)
-                * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
-            ).clamp(min=1e-4)
-            module.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt)))
-
-
-class SurjoModel(SurjoPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList(
-            [SurjoDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
-        )
-        self.norm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary = SurjoRotaryEmbedding(config)
-        self.grad_checkpointing = False
-
-        slot = 0
-        self.xsa_slot_map = {}
-        self.exec_plan = []
-        for li in range(config.prelude_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        for p in range(config.recurrent_passes):
-            for g in range(config.num_groups):
-                base = config.prelude_layers + g * (config.gdn_per_xsa + 1)
-                for j in range(config.gdn_per_xsa):
-                    self.exec_plan.append((self.layers[base + j], base + j, p, -1))
-                li = base + config.gdn_per_xsa
-                self.xsa_slot_map[(li, p)] = slot
-                self.exec_plan.append((self.layers[li], li, p, slot))
-                slot += 1
-        for li in range(config.prelude_layers + config.recurrent_layers,
-                        config.num_hidden_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        self.num_xsa_slots = slot
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, use_cache=None,
-                cache_position=None, return_dict=True,
-                output_hidden_states=False, output_attentions=False,
-                cu_seqlens=None, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        if use_cache is None:
-            use_cache = past_key_values is not None
-        if use_cache and past_key_values is None:
-            past_key_values = SurjoCache()
-        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
-
-        if past_len > 0:
-            if input_ids is not None:
-                n_ids = input_ids.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_ids
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    input_ids = input_ids[:, -n_new:] if n_new <= n_ids else input_ids
-                elif n_ids > past_len:
-                    input_ids = input_ids[:, past_len:]
-                else:
-                    input_ids = input_ids[:, -1:]
-            elif inputs_embeds is not None:
-                n_emb = inputs_embeds.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_emb
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    inputs_embeds = (inputs_embeds[:, -n_new:] if n_new <= n_emb
-                                     else inputs_embeds)
-                elif n_emb > past_len:
-                    inputs_embeds = inputs_embeds[:, past_len:]
-                else:
-                    inputs_embeds = inputs_embeds[:, -1:]
-
-        if inputs_embeds is None:
-            if input_ids is None:
-                raise ValueError("Specify input_ids or inputs_embeds")
-            inputs_embeds = self.embed_tokens(input_ids)
-        B, T = inputs_embeds.shape[0], inputs_embeds.shape[1]
-
-        cache_position = torch.arange(past_len, past_len + T,
-                                      device=inputs_embeds.device)
-        if position_ids is not None and position_ids.shape[-1] != T:
-            position_ids = None
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0).expand(B, -1)
-        cos, sin = self.rotary(position_ids)
-
-        gdn_states = dict(past_key_values.gdn_state) if past_key_values is not None else {}
-        conv_states = dict(past_key_values.gdn_conv) if past_key_values is not None else {}
-
-        noise_std = self.config.loop_noise_std
-        hidden = inputs_embeds
-        prev_p = 0
-        use_ckpt = self.training and self.grad_checkpointing
-        all_hidden_states = () if output_hidden_states else None
-
-        for (layer, li, p, slot) in self.exec_plan:
-            if output_hidden_states:
-                all_hidden_states += (hidden,)
-            if p != prev_p and p > 0 and self.training and noise_std > 0.0:
-                hidden = hidden + torch.randn_like(hidden) * noise_std
-            prev_p = p
-
-            if use_ckpt:
-                hidden = torch.utils.checkpoint.checkpoint(
-                    self._train_layer, layer, hidden, attention_mask, cos, sin,
-                    cu_seqlens, use_reentrant=False,
-                )
-            elif use_cache:
-                h = layer.input_layernorm(hidden)
-                if layer.layer_type == "xsa":
-                    out = layer.self_attn(h, attention_mask, past_key_values,
-                                          slot, use_cache, cache_position,
-                                          cos, sin, cu_seqlens)
-                else:
-                    gs = gdn_states.get((li, p))
-                    cs = conv_states.get((li, p))
-                    out, ngs, ncs = layer.linear_attn(h, attention_mask, gs, cs,
-                                                      use_cache, cu_seqlens)
-                    if ngs is not None:
-                        gdn_states[(li, p)] = ngs
-                    if ncs is not None:
-                        conv_states[(li, p)] = ncs
-                hidden = hidden + out
-                h2 = layer.post_attention_layernorm(hidden)
-                hidden = hidden + layer.mlp(h2)
-            else:
-                hidden = layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-        hidden = self.norm(hidden)
-        if output_hidden_states:
-            all_hidden_states += (hidden,)
-
-        if use_cache and past_key_values is not None:
-            past_key_values.gdn_state.update(gdn_states)
-            past_key_values.gdn_conv.update(conv_states)
-
-        if not return_dict:
-            out = tuple(v for v in [hidden, past_key_values if use_cache else None,
-                                    all_hidden_states] if v is not None)
-            return out
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=None,
-        )
-
-    def _train_layer(self, layer, hidden, attention_mask, cos, sin, cu_seqlens):
-        return layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-
-class SurjoForCausalLM(SurjoPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = SurjoModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.post_init()
-
-    def tie_weights(self, **kwargs):
-        try:
-            super().tie_weights(**kwargs)
-        except TypeError:
-            pass
-        except Exception:
-            pass
-        if (getattr(self.config, "tie_word_embeddings", False)
-                and getattr(self, "model", None) is not None
-                and getattr(self, "lm_head", None) is not None):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-    def save_pretrained(self, *args, **kwargs):
-        was_tied = (getattr(self, "lm_head", None) is not None
-                    and getattr(self, "model", None) is not None
-                    and self.lm_head.weight is self.model.embed_tokens.weight)
-        if was_tied and self.config.tie_word_embeddings:
-            self.lm_head.weight = nn.Parameter(
-                self.model.embed_tokens.weight.detach().clone(),
-                requires_grad=self.model.embed_tokens.weight.requires_grad,
-            )
-        try:
-            return super().save_pretrained(*args, **kwargs)
-        finally:
-            if was_tied and self.config.tie_word_embeddings:
-                self.tie_weights()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        self.lm_head = value
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
-        self.model.grad_checkpointing = True
-
-    def gradient_checkpointing_disable(self):
-        self.model.grad_checkpointing = False
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, labels=None,
-                use_cache=None, output_hidden_states=None, output_attentions=None,
-                return_dict=None, logits_to_keep=0, cache_position=None,
-                cu_seqlens=None, **kwargs):
-        if logits_to_keep in (None, 0) and kwargs.get("num_logits_to_keep"):
-            logits_to_keep = kwargs.pop("num_logits_to_keep")
-        return_dict = True if return_dict is None else return_dict
-        if use_cache is None:
-            use_cache = not (self.training or labels is not None)
-
-        outputs = self.model(
-            input_ids=input_ids, attention_mask=attention_mask,
-            position_ids=position_ids, past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds, use_cache=use_cache,
-            cache_position=cache_position,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            cu_seqlens=cu_seqlens,
-            **kwargs,
-        )
-        hidden_states = outputs[0]
-
-        loss, logits = None, None
-        if labels is not None:
-            sh = hidden_states[..., :-1, :].contiguous()
-            sl = labels[..., 1:].contiguous()
-            if _USE_FLA_FUSED_CE and _fla_flce is not None:
-                loss = _fla_flce(
-                    sh.view(-1, sh.shape[-1]),
-                    sl.view(-1),
-                    self.lm_head.weight,
-                    num_chunks=8,
-                    reduction="mean",
-                )
-            else:
-                n_chunks = max(8, sh.shape[0] // 8)
-                tot = hidden_states.new_zeros((), dtype=torch.float32)
-                n_tok = 0
-                for h_c, l_c in zip(sh.chunk(n_chunks, dim=0), sl.chunk(n_chunks, dim=0)):
-                    if l_c.numel() == 0:
-                        continue
-                    lg = self.lm_head(h_c)
-                    tot = tot + F.cross_entropy(
-                        lg.float().view(-1, lg.size(-1)), l_c.view(-1), reduction="sum"
-                    )
-                    n_tok += l_c.numel()
-                loss = tot / max(n_tok, 1)
-        else:
-            if isinstance(logits_to_keep, int) and logits_to_keep > 0:
-                sl_hidden = hidden_states[:, -logits_to_keep:, :]
-            elif torch.is_tensor(logits_to_keep):
-                sl_hidden = hidden_states[:, logits_to_keep, :]
-            else:
-                sl_hidden = hidden_states
-            logits = self.lm_head(sl_hidden)
-
-        if not return_dict:
-            out = (logits,) + outputs[1:]
-            return (loss,) + out if loss is not None else out
-        return CausalLMOutputWithPast(
-            loss=loss, logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
-                                      attention_mask=None, inputs_embeds=None,
-                                      cache_position=None, use_cache=True, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        has_state = past_key_values is not None and len(past_key_values.kv) > 0
-        if has_state:
-            input_ids = input_ids[:, -1:]
-            if kwargs.get("logits_to_keep") is None:
-                kwargs["logits_to_keep"] = 1
-        past_len = past_key_values.get_seq_length() if has_state else 0
-        cache_position = torch.arange(past_len, past_len + input_ids.shape[1],
-                                      device=input_ids.device)
-        model_inputs = {
-            "input_ids": input_ids.contiguous(),
-            "past_key_values": past_key_values,
-            "use_cache": use_cache,
-            "cache_position": cache_position,
-            "attention_mask": attention_mask,
-        }
-        model_inputs.update(kwargs)
-        return model_inputs
-
-    def _reorder_cache(self, past_key_values: SurjoCache, beam_idx):
-        for slot, (k, v) in past_key_values.kv.items():
-            past_key_values.kv[slot] = (k.index_select(0, beam_idx),
-                                        v.index_select(0, beam_idx))
-        for key in list(past_key_values.gdn_state.keys()):
-            s = past_key_values.gdn_state[key]
-            if s is not None:
-                past_key_values.gdn_state[key] = s.index_select(0, beam_idx)
-        for key in list(past_key_values.gdn_conv.keys()):
-            cs = past_key_values.gdn_conv[key]
-            if cs is not None:
-                past_key_values.gdn_conv[key] = tuple(
-                    c.index_select(0, beam_idx) if c is not None else None for c in cs
-                )
-        return past_key_values
-
-
-SurjoConfig.register_for_auto_class()
-SurjoForCausalLM.register_for_auto_class()

checkpoint-10500/optimizer.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b48fb49509808c6074512892098a39c2be225b1f7964d138e8c71afa64d320e8
-size 148034507

checkpoint-10500/optimizer_1.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a26307eb5c035abcfac3afab28155763576f921db5eb919a341f98dcfaa13958
-size 134405851

checkpoint-10500/random_states_0.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:5dfe2537b8f7e81a5971d1e93fe3fb2f280ed043ff3409bca61cdc2e66aa6cd6
-size 14821

checkpoint-10500/tokenizer_config.json

@@ -1,85 +0,0 @@
-{
-  "backend": "tokenizers",
-  "bos_token": "<|bos|>",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|eos|>",
-  "extra_special_tokens": [
-    "<|unk|>",
-    "<|pad|>",
-    "<|bos|>",
-    "<|eos|>",
-    "<|mask|>",
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|system|>",
-    "<|user|>",
-    "<|assistant|>",
-    "<think>",
-    "</think>",
-    "<|begin_of_thought|>",
-    "<|end_of_thought|>",
-    "<answer>",
-    "</answer>",
-    "<|step|>",
-    "<|/step|>",
-    "<context>",
-    "</context>",
-    "<|doc_start|>",
-    "<|doc_end|>",
-    "<|search|>",
-    "<|search_results|>",
-    "<|tool_list_start|>",
-    "<|tool_list_end|>",
-    "<tools>",
-    "</tools>",
-    "<|tool_call_start|>",
-    "<|tool_call_end|>",
-    "<|tool_call|>",
-    "<|/tool_call|>",
-    "<|tool_response_start|>",
-    "<|tool_response_end|>",
-    "<|tool_response|>",
-    "<|/tool_response|>",
-    "<|image|>",
-    "<|image_pad|>",
-    "<|image_placeholder|>",
-    "<|audio|>",
-    "<|audio_pad|>",
-    "<|audio_placeholder|>",
-    "<|video|>",
-    "<|video_pad|>",
-    "<|fim_prefix|>",
-    "<|fim_suffix|>",
-    "<|fim_middle|>",
-    "<|repo_name|>",
-    "<|file_separator|>",
-    "<|reward|>",
-    "<|reserved_0|>",
-    "<|reserved_1|>",
-    "<|reserved_2|>",
-    "<|reserved_3|>",
-    "<|reserved_4|>",
-    "<|reserved_5|>",
-    "<|reserved_6|>",
-    "<|reserved_7|>",
-    "<|reserved_8|>",
-    "<|reserved_9|>",
-    "<|reserved_10|>",
-    "<|reserved_11|>",
-    "<|reserved_12|>",
-    "<|reserved_13|>",
-    "<|reserved_14|>",
-    "<|reserved_15|>",
-    "<|reserved_16|>",
-    "<|reserved_17|>",
-    "<|reserved_18|>",
-    "<|reserved_19|>"
-  ],
-  "is_local": false,
-  "local_files_only": false,
-  "mask_token": "<|mask|>",
-  "model_max_length": 10000000,
-  "pad_token": "<|pad|>",
-  "tokenizer_class": "TokenizersBackend",
-  "unk_token": "<|unk|>"
-}

checkpoint-11000/chat_template.jinja

@@ -1,4 +0,0 @@
-{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
-' + message['content'] + '<|im_end|>
-' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
-' }}{% endif %}

checkpoint-11000/config.json

@@ -1,62 +0,0 @@
-{
-  "_use_flash_xsa": true,
-  "architectures": [
-    "SurjoForCausalLM"
-  ],
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_surjo.SurjoConfig",
-    "AutoModel": "modeling_surjo.SurjoForCausalLM",
-    "AutoModelForCausalLM": "modeling_surjo.SurjoForCausalLM"
-  },
-  "coda_layers": 1,
-  "dtype": "float32",
-  "gdn_allow_neg_eigval": false,
-  "gdn_conv_kernel_size": 4,
-  "gdn_key_dim": 64,
-  "gdn_num_v_heads": 8,
-  "gdn_per_xsa": 3,
-  "gdn_value_dim": 64,
-  "gradient_checkpointing": false,
-  "head_dim": 64,
-  "hidden_act": "silu",
-  "hidden_size": 512,
-  "initializer_range": 0.02,
-  "intermediate_size": 1536,
-  "layer_types": [
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "full_attention"
-  ],
-  "loop_noise_std": 0.02,
-  "max_position_embeddings": 2048,
-  "model_type": "surjo",
-  "num_attention_heads": 8,
-  "num_groups": 2,
-  "num_hidden_layers": 10,
-  "num_key_value_heads": 4,
-  "partial_rotary_factor": 1.0,
-  "prelude_layers": 1,
-  "recurrent_layers": 8,
-  "recurrent_passes": 2,
-  "rms_norm_eps": 1e-05,
-  "rope_parameters": {
-    "partial_rotary_factor": 1.0,
-    "rope_theta": 10000.0,
-    "rope_type": "default"
-  },
-  "rope_theta": 10000.0,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.16.1",
-  "use_cache": false,
-  "vocab_size": 32768,
-  "xsa_projection": true
-}

checkpoint-11000/configuration_surjo.py

@@ -1,148 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""Surjo gen-3 model configuration."""
-
-from typing import List, Optional
-from transformers import PretrainedConfig
-
-
-class SurjoConfig(PretrainedConfig):
-    r"""Configuration for the Surjo gen-3 hybrid architecture."""
-
-    model_type = "surjo"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size: int = 50257,
-        hidden_size: int = 384,
-        intermediate_size: int = 1024,
-        num_attention_heads: int = 6,
-        num_key_value_heads: int = 2,
-        head_dim: int = 64,
-        hidden_act: str = "silu",
-        max_position_embeddings: int = 2048,
-        initializer_range: float = 0.02,
-        rms_norm_eps: float = 1e-6,
-        tie_word_embeddings: bool = True,
-        rope_theta: float = 10000.0,
-        attention_dropout: float = 0.0,
-        attention_bias: bool = False,
-        prelude_layers: int = 1,
-        recurrent_layers: int = 8,
-        coda_layers: int = 1,
-        recurrent_passes: int = 2,
-        loop_noise_std: float = 0.02,
-        gdn_per_xsa: int = 3,
-        gdn_key_dim: Optional[int] = None,
-        gdn_value_dim: Optional[int] = None,
-        gdn_num_v_heads: Optional[int] = None,
-        gdn_conv_kernel_size: int = 4,
-        gdn_allow_neg_eigval: bool = False,
-        xsa_projection: bool = True,
-        rope_parameters: Optional[dict] = None,
-        rope_scaling: Optional[dict] = None,
-        partial_rotary_factor: float = 1.0,
-        gradient_checkpointing: bool = False,
-        use_cache: bool = True,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.tie_word_embeddings = tie_word_embeddings
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.attention_bias = attention_bias
-
-        self.prelude_layers = prelude_layers
-        self.recurrent_layers = recurrent_layers
-        self.coda_layers = coda_layers
-        self.recurrent_passes = recurrent_passes
-        self.loop_noise_std = loop_noise_std
-        self.gdn_per_xsa = gdn_per_xsa
-
-        self.gdn_key_dim = gdn_key_dim if gdn_key_dim is not None else self.head_dim
-        self.gdn_value_dim = gdn_value_dim if gdn_value_dim is not None else self.head_dim
-        self.gdn_num_v_heads = (
-            gdn_num_v_heads if gdn_num_v_heads is not None else num_attention_heads
-        )
-        self.gdn_conv_kernel_size = gdn_conv_kernel_size
-        self.gdn_allow_neg_eigval = gdn_allow_neg_eigval
-
-        self.xsa_projection = xsa_projection
-
-        if rope_parameters is not None:
-            self.rope_parameters = dict(rope_parameters)
-        elif rope_scaling is not None:
-            self.rope_parameters = dict(rope_scaling)
-        else:
-            self.rope_parameters = {"rope_type": "default"}
-        self.rope_parameters.setdefault("rope_type", "default")
-        self.rope_parameters.setdefault("rope_theta", rope_theta)
-        if self.rope_parameters.get("rope_type") != "default":
-            self.rope_parameters.setdefault("factor", 1.0)
-        self.rope_scaling = dict(self.rope_parameters)
-        self.partial_rotary_factor = partial_rotary_factor
-
-        self.gradient_checkpointing = gradient_checkpointing
-        self.use_cache = use_cache
-
-        group = gdn_per_xsa + 1
-        if recurrent_layers % group != 0:
-            raise ValueError(
-                f"recurrent_layers ({recurrent_layers}) must be divisible by "
-                f"gdn_per_xsa + 1 = {group}"
-            )
-        if num_attention_heads % num_key_value_heads != 0:
-            raise ValueError(
-                f"num_attention_heads ({num_attention_heads}) must be divisible by "
-                f"num_key_value_heads ({num_key_value_heads})"
-            )
-        if (
-            self.gdn_num_v_heads < num_attention_heads
-            or self.gdn_num_v_heads % num_attention_heads != 0
-        ):
-            raise ValueError(
-                f"gdn_num_v_heads ({self.gdn_num_v_heads}) must be >= "
-                f"num_attention_heads ({num_attention_heads}) and divisible by it"
-            )
-        valid_rope_types = {"default", "linear", "dynamic", "yarn", "longrope", "llama3"}
-        rt = self.rope_parameters.get("rope_type", "default")
-        if rt not in valid_rope_types:
-            raise ValueError(f"rope_type must be one of {valid_rope_types}, got {rt}")
-
-        self.num_groups = recurrent_layers // group
-        self.num_hidden_layers = prelude_layers + recurrent_layers + coda_layers
-        self.layer_types = self._build_layer_types()
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    def _build_layer_types(self) -> List[str]:
-        types = ["full_attention"] * self.prelude_layers
-        for _ in range(self.num_groups):
-            types.extend(["linear_attention"] * self.gdn_per_xsa)
-            types.append("full_attention")
-        types.extend(["full_attention"] * self.coda_layers)
-        assert len(types) == self.num_hidden_layers
-        return types
-
-    @property
-    def surjo_layer_types(self) -> List[str]:
-        return ["xsa" if t == "full_attention" else "gdn" for t in self.layer_types]
-
-    @property
-    def effective_num_layers(self) -> int:
-        return (
-            self.prelude_layers
-            + self.recurrent_layers * self.recurrent_passes
-            + self.coda_layers
-        )

checkpoint-11000/generation_config.json

@@ -1,7 +0,0 @@
-{
-  "_from_model_config": true,
-  "output_attentions": false,
-  "output_hidden_states": false,
-  "transformers_version": "5.16.1",
-  "use_cache": false
-}

checkpoint-11000/model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4947b5f518852154f008f7311c50e54b07859b5ede121424ec7fd2ed5e9aefcb
-size 215196728

checkpoint-11000/modeling_surjo.py

@@ -1,926 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""SurjoLabs Surjo Architecture"""
-
-import math
-from typing import Dict, List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-
-from transformers import PreTrainedModel
-from transformers.activations import ACT2FN
-
-try:
-    from transformers import GenerationMixin
-except ImportError:
-    from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-
-try:
-    from .configuration_surjo import SurjoConfig
-except (ImportError, ValueError):
-    from configuration_surjo import SurjoConfig
-
-try:
-    from fla.modules.fused_linear_cross_entropy import (
-        fused_linear_cross_entropy_loss as _fla_flce,
-    )
-except ImportError:
-    try:
-        from fla.modules import fused_linear_cross_entropy_loss as _fla_flce
-    except ImportError:
-        _fla_flce = None
-
-try:
-    from fla.modules.norm import rms_norm as _fla_rms_norm
-except ImportError:
-    try:
-        from fla.modules import rms_norm as _fla_rms_norm
-    except ImportError:
-        _fla_rms_norm = None
-
-try:
-    from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
-    FLA_AVAILABLE = True
-except ImportError:
-    chunk_gdn2 = None
-    fused_recurrent_gdn2 = None
-    FLA_AVAILABLE = False
-
-_USE_FLA_FUSED_CE = True
-_USE_FLA_RMS_NORM = False
-
-if FLA_AVAILABLE:
-
-    @torch.compiler.disable
-    def _fla_chunk(q, k, v, g, b, w, initial_state=None, output_final_state=False,
-                   cu_seqlens=None):
-        return chunk_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
-        )
-
-    @torch.compiler.disable
-    def _fla_recurrent(q, k, v, g, b, w, initial_state=None, output_final_state=False):
-        return fused_recurrent_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-        )
-
-try:
-    from flash_attn import flash_attn_varlen_func
-    FLASH_ATTN_AVAILABLE = True
-except ImportError:
-    flash_attn_varlen_func = None
-    FLASH_ATTN_AVAILABLE = False
-
-
-def _mask_padding_states(hidden_states, padding_mask):
-    if padding_mask is not None:
-        return hidden_states * padding_mask.unsqueeze(-1).to(hidden_states.dtype)
-    return hidden_states
-
-
-def naive_gdn2_reference(q, k, v, g, b, w, initial_state=None):
-    """fp32 reference for GDN-2."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    S = (
-        torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
-        if initial_state is None
-        else initial_state.float()
-    )
-    scale = K ** -0.5
-    outs = []
-    for t in range(T):
-        q_t = F.normalize(q[:, t], dim=-1)
-        k_t = F.normalize(k[:, t], dim=-1)
-        v_t, g_t, b_t, w_t = v[:, t], g[:, t], b[:, t], w[:, t]
-        S = S * torch.exp(g_t).unsqueeze(-1)
-        e = b_t * k_t
-        r = torch.einsum("bhkv,bhk->bhv", S, e)
-        z = w_t * v_t
-        S = S + torch.einsum("bhk,bhv->bhkv", k_t, z - r)
-        outs.append(scale * torch.einsum("bhkv,bhk->bhv", S, q_t))
-    return torch.stack(outs, dim=1), S
-
-
-class SurjoRMSNorm(nn.Module):
-    """FLA Triton RMSNorm when available; else fused ATen F.rms_norm."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, x):
-        if x.dtype != self.weight.dtype:
-            x = x.to(self.weight.dtype)
-        if _USE_FLA_RMS_NORM and _fla_rms_norm is not None and x.is_cuda:
-            return _fla_rms_norm(x, self.weight, None, eps=self.eps)
-        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
-
-
-class ClampedMLP(nn.Module):
-    def __init__(self, config):
-        hidden_size = config.hidden_size
-        intermediate_size = config.intermediate_size
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        return self.down_proj(
-            self.act_fn(self.gate_proj(x).clamp(-15.0, 15.0)) * self.up_proj(x)
-        )
-
-
-class SurjoRotaryEmbedding(nn.Module):
-    """RoPE with frequencies computed in forward() (no stored buffer)."""
-
-    def __init__(self, config, device=None):
-        super().__init__()
-        self.config = config
-        self.rope_type = config.rope_parameters.get("rope_type", "default")
-        self.max_position_embeddings = config.max_position_embeddings
-        self.base = config.rope_parameters.get("rope_theta", config.rope_theta)
-        self.factor = config.rope_parameters.get("factor", 1.0)
-        self.dim = int(config.head_dim * config.partial_rotary_factor)
-
-    def _inv_freq(self, device):
-        inv = 1.0 / (
-            self.base
-            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
-        )
-        if self.rope_type == "linear":
-            inv = inv / self.factor
-        return inv
-
-    @torch.no_grad()
-    def forward(self, position_ids):
-        inv = self._inv_freq(position_ids.device)
-        freqs = torch.einsum("bt,f->btf", position_ids.float(), inv)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        return emb.cos(), emb.sin()
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rope(q, k, cos, sin):
-    cos = cos[:, None, :, :].to(q.dtype)
-    sin = sin[:, None, :, :].to(q.dtype)
-    q = q * cos + rotate_half(q) * sin
-    k = k * cos + rotate_half(k) * sin
-    return q, k
-
-
-def apply_rope_flat(q, k, cos, sin):
-    """RoPE on flat varlen tensors. q, k: [total, H, D]; cos, sin: [B, T, D]."""
-    total = q.shape[0]
-    cos_f = cos.reshape(total, -1)[:, None, :].to(q.dtype)
-    sin_f = sin.reshape(total, -1)[:, None, :].to(q.dtype)
-    q = q * cos_f + rotate_half(q) * sin_f
-    k = k * cos_f + rotate_half(k) * sin_f
-    return q, k
-
-
-@torch.compiler.disable
-def _flash_varlen(q, k, v, cu_seqlens, max_seqlen):
-    """FA2 varlen. cu_seqlens device int32; max_seqlen a constant (no sync)."""
-    return flash_attn_varlen_func(
-        q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
-        dropout_p=0.0, causal=True,
-    )
-
-
-class ShortConvolution(nn.Module):
-    def __init__(self, dim: int, kernel_size: int = 4):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.conv = nn.Conv1d(dim, dim, kernel_size, groups=dim, bias=False)
-
-    def forward(self, x, cache=None, output_state=False):
-        B, T, C = x.shape
-        xt = x.transpose(1, 2)
-        if T == 1 and cache is not None:
-            comb = torch.cat([cache, xt], dim=2)
-            w = self.conv.weight.squeeze(1)
-            out = F.silu(torch.einsum("bck,ck->bc", comb, w).unsqueeze(-1))
-            new = comb[:, :, 1:] if output_state else None
-            return out.transpose(1, 2), new
-        if cache is not None:
-            xt_in = torch.cat([cache, xt], dim=2)
-        else:
-            xt_in = F.pad(xt, (self.kernel_size - 1, 0))
-        out = F.silu(self.conv(xt_in))
-        new = xt_in[:, :, -(self.kernel_size - 1):] if output_state else None
-        return out.transpose(1, 2), new
-
-
-class SurjoCache:
-    def __init__(self):
-        self.kv = {}
-        self.gdn_state = {}
-        self.gdn_conv = {}
-
-    def get_seq_length(self):
-        if not self.kv:
-            return 0
-        return next(iter(self.kv.values()))[0].shape[2]
-
-    @staticmethod
-    def num_xsa_slots(config):
-        return (
-            config.prelude_layers
-            + config.num_groups * config.recurrent_passes
-            + config.coda_layers
-        )
-
-
-class XSAAttention(nn.Module):
-    """Full attention with XSA value-projection subtraction."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.num_kv_groups = self.num_heads // self.num_kv_heads
-        self.head_dim = config.head_dim
-        self.xsa_projection = config.xsa_projection
-        self.dropout = config.attention_dropout
-        self.hidden_size = config.hidden_size
-        self.use_flash = getattr(config, "_use_flash_xsa", False)
-        self.q_proj = nn.Linear(
-            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = SurjoRMSNorm(self.head_dim)
-        self.k_norm = SurjoRMSNorm(self.head_dim)
-
-    def _xsa_projection_step(self, attn, v):
-        dv = (attn.float() * v.float()).sum(dim=-1, keepdim=True)
-        vv = (v.float() * v.float()).sum(dim=-1, keepdim=True).clamp_min(1e-4)
-        scale = (dv / vv).to(attn.dtype)
-        return attn - scale * v
-
-    def forward(self, hidden_states, attention_mask, cache, slot, use_cache,
-                cache_position, cos, sin, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        H, Hkv, D = self.num_heads, self.num_kv_heads, self.head_dim
-
-        q = self.q_proj(hidden_states).view(B, T, H, D).transpose(1, 2)
-        k = self.k_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        use_varlen = (
-            cu_seqlens is not None
-            and cache is None
-            and self.use_flash
-            and FLASH_ATTN_AVAILABLE
-            and hidden_states.is_cuda
-        )
-        if use_varlen:
-            total = B * T
-            q_f = q.transpose(1, 2).reshape(total, H, D).to(torch.bfloat16)
-            k_f = k.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            v_f = v.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            if Hkv != H:
-                k_f = k_f.repeat_interleave(self.num_kv_groups, dim=1)
-                v_f = v_f.repeat_interleave(self.num_kv_groups, dim=1)
-            q_f, k_f = apply_rope_flat(q_f, k_f, cos, sin)
-            cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-            o_f = _flash_varlen(q_f, k_f, v_f, cu_dev, T)
-            attn_out = o_f.view(B, T, H, D).to(hidden_states.dtype)
-            v_xsa = v_f.view(B, T, H, D).to(hidden_states.dtype)
-            if self.xsa_projection:
-                attn_out = self._xsa_projection_step(attn_out, v_xsa)
-            out = attn_out.reshape(B, T, H * D)
-            return self.o_proj(out)
-
-        q, k = apply_rope(q, k, cos, sin)
-        v_cur = v
-        if cache is not None and use_cache:
-            if slot in cache.kv:
-                pk, pv = cache.kv[slot]
-                k = torch.cat([pk, k.to(pk.dtype)], dim=2)
-                v = torch.cat([pv, v.to(pv.dtype)], dim=2)
-            cache.kv[slot] = (k, v)
-        S = k.shape[2]
-        if Hkv != H:
-            k_e = k.repeat_interleave(self.num_kv_groups, dim=1)
-            v_e = v.repeat_interleave(self.num_kv_groups, dim=1)
-            v_cur_e = v_cur.repeat_interleave(self.num_kv_groups, dim=1)
-        else:
-            k_e, v_e, v_cur_e = k, v, v_cur
-
-        is_causal = False
-        mask4 = None
-        if attention_mask is not None:
-            L = attention_mask.shape[-1]
-            if L < S:
-                attention_mask = F.pad(attention_mask, (S - L, 0), value=1)
-            elif L > S:
-                attention_mask = attention_mask[:, -S:]
-            neg = torch.finfo(q.dtype).min
-            m = attention_mask[:, None, None, :].to(q.dtype)
-            if T == S:
-                if T > 1:
-                    causal = torch.ones(T, S, device=q.device, dtype=q.dtype).tril()
-                    allow = causal[None, None] * m
-                    mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                        torch.full_like(allow, neg))
-                else:
-                    mask4 = torch.where(m > 0, torch.zeros_like(m),
-                                        torch.full_like(m, neg))
-            else:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                causal = kv_pos[None, :] <= cache_position[:, None].to(q.device)
-                allow = causal[None, None].to(q.dtype) * m
-                mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                    torch.full_like(allow, neg))
-        elif cu_seqlens is not None and T > 1:
-            neg = torch.finfo(q.dtype).min
-            cu_dev = cu_seqlens.to(q.device, torch.int32)
-            gpos = torch.arange(B * T, device=q.device)
-            doc_id = torch.searchsorted(cu_dev, gpos, right=True) - 1
-            doc_id = doc_id.clamp(min=0).view(B, T)
-            same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
-            idx = torch.arange(T, device=q.device)
-            causal = idx.unsqueeze(0) <= idx.unsqueeze(1)
-            allow = (same & causal).unsqueeze(1)
-            mask4 = torch.where(
-                allow,
-                torch.zeros(B, 1, T, T, dtype=q.dtype, device=q.device),
-                torch.full((B, 1, T, T), neg, dtype=q.dtype, device=q.device),
-            )
-        else:
-            if T == S and T > 1:
-                is_causal = True
-            elif T > 1:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                neg = torch.finfo(q.dtype).min
-                causal = torch.zeros(T, S, device=q.device, dtype=q.dtype)
-                causal = causal.masked_fill(
-                    kv_pos[None, :] > cache_position[:, None], neg
-                )
-                mask4 = causal[None, None, :, :]
-
-        attn = F.scaled_dot_product_attention(
-            q, k_e, v_e, attn_mask=mask4, is_causal=is_causal,
-            dropout_p=self.dropout if self.training else 0.0,
-        )
-
-        if self.xsa_projection:
-            attn = self._xsa_projection_step(attn, v_cur_e)
-
-        out = attn.transpose(1, 2).reshape(B, T, H * D)
-        return self.o_proj(out)
-
-
-class GDN2Attention(nn.Module):
-    """Gated DeltaNet-2 linear attention."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        H = config.num_attention_heads
-        Hv = config.gdn_num_v_heads
-        self.num_heads = H
-        self.num_v_heads = Hv
-        self.gva = Hv // H
-        self.head_k_dim = config.gdn_key_dim
-        self.head_v_dim = config.gdn_value_dim
-        self.key_dim = H * self.head_k_dim
-        self.value_dim = Hv * self.head_v_dim
-        self.conv_kernel = config.gdn_conv_kernel_size
-        self.allow_neg_eigval = config.gdn_allow_neg_eigval
-        hidden = config.hidden_size
-        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.q_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.k_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.v_conv = ShortConvolution(self.value_dim, self.conv_kernel)
-        self.f_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.key_dim, bias=False),
-        )
-        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.A_log = nn.Parameter(
-            torch.log(torch.empty(H, dtype=torch.float32).uniform_(1, 16))
-        )
-        self.A_log._no_weight_decay = True
-        dt = torch.exp(
-            torch.rand(self.key_dim, dtype=torch.float32)
-            * (math.log(0.1) - math.log(0.001))
-            + math.log(0.001)
-        ).clamp(min=1e-4)
-        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
-        self.dt_bias._no_weight_decay = True
-        self.g_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
-        )
-        self.o_norm = SurjoRMSNorm(self.head_v_dim, eps=config.rms_norm_eps)
-        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)
-
-    def forward(self, hidden_states, attention_mask, gdn_state, conv_state,
-                output_state, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != T:
-                attention_mask = (
-                    attention_mask[:, -T:] if attention_mask.shape[-1] > T else None
-                )
-            if attention_mask is not None:
-                hidden_states = _mask_padding_states(hidden_states, attention_mask)
-        cq = conv_state[0] if conv_state is not None else None
-        ck = conv_state[1] if conv_state is not None else None
-        cvs = conv_state[2] if conv_state is not None else None
-        q, nq = self.q_conv(self.q_proj(hidden_states), cq, output_state)
-        k, nk = self.k_conv(self.k_proj(hidden_states), ck, output_state)
-        v, nv = self.v_conv(self.v_proj(hidden_states), cvs, output_state)
-        g = F.softplus(self.f_proj(hidden_states).float() + self.dt_bias)
-        b = self.b_proj(hidden_states).sigmoid()
-        w = self.w_proj(hidden_states).sigmoid()
-        q = q.reshape(B, T, self.num_heads, self.head_k_dim)
-        k = k.reshape(B, T, self.num_heads, self.head_k_dim)
-        v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
-        g = g.view(B, T, self.num_heads, self.head_k_dim)
-        b = b.view(B, T, self.num_heads, self.head_k_dim)
-        w = w.view(B, T, self.num_v_heads, self.head_v_dim)
-        g = -self.A_log.float().exp().view(1, 1, self.num_heads, 1) * g
-        if self.allow_neg_eigval:
-            b = b * 2.0
-        if self.gva > 1:
-            q = q.repeat_interleave(self.gva, dim=2)
-            k = k.repeat_interleave(self.gva, dim=2)
-            g = g.repeat_interleave(self.gva, dim=2)
-            b = b.repeat_interleave(self.gva, dim=2)
-        init = gdn_state.float() if gdn_state is not None else None
-
-        if FLA_AVAILABLE and hidden_states.is_cuda:
-            if T == 1 and init is not None:
-                o, fs = _fla_recurrent(q, k, v, g, b, w, init, output_state)
-            elif cu_seqlens is not None:
-                total = B * T
-                cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-                q_f = q.reshape(1, total, self.num_heads, self.head_k_dim)
-                k_f = k.reshape(1, total, self.num_heads, self.head_k_dim)
-                v_f = v.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                g_f = g.reshape(1, total, self.num_heads, self.head_k_dim)
-                b_f = b.reshape(1, total, self.num_heads, self.head_k_dim)
-                w_f = w.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                o, fs = _fla_chunk(q_f, k_f, v_f, g_f, b_f, w_f, init,
-                                   output_state, cu_dev)
-                o = o.reshape(B, T, self.num_v_heads, self.head_v_dim)
-            else:
-                o, fs = _fla_chunk(q, k, v, g, b, w, init, output_state, None)
-        else:
-            o, fs = naive_gdn2_reference(q.float(), k.float(), v.float(), g,
-                                         b.float(), w.float(), init)
-
-        # Cast back to activation precision
-        o = o.to(hidden_states.dtype)
-
-        fs = fs if output_state else None
-        gate = self.g_proj(hidden_states).view(B, T, self.num_v_heads, self.head_v_dim)
-        o = self.o_norm(o)
-        o = o * torch.sigmoid(gate)
-        o = o.reshape(B, T, self.value_dim)
-        conv_out = (nq, nk, nv) if output_state else None
-        return self.o_proj(o), fs, conv_out
-
-
-def _is_xsa(config, li: int) -> bool:
-    if li < config.prelude_layers:
-        return True
-    if li >= config.prelude_layers + config.recurrent_layers:
-        return True
-    return (li - config.prelude_layers) % (config.gdn_per_xsa + 1) == config.gdn_per_xsa
-
-
-class SurjoDecoderLayer(nn.Module):
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_type = "xsa" if _is_xsa(config, layer_idx) else "gdn"
-        if self.layer_type == "xsa":
-            self.self_attn = XSAAttention(config, layer_idx)
-        else:
-            self.linear_attn = GDN2Attention(config, layer_idx)
-        self.mlp = ClampedMLP(config)
-        self.input_layernorm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SurjoRMSNorm(config.hidden_size,
-                                                     eps=config.rms_norm_eps)
-
-    def forward(self, hidden_states, attention_mask=None, cos=None, sin=None,
-                cu_seqlens=None):
-        h = self.input_layernorm(hidden_states)
-        if self.layer_type == "xsa":
-            out = self.self_attn(h, attention_mask, None, -1, False, None,
-                                 cos, sin, cu_seqlens)
-        else:
-            out, _, _ = self.linear_attn(h, attention_mask, None, None, False,
-                                         cu_seqlens)
-        hidden_states = hidden_states + out
-        h2 = self.post_attention_layernorm(hidden_states)
-        return hidden_states + self.mlp(h2)
-
-
-class SurjoPreTrainedModel(PreTrainedModel):
-    config_class = SurjoConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["SurjoDecoderLayer"]
-    _supports_flash_attn = False
-    _supports_flash_attn_2 = False
-    _supports_sdpa = True
-    _is_stateful = True
-    _can_compile_fullgraph = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, GDN2Attention):
-            module.A_log.data.copy_(
-                torch.log(
-                    torch.empty(module.num_heads, dtype=torch.float32,
-                                device=module.A_log.device).uniform_(1, 16)
-                )
-            )
-            dt = torch.exp(
-                torch.rand(module.key_dim, dtype=torch.float32,
-                           device=module.dt_bias.device)
-                * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
-            ).clamp(min=1e-4)
-            module.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt)))
-
-
-class SurjoModel(SurjoPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList(
-            [SurjoDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
-        )
-        self.norm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary = SurjoRotaryEmbedding(config)
-        self.grad_checkpointing = False
-
-        slot = 0
-        self.xsa_slot_map = {}
-        self.exec_plan = []
-        for li in range(config.prelude_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        for p in range(config.recurrent_passes):
-            for g in range(config.num_groups):
-                base = config.prelude_layers + g * (config.gdn_per_xsa + 1)
-                for j in range(config.gdn_per_xsa):
-                    self.exec_plan.append((self.layers[base + j], base + j, p, -1))
-                li = base + config.gdn_per_xsa
-                self.xsa_slot_map[(li, p)] = slot
-                self.exec_plan.append((self.layers[li], li, p, slot))
-                slot += 1
-        for li in range(config.prelude_layers + config.recurrent_layers,
-                        config.num_hidden_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        self.num_xsa_slots = slot
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, use_cache=None,
-                cache_position=None, return_dict=True,
-                output_hidden_states=False, output_attentions=False,
-                cu_seqlens=None, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        if use_cache is None:
-            use_cache = past_key_values is not None
-        if use_cache and past_key_values is None:
-            past_key_values = SurjoCache()
-        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
-
-        if past_len > 0:
-            if input_ids is not None:
-                n_ids = input_ids.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_ids
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    input_ids = input_ids[:, -n_new:] if n_new <= n_ids else input_ids
-                elif n_ids > past_len:
-                    input_ids = input_ids[:, past_len:]
-                else:
-                    input_ids = input_ids[:, -1:]
-            elif inputs_embeds is not None:
-                n_emb = inputs_embeds.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_emb
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    inputs_embeds = (inputs_embeds[:, -n_new:] if n_new <= n_emb
-                                     else inputs_embeds)
-                elif n_emb > past_len:
-                    inputs_embeds = inputs_embeds[:, past_len:]
-                else:
-                    inputs_embeds = inputs_embeds[:, -1:]
-
-        if inputs_embeds is None:
-            if input_ids is None:
-                raise ValueError("Specify input_ids or inputs_embeds")
-            inputs_embeds = self.embed_tokens(input_ids)
-        B, T = inputs_embeds.shape[0], inputs_embeds.shape[1]
-
-        cache_position = torch.arange(past_len, past_len + T,
-                                      device=inputs_embeds.device)
-        if position_ids is not None and position_ids.shape[-1] != T:
-            position_ids = None
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0).expand(B, -1)
-        cos, sin = self.rotary(position_ids)
-
-        gdn_states = dict(past_key_values.gdn_state) if past_key_values is not None else {}
-        conv_states = dict(past_key_values.gdn_conv) if past_key_values is not None else {}
-
-        noise_std = self.config.loop_noise_std
-        hidden = inputs_embeds
-        prev_p = 0
-        use_ckpt = self.training and self.grad_checkpointing
-        all_hidden_states = () if output_hidden_states else None
-
-        for (layer, li, p, slot) in self.exec_plan:
-            if output_hidden_states:
-                all_hidden_states += (hidden,)
-            if p != prev_p and p > 0 and self.training and noise_std > 0.0:
-                hidden = hidden + torch.randn_like(hidden) * noise_std
-            prev_p = p
-
-            if use_ckpt:
-                hidden = torch.utils.checkpoint.checkpoint(
-                    self._train_layer, layer, hidden, attention_mask, cos, sin,
-                    cu_seqlens, use_reentrant=False,
-                )
-            elif use_cache:
-                h = layer.input_layernorm(hidden)
-                if layer.layer_type == "xsa":
-                    out = layer.self_attn(h, attention_mask, past_key_values,
-                                          slot, use_cache, cache_position,
-                                          cos, sin, cu_seqlens)
-                else:
-                    gs = gdn_states.get((li, p))
-                    cs = conv_states.get((li, p))
-                    out, ngs, ncs = layer.linear_attn(h, attention_mask, gs, cs,
-                                                      use_cache, cu_seqlens)
-                    if ngs is not None:
-                        gdn_states[(li, p)] = ngs
-                    if ncs is not None:
-                        conv_states[(li, p)] = ncs
-                hidden = hidden + out
-                h2 = layer.post_attention_layernorm(hidden)
-                hidden = hidden + layer.mlp(h2)
-            else:
-                hidden = layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-        hidden = self.norm(hidden)
-        if output_hidden_states:
-            all_hidden_states += (hidden,)
-
-        if use_cache and past_key_values is not None:
-            past_key_values.gdn_state.update(gdn_states)
-            past_key_values.gdn_conv.update(conv_states)
-
-        if not return_dict:
-            out = tuple(v for v in [hidden, past_key_values if use_cache else None,
-                                    all_hidden_states] if v is not None)
-            return out
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=None,
-        )
-
-    def _train_layer(self, layer, hidden, attention_mask, cos, sin, cu_seqlens):
-        return layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-
-class SurjoForCausalLM(SurjoPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = SurjoModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.post_init()
-
-    def tie_weights(self, **kwargs):
-        try:
-            super().tie_weights(**kwargs)
-        except TypeError:
-            pass
-        except Exception:
-            pass
-        if (getattr(self.config, "tie_word_embeddings", False)
-                and getattr(self, "model", None) is not None
-                and getattr(self, "lm_head", None) is not None):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-    def save_pretrained(self, *args, **kwargs):
-        was_tied = (getattr(self, "lm_head", None) is not None
-                    and getattr(self, "model", None) is not None
-                    and self.lm_head.weight is self.model.embed_tokens.weight)
-        if was_tied and self.config.tie_word_embeddings:
-            self.lm_head.weight = nn.Parameter(
-                self.model.embed_tokens.weight.detach().clone(),
-                requires_grad=self.model.embed_tokens.weight.requires_grad,
-            )
-        try:
-            return super().save_pretrained(*args, **kwargs)
-        finally:
-            if was_tied and self.config.tie_word_embeddings:
-                self.tie_weights()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        self.lm_head = value
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
-        self.model.grad_checkpointing = True
-
-    def gradient_checkpointing_disable(self):
-        self.model.grad_checkpointing = False
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, labels=None,
-                use_cache=None, output_hidden_states=None, output_attentions=None,
-                return_dict=None, logits_to_keep=0, cache_position=None,
-                cu_seqlens=None, **kwargs):
-        if logits_to_keep in (None, 0) and kwargs.get("num_logits_to_keep"):
-            logits_to_keep = kwargs.pop("num_logits_to_keep")
-        return_dict = True if return_dict is None else return_dict
-        if use_cache is None:
-            use_cache = not (self.training or labels is not None)
-
-        outputs = self.model(
-            input_ids=input_ids, attention_mask=attention_mask,
-            position_ids=position_ids, past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds, use_cache=use_cache,
-            cache_position=cache_position,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            cu_seqlens=cu_seqlens,
-            **kwargs,
-        )
-        hidden_states = outputs[0]
-
-        loss, logits = None, None
-        if labels is not None:
-            sh = hidden_states[..., :-1, :].contiguous()
-            sl = labels[..., 1:].contiguous()
-            if _USE_FLA_FUSED_CE and _fla_flce is not None:
-                loss = _fla_flce(
-                    sh.view(-1, sh.shape[-1]),
-                    sl.view(-1),
-                    self.lm_head.weight,
-                    num_chunks=8,
-                    reduction="mean",
-                )
-            else:
-                n_chunks = max(8, sh.shape[0] // 8)
-                tot = hidden_states.new_zeros((), dtype=torch.float32)
-                n_tok = 0
-                for h_c, l_c in zip(sh.chunk(n_chunks, dim=0), sl.chunk(n_chunks, dim=0)):
-                    if l_c.numel() == 0:
-                        continue
-                    lg = self.lm_head(h_c)
-                    tot = tot + F.cross_entropy(
-                        lg.float().view(-1, lg.size(-1)), l_c.view(-1), reduction="sum"
-                    )
-                    n_tok += l_c.numel()
-                loss = tot / max(n_tok, 1)
-        else:
-            if isinstance(logits_to_keep, int) and logits_to_keep > 0:
-                sl_hidden = hidden_states[:, -logits_to_keep:, :]
-            elif torch.is_tensor(logits_to_keep):
-                sl_hidden = hidden_states[:, logits_to_keep, :]
-            else:
-                sl_hidden = hidden_states
-            logits = self.lm_head(sl_hidden)
-
-        if not return_dict:
-            out = (logits,) + outputs[1:]
-            return (loss,) + out if loss is not None else out
-        return CausalLMOutputWithPast(
-            loss=loss, logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
-                                      attention_mask=None, inputs_embeds=None,
-                                      cache_position=None, use_cache=True, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        has_state = past_key_values is not None and len(past_key_values.kv) > 0
-        if has_state:
-            input_ids = input_ids[:, -1:]
-            if kwargs.get("logits_to_keep") is None:
-                kwargs["logits_to_keep"] = 1
-        past_len = past_key_values.get_seq_length() if has_state else 0
-        cache_position = torch.arange(past_len, past_len + input_ids.shape[1],
-                                      device=input_ids.device)
-        model_inputs = {
-            "input_ids": input_ids.contiguous(),
-            "past_key_values": past_key_values,
-            "use_cache": use_cache,
-            "cache_position": cache_position,
-            "attention_mask": attention_mask,
-        }
-        model_inputs.update(kwargs)
-        return model_inputs
-
-    def _reorder_cache(self, past_key_values: SurjoCache, beam_idx):
-        for slot, (k, v) in past_key_values.kv.items():
-            past_key_values.kv[slot] = (k.index_select(0, beam_idx),
-                                        v.index_select(0, beam_idx))
-        for key in list(past_key_values.gdn_state.keys()):
-            s = past_key_values.gdn_state[key]
-            if s is not None:
-                past_key_values.gdn_state[key] = s.index_select(0, beam_idx)
-        for key in list(past_key_values.gdn_conv.keys()):
-            cs = past_key_values.gdn_conv[key]
-            if cs is not None:
-                past_key_values.gdn_conv[key] = tuple(
-                    c.index_select(0, beam_idx) if c is not None else None for c in cs
-                )
-        return past_key_values
-
-
-SurjoConfig.register_for_auto_class()
-SurjoForCausalLM.register_for_auto_class()

checkpoint-11000/optimizer.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ec2f2c7498c5486b6a9070389bf8004fc2664292455b891c2b66ebebef9a1d22
-size 148034507

checkpoint-11000/optimizer_1.bin

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:30d8b290568d14613f48a45c46ea1e1b1951e7d4f229533ef02526f2c35647ee
-size 134405851

checkpoint-11000/random_states_0.pkl

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:4d0af987318d035cbf3b8ebc860997b833d6d9ae9456c5fc8fd84590e13d5040
-size 14821

checkpoint-11000/tokenizer_config.json

@@ -1,85 +0,0 @@
-{
-  "backend": "tokenizers",
-  "bos_token": "<|bos|>",
-  "clean_up_tokenization_spaces": false,
-  "eos_token": "<|eos|>",
-  "extra_special_tokens": [
-    "<|unk|>",
-    "<|pad|>",
-    "<|bos|>",
-    "<|eos|>",
-    "<|mask|>",
-    "<|im_start|>",
-    "<|im_end|>",
-    "<|system|>",
-    "<|user|>",
-    "<|assistant|>",
-    "<think>",
-    "</think>",
-    "<|begin_of_thought|>",
-    "<|end_of_thought|>",
-    "<answer>",
-    "</answer>",
-    "<|step|>",
-    "<|/step|>",
-    "<context>",
-    "</context>",
-    "<|doc_start|>",
-    "<|doc_end|>",
-    "<|search|>",
-    "<|search_results|>",
-    "<|tool_list_start|>",
-    "<|tool_list_end|>",
-    "<tools>",
-    "</tools>",
-    "<|tool_call_start|>",
-    "<|tool_call_end|>",
-    "<|tool_call|>",
-    "<|/tool_call|>",
-    "<|tool_response_start|>",
-    "<|tool_response_end|>",
-    "<|tool_response|>",
-    "<|/tool_response|>",
-    "<|image|>",
-    "<|image_pad|>",
-    "<|image_placeholder|>",
-    "<|audio|>",
-    "<|audio_pad|>",
-    "<|audio_placeholder|>",
-    "<|video|>",
-    "<|video_pad|>",
-    "<|fim_prefix|>",
-    "<|fim_suffix|>",
-    "<|fim_middle|>",
-    "<|repo_name|>",
-    "<|file_separator|>",
-    "<|reward|>",
-    "<|reserved_0|>",
-    "<|reserved_1|>",
-    "<|reserved_2|>",
-    "<|reserved_3|>",
-    "<|reserved_4|>",
-    "<|reserved_5|>",
-    "<|reserved_6|>",
-    "<|reserved_7|>",
-    "<|reserved_8|>",
-    "<|reserved_9|>",
-    "<|reserved_10|>",
-    "<|reserved_11|>",
-    "<|reserved_12|>",
-    "<|reserved_13|>",
-    "<|reserved_14|>",
-    "<|reserved_15|>",
-    "<|reserved_16|>",
-    "<|reserved_17|>",
-    "<|reserved_18|>",
-    "<|reserved_19|>"
-  ],
-  "is_local": false,
-  "local_files_only": false,
-  "mask_token": "<|mask|>",
-  "model_max_length": 10000000,
-  "pad_token": "<|pad|>",
-  "tokenizer_class": "TokenizersBackend",
-  "unk_token": "<|unk|>"
-}

checkpoint-11500/chat_template.jinja

@@ -1,4 +0,0 @@
-{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
-' + message['content'] + '<|im_end|>
-' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
-' }}{% endif %}

checkpoint-11500/config.json

@@ -1,62 +0,0 @@
-{
-  "_use_flash_xsa": true,
-  "architectures": [
-    "SurjoForCausalLM"
-  ],
-  "attention_bias": false,
-  "attention_dropout": 0.0,
-  "auto_map": {
-    "AutoConfig": "configuration_surjo.SurjoConfig",
-    "AutoModel": "modeling_surjo.SurjoForCausalLM",
-    "AutoModelForCausalLM": "modeling_surjo.SurjoForCausalLM"
-  },
-  "coda_layers": 1,
-  "dtype": "float32",
-  "gdn_allow_neg_eigval": false,
-  "gdn_conv_kernel_size": 4,
-  "gdn_key_dim": 64,
-  "gdn_num_v_heads": 8,
-  "gdn_per_xsa": 3,
-  "gdn_value_dim": 64,
-  "gradient_checkpointing": false,
-  "head_dim": 64,
-  "hidden_act": "silu",
-  "hidden_size": 512,
-  "initializer_range": 0.02,
-  "intermediate_size": 1536,
-  "layer_types": [
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "linear_attention",
-    "linear_attention",
-    "linear_attention",
-    "full_attention",
-    "full_attention"
-  ],
-  "loop_noise_std": 0.02,
-  "max_position_embeddings": 2048,
-  "model_type": "surjo",
-  "num_attention_heads": 8,
-  "num_groups": 2,
-  "num_hidden_layers": 10,
-  "num_key_value_heads": 4,
-  "partial_rotary_factor": 1.0,
-  "prelude_layers": 1,
-  "recurrent_layers": 8,
-  "recurrent_passes": 2,
-  "rms_norm_eps": 1e-05,
-  "rope_parameters": {
-    "partial_rotary_factor": 1.0,
-    "rope_theta": 10000.0,
-    "rope_type": "default"
-  },
-  "rope_theta": 10000.0,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.16.1",
-  "use_cache": false,
-  "vocab_size": 32768,
-  "xsa_projection": true
-}

checkpoint-11500/configuration_surjo.py

@@ -1,148 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""Surjo gen-3 model configuration."""
-
-from typing import List, Optional
-from transformers import PretrainedConfig
-
-
-class SurjoConfig(PretrainedConfig):
-    r"""Configuration for the Surjo gen-3 hybrid architecture."""
-
-    model_type = "surjo"
-    keys_to_ignore_at_inference = ["past_key_values"]
-
-    def __init__(
-        self,
-        vocab_size: int = 50257,
-        hidden_size: int = 384,
-        intermediate_size: int = 1024,
-        num_attention_heads: int = 6,
-        num_key_value_heads: int = 2,
-        head_dim: int = 64,
-        hidden_act: str = "silu",
-        max_position_embeddings: int = 2048,
-        initializer_range: float = 0.02,
-        rms_norm_eps: float = 1e-6,
-        tie_word_embeddings: bool = True,
-        rope_theta: float = 10000.0,
-        attention_dropout: float = 0.0,
-        attention_bias: bool = False,
-        prelude_layers: int = 1,
-        recurrent_layers: int = 8,
-        coda_layers: int = 1,
-        recurrent_passes: int = 2,
-        loop_noise_std: float = 0.02,
-        gdn_per_xsa: int = 3,
-        gdn_key_dim: Optional[int] = None,
-        gdn_value_dim: Optional[int] = None,
-        gdn_num_v_heads: Optional[int] = None,
-        gdn_conv_kernel_size: int = 4,
-        gdn_allow_neg_eigval: bool = False,
-        xsa_projection: bool = True,
-        rope_parameters: Optional[dict] = None,
-        rope_scaling: Optional[dict] = None,
-        partial_rotary_factor: float = 1.0,
-        gradient_checkpointing: bool = False,
-        use_cache: bool = True,
-        **kwargs,
-    ):
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_attention_heads = num_attention_heads
-        self.num_key_value_heads = num_key_value_heads
-        self.head_dim = head_dim if head_dim is not None else hidden_size // num_attention_heads
-        self.hidden_act = hidden_act
-        self.max_position_embeddings = max_position_embeddings
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.tie_word_embeddings = tie_word_embeddings
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.attention_bias = attention_bias
-
-        self.prelude_layers = prelude_layers
-        self.recurrent_layers = recurrent_layers
-        self.coda_layers = coda_layers
-        self.recurrent_passes = recurrent_passes
-        self.loop_noise_std = loop_noise_std
-        self.gdn_per_xsa = gdn_per_xsa
-
-        self.gdn_key_dim = gdn_key_dim if gdn_key_dim is not None else self.head_dim
-        self.gdn_value_dim = gdn_value_dim if gdn_value_dim is not None else self.head_dim
-        self.gdn_num_v_heads = (
-            gdn_num_v_heads if gdn_num_v_heads is not None else num_attention_heads
-        )
-        self.gdn_conv_kernel_size = gdn_conv_kernel_size
-        self.gdn_allow_neg_eigval = gdn_allow_neg_eigval
-
-        self.xsa_projection = xsa_projection
-
-        if rope_parameters is not None:
-            self.rope_parameters = dict(rope_parameters)
-        elif rope_scaling is not None:
-            self.rope_parameters = dict(rope_scaling)
-        else:
-            self.rope_parameters = {"rope_type": "default"}
-        self.rope_parameters.setdefault("rope_type", "default")
-        self.rope_parameters.setdefault("rope_theta", rope_theta)
-        if self.rope_parameters.get("rope_type") != "default":
-            self.rope_parameters.setdefault("factor", 1.0)
-        self.rope_scaling = dict(self.rope_parameters)
-        self.partial_rotary_factor = partial_rotary_factor
-
-        self.gradient_checkpointing = gradient_checkpointing
-        self.use_cache = use_cache
-
-        group = gdn_per_xsa + 1
-        if recurrent_layers % group != 0:
-            raise ValueError(
-                f"recurrent_layers ({recurrent_layers}) must be divisible by "
-                f"gdn_per_xsa + 1 = {group}"
-            )
-        if num_attention_heads % num_key_value_heads != 0:
-            raise ValueError(
-                f"num_attention_heads ({num_attention_heads}) must be divisible by "
-                f"num_key_value_heads ({num_key_value_heads})"
-            )
-        if (
-            self.gdn_num_v_heads < num_attention_heads
-            or self.gdn_num_v_heads % num_attention_heads != 0
-        ):
-            raise ValueError(
-                f"gdn_num_v_heads ({self.gdn_num_v_heads}) must be >= "
-                f"num_attention_heads ({num_attention_heads}) and divisible by it"
-            )
-        valid_rope_types = {"default", "linear", "dynamic", "yarn", "longrope", "llama3"}
-        rt = self.rope_parameters.get("rope_type", "default")
-        if rt not in valid_rope_types:
-            raise ValueError(f"rope_type must be one of {valid_rope_types}, got {rt}")
-
-        self.num_groups = recurrent_layers // group
-        self.num_hidden_layers = prelude_layers + recurrent_layers + coda_layers
-        self.layer_types = self._build_layer_types()
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    def _build_layer_types(self) -> List[str]:
-        types = ["full_attention"] * self.prelude_layers
-        for _ in range(self.num_groups):
-            types.extend(["linear_attention"] * self.gdn_per_xsa)
-            types.append("full_attention")
-        types.extend(["full_attention"] * self.coda_layers)
-        assert len(types) == self.num_hidden_layers
-        return types
-
-    @property
-    def surjo_layer_types(self) -> List[str]:
-        return ["xsa" if t == "full_attention" else "gdn" for t in self.layer_types]
-
-    @property
-    def effective_num_layers(self) -> int:
-        return (
-            self.prelude_layers
-            + self.recurrent_layers * self.recurrent_passes
-            + self.coda_layers
-        )

checkpoint-11500/generation_config.json

@@ -1,7 +0,0 @@
-{
-  "_from_model_config": true,
-  "output_attentions": false,
-  "output_hidden_states": false,
-  "transformers_version": "5.16.1",
-  "use_cache": false
-}

checkpoint-11500/model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:117dcbdaf78bdc4fe621f314821ef375e693f6b47f31cce978eebd2f1d64abd0
-size 215196728

checkpoint-11500/modeling_surjo.py

@@ -1,926 +0,0 @@
-# Copyright 2026 SurjoLabs. All rights reserved.
-# Licensed under the Apache License, Version 2.0
-
-"""SurjoLabs Surjo Architecture"""
-
-import math
-from typing import Dict, List, Optional, Tuple
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.utils.checkpoint
-
-from transformers import PreTrainedModel
-from transformers.activations import ACT2FN
-
-try:
-    from transformers import GenerationMixin
-except ImportError:
-    from transformers.generation import GenerationMixin
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-
-try:
-    from .configuration_surjo import SurjoConfig
-except (ImportError, ValueError):
-    from configuration_surjo import SurjoConfig
-
-try:
-    from fla.modules.fused_linear_cross_entropy import (
-        fused_linear_cross_entropy_loss as _fla_flce,
-    )
-except ImportError:
-    try:
-        from fla.modules import fused_linear_cross_entropy_loss as _fla_flce
-    except ImportError:
-        _fla_flce = None
-
-try:
-    from fla.modules.norm import rms_norm as _fla_rms_norm
-except ImportError:
-    try:
-        from fla.modules import rms_norm as _fla_rms_norm
-    except ImportError:
-        _fla_rms_norm = None
-
-try:
-    from fla.ops.gdn2 import chunk_gdn2, fused_recurrent_gdn2
-    FLA_AVAILABLE = True
-except ImportError:
-    chunk_gdn2 = None
-    fused_recurrent_gdn2 = None
-    FLA_AVAILABLE = False
-
-_USE_FLA_FUSED_CE = True
-_USE_FLA_RMS_NORM = False
-
-if FLA_AVAILABLE:
-
-    @torch.compiler.disable
-    def _fla_chunk(q, k, v, g, b, w, initial_state=None, output_final_state=False,
-                   cu_seqlens=None):
-        return chunk_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
-        )
-
-    @torch.compiler.disable
-    def _fla_recurrent(q, k, v, g, b, w, initial_state=None, output_final_state=False):
-        return fused_recurrent_gdn2(
-            q=q, k=k, v=v, g=g, b=b, w=w,
-            initial_state=initial_state,
-            output_final_state=output_final_state,
-            use_qk_l2norm_in_kernel=True,
-        )
-
-try:
-    from flash_attn import flash_attn_varlen_func
-    FLASH_ATTN_AVAILABLE = True
-except ImportError:
-    flash_attn_varlen_func = None
-    FLASH_ATTN_AVAILABLE = False
-
-
-def _mask_padding_states(hidden_states, padding_mask):
-    if padding_mask is not None:
-        return hidden_states * padding_mask.unsqueeze(-1).to(hidden_states.dtype)
-    return hidden_states
-
-
-def naive_gdn2_reference(q, k, v, g, b, w, initial_state=None):
-    """fp32 reference for GDN-2."""
-    B, T, H, K = q.shape
-    V = v.shape[-1]
-    S = (
-        torch.zeros(B, H, K, V, dtype=torch.float32, device=q.device)
-        if initial_state is None
-        else initial_state.float()
-    )
-    scale = K ** -0.5
-    outs = []
-    for t in range(T):
-        q_t = F.normalize(q[:, t], dim=-1)
-        k_t = F.normalize(k[:, t], dim=-1)
-        v_t, g_t, b_t, w_t = v[:, t], g[:, t], b[:, t], w[:, t]
-        S = S * torch.exp(g_t).unsqueeze(-1)
-        e = b_t * k_t
-        r = torch.einsum("bhkv,bhk->bhv", S, e)
-        z = w_t * v_t
-        S = S + torch.einsum("bhk,bhv->bhkv", k_t, z - r)
-        outs.append(scale * torch.einsum("bhkv,bhk->bhv", S, q_t))
-    return torch.stack(outs, dim=1), S
-
-
-class SurjoRMSNorm(nn.Module):
-    """FLA Triton RMSNorm when available; else fused ATen F.rms_norm."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.eps = eps
-
-    def forward(self, x):
-        if x.dtype != self.weight.dtype:
-            x = x.to(self.weight.dtype)
-        if _USE_FLA_RMS_NORM and _fla_rms_norm is not None and x.is_cuda:
-            return _fla_rms_norm(x, self.weight, None, eps=self.eps)
-        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
-
-
-class ClampedMLP(nn.Module):
-    def __init__(self, config):
-        hidden_size = config.hidden_size
-        intermediate_size = config.intermediate_size
-        super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
-        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-        self.act_fn = ACT2FN[config.hidden_act]
-
-    def forward(self, x):
-        return self.down_proj(
-            self.act_fn(self.gate_proj(x).clamp(-15.0, 15.0)) * self.up_proj(x)
-        )
-
-
-class SurjoRotaryEmbedding(nn.Module):
-    """RoPE with frequencies computed in forward() (no stored buffer)."""
-
-    def __init__(self, config, device=None):
-        super().__init__()
-        self.config = config
-        self.rope_type = config.rope_parameters.get("rope_type", "default")
-        self.max_position_embeddings = config.max_position_embeddings
-        self.base = config.rope_parameters.get("rope_theta", config.rope_theta)
-        self.factor = config.rope_parameters.get("factor", 1.0)
-        self.dim = int(config.head_dim * config.partial_rotary_factor)
-
-    def _inv_freq(self, device):
-        inv = 1.0 / (
-            self.base
-            ** (torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim)
-        )
-        if self.rope_type == "linear":
-            inv = inv / self.factor
-        return inv
-
-    @torch.no_grad()
-    def forward(self, position_ids):
-        inv = self._inv_freq(position_ids.device)
-        freqs = torch.einsum("bt,f->btf", position_ids.float(), inv)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        return emb.cos(), emb.sin()
-
-
-def rotate_half(x):
-    x1, x2 = x.chunk(2, dim=-1)
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rope(q, k, cos, sin):
-    cos = cos[:, None, :, :].to(q.dtype)
-    sin = sin[:, None, :, :].to(q.dtype)
-    q = q * cos + rotate_half(q) * sin
-    k = k * cos + rotate_half(k) * sin
-    return q, k
-
-
-def apply_rope_flat(q, k, cos, sin):
-    """RoPE on flat varlen tensors. q, k: [total, H, D]; cos, sin: [B, T, D]."""
-    total = q.shape[0]
-    cos_f = cos.reshape(total, -1)[:, None, :].to(q.dtype)
-    sin_f = sin.reshape(total, -1)[:, None, :].to(q.dtype)
-    q = q * cos_f + rotate_half(q) * sin_f
-    k = k * cos_f + rotate_half(k) * sin_f
-    return q, k
-
-
-@torch.compiler.disable
-def _flash_varlen(q, k, v, cu_seqlens, max_seqlen):
-    """FA2 varlen. cu_seqlens device int32; max_seqlen a constant (no sync)."""
-    return flash_attn_varlen_func(
-        q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen,
-        dropout_p=0.0, causal=True,
-    )
-
-
-class ShortConvolution(nn.Module):
-    def __init__(self, dim: int, kernel_size: int = 4):
-        super().__init__()
-        self.kernel_size = kernel_size
-        self.conv = nn.Conv1d(dim, dim, kernel_size, groups=dim, bias=False)
-
-    def forward(self, x, cache=None, output_state=False):
-        B, T, C = x.shape
-        xt = x.transpose(1, 2)
-        if T == 1 and cache is not None:
-            comb = torch.cat([cache, xt], dim=2)
-            w = self.conv.weight.squeeze(1)
-            out = F.silu(torch.einsum("bck,ck->bc", comb, w).unsqueeze(-1))
-            new = comb[:, :, 1:] if output_state else None
-            return out.transpose(1, 2), new
-        if cache is not None:
-            xt_in = torch.cat([cache, xt], dim=2)
-        else:
-            xt_in = F.pad(xt, (self.kernel_size - 1, 0))
-        out = F.silu(self.conv(xt_in))
-        new = xt_in[:, :, -(self.kernel_size - 1):] if output_state else None
-        return out.transpose(1, 2), new
-
-
-class SurjoCache:
-    def __init__(self):
-        self.kv = {}
-        self.gdn_state = {}
-        self.gdn_conv = {}
-
-    def get_seq_length(self):
-        if not self.kv:
-            return 0
-        return next(iter(self.kv.values()))[0].shape[2]
-
-    @staticmethod
-    def num_xsa_slots(config):
-        return (
-            config.prelude_layers
-            + config.num_groups * config.recurrent_passes
-            + config.coda_layers
-        )
-
-
-class XSAAttention(nn.Module):
-    """Full attention with XSA value-projection subtraction."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
-        self.num_kv_groups = self.num_heads // self.num_kv_heads
-        self.head_dim = config.head_dim
-        self.xsa_projection = config.xsa_projection
-        self.dropout = config.attention_dropout
-        self.hidden_size = config.hidden_size
-        self.use_flash = getattr(config, "_use_flash_xsa", False)
-        self.q_proj = nn.Linear(
-            self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.k_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.v_proj = nn.Linear(
-            self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias
-        )
-        self.o_proj = nn.Linear(
-            self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
-        )
-        self.q_norm = SurjoRMSNorm(self.head_dim)
-        self.k_norm = SurjoRMSNorm(self.head_dim)
-
-    def _xsa_projection_step(self, attn, v):
-        dv = (attn.float() * v.float()).sum(dim=-1, keepdim=True)
-        vv = (v.float() * v.float()).sum(dim=-1, keepdim=True).clamp_min(1e-4)
-        scale = (dv / vv).to(attn.dtype)
-        return attn - scale * v
-
-    def forward(self, hidden_states, attention_mask, cache, slot, use_cache,
-                cache_position, cos, sin, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        H, Hkv, D = self.num_heads, self.num_kv_heads, self.head_dim
-
-        q = self.q_proj(hidden_states).view(B, T, H, D).transpose(1, 2)
-        k = self.k_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        v = self.v_proj(hidden_states).view(B, T, Hkv, D).transpose(1, 2)
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        use_varlen = (
-            cu_seqlens is not None
-            and cache is None
-            and self.use_flash
-            and FLASH_ATTN_AVAILABLE
-            and hidden_states.is_cuda
-        )
-        if use_varlen:
-            total = B * T
-            q_f = q.transpose(1, 2).reshape(total, H, D).to(torch.bfloat16)
-            k_f = k.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            v_f = v.transpose(1, 2).reshape(total, Hkv, D).to(torch.bfloat16)
-            if Hkv != H:
-                k_f = k_f.repeat_interleave(self.num_kv_groups, dim=1)
-                v_f = v_f.repeat_interleave(self.num_kv_groups, dim=1)
-            q_f, k_f = apply_rope_flat(q_f, k_f, cos, sin)
-            cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-            o_f = _flash_varlen(q_f, k_f, v_f, cu_dev, T)
-            attn_out = o_f.view(B, T, H, D).to(hidden_states.dtype)
-            v_xsa = v_f.view(B, T, H, D).to(hidden_states.dtype)
-            if self.xsa_projection:
-                attn_out = self._xsa_projection_step(attn_out, v_xsa)
-            out = attn_out.reshape(B, T, H * D)
-            return self.o_proj(out)
-
-        q, k = apply_rope(q, k, cos, sin)
-        v_cur = v
-        if cache is not None and use_cache:
-            if slot in cache.kv:
-                pk, pv = cache.kv[slot]
-                k = torch.cat([pk, k.to(pk.dtype)], dim=2)
-                v = torch.cat([pv, v.to(pv.dtype)], dim=2)
-            cache.kv[slot] = (k, v)
-        S = k.shape[2]
-        if Hkv != H:
-            k_e = k.repeat_interleave(self.num_kv_groups, dim=1)
-            v_e = v.repeat_interleave(self.num_kv_groups, dim=1)
-            v_cur_e = v_cur.repeat_interleave(self.num_kv_groups, dim=1)
-        else:
-            k_e, v_e, v_cur_e = k, v, v_cur
-
-        is_causal = False
-        mask4 = None
-        if attention_mask is not None:
-            L = attention_mask.shape[-1]
-            if L < S:
-                attention_mask = F.pad(attention_mask, (S - L, 0), value=1)
-            elif L > S:
-                attention_mask = attention_mask[:, -S:]
-            neg = torch.finfo(q.dtype).min
-            m = attention_mask[:, None, None, :].to(q.dtype)
-            if T == S:
-                if T > 1:
-                    causal = torch.ones(T, S, device=q.device, dtype=q.dtype).tril()
-                    allow = causal[None, None] * m
-                    mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                        torch.full_like(allow, neg))
-                else:
-                    mask4 = torch.where(m > 0, torch.zeros_like(m),
-                                        torch.full_like(m, neg))
-            else:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                causal = kv_pos[None, :] <= cache_position[:, None].to(q.device)
-                allow = causal[None, None].to(q.dtype) * m
-                mask4 = torch.where(allow > 0, torch.zeros_like(allow),
-                                    torch.full_like(allow, neg))
-        elif cu_seqlens is not None and T > 1:
-            neg = torch.finfo(q.dtype).min
-            cu_dev = cu_seqlens.to(q.device, torch.int32)
-            gpos = torch.arange(B * T, device=q.device)
-            doc_id = torch.searchsorted(cu_dev, gpos, right=True) - 1
-            doc_id = doc_id.clamp(min=0).view(B, T)
-            same = doc_id.unsqueeze(2) == doc_id.unsqueeze(1)
-            idx = torch.arange(T, device=q.device)
-            causal = idx.unsqueeze(0) <= idx.unsqueeze(1)
-            allow = (same & causal).unsqueeze(1)
-            mask4 = torch.where(
-                allow,
-                torch.zeros(B, 1, T, T, dtype=q.dtype, device=q.device),
-                torch.full((B, 1, T, T), neg, dtype=q.dtype, device=q.device),
-            )
-        else:
-            if T == S and T > 1:
-                is_causal = True
-            elif T > 1:
-                if cache_position is None:
-                    cache_position = torch.arange(S - T, S, device=q.device)
-                kv_pos = torch.arange(S, device=q.device)
-                neg = torch.finfo(q.dtype).min
-                causal = torch.zeros(T, S, device=q.device, dtype=q.dtype)
-                causal = causal.masked_fill(
-                    kv_pos[None, :] > cache_position[:, None], neg
-                )
-                mask4 = causal[None, None, :, :]
-
-        attn = F.scaled_dot_product_attention(
-            q, k_e, v_e, attn_mask=mask4, is_causal=is_causal,
-            dropout_p=self.dropout if self.training else 0.0,
-        )
-
-        if self.xsa_projection:
-            attn = self._xsa_projection_step(attn, v_cur_e)
-
-        out = attn.transpose(1, 2).reshape(B, T, H * D)
-        return self.o_proj(out)
-
-
-class GDN2Attention(nn.Module):
-    """Gated DeltaNet-2 linear attention."""
-
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        H = config.num_attention_heads
-        Hv = config.gdn_num_v_heads
-        self.num_heads = H
-        self.num_v_heads = Hv
-        self.gva = Hv // H
-        self.head_k_dim = config.gdn_key_dim
-        self.head_v_dim = config.gdn_value_dim
-        self.key_dim = H * self.head_k_dim
-        self.value_dim = Hv * self.head_v_dim
-        self.conv_kernel = config.gdn_conv_kernel_size
-        self.allow_neg_eigval = config.gdn_allow_neg_eigval
-        hidden = config.hidden_size
-        self.q_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.k_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.v_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.q_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.k_conv = ShortConvolution(self.key_dim, self.conv_kernel)
-        self.v_conv = ShortConvolution(self.value_dim, self.conv_kernel)
-        self.f_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.key_dim, bias=False),
-        )
-        self.b_proj = nn.Linear(hidden, self.key_dim, bias=False)
-        self.w_proj = nn.Linear(hidden, self.value_dim, bias=False)
-        self.A_log = nn.Parameter(
-            torch.log(torch.empty(H, dtype=torch.float32).uniform_(1, 16))
-        )
-        self.A_log._no_weight_decay = True
-        dt = torch.exp(
-            torch.rand(self.key_dim, dtype=torch.float32)
-            * (math.log(0.1) - math.log(0.001))
-            + math.log(0.001)
-        ).clamp(min=1e-4)
-        self.dt_bias = nn.Parameter(dt + torch.log(-torch.expm1(-dt)))
-        self.dt_bias._no_weight_decay = True
-        self.g_proj = nn.Sequential(
-            nn.Linear(hidden, self.head_v_dim, bias=False),
-            nn.Linear(self.head_v_dim, self.value_dim, bias=True),
-        )
-        self.o_norm = SurjoRMSNorm(self.head_v_dim, eps=config.rms_norm_eps)
-        self.o_proj = nn.Linear(self.value_dim, hidden, bias=False)
-
-    def forward(self, hidden_states, attention_mask, gdn_state, conv_state,
-                output_state, cu_seqlens=None):
-        B, T, _ = hidden_states.shape
-        if attention_mask is not None:
-            if attention_mask.shape[-1] != T:
-                attention_mask = (
-                    attention_mask[:, -T:] if attention_mask.shape[-1] > T else None
-                )
-            if attention_mask is not None:
-                hidden_states = _mask_padding_states(hidden_states, attention_mask)
-        cq = conv_state[0] if conv_state is not None else None
-        ck = conv_state[1] if conv_state is not None else None
-        cvs = conv_state[2] if conv_state is not None else None
-        q, nq = self.q_conv(self.q_proj(hidden_states), cq, output_state)
-        k, nk = self.k_conv(self.k_proj(hidden_states), ck, output_state)
-        v, nv = self.v_conv(self.v_proj(hidden_states), cvs, output_state)
-        g = F.softplus(self.f_proj(hidden_states).float() + self.dt_bias)
-        b = self.b_proj(hidden_states).sigmoid()
-        w = self.w_proj(hidden_states).sigmoid()
-        q = q.reshape(B, T, self.num_heads, self.head_k_dim)
-        k = k.reshape(B, T, self.num_heads, self.head_k_dim)
-        v = v.reshape(B, T, self.num_v_heads, self.head_v_dim)
-        g = g.view(B, T, self.num_heads, self.head_k_dim)
-        b = b.view(B, T, self.num_heads, self.head_k_dim)
-        w = w.view(B, T, self.num_v_heads, self.head_v_dim)
-        g = -self.A_log.float().exp().view(1, 1, self.num_heads, 1) * g
-        if self.allow_neg_eigval:
-            b = b * 2.0
-        if self.gva > 1:
-            q = q.repeat_interleave(self.gva, dim=2)
-            k = k.repeat_interleave(self.gva, dim=2)
-            g = g.repeat_interleave(self.gva, dim=2)
-            b = b.repeat_interleave(self.gva, dim=2)
-        init = gdn_state.float() if gdn_state is not None else None
-
-        if FLA_AVAILABLE and hidden_states.is_cuda:
-            if T == 1 and init is not None:
-                o, fs = _fla_recurrent(q, k, v, g, b, w, init, output_state)
-            elif cu_seqlens is not None:
-                total = B * T
-                cu_dev = cu_seqlens.to(q.device, torch.int32).contiguous()
-                q_f = q.reshape(1, total, self.num_heads, self.head_k_dim)
-                k_f = k.reshape(1, total, self.num_heads, self.head_k_dim)
-                v_f = v.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                g_f = g.reshape(1, total, self.num_heads, self.head_k_dim)
-                b_f = b.reshape(1, total, self.num_heads, self.head_k_dim)
-                w_f = w.reshape(1, total, self.num_v_heads, self.head_v_dim)
-                o, fs = _fla_chunk(q_f, k_f, v_f, g_f, b_f, w_f, init,
-                                   output_state, cu_dev)
-                o = o.reshape(B, T, self.num_v_heads, self.head_v_dim)
-            else:
-                o, fs = _fla_chunk(q, k, v, g, b, w, init, output_state, None)
-        else:
-            o, fs = naive_gdn2_reference(q.float(), k.float(), v.float(), g,
-                                         b.float(), w.float(), init)
-
-        # Cast back to activation precision
-        o = o.to(hidden_states.dtype)
-
-        fs = fs if output_state else None
-        gate = self.g_proj(hidden_states).view(B, T, self.num_v_heads, self.head_v_dim)
-        o = self.o_norm(o)
-        o = o * torch.sigmoid(gate)
-        o = o.reshape(B, T, self.value_dim)
-        conv_out = (nq, nk, nv) if output_state else None
-        return self.o_proj(o), fs, conv_out
-
-
-def _is_xsa(config, li: int) -> bool:
-    if li < config.prelude_layers:
-        return True
-    if li >= config.prelude_layers + config.recurrent_layers:
-        return True
-    return (li - config.prelude_layers) % (config.gdn_per_xsa + 1) == config.gdn_per_xsa
-
-
-class SurjoDecoderLayer(nn.Module):
-    def __init__(self, config, layer_idx: int):
-        super().__init__()
-        self.layer_idx = layer_idx
-        self.layer_type = "xsa" if _is_xsa(config, layer_idx) else "gdn"
-        if self.layer_type == "xsa":
-            self.self_attn = XSAAttention(config, layer_idx)
-        else:
-            self.linear_attn = GDN2Attention(config, layer_idx)
-        self.mlp = ClampedMLP(config)
-        self.input_layernorm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = SurjoRMSNorm(config.hidden_size,
-                                                     eps=config.rms_norm_eps)
-
-    def forward(self, hidden_states, attention_mask=None, cos=None, sin=None,
-                cu_seqlens=None):
-        h = self.input_layernorm(hidden_states)
-        if self.layer_type == "xsa":
-            out = self.self_attn(h, attention_mask, None, -1, False, None,
-                                 cos, sin, cu_seqlens)
-        else:
-            out, _, _ = self.linear_attn(h, attention_mask, None, None, False,
-                                         cu_seqlens)
-        hidden_states = hidden_states + out
-        h2 = self.post_attention_layernorm(hidden_states)
-        return hidden_states + self.mlp(h2)
-
-
-class SurjoPreTrainedModel(PreTrainedModel):
-    config_class = SurjoConfig
-    base_model_prefix = "model"
-    _no_split_modules = ["SurjoDecoderLayer"]
-    _supports_flash_attn = False
-    _supports_flash_attn_2 = False
-    _supports_sdpa = True
-    _is_stateful = True
-    _can_compile_fullgraph = True
-
-    def _init_weights(self, module):
-        std = self.config.initializer_range
-        if isinstance(module, nn.Linear):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.Embedding):
-            module.weight.data.normal_(mean=0.0, std=std)
-            if module.padding_idx is not None:
-                module.weight.data[module.padding_idx].zero_()
-        elif isinstance(module, GDN2Attention):
-            module.A_log.data.copy_(
-                torch.log(
-                    torch.empty(module.num_heads, dtype=torch.float32,
-                                device=module.A_log.device).uniform_(1, 16)
-                )
-            )
-            dt = torch.exp(
-                torch.rand(module.key_dim, dtype=torch.float32,
-                           device=module.dt_bias.device)
-                * (math.log(0.1) - math.log(0.001)) + math.log(0.001)
-            ).clamp(min=1e-4)
-            module.dt_bias.data.copy_(dt + torch.log(-torch.expm1(-dt)))
-
-
-class SurjoModel(SurjoPreTrainedModel):
-    def __init__(self, config):
-        super().__init__(config)
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = nn.ModuleList(
-            [SurjoDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
-        )
-        self.norm = SurjoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.rotary = SurjoRotaryEmbedding(config)
-        self.grad_checkpointing = False
-
-        slot = 0
-        self.xsa_slot_map = {}
-        self.exec_plan = []
-        for li in range(config.prelude_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        for p in range(config.recurrent_passes):
-            for g in range(config.num_groups):
-                base = config.prelude_layers + g * (config.gdn_per_xsa + 1)
-                for j in range(config.gdn_per_xsa):
-                    self.exec_plan.append((self.layers[base + j], base + j, p, -1))
-                li = base + config.gdn_per_xsa
-                self.xsa_slot_map[(li, p)] = slot
-                self.exec_plan.append((self.layers[li], li, p, slot))
-                slot += 1
-        for li in range(config.prelude_layers + config.recurrent_layers,
-                        config.num_hidden_layers):
-            self.xsa_slot_map[(li, 0)] = slot
-            self.exec_plan.append((self.layers[li], li, 0, slot))
-            slot += 1
-        self.num_xsa_slots = slot
-        self.post_init()
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = value
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, use_cache=None,
-                cache_position=None, return_dict=True,
-                output_hidden_states=False, output_attentions=False,
-                cu_seqlens=None, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        if use_cache is None:
-            use_cache = past_key_values is not None
-        if use_cache and past_key_values is None:
-            past_key_values = SurjoCache()
-        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
-
-        if past_len > 0:
-            if input_ids is not None:
-                n_ids = input_ids.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_ids
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    input_ids = input_ids[:, -n_new:] if n_new <= n_ids else input_ids
-                elif n_ids > past_len:
-                    input_ids = input_ids[:, past_len:]
-                else:
-                    input_ids = input_ids[:, -1:]
-            elif inputs_embeds is not None:
-                n_emb = inputs_embeds.shape[1]
-                if (attention_mask is not None
-                        and attention_mask.shape[-1] > n_emb
-                        and attention_mask.shape[-1] > past_len):
-                    n_new = attention_mask.shape[-1] - past_len
-                    inputs_embeds = (inputs_embeds[:, -n_new:] if n_new <= n_emb
-                                     else inputs_embeds)
-                elif n_emb > past_len:
-                    inputs_embeds = inputs_embeds[:, past_len:]
-                else:
-                    inputs_embeds = inputs_embeds[:, -1:]
-
-        if inputs_embeds is None:
-            if input_ids is None:
-                raise ValueError("Specify input_ids or inputs_embeds")
-            inputs_embeds = self.embed_tokens(input_ids)
-        B, T = inputs_embeds.shape[0], inputs_embeds.shape[1]
-
-        cache_position = torch.arange(past_len, past_len + T,
-                                      device=inputs_embeds.device)
-        if position_ids is not None and position_ids.shape[-1] != T:
-            position_ids = None
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0).expand(B, -1)
-        cos, sin = self.rotary(position_ids)
-
-        gdn_states = dict(past_key_values.gdn_state) if past_key_values is not None else {}
-        conv_states = dict(past_key_values.gdn_conv) if past_key_values is not None else {}
-
-        noise_std = self.config.loop_noise_std
-        hidden = inputs_embeds
-        prev_p = 0
-        use_ckpt = self.training and self.grad_checkpointing
-        all_hidden_states = () if output_hidden_states else None
-
-        for (layer, li, p, slot) in self.exec_plan:
-            if output_hidden_states:
-                all_hidden_states += (hidden,)
-            if p != prev_p and p > 0 and self.training and noise_std > 0.0:
-                hidden = hidden + torch.randn_like(hidden) * noise_std
-            prev_p = p
-
-            if use_ckpt:
-                hidden = torch.utils.checkpoint.checkpoint(
-                    self._train_layer, layer, hidden, attention_mask, cos, sin,
-                    cu_seqlens, use_reentrant=False,
-                )
-            elif use_cache:
-                h = layer.input_layernorm(hidden)
-                if layer.layer_type == "xsa":
-                    out = layer.self_attn(h, attention_mask, past_key_values,
-                                          slot, use_cache, cache_position,
-                                          cos, sin, cu_seqlens)
-                else:
-                    gs = gdn_states.get((li, p))
-                    cs = conv_states.get((li, p))
-                    out, ngs, ncs = layer.linear_attn(h, attention_mask, gs, cs,
-                                                      use_cache, cu_seqlens)
-                    if ngs is not None:
-                        gdn_states[(li, p)] = ngs
-                    if ncs is not None:
-                        conv_states[(li, p)] = ncs
-                hidden = hidden + out
-                h2 = layer.post_attention_layernorm(hidden)
-                hidden = hidden + layer.mlp(h2)
-            else:
-                hidden = layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-        hidden = self.norm(hidden)
-        if output_hidden_states:
-            all_hidden_states += (hidden,)
-
-        if use_cache and past_key_values is not None:
-            past_key_values.gdn_state.update(gdn_states)
-            past_key_values.gdn_conv.update(conv_states)
-
-        if not return_dict:
-            out = tuple(v for v in [hidden, past_key_values if use_cache else None,
-                                    all_hidden_states] if v is not None)
-            return out
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden,
-            past_key_values=past_key_values if use_cache else None,
-            hidden_states=all_hidden_states,
-            attentions=None,
-        )
-
-    def _train_layer(self, layer, hidden, attention_mask, cos, sin, cu_seqlens):
-        return layer(hidden, attention_mask, cos, sin, cu_seqlens)
-
-
-class SurjoForCausalLM(SurjoPreTrainedModel, GenerationMixin):
-    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
-
-    def __init__(self, config):
-        super().__init__(config)
-        self.model = SurjoModel(config)
-        self.vocab_size = config.vocab_size
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.post_init()
-
-    def tie_weights(self, **kwargs):
-        try:
-            super().tie_weights(**kwargs)
-        except TypeError:
-            pass
-        except Exception:
-            pass
-        if (getattr(self.config, "tie_word_embeddings", False)
-                and getattr(self, "model", None) is not None
-                and getattr(self, "lm_head", None) is not None):
-            self.lm_head.weight = self.model.embed_tokens.weight
-
-    def save_pretrained(self, *args, **kwargs):
-        was_tied = (getattr(self, "lm_head", None) is not None
-                    and getattr(self, "model", None) is not None
-                    and self.lm_head.weight is self.model.embed_tokens.weight)
-        if was_tied and self.config.tie_word_embeddings:
-            self.lm_head.weight = nn.Parameter(
-                self.model.embed_tokens.weight.detach().clone(),
-                requires_grad=self.model.embed_tokens.weight.requires_grad,
-            )
-        try:
-            return super().save_pretrained(*args, **kwargs)
-        finally:
-            if was_tied and self.config.tie_word_embeddings:
-                self.tie_weights()
-
-    def get_input_embeddings(self):
-        return self.model.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.model.embed_tokens = value
-
-    def get_output_embeddings(self):
-        return self.lm_head
-
-    def set_output_embeddings(self, value):
-        self.lm_head = value
-
-    def set_decoder(self, decoder):
-        self.model = decoder
-
-    def get_decoder(self):
-        return self.model
-
-    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None, **kwargs):
-        self.model.grad_checkpointing = True
-
-    def gradient_checkpointing_disable(self):
-        self.model.grad_checkpointing = False
-
-    def forward(self, input_ids=None, attention_mask=None, position_ids=None,
-                past_key_values=None, inputs_embeds=None, labels=None,
-                use_cache=None, output_hidden_states=None, output_attentions=None,
-                return_dict=None, logits_to_keep=0, cache_position=None,
-                cu_seqlens=None, **kwargs):
-        if logits_to_keep in (None, 0) and kwargs.get("num_logits_to_keep"):
-            logits_to_keep = kwargs.pop("num_logits_to_keep")
-        return_dict = True if return_dict is None else return_dict
-        if use_cache is None:
-            use_cache = not (self.training or labels is not None)
-
-        outputs = self.model(
-            input_ids=input_ids, attention_mask=attention_mask,
-            position_ids=position_ids, past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds, use_cache=use_cache,
-            cache_position=cache_position,
-            output_hidden_states=output_hidden_states,
-            output_attentions=output_attentions,
-            cu_seqlens=cu_seqlens,
-            **kwargs,
-        )
-        hidden_states = outputs[0]
-
-        loss, logits = None, None
-        if labels is not None:
-            sh = hidden_states[..., :-1, :].contiguous()
-            sl = labels[..., 1:].contiguous()
-            if _USE_FLA_FUSED_CE and _fla_flce is not None:
-                loss = _fla_flce(
-                    sh.view(-1, sh.shape[-1]),
-                    sl.view(-1),
-                    self.lm_head.weight,
-                    num_chunks=8,
-                    reduction="mean",
-                )
-            else:
-                n_chunks = max(8, sh.shape[0] // 8)
-                tot = hidden_states.new_zeros((), dtype=torch.float32)
-                n_tok = 0
-                for h_c, l_c in zip(sh.chunk(n_chunks, dim=0), sl.chunk(n_chunks, dim=0)):
-                    if l_c.numel() == 0:
-                        continue
-                    lg = self.lm_head(h_c)
-                    tot = tot + F.cross_entropy(
-                        lg.float().view(-1, lg.size(-1)), l_c.view(-1), reduction="sum"
-                    )
-                    n_tok += l_c.numel()
-                loss = tot / max(n_tok, 1)
-        else:
-            if isinstance(logits_to_keep, int) and logits_to_keep > 0:
-                sl_hidden = hidden_states[:, -logits_to_keep:, :]
-            elif torch.is_tensor(logits_to_keep):
-                sl_hidden = hidden_states[:, logits_to_keep, :]
-            else:
-                sl_hidden = hidden_states
-            logits = self.lm_head(sl_hidden)
-
-        if not return_dict:
-            out = (logits,) + outputs[1:]
-            return (loss,) + out if loss is not None else out
-        return CausalLMOutputWithPast(
-            loss=loss, logits=logits,
-            past_key_values=outputs.past_key_values,
-            hidden_states=outputs.hidden_states,
-            attentions=outputs.attentions,
-        )
-
-    def prepare_inputs_for_generation(self, input_ids, past_key_values=None,
-                                      attention_mask=None, inputs_embeds=None,
-                                      cache_position=None, use_cache=True, **kwargs):
-        if past_key_values is not None and not isinstance(past_key_values, SurjoCache):
-            past_key_values = None
-        has_state = past_key_values is not None and len(past_key_values.kv) > 0
-        if has_state:
-            input_ids = input_ids[:, -1:]
-            if kwargs.get("logits_to_keep") is None:
-                kwargs["logits_to_keep"] = 1
-        past_len = past_key_values.get_seq_length() if has_state else 0
-        cache_position = torch.arange(past_len, past_len + input_ids.shape[1],
-                                      device=input_ids.device)
-        model_inputs = {
-            "input_ids": input_ids.contiguous(),
-            "past_key_values": past_key_values,
-            "use_cache": use_cache,
-            "cache_position": cache_position,
-            "attention_mask": attention_mask,
-        }
-        model_inputs.update(kwargs)
-        return model_inputs
-
-    def _reorder_cache(self, past_key_values: SurjoCache, beam_idx):
-        for slot, (k, v) in past_key_values.kv.items():
-            past_key_values.kv[slot] = (k.index_select(0, beam_idx),
-                                        v.index_select(0, beam_idx))
-        for key in list(past_key_values.gdn_state.keys()):
-            s = past_key_values.gdn_state[key]
-            if s is not None:
-                past_key_values.gdn_state[key] = s.index_select(0, beam_idx)
-        for key in list(past_key_values.gdn_conv.keys()):
-            cs = past_key_values.gdn_conv[key]
-            if cs is not None:
-                past_key_values.gdn_conv[key] = tuple(
-                    c.index_select(0, beam_idx) if c is not None else None for c in cs
-                )
-        return past_key_values
-
-
-SurjoConfig.register_for_auto_class()
-SurjoForCausalLM.register_for_auto_class()