########################################################################################################
# RWKV-7 "Goose" (x070 / g1d) HuggingFace modeling code
# Based on the reference implementation from https://github.com/BlinkDL/RWKV-LM
#
# This file provides a `trust_remote_code=True` compatible RWKV-7 model that:
#   * uses a fused CUDA kernel (fwd + bwd, "wind_backstepping" bf16) when a GPU
#     + a working CUDA toolchain are available (training / no-cache path), and
#   * transparently falls back to a pure-PyTorch implementation otherwise
#     (autograd handles the backward pass automatically in the fallback).
#
# Generation now uses a recurrent state cache (RNN mode) instead of recomputing
# the full sequence every step.  Stateful inference calls run on a second,
# forward-only CUDA kernel ("wkv7s", state in/out, no chunk padding) when
# available, so both prefill and decode are kernel-speed under no_grad.
# The per-layer state is:
#   (att_x_prev, wkv_state, ffn_x_prev)
#     att_x_prev : (B, C)        last post-ln1 token, for time-shift
#     wkv_state  : (B, H, N, N)  float32 WKV matrix state
#     ffn_x_prev : (B, C)        last post-ln2 token, for time-shift
# The state is carried through `state=` / `outputs.state`, which HF's
# GenerationMixin propagates between steps (same convention as Rwkv/Mamba).
########################################################################################################

import os
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.modeling_utils import PreTrainedModel
from transformers.utils import ModelOutput
from transformers.generation import GenerationMixin

try:
    from .configuration_rwkv7 import RWKV7Config
except ImportError:  # allow running as a plain script (e.g. conversion)
    from configuration_rwkv7 import RWKV7Config


# Per-layer recurrent state: (att_x_prev, wkv_state, ffn_x_prev)
LayerState = Tuple[torch.Tensor, torch.Tensor, torch.Tensor]


########################################################################################################
# HF output dataclasses (carry `state` so GenerationMixin can propagate it)
########################################################################################################


