modeling_evo1.py

"""Minimal Evo1 (StripedHyena) HuggingFace port.

This module is a refactor of togethercomputer/evo-1-131k-base@1.1_fix's
``modeling_hyena.py`` + ``model.py`` into a single self-contained file with:

  * ``output_hidden_states`` and ``output_attentions`` plumbed end-to-end,
  * ``attn_implementation`` switch (``eager`` / ``sdpa`` / ``flash_attention_2``),
  * ``Evo1Model`` (no LM head, ``BaseModelOutputWithPast``) for ``AutoModel``,
  * ``Evo1ForCausalLM`` (with logits, ``CausalLMOutputWithPast``)
    for ``AutoModelForCausalLM``,
  * minimal external imports (only ``torch`` + ``transformers``; ``flash-attn``
    is loaded lazily and only when ``attn_implementation='flash_attention_2'``).

Hyena blocks have no attention matrix by construction, so they always emit
``None`` in the per-layer ``attentions`` tuple. Attention blocks (layers 8,
16, 24 for Evo1) emit the (B, H, T, T) softmax matrix when
``output_attentions=True`` (this triggers a one-time fallback from sdpa /
flash_attention_2 to the eager backend).
"""

from __future__ import annotations

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.generation import GenerationMixin
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)
from transformers.utils import logging

from .attention import MHA
from .cache import Evo1Cache, InferenceParams, RecurrentInferenceParams
from .configuration_evo1 import Evo1Config
from .engine import HyenaInferenceEngine
from .layers import ParallelGatedMLP, RMSNorm, VocabParallelEmbedding
from .rotary import swap_mha_rope
# dummy import so that trust_remote_code bundles the tokenizer file
from .tokenization_evo1 import ByteTokenizer  # noqa: F401

logger = logging.get_logger(__name__)


# =============================================================================
# Block: attention (used at layers config.attn_layer_idxs)
# =============================================================================


class AttentionBlock(nn.Module):
    """Pre-norm Transformer block: norm -> MHA -> residual -> norm -> MLP -> residual."""

    def __init__(self, config, layer_idx) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.pre_norm, self.post_norm = RMSNorm(config), RMSNorm(config)
        self.proj_groups = config.get("proj_groups", 1)
        dtype = config.get("attn_block_dtype", torch.bfloat16)
        mlp_dtype = config.get("mlp_dtype", torch.bfloat16)
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size_per_attention_head = (
            config.hidden_size // config.num_attention_heads
        )

        attn_impl = getattr(config, "_attn_implementation", "eager")

        self.inner_mha_cls = MHA(
            embed_dim=config.hidden_size,
            num_heads=config.num_attention_heads,
            num_heads_kv=config.num_attention_heads // self.proj_groups,
            rotary_emb_dim=config.hidden_size // config.num_attention_heads,
            qkv_proj_bias=config.get("qkv_proj_bias", True),
            rotary_emb_base=config.get("rotary_emb_base", 10000),
            causal=True,
            layer_idx=layer_idx,
            out_proj_bias=config.get("mha_out_proj_bias", True),
            attn_implementation=attn_impl,
        ).to(dtype=dtype)

        if config.get("use_interpolated_rotary_pos_emb", False):
            swap_mha_rope(
                mha=self.inner_mha_cls,
                kwargs_new_rope={
                    "scaling_factor": config.get("rotary_emb_scaling_factor", 1.0)
                },
            )

        if self.config.get("smeared_gqa", False):
            self.inner_mha_cls.num_heads_kv = self.inner_mha_cls.num_heads
        # Make sure the inv_freq buffer round-trips through to_bfloat16/state_dict.
        self.inner_mha_cls.rotary_emb.register_buffer(
            "inv_freq", self.inner_mha_cls.rotary_emb.inv_freq
        )

        self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)

    def forward(
        self,
        u: torch.Tensor,
        inference_params=None,
        padding_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        *args,
        **kwargs,
    ):
        if isinstance(padding_mask, torch.Tensor):
            # Workaround for masking with no qkv bias: this zeros the attended
            # values at pad positions so they don't leak via attention.
            u = u * padding_mask[..., None]

        attn_out, attn_weights = self.inner_mha_cls(
            self.pre_norm(u),
            inference_params=inference_params,
            output_attentions=output_attentions,
        )
        u = attn_out + u

        if isinstance(padding_mask, torch.Tensor):
            u = u * padding_mask[..., None]
        u = self.mlp(self.post_norm(u)) + u
        return u, attn_weights


