# 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