Upload folder using huggingface_hub

chat_template.jinja

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

config.json

@@ -0,0 +1,202 @@
+{
+  "architectures": [
+    "RicursiveForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_ricursive.RicursiveConfig",
+    "AutoModel": "modeling_ricursive.RicursiveModel",
+    "AutoModelForCausalLM": "modeling_ricursive.RicursiveForCausalLM"
+  },
+  "bos_token_id": 100257,
+  "embed_scale": 71.55417527999327,
+  "eos_token_id": [
+    100265,
+    100257
+  ],
+  "ffn_clamp_threshold": 20.0,
+  "gate_full_precision": true,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 5120,
+  "hybrid_layer_pattern": [
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0,
+    1,
+    1,
+    1,
+    1,
+    0
+  ],
+  "intermediate_size": 40960,
+  "is_hybrid_swa": true,
+  "layer_types": [
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 2097152,
+  "model_type": "ricursive",
+  "num_attention_heads": 80,
+  "num_hidden_layers": 80,
+  "num_key_value_heads": 8,
+  "pad_token_id": 100277,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 10000.0,
+  "sliding_window": 1024,
+  "tie_word_embeddings": false,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "use_gated_attention": true,
+  "use_valres": true,
+  "vocab_size": 100352
+}

configuration_ricursive.py

@@ -0,0 +1,123 @@
+"""HuggingFace PretrainedConfig for the Ricursive transformer architecture."""
+
+from transformers import PretrainedConfig
+
+
+class RicursiveConfig(PretrainedConfig):
+    """
+    Configuration for the Ricursive transformer.
+
+    Key architectural features:
+    - PeriNorm blocks (4 layer norms per block instead of 2)
+    - Gated attention with elementwise sigmoid gate
+    - Mixed local/global layers: local layers use sliding window + RoPE,
+      global layers (every 5th) use full attention with no RoPE
+    - Per-head Q/K RMSNorm
+    """
+
+    model_type = "ricursive"
+
+    def __init__(
+        self,
+        vocab_size=100352,
+        hidden_size=3072,
+        intermediate_size=24576,
+        num_hidden_layers=50,
+        num_attention_heads=48,
+        num_key_value_heads=8,
+        head_dim=128,
+        max_position_embeddings=2048,
+        hidden_act="silu",
+        rms_norm_eps=1e-6,
+        rope_theta=10000.0,
+        sliding_window=1024,
+        layer_types=None,
+        use_gated_attention=True,
+        gate_full_precision=True,
+        use_valres=False,
+        ffn_clamp_threshold=None,
+        embed_scale=None,
+        tie_word_embeddings=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_act = hidden_act
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.sliding_window = sliding_window
+        self.use_gated_attention = use_gated_attention
+        self.gate_full_precision = gate_full_precision
+        self.use_valres = use_valres
+        self.ffn_clamp_threshold = ffn_clamp_threshold
+        # Hybrid-SWA metadata is derived below. Serialized copies must not override the
+        # architecture encoded by layer_types.
+        kwargs.pop("is_hybrid_swa", None)
+        kwargs.pop("hybrid_layer_pattern", None)
+        # Embedding scale convention (see modeling_ricursive.RicursiveModel.forward):
+        #   - non-folded checkpoints set embed_scale (e.g. sqrt(hidden_size)) and the model
+        #     scales the embeddings at runtime;
+        #   - folded checkpoints bake the scale into embed_tokens.weight and leave this None.
+        # Either way the SGLang native port (sglang_ricursive) reads this same key, so the
+        # served and trained forwards agree.
+        self.embed_scale = embed_scale
+
+        # Per-layer type: "sliding_attention" (window + RoPE) or "full_attention" (no RoPE)
+        # These names follow the HuggingFace PretrainedConfig layer_types validation.
+        if layer_types is not None:
+            self.layer_types = layer_types
+        else:
+            # Default: every 5th layer is full_attention (indices 4, 9, 14, ...)
+            self.layer_types = [
+                "full_attention" if i % 5 == 4 else "sliding_attention"
+                for i in range(num_hidden_layers)
+            ]
+        if len(self.layer_types) != num_hidden_layers:
+            raise ValueError(
+                f"layer_types has {len(self.layer_types)} entries, expected {num_hidden_layers}"
+            )
+        invalid_layer_types = set(self.layer_types) - {
+            "sliding_attention",
+            "full_attention",
+        }
+        if invalid_layer_types:
+            raise ValueError(f"unsupported Ricursive layer_types: {sorted(invalid_layer_types)}")
+
+        # SGLang's compact hybrid KV allocator uses 1 for sliding attention and 0 for full
+        # attention. Keep these fields fully derived so they cannot drift from model semantics.
+        self.hybrid_layer_pattern = [
+            1 if layer_type == "sliding_attention" else 0 for layer_type in self.layer_types
+        ]
+        self.is_hybrid_swa = 0 in self.hybrid_layer_pattern and 1 in self.hybrid_layer_pattern
+
+        # Standard HF config fields expected by modeling code
+        self.use_cache = kwargs.pop("use_cache", True)
+
+        self.base_model_tp_plan = {
+            "layers.*.self_attn.q_proj": "colwise",
+            "layers.*.self_attn.k_proj": "colwise",
+            "layers.*.self_attn.v_proj": "colwise",
+            "layers.*.self_attn.o_proj": "rowwise",
+            "layers.*.self_attn.g_proj": "colwise",
+            "layers.*.mlp.gate_proj": "colwise",
+            "layers.*.mlp.up_proj": "colwise",
+            "layers.*.mlp.down_proj": "rowwise",
+        }
+        self.auto_map = {
+            "AutoConfig": "configuration_ricursive.RicursiveConfig",
+            "AutoModel": "modeling_ricursive.RicursiveModel",
+            "AutoModelForCausalLM": "modeling_ricursive.RicursiveForCausalLM",
+        }
+
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+
+        # Must be set AFTER super().__init__ — HF's PretrainedConfig.__init__
+        # pops "architectures" from kwargs (default None) and assigns it to
+        # self.architectures, which would otherwise overwrite our value.
+        self.architectures = ["RicursiveForCausalLM"]

model-00001-of-00016.safetensors

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

model-00002-of-00016.safetensors

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

model-00003-of-00016.safetensors

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

model-00004-of-00016.safetensors

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

model-00005-of-00016.safetensors

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

model-00006-of-00016.safetensors

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

model-00007-of-00016.safetensors

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

model-00008-of-00016.safetensors

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

model-00009-of-00016.safetensors

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

model-00010-of-00016.safetensors

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

model-00011-of-00016.safetensors

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

model-00012-of-00016.safetensors

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

model-00013-of-00016.safetensors

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

model-00014-of-00016.safetensors

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

model-00015-of-00016.safetensors

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

model-00016-of-00016.safetensors

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

modeling_ricursive.py

@@ -0,0 +1,606 @@
+"""HuggingFace modeling implementation for the Ricursive transformer architecture.
+
+Designed for inference with vLLM (via --trust-remote-code). Key differences from
+standard Llama/Mistral:
+- PeriNorm blocks: 4 layer norms per block (pre/post attention + pre/post FFN)
+- Gated attention: elementwise sigmoid gate on attention output
+- Mixed local/global layers: local = sliding window + RoPE, global = full attention + no RoPE
+- Per-head Q/K RMSNorm
+"""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+
+from .configuration_ricursive import RicursiveConfig
+
+
+class RicursiveRMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+class RicursiveRotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        # NOTE: inv_freq is computed on-the-fly in forward() rather than stored
+        # as a buffer, because non-persistent buffers get corrupted during
+        # from_pretrained() loading.
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
+        inv_freq = 1.0 / (
+            self.theta ** (torch.arange(0, self.dim, 2, device=x.device).float() / self.dim)
+        )
+        inv_freq_expanded = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos()
+            sin = emb.sin()
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    cos = cos.unsqueeze(1)  # (B, 1, T, dim)
+    sin = sin.unsqueeze(1)  # (B, 1, T, dim)
+    # OLMo-core applies RoPE in full precision (float32) for numerical stability
+    orig_dtype = q.dtype
+    q_float, k_float = q.float(), k.float()
+    cos_float, sin_float = cos.float(), sin.float()
+    q_embed = (q_float * cos_float) + (rotate_half(q_float) * sin_float)
+    k_embed = (k_float * cos_float) + (rotate_half(k_float) * sin_float)
+    return q_embed.to(orig_dtype), k_embed.to(orig_dtype)
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    if n_rep == 1:
+        return hidden_states
+    batch, num_kv_heads, slen, head_dim = hidden_states.shape
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_kv_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)
+
+
+class RicursiveAttention(nn.Module):
+    def __init__(self, config: RicursiveConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.layer_type = config.layer_types[layer_idx]
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+        if config.use_valres:
+            self.value_res_lambda = nn.Parameter(torch.zeros(1))
+
+        # Per-head Q/K norms (RMSNorm on head_dim)
+        self.q_norm = RicursiveRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RicursiveRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+        # Gated attention
+        if config.use_gated_attention:
+            self.g_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+
+        # RoPE only for local (sliding_attention) layers
+        if self.layer_type == "sliding_attention":
+            self.rotary_emb = RicursiveRotaryEmbedding(self.head_dim, theta=config.rope_theta)
+
+    def _apply_sliding_window_mask(
+        self,
+        attention_mask: Optional[torch.Tensor],
+        q_len: int,
+        kv_len: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_position: Optional[torch.LongTensor],
+    ) -> Optional[torch.Tensor]:
+        if self.layer_type != "sliding_attention":
+            return attention_mask
+
+        key_positions = torch.arange(kv_len, device=device)
+        if cache_position is not None:
+            query_positions = cache_position.to(device=device)
+        else:
+            query_positions = torch.arange(kv_len - q_len, kv_len, device=device)
+
+        window = self.config.sliding_window
+        allowed = (
+            (key_positions[None, :] <= query_positions[:, None])
+            & (key_positions[None, :] >= query_positions[:, None] - (window - 1))
+        )[None, None, :, :]
+
+        if attention_mask is None:
+            return allowed
+        if attention_mask.dtype == torch.bool:
+            return attention_mask & allowed
+        return attention_mask.masked_fill(~allowed, torch.finfo(dtype).min)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        value_residual: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.shape
+
+        # Project Q, K, V
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        value_states = hidden_states
+        if value_residual is not None:
+            if not self.config.use_valres:
+                raise RuntimeError("value_residual was supplied but config.use_valres is false")
+            blend = torch.sigmoid(self.value_res_lambda[0]).to(dtype=hidden_states.dtype)
+            value_states = torch.lerp(
+                value_residual.to(dtype=hidden_states.dtype), hidden_states, blend
+            )
+        v = self.v_proj(value_states)
+
+        # Reshape to (B, T, n_heads, head_dim) then transpose to (B, n_heads, T, head_dim)
+        q = q.view(bsz, q_len, self.num_heads, self.head_dim)
+        k = k.view(bsz, q_len, self.num_kv_heads, self.head_dim)
+        v = v.view(bsz, q_len, self.num_kv_heads, self.head_dim)
+
+        # Per-head Q/K norm (applied before transpose, on head_dim)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Transpose to (B, n_heads, T, head_dim) for attention
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        # RoPE only for local layers
+        if self.layer_type == "sliding_attention":
+            cos, sin = self.rotary_emb(v, position_ids)
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        # KV cache
+        if past_key_value is not None:
+            k, v = past_key_value.update(k, v, self.layer_idx)
+
+        # GQA: expand KV heads
+        k = repeat_kv(k, self.num_kv_groups)
+        v = repeat_kv(v, self.num_kv_groups)
+
+        attention_mask = self._apply_sliding_window_mask(
+            attention_mask=attention_mask,
+            q_len=q.shape[2],
+            kv_len=k.shape[2],
+            dtype=q.dtype,
+            device=q.device,
+            cache_position=cache_position,
+        )
+
+        # Attention (use is_causal when no explicit mask for correct causal behavior)
+        # Cast to plain bool to avoid SymBool issues with torch.compile
+        is_causal = bool(attention_mask is None and q.shape[2] > 1)
+        attn_output = F.scaled_dot_product_attention(
+            q, k, v, attn_mask=attention_mask, dropout_p=0.0, is_causal=is_causal
+        )
+
+        # (B, n_heads, T, head_dim) -> (B, T, n_heads * head_dim)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+
+        # Gated attention: output * sigmoid(g_proj(input))
+        if self.config.use_gated_attention:
+            gate = self.g_proj(hidden_states)
+            if self.config.gate_full_precision:
+                gate = gate.float()
+            gate = torch.sigmoid(gate).to(attn_output.dtype)
+            attn_output = attn_output.view(bsz, q_len, -1) * gate
+        else:
+            attn_output = attn_output.view(bsz, q_len, -1)
+
+        # Output projection
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+class RicursiveMLP(nn.Module):
+    def __init__(self, config: RicursiveConfig):
+        super().__init__()
+        self.clamp_threshold = config.ffn_clamp_threshold
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        if self.clamp_threshold:
+            gate = gate.clamp(min=None, max=self.clamp_threshold)
+            up = up.clamp(min=-self.clamp_threshold, max=self.clamp_threshold)
+        return self.down_proj(F.silu(gate) * up)
+
+
+class RicursiveDecoderLayer(nn.Module):
+    """PeriNorm decoder layer with 4 layer norms."""
+
+    def __init__(self, config: RicursiveConfig, layer_idx: int):
+        super().__init__()
+        self.self_attn = RicursiveAttention(config, layer_idx)
+        self.mlp = RicursiveMLP(config)
+
+        # PeriNorm: 4 norms
+        self.input_layernorm = RicursiveRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RicursiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.pre_feedforward_layernorm = RicursiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = RicursiveRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        value_residual: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, ...]:
+        # Attention sub-block (PeriNorm: pre-norm -> attn -> post-norm -> residual)
+        normed = self.input_layernorm(hidden_states)
+        attn_output, attn_weights, present_key_value = self.self_attn(
+            normed,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            value_residual=value_residual,
+            **kwargs,
+        )
+        hidden_states = hidden_states + self.post_attention_layernorm(attn_output)
+
+        # FFN sub-block (PeriNorm: pre-norm -> ffn -> post-norm -> residual)
+        normed = self.pre_feedforward_layernorm(hidden_states)
+        ff_output = self.mlp(normed)
+        hidden_states = hidden_states + self.post_feedforward_layernorm(ff_output)
+
+        outputs = (hidden_states,)
+        if output_attentions:
+            outputs += (attn_weights,)
+        if use_cache:
+            outputs += (present_key_value,)
+        return outputs
+
+
+class RicursivePreTrainedModel(PreTrainedModel):
+    config_class = RicursiveConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["RicursiveDecoderLayer"]
+    _supports_sdpa = True
+    _supports_cache_class = True
+
+
+class RicursiveModel(RicursivePreTrainedModel):
+    def __init__(self, config: RicursiveConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [RicursiveDecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = RicursiveRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # embed_scale: non-folded checkpoints carry the scale (e.g. sqrt(hidden_size)) in the
+        # config and apply it here at runtime; folded checkpoints set embed_scale=None and bake
+        # the scale into embed_tokens.weight, so this is a no-op. This mirrors the SGLang native
+        # port (sglang_ricursive) so served and trained forwards agree. It must run BEFORE the
+        # value-residual seed below, whose source is the (scaled) layer-0 embeddings.
+        embed_scale = getattr(self.config, "embed_scale", None)
+        if embed_scale is not None:
+            inputs_embeds = inputs_embeds * embed_scale
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Build causal mask
+        causal_mask = self._update_causal_mask(
+            attention_mask, inputs_embeds, cache_position, past_key_values
+        )
+
+        hidden_states = inputs_embeds
+        value_residual = hidden_states if self.config.use_valres else None
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+
+        for layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            layer_outputs = layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_values,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                value_residual=value_residual,
+                **kwargs,
+            )
+            hidden_states = layer_outputs[0]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = past_key_values if use_cache else None
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
+                if v is not None
+            )
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    def _update_causal_mask(
+        self,
+        attention_mask: Optional[torch.Tensor],
+        input_tensor: torch.Tensor,
+        cache_position: torch.Tensor,
+        past_key_values: Optional[Cache],
+    ) -> Optional[torch.Tensor]:
+        # When there's no explicit attention_mask (no padding) and we're doing
+        # standard prefill or cached decoding, return None and let SDPA use
+        # its efficient is_causal path. This also avoids numerical issues
+        # from float min mask values with large activations.
+        if attention_mask is None:
+            return None
+        if torch.all(attention_mask == 1):
+            return None
+
+        # Only build an explicit mask when we have a padding mask
+        dtype, device = input_tensor.dtype, input_tensor.device
+        sequence_length = input_tensor.shape[1]
+
+        if past_key_values is not None:
+            target_length = past_key_values.get_max_cache_shape()
+            if target_length is None or target_length < 0:
+                target_length = (
+                    cache_position[-1].item() + 1 if cache_position is not None else sequence_length
+                )
+        else:
+            target_length = sequence_length
+
+        causal_mask = torch.full(
+            (sequence_length, target_length),
+            fill_value=torch.finfo(dtype).min,
+            dtype=dtype,
+            device=device,
+        )
+        if sequence_length != 1:
+            causal_mask = torch.triu(causal_mask, diagonal=target_length - sequence_length + 1)
+
+        causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
+
+        if attention_mask.dim() == 2:
+            mask_length = attention_mask.shape[-1]
+            padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(
+                causal_mask.dtype
+            ).masked_fill(attention_mask[:, None, None, :] == 0, torch.finfo(dtype).min)
+            causal_mask = (
+                torch.cat([padding_mask, causal_mask[:, :, :, mask_length:]], dim=-1)
+                if mask_length < target_length
+                else padding_mask
+            )
+
+        return causal_mask
+
+
+class RicursiveForCausalLM(RicursivePreTrainedModel):
+    _tied_weights_keys = []
+
+    def __init__(self, config: RicursiveConfig):
+        super().__init__(config)
+        self.model = RicursiveModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    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, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        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,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        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: torch.LongTensor,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            if input_ids.shape[1] != 1:
+                input_ids = input_ids[:, cache_position]
+
+        if inputs_embeds is not None and cache_position[0] == 0:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": cache_position.unsqueeze(0) if cache_position is not None else None,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "cache_position": cache_position,
+            }
+        )
+        return model_inputs

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": {"content": "<|endoftext|>", "lstrip": false, "normalized": false, "rstrip": false, "single_word": false}, "eos_token": "<|im_end|>", "pad_token": {"content": "<|pad|>", "lstrip": false, "normalized": false, "rstrip": false, "single_word": false}, "unk_token": {"content": "<|endoftext|>", "lstrip": false, "normalized": false, "rstrip": false, "single_word": false}}

tokenizer_config.json

@@ -0,0 +1,190 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "100256": {
+      "content": "<|extra_id_0|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100257": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100258": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100259": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100260": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100261": {
+      "content": "|||PHONE_NUMBER|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100262": {
+      "content": "|||EMAIL_ADDRESS|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100263": {
+      "content": "|||IP_ADDRESS|||",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100264": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100265": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100266": {
+      "content": "<|extra_id_1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100267": {
+      "content": "<|extra_id_2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100268": {
+      "content": "<|extra_id_3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100269": {
+      "content": "<|extra_id_4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100270": {
+      "content": "<|extra_id_5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100271": {
+      "content": "<|extra_id_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100272": {
+      "content": "<|extra_id_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100273": {
+      "content": "<|extra_id_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100274": {
+      "content": "<|extra_id_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100275": {
+      "content": "<|extra_id_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "100276": {
+      "content": "<|endofprompt|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100277": {
+      "content": "<|pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "chat_template": "{%- for message in messages %}{{- '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n' }}{%- endfor %}{%- if add_generation_prompt %}{{- '<|im_start|>assistant\n' }}{%- endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": "<|endoftext|>"
+}