# =============================================================================
# Block: Hyena (used at all other layers)
# =============================================================================


class ParallelHyenaFilter(nn.Module):
    def __init__(self, config, layer_idx) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hyena_filter_groups = config.get(
            "hyena_filter_groups", self.config.hidden_size
        )

        self.use_flashfft = config.get("use_flashfft", False)
        self.state_size = config.state_size
        self.hidden_size = config.hidden_size
        self.num_filters = config.num_filters
        self.inference_mode = config.get("inference_mode", True)
        self.column_split_hyena = config.get("column_split_hyena", True)

        assert self.hidden_size % self.num_filters == 0
        assert self.num_filters <= self.hidden_size

        self.D = nn.Parameter(torch.zeros(self.hidden_size))

        # heads only used to slice post-FIR projections like the checkpoint
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size_per_attention_head = (
            self.hidden_size // self.num_attention_heads
        )

        self.short_filter_length = config.short_filter_length
        self.short_filter_weight = nn.Parameter(
            torch.randn(3 * config.hidden_size, 1, config.short_filter_length)
        )
        self.short_filter_bias = (
            nn.Parameter(torch.randn(3 * config.hidden_size))
            if config.short_filter_bias
            else None
        )

        self.engine = HyenaInferenceEngine(layer_idx=layer_idx)
        self.use_flash_depthwise = config.get("use_flash_depthwise", False)
        self.data_dtype = None

        if self.use_flash_depthwise:
            # importlib avoids the top-level static-import check that HF's
            # dynamic_module_utils.check_imports performs against the file.
            import importlib
            FlashDepthwiseConv1d = importlib.import_module("flashfftconv").FlashDepthwiseConv1d
            self.fir_fn = FlashDepthwiseConv1d(
                channels=3 * self.hidden_size,
                kernel_size=self.short_filter_length,
                padding=self.short_filter_length - 1,
                weights=self.short_filter_weight,
                bias=self.short_filter_bias,
                device=None,
                dtype=self.config.get("depthwise_dtype", torch.bfloat16),
            )
        else:
            self.fir_fn = F.conv1d

        self.fftconv_fn = None
        self.long_fir_threshold = config.get("long_fir_threshold", None)
        if self.long_fir_threshold is not None:
            assert self.use_flashfft is False, (
                "long_fir_threshold not compatible with fused flashfft"
            )

        self.num_systems = self.hidden_size // self.hyena_filter_groups
        poles = torch.randn(self.num_systems, self.state_size, 1, 2)
        poles[..., 0] = 1e-2 * torch.randn(self.num_systems, self.state_size, 1)
        poles[..., 1] = 1e-3 * torch.randn(self.num_systems, self.state_size, 1)
        self.poles = nn.Parameter(poles)
        self.residues = nn.Parameter(
            torch.randn(self.num_systems, self.state_size, 1, 2)
        )
        self.h = None

    def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):
        if (
            inference_params is not None
            and self.layer_idx in inference_params.fir_state_dict.keys()
        ):
            return self.sequential_forward(u, inference_params)
        return self.parallel_forward(u, inference_params, padding_mask)

    def parallel_forward(self, u, inference_params=None, padding_mask=None):
        L = u.shape[1]
        z_pre, fir_state = self.engine.parallel_fir(
            self.fir_fn,
            u,
            self.short_filter_weight,
            self.short_filter_bias,
            L,
            fir_length=self.short_filter_length,
            inference_params=inference_params,
            padding_mask=padding_mask,
        )
        if inference_params:
            inference_params.fir_state_dict[self.layer_idx] = fir_state

        if self.h is None:
            h, _, _, _ = self.compute_filter(L, u.device)
        else:
            h = self.h

        if self.hyena_filter_groups > 1:
            h = h.repeat_interleave(self.hidden_size // self.hyena_filter_groups, 1)

        dims = (
            self.hidden_size,
            self.num_attention_heads,
            self.hidden_size_per_attention_head,
            self.state_size,
            self.hyena_filter_groups,
        )
        y = self.engine.parallel_iir(
            z_pre,
            h,
            self.D,
            L,
            t=self.t,
            poles=self.poles,
            residues=self.residues,
            dims=dims,
            inference_params=inference_params,
            layer_idx=self.layer_idx,
            prefill_style=self.config.get("prefill_style", "fft"),
            use_flashfft=self.use_flashfft,
            fftconv_fn=self.fftconv_fn,
            column_split_hyena=self.column_split_hyena,
            long_fir_threshold=self.long_fir_threshold,
            padding_mask=padding_mask,
        )
        return y, inference_params

    def sequential_forward(self, u, inference_params):
        if self.data_dtype is None:
            self.data_dtype = u.dtype
        if len(u.shape) > 2:
            u = u[:, -1]

        fir_state = inference_params.fir_state_dict[self.layer_idx]
        iir_state = inference_params.state_dict[self.layer_idx]

        z_pre, fir_state = self.engine.step_fir(
            u, fir_state,
            weight=self.short_filter_weight, bias=self.short_filter_bias,
        )
        if self.column_split_hyena:
            x_reshaped = z_pre.reshape(
                z_pre.shape[0],
                self.num_attention_heads,
                3 * self.hidden_size_per_attention_head,
            )
            head = self.hidden_size_per_attention_head
            x2 = x_reshaped[:, :, :head].reshape(z_pre.shape[0], -1)
            x1 = x_reshaped[:, :, head : 2 * head].reshape(z_pre.shape[0], -1)
            v = x_reshaped[:, :, 2 * head:].reshape(z_pre.shape[0], -1)
        else:
            x2, x1, v = z_pre.split(
                [self.hidden_size, self.hidden_size, self.hidden_size], dim=1
            )

        y, iir_state = self.engine.step_iir(
            x2, x1, v, self.D, self.residues, self.poles, iir_state,
            iir_groups=self.hyena_filter_groups,
        )
        inference_params.fir_state_dict[self.layer_idx] = fir_state
        inference_params.state_dict[self.layer_idx] = iir_state
        y = y.to(dtype=self.data_dtype)
        return y[:, None], inference_params

    def update_time(self, L, device):
        if not hasattr(self, "t"):
            self.t = torch.arange(L, device=device)[None, None]
        elif self.t.shape[-1] < L:
            self.t = torch.arange(L, device=device)[None, None]
        else:
            self.t = self.t[..., :L]

    def compute_filter(self, L, device):
        self.update_time(L, device)
        filter_dtype = torch.float32
        residues = torch.view_as_complex(self.residues.to(filter_dtype))
        log_poles = torch.view_as_complex(self.poles.to(filter_dtype)).log()
        h = (residues * (log_poles * self.t).exp()).real.sum(1)[None]
        return h, filter_dtype, log_poles, residues


class ParallelGatedConvBlock(nn.Module):
    def __init__(self, config, layer_idx) -> None:
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.low_mem_mode = config.get("low_mem_mode", False)
        dtype = config.get("hyena_block_dtype", torch.float32)
        mlp_dtype = config.get("mlp_dtype", torch.bfloat16)
        self.pre_norm = RMSNorm(config).to(dtype=dtype)
        self.post_norm = RMSNorm(config).to(dtype=dtype)
        self.filter = ParallelHyenaFilter(config, layer_idx).to(dtype=dtype)
        self.projections = nn.Linear(config.hidden_size, 3 * config.hidden_size)
        self.out_filter_dense = nn.Linear(
            config.hidden_size, config.hidden_size
        ).to(dtype)
        self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)

    def forward(
        self,
        u,
        inference_params=None,
        padding_mask=None,
        output_attentions: bool = False,
        *args,
        **kwargs,
    ):
        z = self.projections(self.pre_norm(u))

        if isinstance(padding_mask, torch.Tensor):
            z = z * padding_mask[..., None]

        z, inference_params = self.filter(
            z, inference_params=inference_params, padding_mask=padding_mask
        )
        z_in = self.out_filter_dense(z) + u

        if isinstance(padding_mask, torch.Tensor):
            z_in = z_in * padding_mask[..., None]

        y = self.mlp(self.post_norm(z_in)) + z_in
        # Hyena blocks have no attention matrix.
        return y, None