@dataclass
class RWKV7Output(ModelOutput):
    last_hidden_state: torch.FloatTensor = None
    state: Optional[List[LayerState]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None


@dataclass
class RWKV7CausalLMOutput(ModelOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    state: Optional[List[LayerState]] = None
    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None


########################################################################################################
# CUDA kernel loading (lazy, best-effort).  If anything goes wrong we silently
# fall back to the pure-PyTorch path.
########################################################################################################

_KERNEL_STATE = {"loaded": False, "ok": False, "op": None, "chunk_len": 16, "head_size": 64}


def _try_load_cuda_kernel(head_size: int, chunk_len: int):
    """Compile & register the wind_backstepping RWKV-7 CUDA op. Returns True on success."""
    if _KERNEL_STATE["loaded"]:
        return _KERNEL_STATE["ok"]
    _KERNEL_STATE["loaded"] = True
    _KERNEL_STATE["head_size"] = head_size
    _KERNEL_STATE["chunk_len"] = chunk_len

    if not torch.cuda.is_available():
        _KERNEL_STATE["ok"] = False
        return False

    try:
        from torch.utils.cpp_extension import load

        this_dir = os.path.dirname(os.path.abspath(__file__))
        cuda_dir = os.path.join(this_dir, "cuda")
        sources = [
            os.path.join(cuda_dir, "wkv7_op.cpp"),
            os.path.join(cuda_dir, "wkv7_cuda.cu"),
        ]
        print(sources)
        print([os.path.exists(s) for s in sources])
        if not all(os.path.exists(s) for s in sources):

            _KERNEL_STATE["ok"] = False
            return False

        flags = [
            "-res-usage",
            f"-D_C_={head_size}",
            f"-D_CHUNK_LEN_={chunk_len}",
            "--use_fast_math",
            "-O3",
            "-Xptxas -O3",
            "--extra-device-vectorization",
        ]
        load(
            name=f"wind_backstepping_c{head_size}_l{chunk_len}",
            sources=sources,
            is_python_module=False,
            verbose=False,
            extra_cuda_cflags=flags,
        )
        _KERNEL_STATE["op"] = torch.ops.wind_backstepping
        _KERNEL_STATE["ok"] = True
        return True
    except Exception as e:  # noqa: BLE001 - any failure -> fallback
        print(f"[RWKV7] CUDA kernel unavailable, using PyTorch fallback ({type(e).__name__}: {e})")
        import traceback

        traceback.print_exc()
        _KERNEL_STATE["ok"] = False
        return False


# Stateful, forward-only inference kernel ("wkv7s"): takes an fp32 WKV state
# and updates it in place -> gives us fast prefill + decode with the cache.
_KERNEL_S_STATE = {"loaded": False, "ok": False, "op": None, "head_size": 64}


def _try_load_cuda_kernel_s(head_size: int):
    """Compile & register the wkv7s stateful RWKV-7 CUDA op. Returns True on success."""
    if _KERNEL_S_STATE["loaded"]:
        return _KERNEL_S_STATE["ok"]
    _KERNEL_S_STATE["loaded"] = True
    _KERNEL_S_STATE["head_size"] = head_size

    if not torch.cuda.is_available():
        _KERNEL_S_STATE["ok"] = False
        return False

    try:
        from torch.utils.cpp_extension import load

        this_dir = os.path.dirname(os.path.abspath(__file__))
        cuda_dir = os.path.join(this_dir, "cuda")
        sources = [
            os.path.join(cuda_dir, "wkv7s_op.cpp"),
            os.path.join(cuda_dir, "wkv7s.cu"),
        ]
        if not all(os.path.exists(s) for s in sources):
            _KERNEL_S_STATE["ok"] = False
            return False

        flags = [
            "-res-usage",
            f"-D_N_={head_size}",
            "--use_fast_math",
            "-O3",
            "-Xptxas -O3",
            "--extra-device-vectorization",
        ]
        load(
            name=f"wkv7s_n{head_size}",
            sources=sources,
            is_python_module=False,
            verbose=False,
            extra_cuda_cflags=flags,
        )
        _KERNEL_S_STATE["op"] = torch.ops.wkv7s
        _KERNEL_S_STATE["ok"] = True
        return True
    except Exception as e:  # noqa: BLE001 - any failure -> fallback
        print(f"[RWKV7] stateful CUDA kernel unavailable, using PyTorch fallback ({type(e).__name__}: {e})")
        _KERNEL_S_STATE["ok"] = False
        return False


class _WindBackstepping(torch.autograd.Function):
    """Fused RWKV-7 kernel wrapper (bf16). Implements both forward and backward.
    Inputs are shaped (B, T, H, C) with T % CHUNK_LEN == 0 and dtype bfloat16.
    NOTE: the kernel always starts from a zero WKV state and does not expose the
    final state, so it is only used on the no-cache (training / plain forward) path.
    """

    @staticmethod
    def forward(ctx, w, q, k, v, z, b):
        op = _KERNEL_STATE["op"]
        chunk_len = _KERNEL_STATE["chunk_len"]
        B, T, H, C = w.shape
        assert T % chunk_len == 0, "pad T to a multiple of CHUNK_LEN"
        assert all(i.dtype == torch.bfloat16 for i in [w, q, k, v, z, b])
        assert all(i.is_contiguous() for i in [w, q, k, v, z, b])
        y = torch.empty_like(v)
        s = torch.empty(B, H, T // chunk_len, C, C, dtype=torch.float32, device=w.device)
        sa = torch.empty(B, T, H, C, dtype=torch.float32, device=w.device)
        op.forward(w, q, k, v, z, b, y, s, sa)
        ctx.save_for_backward(w, q, k, v, z, b, s, sa)
        return y

    @staticmethod
    def backward(ctx, dy):
        op = _KERNEL_STATE["op"]
        assert dy.dtype == torch.bfloat16
        dy = dy.contiguous()
        w, q, k, v, z, b, s, sa = ctx.saved_tensors
        dw, dq, dk, dv, dz, db = [torch.empty_like(x) for x in [w, q, k, v, z, b]]
        op.backward(w, q, k, v, z, b, dy, s, sa, dw, dq, dk, dv, dz, db)
        return dw, dq, dk, dv, dz, db


def _rwkv7_cuda(r, w, k, v, a, b, head_size, chunk_len):
    """CUDA path. r,w,k,v,a,b are (B, T, C) bf16.  a = -kk, b = kk*a_gate."""
    B, T, C = r.shape
    H = C // head_size
    pad = (chunk_len - T % chunk_len) % chunk_len
    if pad:
        r, w, k, v, a, b = [F.pad(x, (0, 0, 0, pad)) for x in (r, w, k, v, a, b)]
    Tp = T + pad
    r, w, k, v, a, b = [x.view(B, Tp, H, head_size).contiguous() for x in (r, w, k, v, a, b)]
    y = _WindBackstepping.apply(w, r, k, v, a, b).view(B, Tp, C)
    if pad:
        y = y[:, :T]
    return y


def _rwkv7_cuda_stateful(r, w, k, v, a, b, head_size, initial_state=None):
    """Stateful CUDA path (wkv7s kernel). Forward-only (no backward), any T,
    no chunk padding needed. The kernel updates the fp32 state in place, so we
    always hand it a fresh tensor and return it as the final state."""
    op = _KERNEL_S_STATE["op"]
    B, T, C = r.shape
    H = C // head_size
    N = head_size
    if initial_state is not None:
        state = initial_state.to(device=r.device, dtype=torch.float32).contiguous().clone()
    else:
        state = torch.zeros(B, H, N, N, dtype=torch.float32, device=r.device)
    r, w, k, v, a, b = [x.contiguous() for x in (r, w, k, v, a, b)]
    y = torch.empty_like(v)
    op.forward(B, T, C, H, state, r, w, k, v, a, b, y)
    return y, state


def _rwkv7_pytorch(r, w, k, v, a, b, head_size, initial_state=None, output_final_state=False):
    """Pure-PyTorch reference (sequential over time). Differentiable via autograd.
    w is the raw (pre-exp) log-decay; the recurrence uses exp(-exp(w)).
    Supports an initial WKV state and optionally returns the final state,
    which is what enables cached (RNN-mode) generation.
    """
    B, T, C = r.size()
    H = C // head_size
    N = head_size
    dtype_in = r.dtype
    r = r.view(B, T, H, N).float()
    k = k.view(B, T, H, N).float()
    v = v.view(B, T, H, N).float()
    a = a.view(B, T, H, N).float()
    b = b.view(B, T, H, N).float()
    w = torch.exp(-torch.exp(w.view(B, T, H, N).float()))

    out = torch.zeros((B, T, H, N), device=r.device, dtype=torch.float32)
    if initial_state is not None:
        state = initial_state.to(device=r.device, dtype=torch.float32)
    else:
        state = torch.zeros((B, H, N, N), device=r.device, dtype=torch.float32)
    for t in range(T):
        kk = k[:, t, :].view(B, H, 1, N)
        rr = r[:, t, :].view(B, H, N, 1)
        vv = v[:, t, :].view(B, H, N, 1)
        aa = a[:, t, :].view(B, H, N, 1)
        bb = b[:, t, :].view(B, H, 1, N)
        state = state * w[:, t, :, None, :] + state @ aa @ bb + vv @ kk
        out[:, t, :] = (state @ rr).view(B, H, N)
    final_state = state if output_final_state else None
    return out.view(B, T, C).to(dtype=dtype_in), final_state


def run_rwkv7(
    r, w, k, v, a, b, config,
    initial_state=None,
    output_final_state=False,
    force_fallback=False,
):

    stateful = initial_state is not None or output_final_state

    common_ok = (
        config.use_cuda_kernel
        and not force_fallback
        and r.is_cuda
        and r.dtype == torch.bfloat16
    )

    print(
        "RWKV7 DEBUG:",
        "stateful=", stateful,
        "common_ok=", common_ok,
        "cuda=", r.is_cuda,
        "dtype=", r.dtype,
        "use_kernel=", config.use_cuda_kernel,
        "force=", force_fallback,
        "kernel_loaded=", _KERNEL_STATE["loaded"],
        "kernel_ok=", _KERNEL_STATE["ok"],
    )

    if not stateful:
        if common_ok and _try_load_cuda_kernel(config.head_size, config.chunk_len):
            print("USING CUDA KERNEL")
            return _rwkv7_cuda(
                r,w,k,v,a,b,
                config.head_size,
                config.chunk_len
            ) ,None

    print("USING PYTORCH FALLBACK")
    return _rwkv7_pytorch(
        r, w, k, v, a, b, config.head_size,
        initial_state=initial_state, output_final_state=output_final_state,
    )


########################################################################################################
# RWKV-7 time-mixing ("attention") block
########################################################################################################


class RWKV7TimeMix(nn.Module):
    def __init__(self, config: RWKV7Config, layer_id: int):
        super().__init__()
        self.config = config
        self.layer_id = layer_id
        self.head_size = config.head_size
        C = config.hidden_size
        self.n_head = C // self.head_size
        H, N = self.n_head, self.head_size

        self.x_r = nn.Parameter(torch.empty(1, 1, C))
        self.x_w = nn.Parameter(torch.empty(1, 1, C))
        self.x_k = nn.Parameter(torch.empty(1, 1, C))
        self.x_v = nn.Parameter(torch.empty(1, 1, C))
        self.x_a = nn.Parameter(torch.empty(1, 1, C))
        self.x_g = nn.Parameter(torch.empty(1, 1, C))

        self.w0 = nn.Parameter(torch.empty(1, 1, C))
        self.w1 = nn.Parameter(torch.empty(C, config.decay_lora))
        self.w2 = nn.Parameter(torch.empty(config.decay_lora, C))

        self.a0 = nn.Parameter(torch.empty(1, 1, C))
        self.a1 = nn.Parameter(torch.empty(C, config.aaa_lora))
        self.a2 = nn.Parameter(torch.empty(config.aaa_lora, C))

        self.v0 = nn.Parameter(torch.empty(1, 1, C))
        self.v1 = nn.Parameter(torch.empty(C, config.mv_lora))
        self.v2 = nn.Parameter(torch.empty(config.mv_lora, C))

        self.g1 = nn.Parameter(torch.empty(C, config.gate_lora))
        self.g2 = nn.Parameter(torch.empty(config.gate_lora, C))

        self.k_k = nn.Parameter(torch.empty(1, 1, C))
        self.k_a = nn.Parameter(torch.empty(1, 1, C))
        self.r_k = nn.Parameter(torch.empty(H, N))

        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
        self.receptance = nn.Linear(C, C, bias=False)
        self.key = nn.Linear(C, C, bias=False)
        self.value = nn.Linear(C, C, bias=False)
        self.output = nn.Linear(C, C, bias=False)
        self.ln_x = nn.GroupNorm(H, C, eps=config.group_norm_epsilon)

    def forward(self, x, v_first, x_prev=None, wkv_state=None, use_cache=False):
        """x_prev: (B, C) last input token from the previous chunk (post-ln1),
        wkv_state: (B, H, N, N) float32 WKV state from the previous chunk."""
        B, T, C = x.size()
        H = self.n_head

        if x_prev is None:
            shifted = self.time_shift(x)  # zero-pad == zero initial shift state
        else:
            shifted = torch.cat([x_prev.to(x.dtype).unsqueeze(1), x[:, :-1]], dim=1)
        xx = shifted - x
        new_x_prev = x[:, -1] if use_cache else None

        xr = x + xx * self.x_r
        xw = x + xx * self.x_w
        xk = x + xx * self.x_k
        xv = x + xx * self.x_v
        xa = x + xx * self.x_a
        xg = x + xx * self.x_g

        r = self.receptance(xr)
        # soft-clamp to (-inf, -0.5); the recurrence applies exp(-exp(w))
        w = -F.softplus(-(self.w0 + torch.tanh(xw @ self.w1) @ self.w2)) - 0.5
        k = self.key(xk)
        v = self.value(xv)
        if self.layer_id == 0:
            v_first = v
        else:
            v = v + (v_first - v) * torch.sigmoid(self.v0 + (xv @ self.v1) @ self.v2)
        a = torch.sigmoid(self.a0 + (xa @ self.a1) @ self.a2)  # in-context learning rate
        g = torch.sigmoid(xg @ self.g1) @ self.g2

        kk = k * self.k_k
        kk = F.normalize(kk.view(B, T, H, -1), dim=-1, p=2.0).view(B, T, C)
        k = k * (1 + (a - 1) * self.k_a)

        x, new_wkv_state = run_rwkv7(
            r, w, k, v, -kk, kk * a, self.config,
            initial_state=wkv_state, output_final_state=use_cache,
        )
        x = self.ln_x(x.view(B * T, C)).view(B, T, C)

        x = x + (
            (r.view(B, T, H, -1) * k.view(B, T, H, -1) * self.r_k).sum(dim=-1, keepdim=True)
            * v.view(B, T, H, -1)
        ).view(B, T, C)
        x = self.output(x * g)
        return x, v_first, new_x_prev, new_wkv_state


########################################################################################################
# RWKV-7 channel-mixing (FFN) block
########################################################################################################


class RWKV7ChannelMix(nn.Module):
    def __init__(self, config: RWKV7Config, layer_id: int):
        super().__init__()
        self.layer_id = layer_id
        self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
        self.x_k = nn.Parameter(torch.empty(1, 1, config.hidden_size))
        self.key = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.value = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

    def forward(self, x, x_prev=None, use_cache=False):
        """x_prev: (B, C) last input token from the previous chunk (post-ln2)."""
        if x_prev is None:
            shifted = self.time_shift(x)
        else:
            shifted = torch.cat([x_prev.to(x.dtype).unsqueeze(1), x[:, :-1]], dim=1)
        xx = shifted - x
        new_x_prev = x[:, -1] if use_cache else None
        k = x + xx * self.x_k
        k = torch.relu(self.key(k)) ** 2
        return self.value(k), new_x_prev


class RWKV7Block(nn.Module):
    def __init__(self, config: RWKV7Config, layer_id: int):
        super().__init__()
        self.layer_id = layer_id
        eps = config.layer_norm_epsilon
        if layer_id == 0:
            self.ln0 = nn.LayerNorm(config.hidden_size, eps=eps)
        self.ln1 = nn.LayerNorm(config.hidden_size, eps=eps)
        self.ln2 = nn.LayerNorm(config.hidden_size, eps=eps)
        self.att = RWKV7TimeMix(config, layer_id)
        self.ffn = RWKV7ChannelMix(config, layer_id)

    def forward(self, x, v_first, layer_state: Optional[LayerState] = None, use_cache: bool = False):
        if self.layer_id == 0:
            x = self.ln0(x)

        if layer_state is not None:
            att_x_prev, wkv_state, ffn_x_prev = layer_state
        else:
            att_x_prev = wkv_state = ffn_x_prev = None

        x_attn, v_first, att_x_prev, wkv_state = self.att(
            self.ln1(x), v_first, x_prev=att_x_prev, wkv_state=wkv_state, use_cache=use_cache
        )
        x = x + x_attn

        ffn_out, ffn_x_prev = self.ffn(self.ln2(x), x_prev=ffn_x_prev, use_cache=use_cache)
        x = x + ffn_out

        new_state = (att_x_prev, wkv_state, ffn_x_prev) if use_cache else None
        return x, v_first, new_state


########################################################################################################
# HuggingFace wrappers
########################################################################################################


class RWKV7PreTrainedModel(PreTrainedModel):
    config_class = RWKV7Config
    base_model_prefix = "rwkv"
    supports_gradient_checkpointing = True
    _no_split_modules = ["RWKV7Block"]

    def _init_weights(self, module):
        # Weights normally come from a pretrained checkpoint; this only covers
        # freshly-created (e.g. re-sized embedding / head) parameters.
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=1e-4)
        elif isinstance(module, nn.LayerNorm):
            module.weight.data.fill_(1.0)
            module.bias.data.zero_()


class RWKV7Model(RWKV7PreTrainedModel):
    def __init__(self, config: RWKV7Config):
        super().__init__(config)
        self.config = config
        self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
        self.blocks = nn.ModuleList(
            [RWKV7Block(config, i) for i in range(config.num_hidden_layers)]
        )
        self.ln_out = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_epsilon)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.emb

    def set_input_embeddings(self, value):
        self.emb = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        state: Optional[List[LayerState]] = None,
        use_cache: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, RWKV7Output]:
        return_dict = return_dict if return_dict is not None else True
        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 getattr(self.config, "use_cache", False)
        )
        if state is not None:
            use_cache = True  # an incoming state implies stateful mode
        if self.gradient_checkpointing and self.training and use_cache:
            use_cache = False

        if inputs_embeds is None:
            inputs_embeds = self.emb(input_ids)
        x = inputs_embeds

        all_hidden_states = () if output_hidden_states else None
        new_states: Optional[List[LayerState]] = [] if use_cache else None
        v_first = torch.empty_like(x)
        for i, block in enumerate(self.blocks):
            if output_hidden_states:
                all_hidden_states += (x,)
            layer_state = state[i] if state is not None else None
            if self.gradient_checkpointing and self.training:
                x, v_first, layer_new_state = self._gradient_checkpointing_func(
                    block.__call__, x, v_first, layer_state, use_cache
                )
            else:
                x, v_first, layer_new_state = block(x, v_first, layer_state, use_cache)
            if use_cache:
                new_states.append(layer_new_state)

        x = self.ln_out(x)
        if output_hidden_states:
            all_hidden_states += (x,)

        if not return_dict:
            return tuple(v for v in [x, new_states, all_hidden_states] if v is not None)
        return RWKV7Output(
            last_hidden_state=x,
            state=new_states,
            hidden_states=all_hidden_states,
        )


