Upload folder using huggingface_hub

RESHARD_INFO.json

@@ -0,0 +1,39 @@
+{
+  "source_dir": "/home/ubuntu/checkpoints/ri-hero-65b-11xC-topsel-3m-1c877500-wc2/hf/step877500",
+  "source_manifest": {
+    "fingerprint": "584628ff3081cd3f19e2d3ddde02d58eb2228fd6f0888e9cdb8baff3cc5fb9f6",
+    "model_safetensors_sha256": "fe0af61d5b75ebd80446457e4f91988cd7811ae0bd492e44069ccf05a47d60ae"
+  },
+  "max_shard_size_bytes": 8000000000,
+  "total_size_bytes": 129565378720,
+  "num_shards": 17,
+  "shard_sha256": {
+    "model-00001-of-00017.safetensors": "2defe6ee993a7a8667f286e2f0d171111601d8020ca69f666c005fd4986a99c9",
+    "model-00002-of-00017.safetensors": "eaf95a952ae1ef67c1190e3fd7160f62fac92cb34aadafca23107ecdbb0f0816",
+    "model-00003-of-00017.safetensors": "f081888035e5119ef338fcc87b716e30fe1b936e5c2128ec3efe80e38d7d74ab",
+    "model-00004-of-00017.safetensors": "33df2d04e5bf1fa248782cc81c77ca19bb5d045e4385ab6ad77bbd326a89abbb",
+    "model-00005-of-00017.safetensors": "07391a73962b13a26c39ea61d849fc4f80853e994d1c153a28e9c6256d4c0290",
+    "model-00006-of-00017.safetensors": "f4fadf975816e1fe33ea9fe3ddbd40c6b1e61761c1677368d8e1b407ccce5594",
+    "model-00007-of-00017.safetensors": "5740d3fb1e1ff8d15a5de7e1e1fe235c1534d41ea44f6640872f5cc733b293c8",
+    "model-00008-of-00017.safetensors": "ebd43c1c0f281879ef5324992298e7432b93390c9a251f5e1f43a1bd3de09484",
+    "model-00009-of-00017.safetensors": "fe21a5ac74f580a0d30f33f6aaeb97c204fd650cf377c77d6dd6d0c30547455b",
+    "model-00010-of-00017.safetensors": "9126a023ec3e2f87da062b97e0f5ee6b61d35cde9154a9f5e24eaf08a381d5d0",
+    "model-00011-of-00017.safetensors": "0edc2746125355cf2dcf041ec99edab4b1645bc4965128645ed15953cfe02b2c",
+    "model-00012-of-00017.safetensors": "0d3c5a82a77a38c4713ff1146c05db9cd353985430a7a7fbf4c6faa4a2655f0c",
+    "model-00013-of-00017.safetensors": "e0269bf7f5829e10661f22e893cb8b477bcfcbd8f250a62cc934c11670174c64",
+    "model-00014-of-00017.safetensors": "9cc7af6165de693dc58b871e7e031298c17e3ae542d2fbaf3d4c713c7e925603",
+    "model-00015-of-00017.safetensors": "1b8711409f6c45f182855d6c77bc4603201dbc2f8bbfb1b0400d14045c210897",
+    "model-00016-of-00017.safetensors": "5c0c5595a615657277925fd6a4ee1d483fa58b42c9c835eb31eb319c6a4916ee",
+    "model-00017-of-00017.safetensors": "e6b0f529781875827e7d6cb71ce441a273c44405f1e449434a7894e2b0a9017b"
+  },
+  "verified_against_source": true,
+  "copied_assets": [
+    "config.json",
+    "configuration_ricursive.py",
+    "modeling_ricursive.py",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "special_tokens_map.json",
+    "chat_template.jinja"
+  ]
+}

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,199 @@
+{
+  "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": 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": 8192,
+  "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.1",
+  "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-00017.safetensors

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

model-00002-of-00017.safetensors

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

model-00003-of-00017.safetensors

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

model-00004-of-00017.safetensors

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

model-00005-of-00017.safetensors

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

model-00006-of-00017.safetensors

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

model-00007-of-00017.safetensors

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

model-00008-of-00017.safetensors

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

model-00009-of-00017.safetensors

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

model-00010-of-00017.safetensors

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

model-00011-of-00017.safetensors

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

model-00012-of-00017.safetensors

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

model-00013-of-00017.safetensors

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

model-00014-of-00017.safetensors

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

model-00015-of-00017.safetensors

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

model-00016-of-00017.safetensors

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

model-00017-of-00017.safetensors

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

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,30 @@
+{
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "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,189 @@
+{
+  "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": "<|endoftext|>",
+  "model_max_length": 8192,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}