def get_block(config, layer_idx, flash_fft=None):
    if layer_idx in config.attn_layer_idxs:
        return AttentionBlock(config, layer_idx)
    if layer_idx in config.hyena_layer_idxs:
        block = ParallelGatedConvBlock(config, layer_idx)
        if config.get("use_flashfft", False):
            block.filter.fftconv_fn = flash_fft
        return block
    raise NotImplementedError(f"layer_idx {layer_idx} not in attn or hyena indices")


# =============================================================================
# Backbone (StripedHyena)
# =============================================================================


class StripedHyena(nn.Module):
    """Pure backbone: token embedding -> N blocks -> RMSNorm.

    The unembed step is owned by the LM head wrapper, not here, so that
    ``Evo1Model`` (no LM head) can return the post-norm hidden state as
    ``last_hidden_state`` cleanly.
    """

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embedding_layer = VocabParallelEmbedding(config)
        self.norm = RMSNorm(config) if config.get("final_norm", True) else None

        if config.get("use_flashfft", False):
            import importlib
            FlashFFTConv = importlib.import_module("flashfftconv").FlashFFTConv
            # NOTE: the original togethercomputer reference had ``config.seqlen``
            # here, which is a typo - that attribute doesn't exist on the
            # config (it's ``max_seqlen``). The bug was unreachable upstream
            # because ``use_flashfft`` defaults to False; we fix it so the
            # path is at least loadable for users who do enable it.
            # FlashFFTConv requires its build-time seqlen to be 2x the
            # longest input it'll ever see (zero-padding for FFT).
            self.flash_fft = FlashFFTConv(2 * config.max_seqlen, dtype=torch.bfloat16)
        else:
            self.flash_fft = None

        self.blocks = nn.ModuleList(
            get_block(config, i, flash_fft=self.flash_fft)
            for i in range(config.num_layers)
        )

    def forward(
        self,
        x: torch.Tensor,
        inference_params_dict=None,
        padding_mask: Optional[torch.Tensor] = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
    ):
        x = self.embedding_layer.embed(x)

        all_hidden_states: list[torch.Tensor] = []
        all_attentions: list[Optional[torch.Tensor]] = []
        if output_hidden_states:
            all_hidden_states.append(x)

        if inference_params_dict is not None:
            x, inference_params_dict_out = self._stateful_forward(
                x, inference_params_dict,
                all_hidden_states=all_hidden_states,
                all_attentions=all_attentions,
                output_hidden_states=output_hidden_states,
                output_attentions=output_attentions,
            )
        else:
            x, inference_params_dict_out = self._stateless_forward(
                x, padding_mask=padding_mask,
                all_hidden_states=all_hidden_states,
                all_attentions=all_attentions,
                output_hidden_states=output_hidden_states,
                output_attentions=output_attentions,
            )

        if self.norm is not None:
            x = self.norm(x)
        if output_hidden_states:
            all_hidden_states.append(x)

        return x, inference_params_dict_out, all_hidden_states, all_attentions

    def _stateful_forward(
        self, x, inference_params_dict,
        all_hidden_states, all_attentions,
        output_hidden_states, output_attentions,
    ):
        for block_idx, block in enumerate(self.blocks):
            block_name = (
                "mha" if block_idx in self.config.attn_layer_idxs else "hyena"
            )
            inference_params = inference_params_dict[block_name]
            x, attn = block(
                x, inference_params=inference_params,
                output_attentions=output_attentions,
            )
            if output_hidden_states:
                all_hidden_states.append(x)
            if output_attentions:
                all_attentions.append(attn)
        return x, inference_params_dict

    def _stateless_forward(
        self, x, padding_mask,
        all_hidden_states, all_attentions,
        output_hidden_states, output_attentions,
    ):
        if isinstance(padding_mask, torch.Tensor):
            x = x * padding_mask[..., None]
        for block in self.blocks:
            x, attn = block(
                x, inference_params=None, padding_mask=padding_mask,
                output_attentions=output_attentions,
            )
            if output_hidden_states:
                all_hidden_states.append(x)
            if output_attentions:
                all_attentions.append(attn)
        return x, None

    def initialize_inference_params(self, max_batch_size: int = 1) -> Evo1Cache:
        return Evo1Cache(
            max_seqlen=self.config.get("max_seqlen", 8192),
            max_batch_size=max_batch_size,
            short_filter_length=self.config.short_filter_length,
            state_size=self.config.state_size,
        )

    def to_bfloat16_except_poles_residues(self):
        """Cast all parameters to bfloat16 except Hyena poles/residues."""
        for k, p in self.named_parameters():
            if "poles" not in k and "residues" not in k:
                p.data = p.data.to(torch.bfloat16)