class RWKV7ForCausalLM(RWKV7PreTrainedModel, GenerationMixin):
    _tied_weights_keys = []

    def __init__(self, config: RWKV7Config):
        super().__init__(config)
        self.rwkv = RWKV7Model(config)
        self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.rwkv.emb

    def set_input_embeddings(self, value):
        self.rwkv.emb = value

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, new_embeddings):
        self.head = new_embeddings

    def get_decoder(self):
        return self.rwkv

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        state: Optional[List[LayerState]] = None,
        use_cache: Optional[bool] = None,
        labels: Optional[torch.LongTensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, RWKV7CausalLMOutput]:
        return_dict = return_dict if return_dict is not None else True

        outputs = self.rwkv(
            input_ids=input_ids,
            inputs_embeds=inputs_embeds,
            state=state,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
            return_dict=True,
        )
        hidden = outputs.last_hidden_state
        logits = self.head(hidden)

        loss = None
        if labels is not None:
            labels = labels.to(logits.device)
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)).float(),
                shift_labels.view(-1),
            )

        if not return_dict:
            output = (logits, outputs.state, outputs.hidden_states)
            output = tuple(v for v in output if v is not None)
            return ((loss,) + output) if loss is not None else output

        return RWKV7CausalLMOutput(
            loss=loss,
            logits=logits,
            state=outputs.state,
            hidden_states=outputs.hidden_states,
        )

    def prepare_inputs_for_generation(
        self, input_ids, state=None, inputs_embeds=None, use_cache=True, **kwargs
    ):
        # RNN mode: once we have a state, only the newly generated token needs
        # to be fed; everything before it is already absorbed into the state.
        if state is not None:
            input_ids = input_ids[:, -1:]
            model_inputs = {"input_ids": input_ids}
        elif inputs_embeds is not None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}
        model_inputs["state"] = state
        model_inputs["use_cache"] = use_cache
        return model_inputs

    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, **kwargs):
        # Recent transformers versions pick up `state` automatically (it is in
        # ALL_CACHE_NAMES); this override keeps older versions working too.
        model_kwargs = super()._update_model_kwargs_for_generation(outputs, model_kwargs, **kwargs)
        model_kwargs["state"] = getattr(outputs, "state", None)
        return model_kwargs