modeling_bolmo.py

import copy
from typing import Callable, Optional, Union, cast
import math

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

from transformers.utils.generic import TransformersKwargs

from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin, GenerationConfig, LogitsProcessorList, StoppingCriteriaList
from transformers.generation.utils import GenerateOutput
from transformers.integrations import use_kernel_forward_from_hub
from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
from transformers.modeling_layers import GradientCheckpointingLayer
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from transformers.processing_utils import Unpack
from transformers.utils import can_return_tuple
from transformers.utils.deprecation import deprecate_kwarg
from transformers.utils.generic import check_model_inputs

from .configuration_bolmo import BolmoConfig
from .tokenization_bolmo import BolmoTokenizerConfig
from .utils_bolmo import compute_boundary_mask, pad_right, pad_left, MaskState

try:
    from xlstm.xlstm_large.model import mLSTMLayer, mLSTMLayerConfig, mLSTMLayerStateType, soft_cap, mLSTMBackendConfig
except ImportError:
    raise ImportError("The `xlstm` package is required to use Bolmo. Please install it via `pip install xlstm`.")


@use_kernel_forward_from_hub("RMSNorm")
class BolmoRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        BolmoRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return (self.weight * hidden_states).to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    scaling: float,
    dropout: float = 0.0,
    **kwargs: Unpack[TransformersKwargs],
):
    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
    if attention_mask is not None:
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    """Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        position_ids (`torch.Tensor`, *optional*):
            Deprecated and unused.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    """
    q_type, k_type = q.dtype, k.dtype
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed.to(q_type), k_embed.to(k_type)


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


class BolmoAttention(nn.Module):
    """Multi-headed attention from 'Attention Is All You Need' paper"""

    def __init__(self, config: BolmoConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.attention_dropout = config.attention_dropout
        self.is_causal = True

        self.q_proj = nn.Linear(
            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
        )
        self.k_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.v_proj = nn.Linear(
            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
        )
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
        )
        self.q_norm = BolmoRMSNorm(config.num_attention_heads * self.head_dim, config.rms_norm_eps)
        self.k_norm = BolmoRMSNorm(config.num_key_value_heads * self.head_dim, config.rms_norm_eps)
        assert config.layer_types is not None
        self.attention_type = config.layer_types[layer_idx]
        self.sliding_window = config.sliding_window if self.attention_type == "sliding_attention" else None

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: Optional[torch.Tensor],
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.Tensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_norm(self.q_proj(hidden_states))
        key_states = self.k_norm(self.k_proj(hidden_states))
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(hidden_shape).transpose(1, 2)
        key_states = key_states.view(hidden_shape).transpose(1, 2)
        value_states = value_states.view(hidden_shape).transpose(1, 2)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        if past_key_values is not None:
            # sin and cos are specific to RoPE models; cache_position needed for the static cache
            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)

        attention_interface: Callable = eager_attention_forward
        if self.config._attn_implementation != "eager":
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            sliding_window=self.sliding_window,
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class BolmoMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