# =============================================================================
# HuggingFace wrappers
# =============================================================================


class Evo1PreTrainedModel(PreTrainedModel):
    config_class = Evo1Config
    base_model_prefix = "backbone"
    supports_gradient_checkpointing = False
    _no_split_modules = ["AttentionBlock", "ParallelGatedConvBlock"]
    _skip_keys_device_placement = "past_key_values"
    _keys_to_ignore_on_load_missing = [r"freq", r"\.t$"]
    _keys_to_ignore_on_load_unexpected = [r"fftconv", r"twiddle_factors"]
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    # Hyena filter SSM parameters (poles / residues) MUST stay in fp32: they
    # parametrize a long-range modal-form filter whose stability collapses
    # in bf16. HF will keep these in fp32 even when the rest of the model is
    # loaded in bf16 (or fp16) via the dtype= kwarg of from_pretrained.
    _keep_in_fp32_modules = ["poles", "residues"]

    @classmethod
    def from_pretrained(cls, *args, **kwargs):
        # Evo1 was trained in bfloat16, with the modal-form filter parameters
        # (Hyena poles / residues) kept in fp32 via _keep_in_fp32_modules.
        # bf16 works correctly for all three attention backends (eager, sdpa,
        # flash_attention_2). Default to bf16 so users don't have to pass it
        # explicitly; this also silences HF's flash_attention_2 dtype warning
        # (which inspects the model dtype before force_dtype() runs in __init__).
        if "dtype" not in kwargs and "torch_dtype" not in kwargs:
            kwargs["dtype"] = torch.bfloat16
        return super().from_pretrained(*args, **kwargs)


