Blaze-Title — FFT on Qyrou 115K + Wild 130K (70/30, wild deduped), lr=0.0003, 1 epoch, seq 2048, assistant-only loss, compile default

configuration_blaze.py

@@ -0,0 +1,23 @@
+
+from transformers import LlamaConfig
+
+class BlazeConfig(LlamaConfig):
+    model_type = "blaze"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=1, recurrent_layers=12, coda_layers=1,
+                 recurrent_passes=2,
+                 gradient_checkpointing=False, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        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.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}

modeling_blaze.py

@@ -0,0 +1,719 @@
+# coding=utf-8
+# Copyright 2024 SurjoLabs and HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Optional, Tuple, Union, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRMSNorm, LlamaMLP
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
+from transformers.cache_utils import DynamicCache, Cache
+
+try:
+    from .configuration_blaze import BlazeConfig
+except ImportError:
+    from configuration_blaze import BlazeConfig
+
+# Safe Flash Attention imports with fallback
+try:
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    try:
+        from flash_attn import flash_attn_varlen_func
+        flash_attn_func = None
+        FLASH_ATTN_AVAILABLE = True
+    except ImportError:
+        flash_attn_func = None
+        flash_attn_varlen_func = None
+        FLASH_ATTN_AVAILABLE = False
+
+
+@torch._dynamo.disable()
+def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
+    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
+    return flash_attn_varlen_func(
+        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
+        dropout_p=dropout_p, causal=True,
+    )
+
+
+@torch._dynamo.disable()
+def _flash_attn(q, k, v, dropout_p, causal):
+    return flash_attn_func(
+        q, k, v, dropout_p=dropout_p, causal=causal,
+    )
+
+
+class ClampedLlamaMLP(LlamaMLP):
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+
+class XSAAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.recurrent_cache_idx = None
+        self._use_recurrent_slot = False
+        self._current_pass = 0
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)
+
+        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Union[Cache, Tuple[torch.Tensor]]] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        expected_batch_size: Optional[int] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[Union[int, torch.Tensor]] = None,
+        has_padding: bool = False,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)
+
+        if hidden_states.ndim == 2:
+            if expected_batch_size is None:
+                raise RuntimeError(
+                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
+                    f"without an expected_batch_size to safely restore the batch dim."
+                )
+            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if expected_batch_size is not None and bsz != expected_batch_size:
+            raise RuntimeError(
+                f"XSAAttention: hidden_states batch size {bsz} does not match "
+                f"expected_batch_size {expected_batch_size}."
+            )
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+
+        # 1. Packed Sequence Varlen Flash Attention
+        use_flash_varlen = (
+            cu_seqlens is not None
+            and past_kv is None
+            and (getattr(self.config, "use_flash_attn", False) or getattr(self.config, "_attn_implementation", "") == "flash_attention_2")
+            and FLASH_ATTN_AVAILABLE
+            and flash_attn_varlen_func is not None
+        )
+
+        if use_flash_varlen:
+            total = bsz * q_len
+            q = query_states.reshape(total, self.num_heads, self.head_dim).to(torch.bfloat16)
+            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim).to(torch.bfloat16)
+            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim).to(torch.bfloat16)
+
+            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
+            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
+            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)
+
+            attn_output = _flash_varlen(
+                q, k, v, cu_seqlens, max_seqlen,
+                self.config.attention_dropout if self.training else 0.0,
+            )
+
+            if getattr(self.config, 'xsa_projection', True):
+                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
+                v_grouped = v.unsqueeze(2)
+                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+                scale = (dot_yv / dot_vv).to(y.dtype)
+                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)
+
+            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
+            return (attn_output, None)
+
+        # 2. Standard Attention & KV Caching
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        current_v = value_states
+
+        target_idx = self.layer_idx
+        if getattr(self, "_use_recurrent_slot", False) and self.recurrent_cache_idx is not None:
+            pass_offset = max(0, getattr(self, "_current_pass", 1) - 1)
+            target_idx = self.recurrent_cache_idx + pass_offset * getattr(self.config, "recurrent_layers", 1)
+
+        # Pre-allocate cache slots in bulk without per-token Python overhead
+        if past_kv is not None:
+            if hasattr(past_kv, "layers"):
+                curr_len = len(past_kv.layers)
+                if curr_len <= target_idx:
+                    layer_cls = getattr(past_kv, "layer_class_to_replicate", None)
+                    if layer_cls is None and curr_len > 0:
+                        layer_cls = past_kv.layers[0].__class__
+                    if layer_cls is None:
+                        from transformers.cache_utils import DynamicLayer
+                        layer_cls = DynamicLayer
+                    past_kv.layers.extend([layer_cls() for _ in range(target_idx - curr_len + 1)])
+            elif hasattr(past_kv, "key_cache"):
+                curr_len = len(past_kv.key_cache)
+                if curr_len <= target_idx:
+                    num_to_add = target_idx - curr_len + 1
+                    past_kv.key_cache.extend([
+                        torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device)
+                        for _ in range(num_to_add)
+                    ])
+                    past_kv.value_cache.extend([
+                        torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device)
+                        for _ in range(num_to_add)
+                    ])
+            else:
+                while len(past_kv) <= target_idx:
+                    past_kv.update(
+                        torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
+                        torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
+                        len(past_kv)
+                    )
+            key_states, value_states = past_kv.update(key_states, value_states, target_idx)
+
+        kv_len = key_states.shape[-2]
+
+        is_flash_enabled = getattr(self.config, "use_flash_attn", False) or getattr(self.config, "_attn_implementation", "") == "flash_attention_2"
+        use_flash_func = (
+            FLASH_ATTN_AVAILABLE
+            and flash_attn_func is not None
+            and is_flash_enabled
+            and query_states.is_cuda
+            and not has_padding
+            and (attention_mask is None or attention_mask.ndim == 2)
+            and (q_len == 1 or kv_len == q_len)
+        )
+
+        attn_output = None
+        if use_flash_func:
+            try:
+                q_fa = query_states.transpose(1, 2)
+                k_fa = key_states.transpose(1, 2)
+                v_fa = value_states.transpose(1, 2)
+
+                orig_dtype = q_fa.dtype
+                if orig_dtype not in (torch.float16, torch.bfloat16):
+                    target_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+                    q_fa = q_fa.to(target_dtype)
+                    k_fa = k_fa.to(target_dtype)
+                    v_fa = v_fa.to(target_dtype)
+
+                causal = (q_len > 1 and kv_len == q_len)
+                drop_p = self.config.attention_dropout if self.training else 0.0
+                out_fa = _flash_attn(q_fa, k_fa, v_fa, drop_p, causal)
+                if orig_dtype not in (torch.float16, torch.bfloat16):
+                    out_fa = out_fa.to(orig_dtype)
+
+                attn_output = out_fa.transpose(1, 2)
+            except Exception:
+                attn_output = None
+
+        if attn_output is None:
+            key_states_sdpa = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+            value_states_sdpa = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+
+            is_causal = False
+            attn_mask = None
+
+            if attention_mask is not None:
+                if attention_mask.ndim == 2:
+                    if has_padding:
+                        if attention_mask.shape[-1] < kv_len:
+                            attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
+                        elif attention_mask.shape[-1] > kv_len:
+                            attention_mask = attention_mask[:, -kv_len:]
+
+                        pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min
+
+                        if q_len > 1:
+                            if cache_position is None:
+                                cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                            kv_positions = torch.arange(kv_len, device=query_states.device)
+
+                            neg_inf = torch.finfo(query_states.dtype).min
+                            causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                            causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
+                            attn_mask = causal_mask[None, None, :, :] + pad_mask
+
+                            diag_idx = torch.arange(q_len, device=attn_mask.device)
+                            start_idx = attn_mask.shape[-1] - q_len
+                            attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
+                        else:
+                            attn_mask = pad_mask
+                        is_causal = False
+                    else:
+                        if q_len > 1 and kv_len == q_len:
+                            is_causal = True
+                            attn_mask = None
+                        elif q_len > 1:
+                            if cache_position is None:
+                                cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                            kv_positions = torch.arange(kv_len, device=query_states.device)
+                            causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                            causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], torch.finfo(query_states.dtype).min)
+                            attn_mask = causal_mask[None, None, :, :]
+                            is_causal = False
+                        else:
+                            is_causal = False
+                            attn_mask = None
+                elif attention_mask.ndim == 4:
+                    attn_mask = attention_mask.to(dtype=query_states.dtype)
+                    is_causal = False
+                elif attention_mask.ndim == 3:
+                    attn_mask = attention_mask.unsqueeze(1).to(dtype=query_states.dtype)
+                    is_causal = False
+            else:
+                if q_len > 1 and kv_len == q_len:
+                    is_causal = True
+                    attn_mask = None
+                elif q_len > 1:
+                    if cache_position is None:
+                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                    kv_positions = torch.arange(kv_len, device=query_states.device)
+                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], torch.finfo(query_states.dtype).min)
+                    attn_mask = causal_mask[None, None, :, :]
+                    is_causal = False
+                else:
+                    # Single-token decode attends to all past tokens without causal truncation
+                    is_causal = False
+                    attn_mask = None
+
+            attn_output = F.scaled_dot_product_attention(
+                query_states, key_states_sdpa, value_states_sdpa, attn_mask=attn_mask,
+                dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
+            )
+
+        if getattr(self.config, 'xsa_projection', True):
+            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
+            v_grouped = current_v.unsqueeze(2)
+            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+            scale = (dot_yv / dot_vv).to(y.dtype)
+            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return (attn_output, None)
+
+
+@torch._dynamo.disable()
+def _checkpointed_layer_forward(
+    layer, hidden_states, attention_mask, position_ids,
+    cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen
+):
+    out = layer(
+        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+        past_key_value=None, use_cache=False,
+        cache_position=cache_position, position_embeddings=(cos, sin),
+        expected_batch_size=expected_batch_size,
+        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+    )
+    hs_out = out[0] if isinstance(out, tuple) else out
+    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
+        raise RuntimeError(
+            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+            f"batch size {expected_batch_size}."
+        )
+    return hs_out
+
+
+class BlazeModel(LlamaModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
+            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"
+
+        p1 = config.prelude_layers
+        r1 = p1 + config.recurrent_layers
+
+        for i, layer in enumerate(self.layers):
+            layer.self_attn = XSAAttention(config, layer_idx=i)
+            layer.mlp = ClampedLlamaMLP(config)
+
+        for i, layer in enumerate(self.layers[p1:r1]):
+            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i
+
+        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", False)
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self.gradient_checkpointing = False
+
+    def _get_cache_seq_length(self, past_key_values) -> int:
+        if past_key_values is None:
+            return 0
+        if hasattr(past_key_values, "get_seq_length"):
+            return past_key_values.get_seq_length(0)
+        return past_key_values[0][0].shape[-2] if len(past_key_values) > 0 else 0
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        past_key_values: Optional[Union[Cache, Tuple[torch.Tensor]]] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = False,
+        output_hidden_states: Optional[bool] = False,
+        return_dict: Optional[bool] = True,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[Union[int, torch.Tensor]] = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        cache_position = kwargs.get("cache_position", None)
+        if use_cache is None:
+            use_cache = getattr(self.config, "use_cache", False)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+        elif past_key_values is not None and not isinstance(past_key_values, DynamicCache):
+            if hasattr(DynamicCache, "from_legacy_cache"):
+                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
+
+        if cache_position is None:
+            past_seen = self._get_cache_seq_length(past_key_values) if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
+        elif cache_position.shape[-1] > seq_len:
+            cache_position = cache_position[-seq_len:]
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)
+        elif position_ids.shape[-1] > seq_len:
+            position_ids = position_ids[:, -seq_len:]
+
+        hidden_states = inputs_embeds
+        try:
+            position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        except TypeError:
+            position_embeddings = self.rotary_emb(hidden_states, seq_len=seq_len)
+        cos, sin = position_embeddings
+
+        p1 = self.config.prelude_layers
+        r1 = p1 + self.config.recurrent_layers
+        c1 = r1 + self.config.coda_layers
+
+        prelude = self.layers[:p1]
+        recurrent = self.layers[p1:r1]
+        coda = self.layers[r1:c1]
+
+        use_ckpt = self.training and self.gradient_checkpointing and not use_cache
+
+        # Check padding ONCE per forward pass to avoid per-layer GPU-to-CPU stalls
+        has_padding = False
+        if attention_mask is not None and bsz > 1 and attention_mask.ndim == 2:
+            has_padding = bool((attention_mask == 0).any())
+
+        def run_layer(layer, hs):
+            if cu_seqlens is not None:
+                torch._dynamo.mark_dynamic(cu_seqlens, 0)
+
+            out = layer(
+                hs, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
+                cache_position=cache_position, position_embeddings=position_embeddings,
+                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+                has_padding=has_padding,
+            )
+            hs_out = out[0] if isinstance(out, tuple) else out
+            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
+                raise RuntimeError(
+                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+                    f"batch size {bsz}."
+                )
+            return hs_out
+
+        def run_layer_maybe_ckpt(layer, hs):
+            if use_ckpt:
+                return torch.utils.checkpoint.checkpoint(
+                    _checkpointed_layer_forward,
+                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
+                    cu_seqlens, max_seqlen,
+                    use_reentrant=False,
+                )
+            return run_layer(layer, hs)
+
+        all_hidden_states = () if output_hidden_states else None
+
+        for layer in prelude:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        recurrent_passes = getattr(self.config, "recurrent_passes", 2)
+        for pass_idx in range(recurrent_passes):
+            if self.training:
+                hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+            is_recurrent_slot = pass_idx > 0
+
+            for layer in recurrent:
+                if output_hidden_states:
+                    all_hidden_states += (hidden_states,)
+                layer.self_attn._use_recurrent_slot = is_recurrent_slot
+                layer.self_attn._current_pass = pass_idx
+                try:
+                    hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+                finally:
+                    layer.self_attn._use_recurrent_slot = False
+                    layer.self_attn._current_pass = 0
+
+        for layer in coda:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, past_key_values if use_cache else None, all_hidden_states] if v is not None)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=all_hidden_states,
+        )
+
+
+class BlazeForCausalLM(LlamaForCausalLM):
+    config_class = BlazeConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = BlazeModel(config)
+        self.vocab_size = config.vocab_size
+        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 gradient_checkpointing_enable(self, **kwargs):
+        self.model.gradient_checkpointing_enable(**kwargs)
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing_disable()
+
+    def _get_cache_seq_length(self, past_key_values) -> int:
+        return self.model._get_cache_seq_length(past_key_values)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        num_logits_to_keep: Optional[int] = 0,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, Tuple[torch.Tensor]]] = None,
+        cu_seqlens: Optional[torch.Tensor] = None,
+        max_seqlen: Optional[Union[int, torch.Tensor]] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        return_dict = return_dict if return_dict is not None else getattr(self.config, "return_dict", True)
+        if use_cache is None:
+            use_cache = False if (self.training or labels is not None) else True
+
+        if num_logits_to_keep is None or num_logits_to_keep == 0:
+            if "logits_to_keep" in kwargs:
+                num_logits_to_keep = kwargs.get("logits_to_keep", 0) or 0
+            else:
+                num_logits_to_keep = 0
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+            return_dict=return_dict,
+            **kwargs,
+        )
+        hidden_states = outputs[0]
+
+        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
+            raise RuntimeError(
+                f"BlazeModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
+                f"expected batch size {expected_bsz}."
+            )
+
+        loss = None
+        logits = None
+
+        if labels is not None:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            num_chunks = 8
+            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
+            l_chunks = shift_labels.chunk(num_chunks, dim=0)
+
+            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+            total_tokens = 0
+            for h_c, l_c in zip(h_chunks, l_chunks):
+                if l_c.numel() == 0:
+                    continue
+                logits_c = self.lm_head(h_c)
+                chunk_loss = F.cross_entropy(
+                    logits_c.view(-1, logits_c.size(-1)).float(),
+                    l_c.view(-1),
+                    reduction="sum",
+                )
+                total_loss = total_loss + chunk_loss
+                total_tokens += l_c.numel()
+            loss = (total_loss / max(total_tokens, 1)).to(hidden_states.dtype)
+        else:
+            if num_logits_to_keep == 0:
+                slice_hidden = hidden_states
+            else:
+                slice_hidden = hidden_states[:, -num_logits_to_keep:, :]
+            logits = self.lm_head(slice_hidden)
+
+        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,
+        position_ids: Optional[torch.LongTensor] = None,
+        use_cache: bool = True,
+        num_logits_to_keep: Optional[int] = None,
+        **kwargs,
+    ) -> dict:
+        cache_position = kwargs.get("cache_position", None)
+        past_length = 0
+
+        if past_key_values is not None:
+            past_length = self._get_cache_seq_length(past_key_values)
+
+            # Nanbeige & Llama strict token slicing contract
+            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length):]
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            else:
+                input_ids = input_ids[:, -1:]
+
+        if inputs_embeds is not None and past_length == 0:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids.contiguous()}
+
+        input_length = input_ids.shape[1]
+
+        if cache_position is None:
+            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
+        else:
+            cache_position = cache_position[-input_length:]
+
+        if position_ids is None and attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None:
+                position_ids = position_ids[:, -input_length:]
+        elif position_ids is not None:
+            position_ids = position_ids[:, -input_length:]
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        if num_logits_to_keep is not None:
+            model_inputs["num_logits_to_keep"] = num_logits_to_keep
+        return model_inputs
+
+    def _reorder_cache(self, past_key_values, beam_idx):
+        if hasattr(past_key_values, "reorder_cache"):
+            return past_key_values.reorder_cache(beam_idx)
+        return past_key_values