class BolmoDecoderLayer(GradientCheckpointingLayer):
    def __init__(self, config: BolmoConfig, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = BolmoAttention(config=config, layer_idx=layer_idx)

        self.mlp = BolmoMLP(config)
        self.post_attention_layernorm = BolmoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_feedforward_layernorm = BolmoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        use_cache: Optional[bool] = False,
        cache_position: Optional[torch.Tensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        residual = hidden_states
        attn_out, _ = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = self.post_attention_layernorm(attn_out)
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        mlp_out = self.mlp(hidden_states)
        hidden_states = self.post_feedforward_layernorm(mlp_out)
        hidden_states = residual + hidden_states

        return hidden_states


class BolmoBoundaryPredictor(nn.Module):
    def __init__(self, config: BolmoConfig):
        super().__init__()

        self.d_model = config.hidden_size
        self.boundary_threshold = config.boundary_threshold
        self.boundary_predictor_lookahead = config.boundary_predictor_lookahead
        self.q_proj_layer = nn.Linear(self.d_model, self.d_model, bias=False)
        self.k_proj_layer = nn.Linear(self.d_model, self.d_model, bias=False)

    def forward(
        self,
        hidden_states: torch.Tensor,
        sequence_start_indices: Optional[torch.Tensor] = None,
        epsilon: float = 1e-3,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if self.boundary_predictor_lookahead == 0:
            # do not use the same rep for k and v, use current and one before as in H-Net + pad with negative to the left
            cos_sim = torch.cat([
                torch.ones((hidden_states.shape[0], 1), device=hidden_states.device, dtype=hidden_states.dtype) * -1,
                torch.einsum(
                    "b l d, b l d -> b l",
                    F.normalize(self.q_proj_layer(hidden_states[:, :-1]), dim=-1),
                    F.normalize(self.k_proj_layer(hidden_states[:, 1:]), dim=-1),
                )
            ], dim=1)
        else:
            cos_sim = torch.einsum(
                "b l d, b l d -> b l",
                F.normalize(self.q_proj_layer(hidden_states[:, :-self.boundary_predictor_lookahead]), dim=-1),
                F.normalize(self.k_proj_layer(hidden_states[:, self.boundary_predictor_lookahead:]), dim=-1),
            )
        boundary_logprobs = torch.log1p(-cos_sim.float().clip(max=1.0 - epsilon)) - math.log(2)
        POSITIVE_LOGPROB = 0.0
        NEGATIVE_LOGPROB = -100_000
        if sequence_start_indices is None:
            boundary_logprobs[:, 0] = POSITIVE_LOGPROB
        else:
            pad_mask = torch.arange(boundary_logprobs.shape[1], device=boundary_logprobs.device)[None, :] < sequence_start_indices[:, None]
            boundary_logprobs = boundary_logprobs.masked_fill(pad_mask, NEGATIVE_LOGPROB)
            boundary_logprobs[torch.arange(len(boundary_logprobs), device=boundary_logprobs.device), sequence_start_indices] = POSITIVE_LOGPROB

        boundary_logprobs = F.pad(boundary_logprobs, (0, self.boundary_predictor_lookahead), "constant", NEGATIVE_LOGPROB)
        boundary_mask = compute_boundary_mask(boundary_logprobs, self.boundary_threshold)

        return boundary_logprobs, boundary_mask


class BolmoXLSTMLayer(mLSTMLayer):
    def __init__(self, config: BolmoConfig):
        super().__init__(mLSTMLayerConfig(
            embedding_dim=config.hidden_size,
            num_heads=config.num_local_heads,
            mlstm_backend=mLSTMBackendConfig(
                chunkwise_kernel="chunkwise--triton_limit_chunk",
                sequence_kernel="native_sequence__triton",
                step_kernel="triton",
                mode="train",
                return_last_states=True,
                autocast_kernel_dtype="float32",
            )
        ))

    # original forward adapted to support sequence_start_indices
    # i.e. set the forget gate to zero at the start of sequence
    def _original_forward(
        self, x: torch.Tensor,
        state: mLSTMLayerStateType | None = None,
        sequence_start_indices: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, mLSTMLayerStateType | None]:
        assert x.ndim == 3, f"Input must have shape [B, S, D], got {x.shape}"
        B, S, _ = x.shape
        if self.config.weight_mode == "single":
            q = self.q(x)
            k = self.k(x)
            v = self.v(x)
            o_preact = self.ogate_preact(x)
            i_preact = soft_cap(
                self.igate_preact(x), cap_value=self.config.gate_soft_cap
            )
            f_preact = soft_cap(
                self.fgate_preact(x), cap_value=self.config.gate_soft_cap
            )
        elif self.config.weight_mode == "fused":
            qkv_opreact = self.qkv_opreact(x)
            q, k, v, o_preact = torch.tensor_split(
                qkv_opreact,
                (
                    self.qk_dim,
                    2 * self.qk_dim,
                    2 * self.qk_dim + self.v_dim,
                ),
                dim=-1,
            )

            if_preact = soft_cap(
                self.ifgate_preact(x), cap_value=self.config.gate_soft_cap
            )
            i_preact, f_preact = torch.tensor_split(
                if_preact, (self.config.num_heads,), dim=-1
            )
        else:
            raise ValueError(f"Unknown weight_mode: {self.config.weight_mode}")

        q = q.reshape(B, S, self.config.num_heads, -1).transpose(1, 2)
        k = k.reshape(B, S, self.config.num_heads, -1).transpose(1, 2)
        v = v.reshape(B, S, self.config.num_heads, -1).transpose(1, 2)

        if sequence_start_indices is not None:
            f_preact[torch.arange(B, device=f_preact.device), sequence_start_indices] = -100_000

        i_preact = i_preact.transpose(1, 2)
        f_preact = f_preact.transpose(1, 2)
        if state is None:
            c_initial, n_initial, m_initial = None, None, None
        else:
            c_initial, n_initial, m_initial = state

        h, state = self.mlstm_backend(
            q=q,
            k=k,
            v=v,
            i=i_preact,
            f=f_preact,
            c_initial=c_initial,
            n_initial=n_initial,
            m_initial=m_initial,
        )
        expected_h_shape = (
            B,
            self.config.num_heads,
            S,
            self.v_dim // self.config.num_heads,
        )
        assert (
            h.shape == expected_h_shape
        ), f"Got {h.shape}, expected {expected_h_shape}"

        h = h.transpose(1, 2)
        h_norm = self.multihead_norm(h)
        h_norm = h_norm.reshape(B, S, -1)

        h_out = self.ogate_act_fn(o_preact) * h_norm

        y = self.out_proj(h_out)
        return y, state

    def forward(  # type: ignore
        self,
        x: torch.Tensor,
        past_key_values: Optional[dict] = None,
        use_cache: bool = False,
        sequence_start_indices: Optional[torch.Tensor] = None,
        cache_mask: Optional[MaskState] = None
    ):
        if self.training:
            self.mlstm_backend.config.mode = "train"
        else:
            self.mlstm_backend.config.mode = "inference"

        if use_cache:
            assert past_key_values is not None

            prev_mode = self.mlstm_backend.config.mode
            state = past_key_values.get("state", None)

            if cache_mask is not None:
                state_for_model = cast(mLSTMLayerStateType, tuple(cache_mask.selective_get(x, inv=True) for x in state) if state is not None else None)
            else:
                state_for_model = state

            h, new_state = self._original_forward(
                x,
                state=state_for_model,
                sequence_start_indices=sequence_start_indices
            )
            assert new_state is not None

            if state is None or cache_mask is None:
                state = new_state
            else:
                if cache_mask is not None:
                    for i in range(len(state)):
                        cache_mask.selective_put(new_state[i], state[i], inv=True)

            past_key_values["state"] = state
            self.mlstm_backend.config.mode = prev_mode

            return h
        else:
            h, _ = super().forward(x)
            return h

class BolmoLocalLayer(nn.Module):
    def __init__(self, config: BolmoConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size

        self.act_fn = ACT2FN[config.hidden_act]

        self.xlstm = BolmoXLSTMLayer(config)

        local_mlp_config = copy.deepcopy(config)
        local_mlp_config.intermediate_size = config.local_intermediate_size
        self.mlp = BolmoMLP(local_mlp_config)

        self.pre_xlstm_layernorm = BolmoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.pre_feedforward_layernorm = BolmoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        sequence_start_indices: Optional[torch.Tensor] = None,
        past_key_values: Optional[dict] = None,
        use_cache: Optional[bool] = False,
        cache_mask: Optional[MaskState] = None,
    ) -> torch.Tensor:
        residual = hidden_states
        xlstm_out = self.xlstm(self.pre_xlstm_layernorm(hidden_states), sequence_start_indices=sequence_start_indices, past_key_values=past_key_values["xlstm"] if past_key_values is not None else None, use_cache=use_cache, cache_mask=cache_mask)
        hidden_states = residual + xlstm_out

        # Fully Connected
        residual = hidden_states
        ffn_out = self.mlp(self.pre_feedforward_layernorm(hidden_states))
        hidden_states = residual + ffn_out

        return hidden_states


class BolmoLocalEncoder(nn.Module):
    def __init__(self, config: BolmoConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.add_expanded_embeddings = config.add_expanded_embeddings

        self.byte_embedding = nn.Embedding(
            config.vocab_size,
            self.hidden_size,
        )
        if self.add_expanded_embeddings:
            self.subword_embedding = nn.Embedding(
                config.subword_vocab_size,
                self.hidden_size,
            )
        else:
            self.subword_embedding = None

        self.layers = nn.ModuleList(
            [BolmoLocalLayer(config) for _ in range(config.num_local_encoder_layers)]
        )

        self.post_last_block_norm = BolmoRMSNorm(
            self.hidden_size,
            config.local_rms_norm_eps,
        )
        self.out_projection = nn.Linear(
            self.hidden_size,
            self.hidden_size,
            bias=True,
        )

        self.boundary_predictor_module = BolmoBoundaryPredictor(config)

        self.has_cache = False

    def prepare_inference_cache(self, batch_size: int):
        device = next(self.parameters()).device
        self.has_cache = True

        self.cache_seqlens = 0
        self.last_h = torch.zeros((batch_size, self.hidden_size), dtype=self.out_projection.weight.dtype, device=device)
        self.layer_states = [{"xlstm": {}} for _ in range(len(self.layers))]

    def free_inference_cache(self):
        self.has_cache = False
        if hasattr(self, "cache_seqlens"):
            del self.cache_seqlens
        if hasattr(self, "last_h"):
            del self.last_h
        if hasattr(self, "layer_states"):
            del self.layer_states

    def _embed(self, tokens, expanded_input_ids: Optional[torch.Tensor] = None):
        embeddings = self.byte_embedding(tokens)
        if self.add_expanded_embeddings:
            assert expanded_input_ids is not None and self.subword_embedding is not None
            embeddings = embeddings + self.subword_embedding(expanded_input_ids)

        return embeddings

    def _pool(
        self,
        h: torch.Tensor,
        boundary_mask: torch.Tensor | None,
        n_patches: int,
        boundary_state: Optional[MaskState] = None,
    ):
        if self.has_cache and self.cache_seqlens > 0:
            assert boundary_state is not None
            if boundary_state.all():
                assert h.shape[1] == 1
                reduced_h = h
            else:
                reduced_h = h[[], :, :]
        else:
            assert boundary_mask is not None

            L = h.shape[1]
            token_idx = (
                torch.arange(L, device=h.device)[None, :] + (~boundary_mask).long() * L  # type: ignore
            )
            seq_sorted_indices = torch.argsort(token_idx, dim=1)
            index = seq_sorted_indices[:, :n_patches, None].expand(
                -1, -1, h.shape[-1]
            )

            reduced_h = torch.gather(
                h,
                dim=1,
                index=index,
            )

        return reduced_h

    def forward(
        self,
        input_ids,
        true_boundary_mask: Optional[torch.Tensor] = None,
        boundary_state: Optional[MaskState] = None,
        pad_state: Optional[MaskState] = None,
        expanded_input_ids: Optional[torch.Tensor] = None,
        sequence_start_indices: Optional[torch.Tensor] = None,
    ):
        embeddings = self._embed(input_ids, expanded_input_ids)

        # pass through encoder layers
        if self.has_cache and self.cache_seqlens > 0:
            assert pad_state is not None

            # step those batch positions which are not currently idle (i.e. at a boundary position)
            # if all batch positions are idle, skip the step entirely
            # all positions being idle only happens if fuse_boundaries=False. In this case, the step where we
            # obtain a new representation from the global model will have all positions for the local encoder being idle.
            if not pad_state.all():
                h = pad_state.selective_get(embeddings, inv=True)

                for i, block in enumerate(self.layers):
                    h = block(h, past_key_values=self.layer_states[i], use_cache=True, cache_mask=pad_state)

                if self.post_last_block_norm is not None:
                    h = self.post_last_block_norm(h)

                pad_state.selective_put(h[:, -1, :], self.last_h, inv=True)

            h = self.last_h.unsqueeze(1)
        else:
            h = embeddings
            for i, block in enumerate(self.layers):
                if self.has_cache:
                    use_cache = True
                    past_key_values = self.layer_states[i]
                else:
                    use_cache = False
                    past_key_values = None

                h = block(h, past_key_values=past_key_values, use_cache=use_cache, sequence_start_indices=sequence_start_indices)

            if self.post_last_block_norm is not None:
                h = self.post_last_block_norm(h)

            if self.has_cache:
                self.last_h.copy_(h[:, -1, :])

        if not self.has_cache or self.cache_seqlens == 0: # only used for prefill
            boundary_logprobs, boundary_mask = self.boundary_predictor_module(
                h,
                sequence_start_indices=sequence_start_indices,
            )
            if boundary_state is not None:
                # can't predict through encoder - must be through prev local decoder step
                boundary_mask[:, -1] = boundary_state.mask
        else:
            boundary_logprobs = boundary_mask = None

        # overwrite with true boundaries
        if true_boundary_mask is not None:
            boundary_mask = true_boundary_mask

        patch_embeddings = self._pool(
            h=h,
            boundary_mask=boundary_mask,
            n_patches=int(cast(torch.Tensor, boundary_mask).sum(-1).max().item()) if boundary_mask is not None else 1,
            boundary_state=boundary_state,
        )
        patch_embeddings = self.out_projection(patch_embeddings)

        if self.has_cache:
            self.cache_seqlens += input_ids.shape[1]

        return h, patch_embeddings, boundary_logprobs, boundary_mask


class BolmoLocalDecoder(nn.Module):
    def __init__(self, config: BolmoConfig):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size

        self.initial_norm = BolmoRMSNorm(
            self.hidden_size,
            eps=config.local_rms_norm_eps,
        )

        self.in_projection = nn.Linear(
            self.hidden_size,
            self.hidden_size,
            bias=True,
        )

        self.layers = nn.ModuleList(
            [BolmoLocalLayer(config) for _ in range(config.num_local_decoder_layers)]
        )

        self.has_cache = False

    def prepare_inference_cache(self, batch_size: int):
        device = next(self.parameters()).device
        self.has_cache = True

        self.cache_seqlens = 0
        self.last_value = torch.zeros((batch_size, self.hidden_size), dtype=self.in_projection.weight.dtype, device=device)
        self.layer_states = [{"xlstm": {}} for _ in range(len(self.layers))]

    def free_inference_cache(self):
        self.has_cache = False
        if hasattr(self, "cache_seqlens"):
            del self.cache_seqlens
        if hasattr(self, "last_value"):
            del self.last_value
        if hasattr(self, "layer_states"):
            del self.layer_states

    def _depool(
        self,
        embeds: torch.Tensor,
        patch_embeds: torch.Tensor,
        boundary_mask: Optional[torch.Tensor],
        boundary_state: Optional[MaskState] = None,
        sequence_start_indices: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        if self.has_cache and self.cache_seqlens > 0:
            assert boundary_state is not None

            if patch_embeds.numel() > 0:
                # we got a new value from the global model, so must be at boundary position
                h_patch = patch_embeds[:, -1:, :]
                h = embeds + h_patch

                self.last_value.copy_(h_patch[:, -1])
            else:
                h = embeds + self.last_value.unsqueeze(1)

            # skip pad positions until we get a new value from the global model
            if patch_embeds.numel() == 0:
                h = boundary_state.selective_get(h, inv=True)
            else:
                boundary_state = None

            if h.shape[0] > 0:
                for i, layer in enumerate(self.layers):
                    h = layer(h, past_key_values=self.layer_states[i], use_cache=True, cache_mask=boundary_state)

            self.cache_seqlens += h.shape[1]

            return h
        else:
            assert boundary_mask is not None

            h_patch = patch_embeds 
            prepool_out = h_patch

            # TODO(benjaminm): clipping is problematic if it happens too much; track clip %.
            plug_back_idx = (torch.cumsum(boundary_mask, dim=1) - 1).clip(min=0, max=prepool_out.shape[1] - 1)
            depool_out = torch.gather(
                prepool_out,
                dim=1,
                index=plug_back_idx.unsqueeze(-1).expand(-1, -1, self.hidden_size),
            )
                
            depool_out_modulated = depool_out
            h = depool_out_modulated + embeds

            for i, layer in enumerate(self.layers):
                if self.has_cache:
                    use_cache = True
                    past_key_values = self.layer_states[i]
                else:
                    use_cache = False
                    past_key_values = None

                h = layer(h, past_key_values=past_key_values, use_cache=use_cache, sequence_start_indices=sequence_start_indices)

            if self.has_cache:
                self.last_value.copy_(prepool_out[:, -1])
                self.cache_seqlens += h.shape[1]

            return h

    def forward(
        self,
        embeds: torch.Tensor,
        patch_embeds: torch.Tensor,
        boundary_state: Optional[MaskState],
        boundary_mask: torch.Tensor | None,
        sequence_start_indices: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        h = self.in_projection(embeds)
        h_patch = self.initial_norm(patch_embeds)

        return self._depool(
            embeds=h,
            patch_embeds=h_patch,
            boundary_mask=boundary_mask,
            boundary_state=boundary_state,
            sequence_start_indices=sequence_start_indices,
        )


class BolmoRotaryEmbedding(nn.Module):
    inv_freq: torch.Tensor  # fix linting for `register_buffer`

    def __init__(self, config: BolmoConfig, device=None, rope_type: Optional[str] = None):
        super().__init__()
        if rope_type is not None:
            self.rope_type = rope_type
        elif hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
            # BC: "rope_type" was originally "type"
            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
        else:
            self.rope_type = "default"
        assert self.rope_type is not None

        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
        position_ids_expanded = position_ids[:, None, :].float()

        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling
            return cos, sin


class BolmoPreTrainedModel(PreTrainedModel):
    config: BolmoConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["BolmoDecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True

    _can_compile_fullgraph = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": BolmoDecoderLayer,
        "attentions": BolmoAttention,
    }


class BolmoModel(BolmoPreTrainedModel):
    def __init__(self, config: BolmoConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.local_encoder = BolmoLocalEncoder(config)
        self.local_decoder = BolmoLocalDecoder(config)

        self.layers = nn.ModuleList(
            [BolmoDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = BolmoRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.rotary_embs = nn.ModuleDict(
            {
                "sliding_attention": BolmoRotaryEmbedding(config=config, rope_type="default"),
                "full_attention": BolmoRotaryEmbedding(config=config),
            }
        )

        self.tokenizer_config = BolmoTokenizerConfig(**config.tokenizer_config)
        self._tokenizer = None

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.local_encoder.byte_embedding

    def set_input_embeddings(self, value: nn.Embedding):  # type: ignore
        self.local_encoder.byte_embedding = value

    @property
    def tokenizer(self):
        if self._tokenizer is None:
            self._tokenizer = self.tokenizer_config.build()

        return self._tokenizer

    def prefill_boundary_prediction_forward(
        self,
        input_ids: torch.Tensor,
        expanded_input_ids: Optional[torch.Tensor] = None,
        sequence_start_indices: Optional[torch.Tensor] = None,
        last_token_is_boundary: bool = False,
        **kwargs,
    ) -> torch.Tensor:
        _, _, _, boundary_mask = self.local_encoder.forward(  # type: ignore
            input_ids,
            expanded_input_ids=expanded_input_ids,
            boundary_state=MaskState(torch.full((input_ids.shape[0],), fill_value=last_token_is_boundary, device=input_ids.device, dtype=torch.bool)),
            pad_state=MaskState(torch.zeros((input_ids.shape[0],), device=input_ids.device, dtype=torch.bool)),
            sequence_start_indices=sequence_start_indices,
        )

        return cast(torch.Tensor, boundary_mask)

    @check_model_inputs()
    def forward(
        self,
        input_ids: torch.Tensor,
        expanded_input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        cache_position: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        boundary_mask: Optional[torch.Tensor] = None,
        boundary_state: Optional[MaskState] = None,
        pad_state: Optional[MaskState] = None,
        sequence_start_indices: Optional[torch.Tensor] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPast:
        batch_size = input_ids.shape[0]
        device = input_ids.device

        if self.local_encoder.add_expanded_embeddings and expanded_input_ids is None and input_ids is not None:
            # not optimized
            expanded_input_ids_list: list[torch.Tensor] = []
            for example_idx in range(batch_size):
                expanded_input_ids_list.append(torch.tensor(self.tokenizer.expand_byte_ids(input_ids[example_idx].tolist()), dtype=torch.long, device=device))
            expanded_input_ids = pad_right(expanded_input_ids_list, value=self.tokenizer.pad_token_id, multiple_of=1)  # type: ignore

        h_byte, h_patch, _, boundary_mask = self.local_encoder(
            input_ids=input_ids,
            expanded_input_ids=expanded_input_ids,
            true_boundary_mask=boundary_mask,
            boundary_state=boundary_state,
            pad_state=pad_state,
        )

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position: torch.Tensor = torch.arange(
                past_seen_tokens, past_seen_tokens + h_patch.shape[1], device=device
            )

        if position_ids is None:
            position_ids = cache_position.unsqueeze(0)  # type: ignore

        # It may already have been prepared by e.g. `generate`
        if not isinstance(causal_mask_mapping := attention_mask, dict):
            # Prepare mask arguments
            mask_kwargs = {
                "config": self.config,
                "input_embeds": h_patch,
                "attention_mask": attention_mask,
                "cache_position": cache_position,
                "past_key_values": past_key_values,
                "position_ids": position_ids,
            }
            # Create the masks
            causal_mask_mapping = {
                "full_attention": create_causal_mask(**mask_kwargs),
                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),
            }

        position_embeddings_mapping = {
            "sliding_attention": self.rotary_embs["sliding_attention"](h_byte, position_ids),
            "full_attention": self.rotary_embs["full_attention"](h_byte, position_ids),
        }

        if h_patch.numel() > 0:
            # we need to convert from right-pad to left-pad and back for prefill
            # since flash attention expects left-pad and local/enc dec expect right-pad global tokens
            # should add better left-pad support but this only affects prefill so OK for now
            # although super inefficient!
            if boundary_mask is not None: # prefill
                n_boundaries = boundary_mask.sum(-1)

                for i, current_n_boundaries in enumerate(n_boundaries):
                    h_patch[i, -current_n_boundaries:] = h_patch[i, :current_n_boundaries].clone()

            h_patch_after_global = h_patch

            for decoder_layer in self.layers[: self.config.num_hidden_layers]:
                h_patch_after_global = decoder_layer(
                    h_patch_after_global,
                    attention_mask=causal_mask_mapping[decoder_layer.self_attn.attention_type],
                    position_ids=position_ids,
                    past_key_values=past_key_values,
                    cache_position=cache_position,
                    position_embeddings=position_embeddings_mapping[decoder_layer.self_attn.attention_type],
                    **kwargs,
                )

            if boundary_mask is not None: # prefill
                n_boundaries = boundary_mask.sum(-1)

                for i, current_n_boundaries in enumerate(n_boundaries):
                    h_patch_after_global[i, :current_n_boundaries] = h_patch_after_global[i, -current_n_boundaries:].clone()
        else:
            h_patch_after_global = h_patch

        h_out = self.local_decoder.forward(  # type: ignore
            embeds=h_byte,
            patch_embeds=h_patch_after_global,
            boundary_mask=boundary_mask,
            boundary_state=boundary_state,
            sequence_start_indices=sequence_start_indices,
        )
        h_out = self.norm(h_out)

        return BaseModelOutputWithPast(
            last_hidden_state=h_out,
            past_key_values=past_key_values,
        )


class BolmoForCausalLM(BolmoPreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]
    _tp_plan = {"lm_head": "colwise_rep"}
    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}

    def __init__(self, config):
        super().__init__(config)
        self.model = BolmoModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        # Initialize weights and apply final processing
        self.post_init()
    
    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings: nn.Linear):
        self.lm_head = new_embeddings

    @can_return_tuple
    def forward(
        self,
        input_ids: torch.Tensor,
        expanded_input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_values: Optional[Cache] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        cache_position: Optional[torch.Tensor] = None,
        use_cache: Optional[bool] = None,
        boundary_mask: Optional[torch.Tensor] = None,
        boundary_state: Optional[MaskState] = None,
        pad_state: Optional[MaskState] = None,
        sequence_start_indices: Optional[torch.Tensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **kwargs: Unpack[TransformersKwargs],
    ) -> CausalLMOutputWithPast:
        r"""
        Example:

        ```python
        >>> from transformers import AutoTokenizer, BolmoForCausalLM

        >>> model = BolmoForCausalLM.from_pretrained("meta-olmo3/Bolmo-2-7b-hf")
        >>> tokenizer = AutoTokenizer.from_pretrained("meta-olmo3/Bolmo-2-7b-hf")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```"""
        outputs: BaseModelOutputWithPast = self.model(
            input_ids=input_ids,
            expanded_input_ids=expanded_input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            use_cache=use_cache,
            boundary_mask=boundary_mask,
            boundary_state=boundary_state,
            pad_state=pad_state,
            sequence_start_indices=sequence_start_indices,
            **kwargs,
        )

        hidden_states = cast(torch.Tensor, outputs.last_hidden_state)
        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :])

        return CausalLMOutputWithPast(
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @torch.no_grad()
    def generate(  # type: ignore
        self,
        inputs: torch.Tensor,
        generation_config: Optional[GenerationConfig] = None,
        logits_processor: Optional[LogitsProcessorList] = None,
        stopping_criteria: Optional[StoppingCriteriaList] = None,
        use_model_defaults: Optional[bool] = None,
        **kwargs,
    ) -> Union[GenerateOutput, torch.Tensor]:
        # generic preprocessing

        generation_config, model_kwargs = self._prepare_generation_config(
            generation_config, use_model_defaults, **kwargs
        )
        self._prepare_special_tokens(generation_config, device=self.model.device)

        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()

        # start of custom generate

        expand_input_ids = self.model.local_encoder.add_expanded_embeddings
        batch_size = len(inputs)

        if expand_input_ids:
            expanded_input_ids = []

            for i in range(len(inputs)):
                expanded_input_ids.append(torch.tensor(self.model.tokenizer.expand_byte_ids(inputs[i].tolist()), device=self.device, dtype=torch.long))

            expanded_input_ids = pad_left(expanded_input_ids, value=self.model.tokenizer.pad_token_id, multiple_of=1)  # type: ignore
        else:
            expanded_input_ids = None

        byte_input_ids = inputs
        sequence_start_indices = (byte_input_ids == self.model.tokenizer.pad_token_id).sum(-1)
        batch_size, prompt_len = byte_input_ids.shape
        finished = torch.zeros(batch_size, dtype=torch.bool, device=self.device)

        boundary_offset = self.model.tokenizer.offset + 256
        eos = self.model.tokenizer.eos_token_id

        self.model.local_encoder.free_inference_cache()
        self.model.local_decoder.free_inference_cache()

        boundary_mask = self.model.prefill_boundary_prediction_forward(  # type: ignore
            byte_input_ids,
            expanded_input_ids=expanded_input_ids,
            sequence_start_indices=sequence_start_indices,
        )

        self.model.local_encoder.prepare_inference_cache(batch_size)
        self.model.local_decoder.prepare_inference_cache(batch_size)

        # roll back by one and force decoding to account for lookahead
        boundary_mask = boundary_mask[:, :-1]
        # need to roll one byte back and force decoding to detect whether the last byte is a boundary
        forced_decoding_ids = byte_input_ids[:, -1].cpu().tolist()
        byte_input_ids = byte_input_ids[:, :-1]
        expanded_input_ids = expanded_input_ids[:, :-1] if expanded_input_ids is not None else None
        # stays the same unless last token is pad.
        sequence_start_indices = (byte_input_ids == self.model.tokenizer.pad_token_id).sum(-1)

        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
        has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
        generation_config = self._prepare_generated_length(
            generation_config=generation_config,
            has_default_max_length=has_default_max_length,
            has_default_min_length=has_default_min_length,
            model_input_name="input_ids",
            inputs_tensor=byte_input_ids,
            input_ids_length=byte_input_ids.shape[1],
        )

        logits_processor = self._get_logits_processor(
            generation_config=generation_config,  # type: ignore
            input_ids_seq_length=byte_input_ids.shape[1],
            encoder_input_ids=byte_input_ids,  # type: ignore
            logits_processor=logits_processor,
            device=byte_input_ids.device,  # type: ignore
            model_kwargs=model_kwargs,
        )
        stopping_criteria = self._get_stopping_criteria(
            generation_config=generation_config,  # type: ignore
            stopping_criteria=stopping_criteria,
            tokenizer=self.model.tokenizer,
        )

        # output container
        generated = byte_input_ids

        max_n_prefill_patches = boundary_mask.sum(-1).max().item()
        tokens_generated_plus_prefilled = max_n_prefill_patches
        bytes_generated = 0

        # generation state
        boundary_state = MaskState(boundary_mask[:, -1].clone())
        pad_state = MaskState(torch.zeros(batch_size, dtype=torch.bool, device=self.device))
        next_tokens = torch.full((batch_size,), self.model.tokenizer.bpe_token_end_id, device=self.device, dtype=torch.long)  # type: ignore
        non_boundary_generated_tokens = [[byte_input_ids[example_idx, -1].item()] for example_idx in range(batch_size)]
        bytes_since_boundary = (boundary_mask.flip(1).cumsum(-1) == 0).sum(-1)
        is_first_forward = True
        global_past_key_values = None

        while not finished.all():
            input_ids_for_model = (
                generated
                if is_first_forward
                else torch.tensor([x[-1] for x in non_boundary_generated_tokens], device=generated.device, dtype=generated.dtype).unsqueeze(1)
            )
            assert not (
                (input_ids_for_model == self.model.tokenizer.bpe_token_end_id) |
                (input_ids_for_model >= boundary_offset)
            ).any().item()  # type: ignore
            if expand_input_ids:
                expanded_input_ids_for_model = torch.zeros_like(input_ids_for_model)
                for i in range(input_ids_for_model.shape[0]):
                    expanded_input_ids_for_model[i, :] = torch.tensor(self.model.tokenizer.expand_byte_ids(
                        generated[i, :].tolist(),
                        n_last=input_ids_for_model.shape[1],
                    ), device=expanded_input_ids_for_model.device, dtype=expanded_input_ids_for_model.dtype)
            else:
                expanded_input_ids_for_model = None

            out = self.forward(  # type: ignore
                input_ids_for_model,
                expanded_input_ids=expanded_input_ids_for_model,
                boundary_mask=boundary_mask if is_first_forward else None,
                boundary_state=boundary_state,
                pad_state=pad_state,
                sequence_start_indices=sequence_start_indices,
                logits_to_keep=1,
                use_cache=True,
                past_key_values=global_past_key_values,
            )
            next_token_logits = cast(torch.Tensor, out.logits)
            global_past_key_values = out.past_key_values

            if boundary_state.all():
                # new token, must not be boundary
                bytes_since_boundary[:] = 0
            else:
                boundary_state.selective_add(1, bytes_since_boundary, inv=True)

            if any(x is not None for x in forced_decoding_ids):
                # only supported for the first token atm, so len(next_token_logits) == batch_size
                assert len(next_token_logits) == batch_size and is_first_forward
                for example_idx in range(batch_size):
                    forced_decoding_id = forced_decoding_ids[example_idx]

                    if forced_decoding_id is not None:
                        no_boundary_logit = next_token_logits[example_idx, 0, forced_decoding_id].item()
                        boundary_logit = next_token_logits[example_idx, 0, forced_decoding_id + boundary_offset].item()

                        next_token_logits[example_idx, 0, :] = -100_000
                        next_token_logits[example_idx, 0, forced_decoding_id] = no_boundary_logit
                        next_token_logits[example_idx, 0, forced_decoding_id + boundary_offset] = boundary_logit

                        forced_decoding_ids[example_idx] = None  # only force once

            # passing input_ids to logit processor not implemented
            next_token_scores = logits_processor(None, next_token_logits[:, -1]) # type: ignore

            if generation_config is not None and generation_config.do_sample:
                probs = nn.functional.softmax(next_token_scores, dim=-1)
                new_next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
            else:
                new_next_tokens = torch.argmax(next_token_scores, dim=-1)

            if boundary_state.all() or is_first_forward:
                tokens_generated_plus_prefilled += 1

                next_tokens = new_next_tokens
                next_tokens_cpu = next_tokens.cpu()
                for example_idx in range(batch_size):
                    if finished[example_idx].item():
                        continue

                    next_token_cpu = next_tokens_cpu[example_idx].item()

                    if next_token_cpu >= boundary_offset:
                        next_token_cpu -= boundary_offset

                    non_boundary_generated_tokens[example_idx].append(next_token_cpu)
            else:
                next_tokens[:] = self.model.tokenizer.bpe_token_end_id # type: ignore
                boundary_state.selective_put(new_next_tokens, next_tokens, inv=True)
                next_tokens_cpu = next_tokens.cpu()

                for example_idx in range(batch_size):
                    if finished[example_idx].item():
                        continue

                    next_token_cpu = next_tokens_cpu[example_idx].item()

                    if not boundary_state.cpu_mask[example_idx].item():
                        if next_token_cpu >= boundary_offset:
                            next_token_cpu -= boundary_offset

                        non_boundary_generated_tokens[example_idx].append(next_token_cpu)

            is_first_forward = False

            boundary_state = MaskState(
                (next_tokens == self.model.tokenizer.bpe_token_end_id) |
                (next_tokens >= boundary_offset) |
                finished
            )  # type: ignore
            pad_state = MaskState(
                (next_tokens == self.model.tokenizer.bpe_token_end_id) |
                finished
            )

            # Force EOS for (previously) finished sequences
            next_tokens = torch.where(finished, torch.full_like(next_tokens, eos), next_tokens)

            # Append next tokens
            generated = torch.cat([generated, next_tokens.unsqueeze(-1)], dim=1)

            # Handle finished sequences
            stop_hit = next_tokens.eq(eos) | next_tokens.eq(eos + boundary_offset)

            for i in range(batch_size):
                # passing `scores` to stopping criteria not implemented
                if stopping_criteria(torch.tensor(non_boundary_generated_tokens[i], dtype=torch.long).unsqueeze(0), None).squeeze(0).item():  # type: ignore
                    stop_hit[i] = True

            finished |= stop_hit
            bytes_generated += 1

        return pad_left([
            torch.cat([byte_input_ids[i, :-1], torch.tensor(x, dtype=torch.long, device=byte_input_ids.device)])
            for i, x in enumerate(non_boundary_generated_tokens)
        ], value=self.model.tokenizer.pad_token_id, multiple_of=1)  # type: ignore

__all__ = ["BolmoForCausalLM", "BolmoModel", "BolmoPreTrainedModel"]