class Evo1Model(Evo1PreTrainedModel):
    """Bare backbone: returns ``BaseModelOutputWithPast``.

    ``last_hidden_state`` is the final (post-RMSNorm) representation, ready
    to be fed into a downstream head or unembed projection.
    """

    def __init__(self, config: Evo1Config):
        super().__init__(config)
        self.backbone = StripedHyena(config)
        self.config = config
        self.post_init()
        self.force_dtype()

    def force_dtype(self):
        # Cast everything except poles/residues to bf16 (the trained dtype).
        # This runs at __init__ time so the model is usable even without an
        # explicit ``dtype=torch.bfloat16`` kwarg to ``from_pretrained``.
        self.backbone.to_bfloat16_except_poles_residues()

    def get_input_embeddings(self):
        return self.backbone.embedding_layer

    def set_input_embeddings(self, value):
        self.backbone.embedding_layer = value

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.LongTensor] = None,
        past_key_values=None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        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
        )
        # Evo1Model is the bare backbone (no LM head). Default to no caching:
        # KV caches and Hyena recurrent state are only useful for autoregressive
        # generation (Evo1ForCausalLM). For embedding extraction the caches
        # have a large per-layer memory footprint with no benefit. The user
        # can still opt-in by passing ``use_cache=True`` explicitly.
        use_cache = use_cache if use_cache is not None else False
        if use_cache and self.training:
            use_cache = False

        inputs = input_ids
        if use_cache and past_key_values is None:
            past_key_values = self.backbone.initialize_inference_params(
                max_batch_size=input_ids.shape[0],
            )

        last_hidden, past_kv, hidden_states, attentions = self.backbone(
            inputs,
            padding_mask=attention_mask,
            inference_params_dict=past_key_values if use_cache else None,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
        )

        if not return_dict:
            outputs = (last_hidden,)
            if use_cache:
                outputs += (past_kv,)
            if output_hidden_states:
                outputs += (tuple(hidden_states),)
            if output_attentions:
                outputs += (tuple(attentions),)
            return outputs

        return BaseModelOutputWithPast(
            last_hidden_state=last_hidden,
            past_key_values=past_kv if use_cache else None,
            hidden_states=tuple(hidden_states) if output_hidden_states else None,
            attentions=tuple(attentions) if output_attentions else None,
        )


class Evo1ForCausalLM(Evo1PreTrainedModel, GenerationMixin):
    """LM head wrapper. Tied to ``backbone.embedding_layer`` (Evo1 ties weights)."""

    def __init__(self, config: Evo1Config, **kwargs):
        super().__init__(config, **kwargs)
        self.backbone = StripedHyena(config)
        self.config = config

        # Pad-to-multiple-of for the vocab (matches togethercomputer config).
        vocab_size = config.vocab_size
        if vocab_size % config.make_vocab_size_divisible_by != 0:
            vocab_size += config.make_vocab_size_divisible_by - (
                vocab_size % config.make_vocab_size_divisible_by
            )
        self.vocab_size = vocab_size
        self.post_init()
        self.force_dtype()

    def force_dtype(self):
        self.backbone.to_bfloat16_except_poles_residues()

    def get_input_embeddings(self):
        return self.backbone.embedding_layer

    def set_input_embeddings(self, value):
        self.backbone.embedding_layer = value

    def get_output_embeddings(self):
        return self.backbone.embedding_layer

    def set_output_embeddings(self, value):
        self.backbone.embedding_layer = value

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        past_key_values=None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )
        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

        if use_cache and labels is not None:
            logger.warning_once(
                "use_cache=True is incompatible with loss computation; "
                "disabling cache."
            )
            use_cache = False

        inputs = input_ids
        if use_cache:
            # If the user (or HF generation) didn't pass our Evo1Cache,
            # initialize a fresh one on the first call.
            if not isinstance(past_key_values, Evo1Cache):
                past_key_values = self.backbone.initialize_inference_params(
                    max_batch_size=input_ids.shape[0],
                )
            else:
                seqlen_offset = past_key_values.seqlen_offset
                if seqlen_offset == 0:
                    # Prefill done; set offset to prompt length minus the one
                    # token we're about to consume (and that we'll keep
                    # consuming one-at-a-time below).
                    past_key_values.set_offset(input_ids.shape[-1] - 1)
                else:
                    past_key_values.advance(1)
                inputs = input_ids[:, -1:]

        last_hidden, past_kv, hidden_states, attentions = self.backbone(
            inputs,
            padding_mask=attention_mask,
            inference_params_dict=past_key_values if use_cache else None,
            output_hidden_states=output_hidden_states,
            output_attentions=output_attentions,
        )

        # Tied unembed: matmul against embedding weights.
        logits = last_hidden @ self.backbone.embedding_layer.weight.T

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

        if not return_dict:
            outputs = (logits,)
            if use_cache:
                outputs += (past_kv,)
            if output_hidden_states:
                outputs += (tuple(hidden_states),)
            if output_attentions:
                outputs += (tuple(attentions),)
            if loss is not None:
                outputs = (loss,) + outputs
            return outputs

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_kv if use_cache else None,
            hidden_states=tuple(hidden_states) if output_hidden_states else None,
            attentions=tuple(attentions) if output_attentions else None,
        )

    @classmethod
    def can_generate(cls) -> bool:
        return True

    def prepare_inputs_for_generation(
        self, input_ids, attention_mask=None, past_key_values=None, **kwargs
    ):
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "past_key_values": past_key_values,
        }