model.py

# model.py
import torch
from torch import Tensor
import torch.nn as nn
from torch.nn import functional as F
from dataclasses import dataclass
from functools import partial
import math
from fast_attnres import (
    FAST_ATTNRES_DTYPES,
    FAST_ATTNRES_MAX_SOURCES,
    FAST_ATTNRES_MAX_WIDTH,
    FAST_ATTNRES_RELEASE_COMMIT,
    FAST_ATTNRES_SOURCE_SHA256,
    FAST_ATTNRES_VERSION,
    FAST_ATTNRES_WHEEL_SHA256,
    FastAttnResDecision,
    fast_attnres_config_decision,
    fast_attnres_package_provenance,
    load_fast_attnres,
)
from attnres_ops import (
    attention_residual_phase1_from_logits,
    attention_residual_phase2,
    attention_residual_phase2_torch,
    attention_residual_phase2_from_logit,
    attention_residual_read,
    lrid_attention_residual_phase2,
    lrid_attention_residual_phase2_torch,
    lrid_attention_residual_read,
)


def rms_norm_eps(x: torch.Tensor, eps: float = None) -> float:
    if eps is not None:
        return eps
    return torch.finfo(x.dtype).eps


def norm(x: Tensor, eps: float = None):
    if hasattr(F, "rms_norm"):
        return F.rms_norm(x, (x.size(-1),), eps=eps)
    return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + rms_norm_eps(x, eps))


def _norm_lrid_key(x: Tensor, num_heads: int):
    x_shape = x.shape
    x = x.reshape(*x_shape[:-1], num_heads, x_shape[-1] // num_heads)
    return norm(x).reshape(*x_shape)


class TokenEmbedding(nn.Module):
    def __init__(self, num_embeddings, embedding_dim):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.weight = nn.Parameter(torch.empty(num_embeddings, embedding_dim))

    def forward(self, idx):
        return F.embedding(idx, self.weight)


@dataclass
class ModelConfig:
    block_size: int = 1024
    vocab_size: int = 100277
    n_layer: int = 12 
    n_head: int = 12
    n_embd: int = 768
    mlp_hidden_dim: int = None
    mlp_ratio: float = 4.0
    weight_tying: bool = False
    rope_theta: float = 500000.0
    norm_pos: str = "after"
    qk_norm: bool = True
    clip_qkv: float = None
    flash_attention: bool = False
    init_std: float = 0.02
    init_cutoff_factor: float = None
    attnres_type: str = None
    use_attnres: bool = False
    use_fused_attnres: bool = False
    attnres_num_blocks: int = 8
    attnres_block_average: bool = True
    attnres_block_average_mode: str = "count"
    attnres_block_count_prior: bool = True
    attnres_block_alpha: object = "legacy"
    attnres_block_beta: object = "legacy"
    attnres_block_alpha_learned: bool = False
    attnres_block_beta_learned: bool = False
    attnres_block_alpha_scope: str = "shared"
    attnres_block_beta_scope: str = "shared"
    attnres_block_split_sublayers: bool = False
    attnres_block_learned_scale: bool = False
    attnres_block_learned_scale_init: str = "count"
    attnres_block_value_norm: bool = False
    attnres_key_norm: bool = True
    attn_res_query_norm: bool = False
    attn_res_query_init: str = "zero"
    attnres_training_cache_phase1: bool = True
    attnres_training_torch_phase2: bool = True
    attnres_fuse_read_norm: bool = True
    use_lrid: bool = False
    lrid_rank: int = 64
    lrid_projection_rank: int = None
    lrid_num_heads: int = 1
    lrid_input_dependent_query: bool = False
    lrid_static_embedding_key: bool = False
    lrid_add_static_embedding_key: bool = False
    lrid_add_static_source_key: bool = False
    lrid_key_from_value: bool = False
    lrid_key_from_value_shared: bool = False
    lrid_key_from_output_tail: bool = False
    lrid_key_value_norm: bool = True
    lrid_query_from_value: bool = False
    lrid_query_from_value_shared: bool = False
    lrid_use_logit_scale: bool = True
    lrid_logit_scale: float = None
    # ``auto`` selects Fast for standard (R=D) and sliced output-tail LRID
    # (R<=D), while leaving projected/experimental LRID on the legacy path.
    # Keep the dataclass default legacy for old direct-library callers and
    # checkpoints; train.py opts into auto by default.
    attnres_backend: str = "legacy"

    def __post_init__(self):
        requested_attnres_backend = str(self.attnres_backend or "legacy").lower()
        if requested_attnres_backend not in {"auto", "legacy", "fast"}:
            raise ValueError("attnres_backend must be one of: auto, legacy, fast")
        self.attnres_type = (self.attnres_type or "block")
        self.attnres_type = self.attnres_type.lower()
        self.attnres_block_average_mode = (self.attnres_block_average_mode or "count").lower()
        if self.attnres_block_average_mode not in {"count", "sqrt"}:
            raise ValueError("attnres_block_average_mode must be one of: count, sqrt")
        self.attnres_block_alpha_scope = self._normalize_block_power_scope(
            self.attnres_block_alpha_scope,
            "attnres_block_alpha_scope",
        )
        self.attnres_block_beta_scope = self._normalize_block_power_scope(
            self.attnres_block_beta_scope,
            "attnres_block_beta_scope",
        )
        self.attnres_block_alpha = self._normalize_block_power_value(
            self.attnres_block_alpha,
            "attnres_block_alpha",
        )
        self.attnres_block_beta = self._normalize_block_power_value(
            self.attnres_block_beta,
            "attnres_block_beta",
        )
        self._validate_block_power_value_length(
            self.attnres_block_alpha,
            self.attnres_block_alpha_scope,
            "attnres_block_alpha",
        )
        self._validate_block_power_value_length(
            self.attnres_block_beta,
            self.attnres_block_beta_scope,
            "attnres_block_beta",
        )
        self.attnres_block_learned_scale_init = self._normalize_block_scale_init(
            self.attnres_block_learned_scale_init
        )
        if self.attnres_block_learned_scale and self.attnres_type != "block":
            raise ValueError("attnres_block_learned_scale requires attnres_type='block'")
        if self.attnres_block_value_norm and self.attnres_type != "block":
            raise ValueError("attnres_block_value_norm requires attnres_type='block'")
        if self.attnres_block_value_norm and self.attnres_block_learned_scale:
            raise ValueError("attnres_block_value_norm and attnres_block_learned_scale are mutually exclusive")
        explicit_alpha_formula = (
            self.attnres_block_alpha != "legacy"
            or self.attnres_block_alpha_learned
        )
        if explicit_alpha_formula and (self.attnres_block_learned_scale or self.attnres_block_value_norm):
            raise ValueError(
                "attnres_block_alpha formula mode is mutually exclusive with "
                "attnres_block_learned_scale and attnres_block_value_norm"
            )
        if (
            self.attnres_type == "block"
            and self.attnres_block_count_prior
            and (
                self.attnres_block_learned_scale
                or self.attnres_block_value_norm
            )
        ):
            raise ValueError(
                "attnres_block_count_prior requires alpha-formula block summaries: "
                "attnres_block_learned_scale=False, and attnres_block_value_norm=False"
            )
        self.attn_res_query_init = (self.attn_res_query_init or "zero").lower()
        if self.attn_res_query_init not in {"zero", "normal", "trunc_normal"}:
            raise ValueError("attn_res_query_init must be one of: zero, normal, trunc_normal")
        if self.lrid_static_embedding_key and self.lrid_add_static_embedding_key:
            raise ValueError("lrid_static_embedding_key and lrid_add_static_embedding_key are mutually exclusive")
        if self.lrid_key_from_value_shared:
            self.lrid_key_from_value = True
        if self.lrid_key_from_output_tail:
            if self.lrid_key_from_value:
                raise ValueError("lrid_key_from_output_tail and lrid_key_from_value are mutually exclusive")
            if self.lrid_static_embedding_key or self.lrid_add_static_embedding_key or self.lrid_add_static_source_key:
                raise ValueError(
                    "lrid_key_from_output_tail cannot be combined with static LRID key additions"
                )
        if self.lrid_query_from_value_shared:
            self.lrid_query_from_value = True
        if self.lrid_rank < 1:
            raise ValueError("lrid_rank must be >= 1")
        if self.lrid_key_from_output_tail and self.lrid_rank > self.n_embd:
            raise ValueError("lrid_rank must be <= n_embd when lrid_key_from_output_tail=True")
        if self.lrid_projection_rank is None:
            self.lrid_projection_rank = self.lrid_rank
        if self.lrid_projection_rank < self.lrid_rank:
            raise ValueError("lrid_projection_rank must be >= lrid_rank")
        if self.lrid_num_heads < 1:
            raise ValueError("lrid_num_heads must be >= 1")
        if self.lrid_rank % self.lrid_num_heads != 0:
            raise ValueError("lrid_rank must be divisible by lrid_num_heads")
        if self.n_embd % self.lrid_num_heads != 0:
            raise ValueError("n_embd must be divisible by lrid_num_heads")
        if not self.lrid_use_logit_scale:
            self.lrid_logit_scale = 1.0
        elif self.lrid_logit_scale is None:
            self.lrid_logit_scale = 1.0 / math.sqrt(self.lrid_rank // self.lrid_num_heads)
        elif self.lrid_logit_scale <= 0.0:
            raise ValueError("lrid_logit_scale must be positive")
        if self.use_lrid:
            self.use_attnres = True
        self._attnres_backend_requested = requested_attnres_backend
        if requested_attnres_backend == "auto":
            self.attnres_backend = (
                "fast"
                if self.use_attnres and (not self.use_lrid or self.lrid_key_from_output_tail)
                else "legacy"
            )
        else:
            self.attnres_backend = requested_attnres_backend

    @staticmethod
    def _normalize_block_scale_init(value):
        value = str(value or "count").lower().replace(" ", "")
        if value in {"count", "1/c", "inv_count", "inverse_count"}:
            return "count"
        if value in {"sqrt", "1/sqrtc", "1/sqrt(c)", "inv_sqrt", "inverse_sqrt"}:
            return "sqrt"
        if value in {"one", "1"}:
            return "one"
        raise ValueError(
            "attnres_block_learned_scale_init must be one of: count, sqrt, one"
        )

    @staticmethod
    def _normalize_block_power_scope(value, name):
        value = str(value or "shared").lower().replace("-", "_")
        aliases = {
            "all": "shared",
            "global": "shared",
            "layer": "per_residual",
            "layers": "per_residual",
            "per_layer": "per_residual",
            "residual": "per_residual",
            "block": "per_block",
            "blocks": "per_block",
        }
        value = aliases.get(value, value)
        if value not in {"shared", "per_residual", "per_block"}:
            raise ValueError(f"{name} must be one of: shared, per_residual, per_block")
        return value

    @staticmethod
    def _normalize_block_power_value(value, name):
        if isinstance(value, str):
            text = value.strip().lower()
            if text == "legacy":
                return "legacy"
            parts = [part.strip() for part in text.split(",")]
            try:
                if len(parts) == 1:
                    return float(parts[0])
                return [float(part) for part in parts]
            except ValueError as exc:
                raise ValueError(f"{name} must be 'legacy', a float, or comma-separated floats") from exc
        if isinstance(value, (int, float)):
            return float(value)
        try:
            return [float(part) for part in value]
        except TypeError as exc:
            raise ValueError(f"{name} must be 'legacy', a float, or comma-separated floats") from exc

    def _block_power_scope_length(self, scope):
        if scope == "shared":
            return 1
        if scope == "per_residual":
            return 2 * self.n_layer
        return int(self.attnres_num_blocks)

    def _validate_block_power_value_length(self, value, scope, name):
        if value == "legacy" or isinstance(value, float):
            return
        expected = self._block_power_scope_length(scope)
        if len(value) != expected:
            raise ValueError(f"{name} list length must be {expected} for {scope} scope")


class LRIDStaticKey(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.key = nn.Parameter(torch.empty(config.lrid_rank))

    def reset_parameters(self, std=0.02, init_cutoff_factor=None):
        if init_cutoff_factor is not None:
            cutoff = init_cutoff_factor * std
            nn.init.trunc_normal_(self.key, mean=0.0, std=std, a=-cutoff, b=cutoff)
        else:
            nn.init.normal_(self.key, mean=0.0, std=std)

    def forward(self, reference):
        return self.key.to(reference.dtype).view(1, 1, -1).expand(reference.size(0), reference.size(1), -1)


class RotaryEmbedding(nn.Module):
    def __init__(self, config):
        super().__init__()
        dim = config.n_embd // config.n_head
        max_seq_len = config.block_size
        base = config.rope_theta
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2) / dim))
        freq = torch.outer(torch.arange(max_seq_len), inv_freq)
        self.register_buffer("sin", freq.sin()[None, None])
        self.register_buffer("cos", freq.cos()[None, None])

    def _forward_single(self, x, offset=0):
        T = x.size(-2)
        sin = self.sin[:, :, offset:offset + T]
        cos = self.cos[:, :, offset:offset + T]
        x1, x2 = x[..., 0::2], x[..., 1::2]
        return torch.stack([cos * x1 - sin * x2, sin * x1 + cos * x2], dim=-1).flatten(-2)

    def forward(self, q, k=None, offset=0):
        if k is None:
            return self._forward_single(q, offset=offset)

        return self._forward_single(q, offset=offset), self._forward_single(k, offset=offset)


class LRIDFusedProjection(nn.Module):
    def __init__(self, config, input_dim, output_dim):
        super().__init__()
        self.output_dim = output_dim
        self.rank = config.lrid_rank
        self.projection_rank = config.lrid_projection_rank
        self.use_query = config.lrid_input_dependent_query
        self.use_output_tail_key = config.lrid_key_from_output_tail
        self.use_value_key = config.lrid_key_from_value
        self.use_shared_value_key = config.lrid_key_from_value_shared
        self.use_local_value_key = self.use_value_key and not self.use_shared_value_key
        self.use_key = not (self.use_value_key or self.use_output_tail_key)
        self.use_value_query = self.use_query and config.lrid_query_from_value
        self.use_shared_value_query = self.use_query and config.lrid_query_from_value_shared
        self.use_local_value_query = self.use_value_query and not self.use_shared_value_query
        self.use_fused_query = self.use_query and not self.use_value_query
        self.key_offset = output_dim
        self.query_offset = output_dim + (self.projection_rank if self.use_key else 0)
        extra_dim = (self.projection_rank if self.use_key else 0) + (self.rank if self.use_fused_query else 0)
        self.proj = nn.Linear(input_dim, output_dim + extra_dim, bias=False)
        self.use_key_norm = self.use_key and config.attnres_key_norm
        self.use_value_norm = (self.use_local_value_key or self.use_local_value_query) and config.lrid_key_value_norm
        self.num_heads = config.lrid_num_heads
        if self.use_local_value_key:
            self.value_key_proj = nn.Linear(output_dim, config.lrid_rank, bias=False)
        else:
            self.value_key_proj = None
        if self.use_local_value_query:
            self.value_query_proj = nn.Linear(output_dim, config.lrid_rank, bias=False)
        else:
            self.value_query_proj = None

    def _prepare_value_projection_input(self, value):
        if self.use_value_norm:
            return norm(value)
        return value

    def project_key_from_value(self, value):
        if not self.use_local_value_key:
            raise RuntimeError("Local value-key projection is only available in unshared lrid_key_from_value mode")
        return self.value_key_proj(self._prepare_value_projection_input(value))

    def project_query_from_value(self, value):
        if not self.use_local_value_query:
            raise RuntimeError("Local value-query projection is only available in unshared lrid_query_from_value mode")
        return self.value_query_proj(self._prepare_value_projection_input(value))

    def forward(self, x, emit_lrid_key=True):
        if not emit_lrid_key:
            output = F.linear(x, self.proj.weight[:self.output_dim], None)
            if self.use_query:
                return output, None, None
            return output, None

        projected = self.proj(x)
        output = projected[..., :self.output_dim].contiguous()
        key = None
        query = None
        if self.use_key:
            key = projected[..., self.key_offset:self.key_offset + self.rank]
            if self.use_key_norm:
                key = _norm_lrid_key(key, self.num_heads)
        elif self.use_output_tail_key:
            key = output[..., -self.rank:].contiguous()
        elif self.use_local_value_key:
            key = self.project_key_from_value(output)
        if self.use_fused_query:
            query = projected[..., self.query_offset:self.query_offset + self.rank]
        elif self.use_local_value_query:
            query = self.project_query_from_value(output)
        if self.use_query:
            return output, key, query
        return output, key


class LRIDSourceKeyProjection(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.rank = config.lrid_rank
        self.num_heads = config.lrid_num_heads
        self.use_output_tail_key = config.lrid_key_from_output_tail
        uses_value_projection = (
            config.lrid_key_from_value
            or (config.lrid_input_dependent_query and config.lrid_query_from_value_shared)
        )
        self.use_value_norm = uses_value_projection and config.lrid_key_value_norm
        self.proj = None if self.use_output_tail_key else nn.Linear(config.n_embd, config.lrid_rank, bias=False)
        if config.lrid_input_dependent_query and config.lrid_query_from_value_shared:
            self.query_proj = nn.Linear(config.n_embd, config.lrid_rank, bias=False)
        else:
            self.query_proj = None
        self.use_key_norm = config.attnres_key_norm and not config.lrid_key_from_value

    def _prepare_value_projection_input(self, x):
        if self.use_value_norm:
            return norm(x)
        return x

    def forward(self, x):
        if self.use_output_tail_key:
            return x[..., -self.rank:].contiguous()
        if self.config.lrid_key_from_value:
            x = self._prepare_value_projection_input(x)
        key = self.proj(x)
        if self.use_key_norm:
            key = _norm_lrid_key(key, self.num_heads)
        return key

    def project_query_from_value(self, x):
        if self.query_proj is None:
            raise RuntimeError("Shared value-query projection is only available when lrid_query_from_value_shared=True")
        return self.query_proj(self._prepare_value_projection_input(x))

class MultiHeadAttention(nn.Module):
    flash_attn_func = None
    flash_attn_varlen_func = None
    flash_tried = False
    
    def __init__(self, config, layer_idx=0):
        super().__init__()
        assert config.n_embd % config.n_head == 0
        self.n_head = config.n_head
        self.n_embd = config.n_embd
        self.head_dim = config.n_embd // config.n_head
        self.rope = RotaryEmbedding(config)
        self.layer_idx = layer_idx
        self.config = config
        
        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        if config.use_lrid:
            self.c_proj = LRIDFusedProjection(config, config.n_embd, config.n_embd)
        else:
            self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        
        self.use_qk_norm = config.qk_norm
        
        self.clip_qkv = config.clip_qkv
        
        if config.flash_attention and not MultiHeadAttention.flash_tried:
            try:
                from flash_attn import flash_attn_func, flash_attn_varlen_func
                MultiHeadAttention.flash_attn_func = flash_attn_func
                MultiHeadAttention.flash_attn_varlen_func = flash_attn_varlen_func
                MultiHeadAttention.flash_tried = True
            except Exception as e:
                print(f"Error with flash-attn {e}.")
                MultiHeadAttention.flash_tried = True

    def _scaled_dot_product_attention(self, q, k, v, attn_mask=None, is_causal=True, 
                                       cu_doc_len=None, max_doc_len=None):
        B, H, T, D = q.size()
        
        if cu_doc_len is not None and max_doc_len is not None and MultiHeadAttention.flash_attn_varlen_func is not None:
            q_flat = q.transpose(1, 2).reshape(B * T, H, D)
            k_flat = k.transpose(1, 2).reshape(B * T, H, D)
            v_flat = v.transpose(1, 2).reshape(B * T, H, D)

            cu_doc_len = cu_doc_len.to(device=q.device, dtype=torch.int32)
            x = MultiHeadAttention.flash_attn_varlen_func(
                q_flat, k_flat, v_flat,
                cu_seqlens_q=cu_doc_len,
                cu_seqlens_k=cu_doc_len,
                max_seqlen_q=max_doc_len,
                max_seqlen_k=max_doc_len,
                causal=is_causal,
            )
            return x.view(B, T, H, D).contiguous().view(B, T, self.n_embd)

        elif cu_doc_len is not None or max_doc_len is not None:
            raise RuntimeError(
                "Document masking requires flash-attn varlen support. "
                "Install flash-attn or disable use_doc_masking."
            )
        
        elif MultiHeadAttention.flash_attn_func is not None and attn_mask is None:
            x = MultiHeadAttention.flash_attn_func(
                q.transpose(1, 2),
                k.transpose(1, 2),
                v.transpose(1, 2),
                causal=is_causal,
            )
            return x.contiguous().view(B, T, self.n_embd)
        
        else:
            x = F.scaled_dot_product_attention(
                q, k, v,
                attn_mask=attn_mask,
                is_causal=is_causal,
            )
            return x.transpose(1, 2).contiguous().view(B, T, self.n_embd)

    def forward(self, x, past_kv=None, use_cache=False, cu_doc_len=None, max_doc_len=None, emit_lrid_key=True):
        B, T, C = x.size()
        
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        
        if self.clip_qkv is not None:
            q.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
            k.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
            v.clamp_(min=-self.clip_qkv, max=self.clip_qkv)
        
        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
        
        if self.use_qk_norm:
            q = norm(q)
            k = norm(k)
        
        if past_kv is not None:
            past_k, past_v = past_kv
            pos_offset = past_k.size(-2)
        else:
            pos_offset = 0
        
        q, k = self.rope(q, k, offset=pos_offset)
        
        if past_kv is not None:
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        is_causal = past_kv is None

        attention_output = self._scaled_dot_product_attention(
            q, k, v,
            is_causal=is_causal,
            cu_doc_len=cu_doc_len,
            max_doc_len=max_doc_len,
        )
        
        if self.config.use_lrid:
            projected = self.c_proj(attention_output, emit_lrid_key=emit_lrid_key)
        else:
            projected = self.c_proj(attention_output)
        if self.config.use_lrid:
            if self.config.lrid_input_dependent_query:
                x, lrid_key, lrid_query = projected
            else:
                x, lrid_key = projected
        else:
            x = projected
        
        if use_cache:
            if self.config.use_lrid:
                if self.config.lrid_input_dependent_query:
                    return x, (k, v), lrid_key, lrid_query
                return x, (k, v), lrid_key
            return x, (k, v)
        if self.config.use_lrid:
            if self.config.lrid_input_dependent_query:
                return x, lrid_key, lrid_query
            return x, lrid_key
        return x


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_dim = config.mlp_hidden_dim if config.mlp_hidden_dim is not None else int(config.n_embd * config.mlp_ratio)
        self.fc1 = nn.Linear(config.n_embd, self.hidden_dim * 2, bias=False)
        if config.use_lrid:
            self.fc2 = LRIDFusedProjection(config, self.hidden_dim, config.n_embd)
        else:
            self.fc2 = nn.Linear(self.hidden_dim, config.n_embd, bias=False)

    def forward(self, x, emit_lrid_key=True):
        x = self.fc1(x)
        x, gate = x.chunk(2, dim=-1)
        x = F.silu(gate) * x
        if self.config.use_lrid:
            return self.fc2(x, emit_lrid_key=emit_lrid_key)
        return self.fc2(x)


class AttentionResidual(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.use_key_norm = config.attnres_key_norm
        self.query = nn.Parameter(torch.empty(config.n_embd))

    def _query(self, dtype):
        query = self.query
        if self.config.attn_res_query_norm:
            query = norm(query.float())
        return query.to(dtype)

    def forward(self, values, source_counts=None, source_logit_biases=None):
        keys = norm(values) if self.use_key_norm else values
        logits = torch.einsum("d,sbtd->sbt", self._query(keys.dtype), keys)
        if source_counts is not None and source_logit_biases is not None:
            raise RuntimeError("source_counts and source_logit_biases are mutually exclusive")
        if source_counts is not None:
            log_counts = torch.as_tensor(source_counts, device=logits.device, dtype=torch.float32).log()
            logits = logits + log_counts.view(-1, 1, 1)
        if source_logit_biases is not None:
            bias_values = [
                bias.to(device=logits.device, dtype=torch.float32).reshape(())
                if torch.is_tensor(bias)
                else torch.tensor(float(bias), device=logits.device, dtype=torch.float32)
                for bias in source_logit_biases
            ]
            logit_bias = torch.stack(bias_values)
            logits = logits + logit_bias.view(-1, 1, 1)
        weights = F.softmax(logits.float(), dim=0).to(values.dtype)
        return torch.einsum("sbt,sbtd->btd", weights, values)


class Block(nn.Module):
    def __init__(self, config, layer_idx=0):
        super().__init__()
        self.norm_pos = config.norm_pos
        self.attn = MultiHeadAttention(config, layer_idx=layer_idx)
        self.mlp = MLP(config)
        self.layer_idx = layer_idx
        self.config = config

    def forward_attention(self, x, past_kv=None, use_cache=False, cu_doc_len=None, max_doc_len=None, x_is_normalized=False, emit_lrid_key=True):
        if self.norm_pos in {"before", "both"} and not x_is_normalized:
            x = norm(x)

        attn_out = self.attn(
            x,
            past_kv=past_kv,
            use_cache=use_cache,
            cu_doc_len=cu_doc_len,
            max_doc_len=max_doc_len,
            emit_lrid_key=emit_lrid_key,
        )

        if self.config.use_lrid:
            if self.config.lrid_input_dependent_query:
                if use_cache:
                    x, new_kv, lrid_key, lrid_query = attn_out
                else:
                    x, lrid_key, lrid_query = attn_out
                    new_kv = None
            else:
                if use_cache:
                    x, new_kv, lrid_key = attn_out
                else:
                    x, lrid_key = attn_out
                    new_kv = None
        elif use_cache:
            x, new_kv = attn_out
        else:
            x = attn_out
            new_kv = None

        if self.norm_pos in {"after", "both"}:
            x = norm(x)

        if use_cache:
            if self.config.use_lrid:
                if self.config.lrid_input_dependent_query:
                    return x, new_kv, lrid_key, lrid_query
                return x, new_kv, lrid_key
            return x, new_kv
        if self.config.use_lrid:
            if self.config.lrid_input_dependent_query:
                return x, lrid_key, lrid_query
            return x, lrid_key
        return x

    def forward_mlp(self, x, x_is_normalized=False, emit_lrid_key=True):
        if self.norm_pos in {"before", "both"} and not x_is_normalized:
            x = norm(x)

        mlp_out = self.mlp(x, emit_lrid_key=emit_lrid_key)
        if self.config.use_lrid:
            if self.config.lrid_input_dependent_query:
                x, lrid_key, lrid_query = mlp_out
            else:
                x, lrid_key = mlp_out
        else:
            x = mlp_out

        if self.norm_pos in {"after", "both"}:
            x = norm(x)

        if self.config.use_lrid:
            if self.config.lrid_input_dependent_query:
                return x, lrid_key, lrid_query
            return x, lrid_key
        return x

    def forward(self, x, past_kv=None, use_cache=False, cu_doc_len=None, max_doc_len=None):
        residual = x

        attn_out = self.forward_attention(
            x,
            past_kv=past_kv,
            use_cache=use_cache,
            cu_doc_len=cu_doc_len,
            max_doc_len=max_doc_len,
        )

        if use_cache:
            x, new_kv = attn_out
        else:
            x = attn_out
            new_kv = None

        x = residual + x

        residual = x

        x = self.forward_mlp(x)

        x = residual + x

        if use_cache:
            return x, new_kv
        else:
            return x


class OBPM(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.use_attnres = config.use_attnres
        self.use_fused_attnres = config.use_fused_attnres
        self.attnres_type = config.attnres_type
        self.use_lrid = config.use_lrid
        self.attnres_block_ends = self._make_attnres_block_ends()
        
        torch.backends.cuda.enable_flash_sdp(True)
        torch.backends.cuda.enable_mem_efficient_sdp(False)
        
        transformer_modules = dict(
            wte=TokenEmbedding(config.vocab_size, config.n_embd),
            layers=nn.ModuleList([Block(config, layer_idx=i) for i in range(config.n_layer)])
        )

        if self.use_attnres:
            if self.attnres_type not in {"full", "block"}:
                raise ValueError("attnres_type must be 'full' or 'block'")
            if self.attnres_type == "block" and config.attnres_num_blocks < 1:
                raise ValueError("attnres_num_blocks must be >= 1 when using block AttnRes")
            if self.use_lrid:
                needs_lrid_source_projection = (
                    (not config.lrid_static_embedding_key and not config.lrid_key_from_output_tail)
                    or config.lrid_key_from_value_shared
                    or (config.lrid_input_dependent_query and config.lrid_query_from_value_shared)
                )
                if needs_lrid_source_projection:
                    transformer_modules["lrid_embedding_key"] = LRIDSourceKeyProjection(config)
                if config.lrid_static_embedding_key or config.lrid_add_static_embedding_key:
                    transformer_modules["lrid_static_embedding_key"] = LRIDStaticKey(config)
                if config.lrid_add_static_source_key:
                    transformer_modules["lrid_static_source_key"] = LRIDStaticKey(config)
                key_head_dim = config.lrid_rank // config.lrid_num_heads
                transformer_modules["lrid_queries"] = nn.ParameterList(
                    [
                        nn.Parameter(torch.empty(config.lrid_num_heads, key_head_dim))
                        for _ in range(2 * config.n_layer)
                    ]
                )
                if config.lrid_input_dependent_query:
                    transformer_modules["lrid_query_gates"] = nn.ParameterList(
                        [
                            nn.Parameter(torch.zeros(config.lrid_num_heads))
                            for _ in range(2 * config.n_layer)
                        ]
                    )
            else:
                transformer_modules["attn_residuals"] = nn.ModuleList(
                    [AttentionResidual(config) for _ in range(2 * config.n_layer)]
                )

        self.transformer = nn.ModuleDict(transformer_modules)
        if self.use_attnres and self.attnres_type == "block" and config.attnres_block_learned_scale:
            self.transformer.register_parameter(
                "attnres_block_scales",
                nn.Parameter(self._make_attnres_block_scale_init()),
            )
        if self.use_attnres and self.attnres_type == "block" and config.attnres_block_alpha_learned:
            self.transformer.register_parameter(
                "attnres_block_alphas",
                nn.Parameter(self._make_attnres_block_power_init("alpha")),
            )
        if self.use_attnres and self.attnres_type == "block" and config.attnres_block_beta_learned:
            self.transformer.register_parameter(
                "attnres_block_betas",
                nn.Parameter(self._make_attnres_block_power_init("beta")),
            )
        
        if not config.weight_tying:
            self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        
        self.apply(partial(self._init_weights, std=config.init_std, init_cutoff_factor=config.init_cutoff_factor))
        if self.use_lrid:
            for query in self.transformer.lrid_queries:
                self._init_attnres_query(query, config.init_std, config.init_cutoff_factor)
        if self.use_lrid and config.lrid_input_dependent_query:
            for module in self.modules():
                if isinstance(module, LRIDFusedProjection):
                    self._init_lrid_dynamic_query_projection(module, config.init_std, config.init_cutoff_factor)
        self._fast_attnres_op = None
        # Fast-AttnRes v2 is CUDA/BF16-only. Construction happens on CPU, so
        # the production route is armed only after the caller moves the full
        # model to its final CUDA/BF16 runtime and calls require_fast_attnres().
        self._fast_attnres_enabled = False
    def to_mixed_precision(self, dtype=torch.bfloat16):
        # These learned block exponents intentionally remain fp32. Preserve
        # their exact values instead of casting fp32 -> bf16 -> fp32 along with
        # the rest of the module, which would silently perturb checkpoints on
        # resume/evaluation.
        preserved_fp32 = {}
        for name in ("attnres_block_alphas", "attnres_block_betas"):
            if hasattr(self.transformer, name):
                preserved_fp32[name] = getattr(self.transformer, name).detach().clone()
        self.to(dtype=dtype)
        for name, value in preserved_fp32.items():
            param = getattr(self.transformer, name)
            param.data = value.to(device=param.device, dtype=torch.float32)
        return self
    
    def get_num_params(self):
        return sum(p.numel() for p in self.parameters())

    def fast_attnres_startup_report(self, validate_package: bool = False):
        """Describe the statically selected Fast route without mutable counters."""
        decision = (
            self._fast_attnres_config_decision()
            if self.config.attnres_backend == "fast"
            else FastAttnResDecision("legacy", "backend_legacy", "attnres_backend is legacy")
        )
        total_reads = 2 * self.config.n_layer if self.use_attnres else 0
        parameter = next(self.parameters(), None)
        operator_width = self.config.n_embd + (
            self.config.lrid_rank if self._fast_lrid_requires_key_payload() else 0
        )
        max_sources = (
            2 * self.config.n_layer + 1
            if self.attnres_type == "full"
            else (
                min(2 * self.config.n_layer, self.config.attnres_num_blocks)
                + (2 if self.config.attnres_block_split_sublayers else 1)
            )
        )

        if decision.eligible and (
            parameter is None
            or parameter.dtype not in FAST_ATTNRES_DTYPES
            or parameter.device.type != "cuda"
            or operator_width > FAST_ATTNRES_MAX_WIDTH
            or max_sources > FAST_ATTNRES_MAX_SOURCES
        ):
            decision = FastAttnResDecision(
                "legacy",
                "unsupported_runtime_envelope",
                "model must be CUDA BF16 and inside the Fast-AttnRes v2.0.1 shape envelope",
            )

        if decision.eligible and self.attnres_type == "block":
            if self._use_attnres_block_count_prior():
                decision = FastAttnResDecision(
                    "legacy",
                    "source_count_prior",
                    "Fast-AttnRes v2.0.1 has no source-logit-prior API",
                )

        active_reads = total_reads if decision.eligible else 0
        package = {"version": None, "source_hashes": {}}
        if validate_package and active_reads:
            self._fast_attnres_op = load_fast_attnres()
            package = fast_attnres_package_provenance()
        reasons = (
            {}
            if active_reads
            else ({decision.reason or "unspecified": total_reads} if total_reads else {})
        )
        return {
            **decision.as_dict(),
            "backend": self.config.attnres_backend,
            "requested_backend": getattr(
                self.config, "_attnres_backend_requested", self.config.attnres_backend
            ),
            "resolved_backend": "fast-attnres" if active_reads else "legacy",
            "active_reads": active_reads,
            "total_reads": total_reads,
            "legacy_fallback_reads": total_reads - active_reads,
            "fast_reads": active_reads,
            "legacy_reads": total_reads - active_reads,
            "fallback_reasons": reasons,
            "value_width": self.config.n_embd,
            "operator_width": operator_width,
            "uses_exact_key_payload": self._fast_lrid_requires_key_payload(),
            **package,
        }

    def require_fast_attnres(self, validate_package: bool = True):
        """Arm Fast only if every multi-source routed read is guaranteed Fast."""
        report = self.fast_attnres_startup_report(validate_package=validate_package)
        total_reads = int(report["total_reads"])
        failures = []
        if self.config.attnres_backend != "fast":
            failures.append("backend did not resolve to fast")
        if total_reads <= 0:
            failures.append("model has no routed reads")
        if int(report["active_reads"]) != total_reads:
            failures.append("not every routed read is active")
        if int(report["fast_reads"]) != total_reads:
            failures.append("not every routed read resolves to Fast-AttnRes")
        if int(report["legacy_reads"]) or int(report["legacy_fallback_reads"]):
            failures.append("a legacy fallback remains")
        if report["resolved_backend"] != "fast-attnres":
            failures.append("resolved backend is not Fast-AttnRes")
        if validate_package:
            if report.get("version") != FAST_ATTNRES_VERSION:
                failures.append(f"Fast-AttnRes {FAST_ATTNRES_VERSION} provenance is missing")
            if report.get("release_commit") != FAST_ATTNRES_RELEASE_COMMIT:
                failures.append("Fast-AttnRes release commit does not match the qualified release")
            if report.get("wheel_sha256") != FAST_ATTNRES_WHEEL_SHA256:
                failures.append("Fast-AttnRes wheel hash does not match the qualified release")
            if report.get("distribution_source_sha256") != FAST_ATTNRES_SOURCE_SHA256:
                failures.append("installed Fast-AttnRes source hash is not the qualified release")
            if not report.get("source_hashes"):
                failures.append("installed Fast-AttnRes source provenance is incomplete")
        if failures:
            self._fast_attnres_enabled = False
            reason = report.get("reason") or "unknown"
            detail = report.get("detail") or ""
            raise RuntimeError(
                "Fast-AttnRes-only contract failed: "
                + "; ".join(failures)
                + f" (reason={reason}: {detail})"
            )
        self._fast_attnres_enabled = True
        return report

    def _assert_fast_runtime_read(
        self,
        sources,
        *,
        source_counts=None,
        source_logit_biases=None,
    ):
        if not self._fast_attnres_enabled:
            raise RuntimeError(
                "Fast-AttnRes model is not armed; call require_fast_attnres() "
                "after moving it to CUDA BF16"
            )
        if source_counts is not None or source_logit_biases is not None:
            raise RuntimeError("Fast-AttnRes cannot use source priors or logit biases")
        first = sources[0]
        if first.device.type != "cuda" or first.dtype not in FAST_ATTNRES_DTYPES:
            raise RuntimeError("Fast-AttnRes residual reads require CUDA BF16 sources")

    def fast_attnres_route_report(self):
        report = self.fast_attnres_startup_report(validate_package=False)
        report["will_execute"] = bool(report["active_reads"])
        return report

    get_fast_attnres_route_report = fast_attnres_route_report

    def _ensure_fast_attnres_op(self):
        if self._fast_attnres_op is None:
            self._fast_attnres_op = load_fast_attnres()
        return self._fast_attnres_op

    @staticmethod
    def _fast_attnres_eps(dtype: torch.dtype, *, lrid: bool = False) -> float:
        # Public-main normalizes both standard and LRID keys through
        # ``F.rms_norm(..., eps=None)``. In the qualified Torch runtime that
        # uses the FP32 accumulation epsilon, including for BF16 inputs.
        return float(torch.finfo(torch.float32).eps)

    def _fast_attnres_config_decision(self):
        config = self.config
        if self.use_lrid and config.lrid_logit_scale != 1.0:
            return FastAttnResDecision(
                "legacy",
                "non_neutral_lrid_logit_scale",
                "the qualified compiled Fast path requires --no-lrid_logit_scale",
            )
        return fast_attnres_config_decision(
            use_attnres=self.use_attnres,
            use_lrid=self.use_lrid,
            attnres_type=self.attnres_type,
            key_norm=config.attnres_key_norm,
            lrid_rank=config.lrid_rank,
            n_embd=config.n_embd,
            lrid_num_heads=config.lrid_num_heads,
            lrid_input_dependent_query=config.lrid_input_dependent_query,
            lrid_key_from_output_tail=config.lrid_key_from_output_tail,
            lrid_key_from_value=config.lrid_key_from_value,
            lrid_key_from_value_shared=config.lrid_key_from_value_shared,
            lrid_query_from_value=config.lrid_query_from_value,
            lrid_query_from_value_shared=config.lrid_query_from_value_shared,
            lrid_static_embedding_key=config.lrid_static_embedding_key,
            lrid_add_static_embedding_key=config.lrid_add_static_embedding_key,
            lrid_add_static_source_key=config.lrid_add_static_source_key,
            n_layer=config.n_layer,
        )

    def _fast_lrid_requires_key_payload(self):
        if not self.use_lrid:
            return False
        if self.config.norm_pos != "before":
            return True
        if self.attnres_type != "block":
            return False
        return (
            self.config.attnres_block_value_norm
            or self.config.attnres_block_learned_scale
            or self.config.attnres_block_alpha_learned
            or any(value != 0.0 for value in self._attnres_block_power_value_list("alpha"))
        )

    def _make_attnres_block_ends(self):
        if not self.use_attnres or self.attnres_type != "block":
            return None
        total_sublayers = 2 * self.config.n_layer
        num_blocks = min(self.config.attnres_num_blocks, total_sublayers)
        return frozenset(
            math.ceil(total_sublayers * i / num_blocks)
            for i in range(1, num_blocks + 1)
        )

    def _attnres_block_partial_count(self, summary_idx):
        if summary_idx < 1 or summary_idx > 2 * self.config.n_layer:
            raise RuntimeError("Block summary index is out of range")
        previous_end = 0
        for block_end in sorted(self.attnres_block_ends):
            if summary_idx <= block_end:
                return summary_idx - previous_end
            previous_end = block_end
        raise RuntimeError("Block summary index does not belong to any block")

    def _attnres_block_bucket_idx(self, summary_idx):
        if summary_idx is None:
            raise RuntimeError("Block bucket scope requires a block summary index")
        if summary_idx < 1 or summary_idx > 2 * self.config.n_layer:
            raise RuntimeError("Block summary index is out of range")
        for bucket_idx, block_end in enumerate(sorted(self.attnres_block_ends)):
            if summary_idx <= block_end:
                return bucket_idx
        raise RuntimeError("Block summary index does not belong to any block")

    def _attnres_block_scope_index(self, scope, summary_idx=None):
        if scope == "shared":
            return 0
        if summary_idx is None:
            raise RuntimeError(f"{scope} scope requires a block summary index")
        if scope == "per_residual":
            if summary_idx < 1 or summary_idx > 2 * self.config.n_layer:
                raise RuntimeError("Block summary index is out of range")
            return summary_idx - 1
        return self._attnres_block_bucket_idx(summary_idx)

    def _legacy_attnres_block_alpha(self):
        if not self.config.attnres_block_average:
            return 0.0
        if self.config.attnres_block_average_mode == "sqrt":
            return 0.5
        return 1.0

    def _legacy_attnres_block_beta(self):
        return 1.0 if self.config.attnres_block_count_prior else 0.0

    def _attnres_block_power_scope_length(self, scope):
        if scope == "shared":
            return 1
        if scope == "per_residual":
            return 2 * self.config.n_layer
        return int(self.config.attnres_num_blocks)

    def _attnres_block_power_value_list(self, kind):
        if kind == "alpha":
            value = self.config.attnres_block_alpha
            scope = self.config.attnres_block_alpha_scope
            legacy_value = self._legacy_attnres_block_alpha()
        elif kind == "beta":
            value = self.config.attnres_block_beta
            scope = self.config.attnres_block_beta_scope
            legacy_value = self._legacy_attnres_block_beta()
        else:
            raise RuntimeError("Unknown AttnRes block power kind")
        length = self._attnres_block_power_scope_length(scope)
        if value == "legacy":
            values = [legacy_value] * length
        elif isinstance(value, float):
            values = [value] * length
        else:
            values = value
        return [float(power) for power in values]

    def _attnres_block_power_values(self, kind):
        values = self._attnres_block_power_value_list(kind)
        return torch.tensor(values, dtype=torch.float32)

    def _make_attnres_block_power_init(self, kind):
        return self._attnres_block_power_values(kind)

    def attnres_block_power_values_for_logging(self, kind):
        if not self.use_attnres or self.attnres_type != "block":
            return None
        if kind == "alpha":
            learned = self.config.attnres_block_alpha_learned
            param_name = "attnres_block_alphas"
        elif kind == "beta":
            learned = self.config.attnres_block_beta_learned
            param_name = "attnres_block_betas"
        else:
            raise RuntimeError("Unknown AttnRes block power kind")
        if learned and hasattr(self.transformer, param_name):
            return getattr(self.transformer, param_name).detach()
        return self._attnres_block_power_values(kind)

    def _attnres_block_power(self, kind, summary_idx=None, dtype=None, device=None):
        if kind == "alpha":
            scope = self.config.attnres_block_alpha_scope
            param_name = "attnres_block_alphas"
            learned = self.config.attnres_block_alpha_learned
        elif kind == "beta":
            scope = self.config.attnres_block_beta_scope
            param_name = "attnres_block_betas"
            learned = self.config.attnres_block_beta_learned
        else:
            raise RuntimeError("Unknown AttnRes block power kind")
        idx = self._attnres_block_scope_index(scope, summary_idx)
        if learned:
            value = getattr(self.transformer, param_name)[idx]
            if device is not None:
                value = value.to(device=device)
            return value
        values = self._attnres_block_power_value_list(kind)
        return values[idx]

    def _attnres_block_alpha(self, summary_idx=None, dtype=None, device=None):
        return self._attnres_block_power("alpha", summary_idx=summary_idx, dtype=dtype, device=device)

    def _attnres_block_beta(self, read_idx=None, source_summary_idx=None, dtype=None, device=None):
        summary_idx = read_idx if self.config.attnres_block_beta_scope == "per_residual" else source_summary_idx
        return self._attnres_block_power("beta", summary_idx=summary_idx, dtype=dtype, device=device)

    def _attnres_block_count_scale(self, count, alpha):
        if torch.is_tensor(alpha):
            log_count = alpha.new_tensor(float(count)).log()
            return torch.exp(-alpha * log_count)
        return float(count) ** (-float(alpha))

    def _make_attnres_block_scale_init(self):
        values = []
        for summary_idx in range(1, 2 * self.config.n_layer + 1):
            count = self._attnres_block_partial_count(summary_idx)
            if self.config.attnres_block_learned_scale_init == "count":
                scale = 1.0 / count
            elif self.config.attnres_block_learned_scale_init == "sqrt":
                scale = 1.0 / math.sqrt(count)
            else:
                scale = 1.0
            values.append(scale)
        return torch.tensor(values, dtype=torch.float32)

    def _attnres_query_idx(self, residual_idx):
        if residual_idx < 1:
            raise RuntimeError("Residual site 0 has no query because it only reads the embedding source")
        return residual_idx - 1

    def _attnres_block_source_scale(self, count, summary_idx=None, dtype=None, device=None):
        if self.config.attnres_block_learned_scale:
            if summary_idx is None:
                raise RuntimeError("Learned block scaling requires a block summary index")
            scale = self.transformer.attnres_block_scales[summary_idx - 1]
            if dtype is not None or device is not None:
                scale = scale.to(dtype=dtype or scale.dtype, device=device or scale.device)
            return scale
        alpha = self._attnres_block_alpha(summary_idx=summary_idx, dtype=dtype, device=device)
        return self._attnres_block_count_scale(count, alpha)

    def _attnres_block_summary(self, value, count, summary_idx=None):
        if self.config.attnres_block_value_norm:
            return norm(value)
        scale = self._attnres_block_source_scale(
            count,
            summary_idx=summary_idx,
            dtype=value.dtype,
            device=value.device,
        )
        if scale is not None:
            return value * scale
        return value

    def _lrid_block_source(self, value, key=None, query=None, count=1, summary_idx=None):
        if self.config.attnres_block_value_norm:
            value = norm(value)
        else:
            scale = self._attnres_block_source_scale(
                count,
                summary_idx=summary_idx,
                dtype=value.dtype,
                device=value.device,
            )
            if scale is not None:
                value = value * scale
                if key is not None and not self.config.lrid_key_from_value and not self.config.attnres_key_norm:
                    if torch.is_tensor(scale):
                        scale = scale.to(dtype=key.dtype, device=key.device)
                    key = key * scale
        key_value = value if self.config.lrid_key_from_value_shared else None
        query_value = value if self.config.lrid_query_from_value_shared else None
        return self._lrid_source(
            value,
            key,
            query,
            key_value=key_value,
            query_value=query_value,
            add_static_key=True,
        )

    def _attnres_block_average_denominator(self, count):
        if self.config.attnres_block_average_mode == "sqrt":
            return math.sqrt(count)
        return count

    def _use_attnres_block_count_prior(self):
        if self.attnres_type != "block":
            return False
        if self.config.attnres_block_beta_learned:
            return True
        if self.config.attnres_block_beta == "legacy":
            return self.config.attnres_block_count_prior
        values = self._attnres_block_power_value_list("beta")
        return any(value != 0.0 for value in values)

    def _attnres_block_count_logit_bias(
        self,
        count,
        read_idx=None,
        source_summary_idx=None,
        dtype=None,
        device=None,
    ):
        if not self._use_attnres_block_count_prior():
            return 0.0
        if count == 1:
            return 0.0
        beta = self._attnres_block_beta(
            read_idx=read_idx,
            source_summary_idx=source_summary_idx,
            dtype=dtype,
            device=device,
        )
        log_count = math.log(float(count))
        if torch.is_tensor(beta):
            return beta * beta.new_tensor(log_count)
        beta = float(beta)
        if beta == 0.0:
            return 0.0
        return beta * log_count

    def _attnres_biases_to_tensor(self, biases, device):
        has_tensor = any(torch.is_tensor(bias) for bias in biases)
        has_nonzero = any((not torch.is_tensor(bias)) and float(bias) != 0.0 for bias in biases)
        if not has_tensor and not has_nonzero:
            return None
        return torch.stack(
            [
                bias.to(device=device, dtype=torch.float32).reshape(())
                if torch.is_tensor(bias)
                else torch.tensor(float(bias), device=device, dtype=torch.float32)
                for bias in biases
            ]
        )

    def _attnres_block_source_logit_biases(
        self,
        read_idx,
        completed_counts,
        completed_summary_idxs,
        partial_count=None,
        partial_summary_idx=None,
    ):
        if not self._use_attnres_block_count_prior():
            return None
        source_counts = completed_counts if partial_count is None else completed_counts + [partial_count]
        source_summary_idxs = (
            completed_summary_idxs
            if partial_count is None
            else completed_summary_idxs + [partial_summary_idx]
        )
        biases = [
            self._attnres_block_count_logit_bias(
                count,
                read_idx=read_idx,
                source_summary_idx=summary_idx,
            )
            for count, summary_idx in zip(source_counts, source_summary_idxs)
        ]
        has_tensor = any(torch.is_tensor(bias) for bias in biases)
        has_nonzero = any((not torch.is_tensor(bias)) and float(bias) != 0.0 for bias in biases)
        if not has_tensor and not has_nonzero:
            return None
        return biases

    def _attnres_block_logit_biases_for_sources(self, read_idx, source_counts, source_summary_idxs):
        if not self._use_attnres_block_count_prior():
            return None
        biases = [
            self._attnres_block_count_logit_bias(
                count,
                read_idx=read_idx,
                source_summary_idx=summary_idx,
            )
            for count, summary_idx in zip(source_counts, source_summary_idxs)
        ]
        has_tensor = any(torch.is_tensor(bias) for bias in biases)
        has_nonzero = any((not torch.is_tensor(bias)) and float(bias) != 0.0 for bias in biases)
        if not has_tensor and not has_nonzero:
            return None
        return biases

    def _attnres_split_partial_entries(self, partial_values, partial_counts, partial_summary_idxs):
        entries = []
        for value, count, summary_idx in zip(partial_values, partial_counts, partial_summary_idxs):
            if value is None:
                continue
            entries.append((value, count, summary_idx))
        return entries

    def _attnres_split_partial_sources(self, partial_values, partial_counts, partial_summary_idxs):
        return [
            self._attnres_block_summary(value, count, summary_idx)
            for value, count, summary_idx in self._attnres_split_partial_entries(
                partial_values,
                partial_counts,
                partial_summary_idxs,
            )
        ]

    def _attnres_split_source_logit_biases(
        self,
        read_idx,
        completed_counts,
        completed_summary_idxs,
        partial_values,
        partial_counts,
        partial_summary_idxs,
    ):
        partial_entries = self._attnres_split_partial_entries(
            partial_values,
            partial_counts,
            partial_summary_idxs,
        )
        source_counts = completed_counts + [count for _, count, _ in partial_entries]
        source_summary_idxs = completed_summary_idxs + [summary_idx for _, _, summary_idx in partial_entries]
        return self._attnres_block_logit_biases_for_sources(
            read_idx,
            source_counts,
            source_summary_idxs,
        )

    def _attnres_split_add_partial(self, partial_values, partial_counts, partial_summary_idxs, slot, value, summary_idx):
        if partial_values[slot] is None:
            partial_values[slot] = value
            partial_counts[slot] = 1
        else:
            partial_values[slot] = partial_values[slot] + value
            partial_counts[slot] += 1
        partial_summary_idxs[slot] = summary_idx

    def _attnres_split_flush_partials(
        self,
        completed_blocks,
        completed_counts,
        completed_summary_idxs,
        partial_values,
        partial_counts,
        partial_summary_idxs,
    ):
        for slot, (value, count, summary_idx) in enumerate(
            zip(partial_values, partial_counts, partial_summary_idxs)
        ):
            if value is None:
                continue
            completed_blocks.append(self._attnres_block_summary(value, count, summary_idx))
            completed_counts.append(count)
            completed_summary_idxs.append(summary_idx)
            partial_values[slot] = None
            partial_counts[slot] = 0
            partial_summary_idxs[slot] = None

    def _lrid_split_add_partial(
        self,
        partial_values,
        partial_keys,
        partial_queries,
        partial_counts,
        partial_summary_idxs,
        slot,
        value,
        lrid_key,
        lrid_query,
        summary_idx,
        key_projector=None,
        query_projector=None,
    ):
        if partial_values[slot] is None:
            partial_values[slot] = value
            partial_counts[slot] = 1
        else:
            partial_values[slot] = partial_values[slot] + value
            partial_counts[slot] += 1
        partial_summary_idxs[slot] = summary_idx

        block_source_value = self._attnres_block_summary(
            partial_values[slot],
            partial_counts[slot],
            summary_idx,
        )
        if self.config.lrid_key_from_value_shared:
            partial_keys[slot] = None
        elif self.config.lrid_key_from_value:
            partial_keys[slot] = key_projector(block_source_value)
        else:
            partial_keys[slot] = (
                lrid_key
                if partial_counts[slot] == 1
                else partial_keys[slot] + lrid_key
            )

        if self.config.lrid_input_dependent_query:
            if self.config.lrid_query_from_value_shared:
                partial_queries[slot] = None
            elif self.config.lrid_query_from_value:
                partial_queries[slot] = query_projector(block_source_value)
            else:
                partial_queries[slot] = lrid_query

    def _lrid_split_flush_partials(
        self,
        completed_blocks,
        completed_counts,
        completed_summary_idxs,
        partial_values,
        partial_keys,
        partial_queries,
        partial_counts,
        partial_summary_idxs,
    ):
        for slot, (value, key, query, count, summary_idx) in enumerate(
            zip(
                partial_values,
                partial_keys,
                partial_queries,
                partial_counts,
                partial_summary_idxs,
            )
        ):
            if value is None:
                continue
            completed_blocks.append(
                self._lrid_block_source(
                    value,
                    key,
                    query if self.config.lrid_input_dependent_query else None,
                    count,
                    summary_idx,
                )
            )
            completed_counts.append(count)
            completed_summary_idxs.append(summary_idx)
            partial_values[slot] = None
            partial_keys[slot] = None
            partial_queries[slot] = None
            partial_counts[slot] = 0
            partial_summary_idxs[slot] = None

    def _attnres_block_read_logit_biases(self, count, first_read, source_summary_idx, device):
        if not self._use_attnres_block_count_prior() or count == 1:
            return None
        biases = [
            self._attnres_block_count_logit_bias(
                count,
                read_idx=read_idx,
                source_summary_idx=source_summary_idx,
                device=device,
            )
            for read_idx in range(first_read, 2 * self.config.n_layer + 1)
        ]
        return self._attnres_biases_to_tensor(biases, device=device)

    def _apply_attnres(
        self,
        residual_idx,
        sources,
        normalize_output=False,
        source_counts=None,
        source_logit_biases=None,
    ):
        if len(sources) == 1:
            return norm(sources[0]) if normalize_output else sources[0]
        residual = self.transformer.attn_residuals[self._attnres_query_idx(residual_idx)]
        if self.config.attnres_backend == "fast":
            self._assert_fast_runtime_read(
                sources,
                source_counts=source_counts,
                source_logit_biases=source_logit_biases,
            )
            if not residual.use_key_norm:
                raise RuntimeError("Fast-AttnRes residual reads require key normalization")
            query = residual._query(sources[0].dtype)
            output = self._ensure_fast_attnres_op()(
                sources,
                query,
                eps=self._fast_attnres_eps(sources[0].dtype),
                scale=1.0,
            )
            return norm(output) if normalize_output else output
        if self.config.use_fused_attnres:
            query = residual._query(sources[0].dtype)
            return attention_residual_read(
                sources,
                query,
                residual.use_key_norm,
                normalize_output=normalize_output,
                source_counts=source_counts,
                source_logit_biases=source_logit_biases,
            )
        values = torch.stack(sources, dim=0)
        output = residual(values, source_counts=source_counts, source_logit_biases=source_logit_biases)
        return norm(output) if normalize_output else output

    def _attnres_query(self, residual_idx, dtype):
        residual = self.transformer.attn_residuals[self._attnres_query_idx(residual_idx)]
        return residual._query(dtype)

    def _apply_training_phase2(
        self,
        partial_source,
        query,
        interblock_output,
        interblock_lse,
        normalize_output=False,
        partial_count=1,
        read_idx=None,
        source_summary_idx=None,
    ):
        phase2 = attention_residual_phase2_torch if self.config.attnres_training_torch_phase2 else attention_residual_phase2
        return phase2(
            partial_source,
            query,
            interblock_output,
            interblock_lse,
            self.config.attnres_key_norm,
            normalize_output=normalize_output,
            logit_bias=self._attnres_block_count_logit_bias(
                partial_count,
                read_idx=read_idx,
                source_summary_idx=source_summary_idx,
            ),
        )

    def _forward_block_attnres_split_sublayers(
        self,
        x,
        past_kv=None,
        use_cache=False,
        cu_doc_len=None,
        max_doc_len=None,
        return_hidden=False,
        fused_read_norm=False,
    ):
        if past_kv is not None or use_cache:
            raise NotImplementedError("KV-cache generation is not supported with attention residuals yet.")
        if self.use_lrid:
            raise NotImplementedError("attnres_block_split_sublayers is currently implemented for non-LRID block AttnRes.")

        completed_blocks = [x]
        completed_block_counts = [1]
        completed_block_summary_idxs = [None]
        partial_values = [None, None]
        partial_counts = [0, 0]
        partial_summary_idxs = [None, None]
        block_end_set = self.attnres_block_ends

        def read_residual(residual_idx):
            sources = completed_blocks + self._attnres_split_partial_sources(
                partial_values,
                partial_counts,
                partial_summary_idxs,
            )
            source_logit_biases = self._attnres_split_source_logit_biases(
                residual_idx,
                completed_block_counts,
                completed_block_summary_idxs,
                partial_values,
                partial_counts,
                partial_summary_idxs,
            )
            return self._apply_attnres(
                residual_idx,
                sources,
                normalize_output=fused_read_norm,
                source_logit_biases=source_logit_biases,
            )

        def add_partial(slot, value, summary_idx):
            self._attnres_split_add_partial(
                partial_values,
                partial_counts,
                partial_summary_idxs,
                slot,
                value,
                summary_idx,
            )

        def flush_if_block_end(residual_idx):
            if residual_idx not in block_end_set:
                return
            self._attnres_split_flush_partials(
                completed_blocks,
                completed_block_counts,
                completed_block_summary_idxs,
                partial_values,
                partial_counts,
                partial_summary_idxs,
            )

        for layer_idx, block in enumerate(self.transformer.layers):
            x = read_residual(2 * layer_idx)
            attn_out = block.forward_attention(
                x,
                past_kv=None,
                use_cache=False,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                x_is_normalized=fused_read_norm,
            )
            add_partial(0, attn_out, 2 * layer_idx + 1)
            flush_if_block_end(2 * layer_idx + 1)

            x = read_residual(2 * layer_idx + 1)
            mlp_out = block.forward_mlp(x, x_is_normalized=fused_read_norm)
            add_partial(1, mlp_out, 2 * layer_idx + 2)
            flush_if_block_end(2 * layer_idx + 2)

        x = read_residual(2 * self.config.n_layer)
        if not fused_read_norm:
            x = norm(x)

        if return_hidden:
            return x

        if self.config.weight_tying:
            return F.linear(x, self.transformer.wte.weight, None)
        return self.lm_head(x)

    def _forward_block_lrid_attnres_split_sublayers(
        self,
        x,
        past_kv=None,
        use_cache=False,
        cu_doc_len=None,
        max_doc_len=None,
        return_hidden=False,
        fused_read_norm=False,
    ):
        if past_kv is not None or use_cache:
            raise NotImplementedError("KV-cache generation is not supported with attention residuals yet.")

        completed_blocks = [self._embedding_lrid_source(x)]
        completed_block_counts = [1]
        completed_block_summary_idxs = [None]
        partial_values = [None, None]
        partial_keys = [None, None]
        partial_queries = [None, None]
        partial_counts = [0, 0]
        partial_summary_idxs = [None, None]
        latest_partial_slot = None
        block_end_set = self.attnres_block_ends

        def read_residual(residual_idx):
            partial_sources = []
            query_override = None
            for slot, (value, key, query, count, summary_idx) in enumerate(
                zip(
                    partial_values,
                    partial_keys,
                    partial_queries,
                    partial_counts,
                    partial_summary_idxs,
                )
            ):
                if value is None:
                    continue
                source = self._lrid_block_source(
                    value,
                    key,
                    query if self.config.lrid_input_dependent_query else None,
                    count,
                    summary_idx,
                )
                partial_sources.append(source)
                if self.config.lrid_input_dependent_query and slot == latest_partial_slot:
                    query_override = source[2]
            sources = completed_blocks + partial_sources
            source_logit_biases = self._attnres_split_source_logit_biases(
                residual_idx,
                completed_block_counts,
                completed_block_summary_idxs,
                partial_values,
                partial_counts,
                partial_summary_idxs,
            )
            return self._apply_lrid_attnres(
                residual_idx,
                sources,
                query_override=query_override,
                normalize_output=fused_read_norm,
                source_logit_biases=source_logit_biases,
            )

        def add_partial(slot, value, lrid_key, lrid_query, summary_idx, key_projector, query_projector=None):
            nonlocal latest_partial_slot
            self._lrid_split_add_partial(
                partial_values,
                partial_keys,
                partial_queries,
                partial_counts,
                partial_summary_idxs,
                slot,
                value,
                lrid_key,
                lrid_query,
                summary_idx,
                key_projector=key_projector,
                query_projector=query_projector,
            )
            latest_partial_slot = slot

        def flush_if_block_end(residual_idx):
            nonlocal latest_partial_slot
            if residual_idx not in block_end_set:
                return
            self._lrid_split_flush_partials(
                completed_blocks,
                completed_block_counts,
                completed_block_summary_idxs,
                partial_values,
                partial_keys,
                partial_queries,
                partial_counts,
                partial_summary_idxs,
            )
            latest_partial_slot = None

        for layer_idx, block in enumerate(self.transformer.layers):
            x = read_residual(2 * layer_idx)
            attn_out = block.forward_attention(
                x,
                past_kv=None,
                use_cache=False,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                x_is_normalized=fused_read_norm,
                emit_lrid_key=True,
            )
            if self.config.lrid_input_dependent_query:
                attn_out, lrid_key, lrid_query = attn_out
            else:
                attn_out, lrid_key = attn_out
                lrid_query = None
            add_partial(
                0,
                attn_out,
                lrid_key,
                lrid_query,
                2 * layer_idx + 1,
                block.attn.c_proj.project_key_from_value,
                block.attn.c_proj.project_query_from_value,
            )
            flush_if_block_end(2 * layer_idx + 1)

            x = read_residual(2 * layer_idx + 1)
            mlp_out = block.forward_mlp(
                x,
                x_is_normalized=fused_read_norm,
                emit_lrid_key=True,
            )
            if self.config.lrid_input_dependent_query:
                mlp_out, lrid_key, lrid_query = mlp_out
            else:
                mlp_out, lrid_key = mlp_out
                lrid_query = None
            add_partial(
                1,
                mlp_out,
                lrid_key,
                lrid_query,
                2 * layer_idx + 2,
                block.mlp.fc2.project_key_from_value,
                block.mlp.fc2.project_query_from_value,
            )
            flush_if_block_end(2 * layer_idx + 2)

        x = read_residual(2 * self.config.n_layer)
        if not fused_read_norm:
            x = norm(x)

        if return_hidden:
            return x

        if self.config.weight_tying:
            return F.linear(x, self.transformer.wte.weight, None)
        return self.lm_head(x)

    def _apply_lrid_training_phase2(
        self,
        partial_value,
        partial_key,
        query,
        interblock_output,
        interblock_lse,
        normalize_output=False,
        partial_count=1,
        read_idx=None,
        source_summary_idx=None,
    ):
        phase2 = lrid_attention_residual_phase2_torch if self.config.attnres_training_torch_phase2 else lrid_attention_residual_phase2
        return phase2(
            partial_value,
            partial_key,
            query,
            interblock_output,
            interblock_lse,
            self.config.lrid_logit_scale,
            self.config.attnres_key_norm,
            normalize_output=normalize_output,
            logit_bias=self._attnres_block_count_logit_bias(
                partial_count,
                read_idx=read_idx,
                source_summary_idx=source_summary_idx,
            ),
        )

    def _forward_block_attnres_fused_training(
        self,
        x,
        past_kv=None,
        use_cache=False,
        cu_doc_len=None,
        max_doc_len=None,
        return_hidden=False,
        fused_read_norm=False,
    ):
        if past_kv is not None or use_cache:
            raise NotImplementedError("KV-cache generation is not supported with attention residuals yet.")

        total_sublayers = 2 * self.config.n_layer
        block_ends = sorted(self.attnres_block_ends)
        block_end_set = self.attnres_block_ends
        query_bank = torch.stack(
            [
                self._attnres_query(read_idx, x.dtype)
                for read_idx in range(1, total_sublayers + 1)
            ],
            dim=0,
        ).contiguous()

        def make_source_logits(source, first_read, source_count=1, source_summary_idx=None):
            key = norm(source) if self.config.attnres_key_norm else source
            queries = query_bank[first_read - 1:].to(key.dtype)
            logits = F.linear(key, queries).permute(2, 0, 1).contiguous()
            logit_biases = self._attnres_block_read_logit_biases(
                source_count,
                first_read,
                source_summary_idx,
                logits.device,
            )
            if logit_biases is not None:
                logits = logits + logit_biases.view(-1, 1, 1)
            return logits

        completed_blocks = [x]
        completed_logits = [make_source_logits(x, 1, 1)]
        completed_logit_first_reads = [1]
        partial_block = None
        partial_count = 0
        phase_first_read = None
        phase_end = None
        phase_outputs = None
        phase_lses = None

        def next_phase_end(residual_idx):
            for block_end in block_ends:
                if block_end > residual_idx:
                    return block_end
            return residual_idx + 1

        def invalidate_phase():
            nonlocal phase_first_read, phase_end, phase_outputs, phase_lses
            phase_first_read = None
            phase_end = None
            phase_outputs = None
            phase_lses = None

        def ensure_phase(residual_idx):
            nonlocal phase_first_read, phase_end, phase_outputs, phase_lses
            if residual_idx < 1:
                return
            if phase_outputs is not None and phase_first_read <= residual_idx < phase_end:
                return
            phase_first_read = residual_idx
            phase_end = next_phase_end(residual_idx)
            if len(completed_blocks) == 1:
                logit_offset = phase_first_read - completed_logit_first_reads[0]
                phase_lse = completed_logits[0][logit_offset:phase_end - completed_logit_first_reads[0]].float()
                phase_outputs = completed_blocks[0].unsqueeze(0).expand(phase_lse.size(0), -1, -1, -1).unbind(0)
                phase_lses = phase_lse.unbind(0)
                return
            phase_logits = torch.stack(
                [
                    source_logits[phase_first_read - first_read:phase_end - first_read]
                    for source_logits, first_read in zip(completed_logits, completed_logit_first_reads)
                ],
                dim=1,
            ).contiguous()
            phase_output, phase_lse = attention_residual_phase1_from_logits(
                completed_blocks,
                phase_logits,
            )
            phase_outputs = phase_output.unbind(0)
            phase_lses = phase_lse.unbind(0)

        def read_residual(residual_idx):
            if residual_idx == 0:
                output = completed_blocks[0]
            else:
                ensure_phase(residual_idx)
                phase_idx = residual_idx - phase_first_read
                interblock_output = phase_outputs[phase_idx]
                if partial_block is None:
                    output = interblock_output
                else:
                    partial_source = self._attnres_block_summary(partial_block, partial_count, residual_idx)
                    query = query_bank[self._attnres_query_idx(residual_idx)].to(partial_source.dtype)
                    output = self._apply_training_phase2(
                        partial_source,
                        query,
                        interblock_output,
                        phase_lses[phase_idx],
                        normalize_output=fused_read_norm,
                        partial_count=partial_count,
                        read_idx=residual_idx,
                        source_summary_idx=residual_idx,
                    )
            return norm(output) if fused_read_norm and (residual_idx == 0 or partial_block is None) else output

        def append_partial_if_block_end(residual_idx):
            nonlocal completed_blocks, completed_logit_first_reads, partial_block, partial_count
            if residual_idx not in block_end_set:
                return
            completed_source = self._attnres_block_summary(partial_block, partial_count, residual_idx)
            completed_blocks.append(completed_source)
            completed_logits.append(make_source_logits(completed_source, residual_idx, partial_count, residual_idx))
            completed_logit_first_reads.append(residual_idx)
            partial_block = None
            partial_count = 0
            invalidate_phase()

        def add_partial(layer_output):
            nonlocal partial_block, partial_count
            if partial_block is None:
                partial_block = layer_output
                partial_count = 1
            else:
                partial_block = partial_block + layer_output
                partial_count += 1

        for layer_idx, block in enumerate(self.transformer.layers):
            x = read_residual(2 * layer_idx)
            attn_out = block.forward_attention(
                x,
                past_kv=None,
                use_cache=False,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                x_is_normalized=fused_read_norm,
            )
            add_partial(attn_out)
            append_partial_if_block_end(2 * layer_idx + 1)

            x = read_residual(2 * layer_idx + 1)
            mlp_out = block.forward_mlp(x, x_is_normalized=fused_read_norm)
            add_partial(mlp_out)
            append_partial_if_block_end(2 * layer_idx + 2)

        x = read_residual(total_sublayers)
        if not fused_read_norm:
            x = norm(x)

        if return_hidden:
            return x

        if self.config.weight_tying:
            return F.linear(x, self.transformer.wte.weight, None)
        return self.lm_head(x)

    def _forward_block_lrid_attnres_fused_training(
        self,
        x,
        past_kv=None,
        use_cache=False,
        cu_doc_len=None,
        max_doc_len=None,
        return_hidden=False,
        fused_read_norm=False,
    ):
        if past_kv is not None or use_cache:
            raise NotImplementedError("KV-cache generation is not supported with attention residuals yet.")

        total_sublayers = 2 * self.config.n_layer
        block_ends = sorted(self.attnres_block_ends)
        block_end_set = self.attnres_block_ends
        key_dim = self.config.lrid_rank
        query_bank = torch.stack(
            [
                self.transformer.lrid_queries[self._attnres_query_idx(read_idx)].reshape(key_dim)
                for read_idx in range(1, total_sublayers + 1)
            ],
            dim=0,
        )
        if self.config.attn_res_query_norm:
            query_bank = norm(query_bank.float()).to(x.dtype)
        else:
            query_bank = query_bank.to(x.dtype)
        query_bank = query_bank.contiguous()

        def source_key(source):
            key = source[1].reshape(source[1].size(0), source[1].size(1), key_dim)
            if self.config.attnres_key_norm:
                key = norm(key.float()).to(source[0].dtype)
            return key

        def make_source_logits(source, first_read, source_count=1, source_summary_idx=None):
            key = source_key(source)
            queries = query_bank[first_read - 1:].to(key.dtype)
            logits = (
                F.linear(key, queries)
                .permute(2, 0, 1)
                .contiguous()
                * self.config.lrid_logit_scale
            )
            logit_biases = self._attnres_block_read_logit_biases(
                source_count,
                first_read,
                source_summary_idx,
                logits.device,
            )
            if logit_biases is not None:
                logits = logits + logit_biases.view(-1, 1, 1)
            return logits

        embedding_source = self._embedding_lrid_source(x)
        completed_values = [embedding_source[0]]
        completed_logits = [make_source_logits(embedding_source, 1, 1)]
        completed_logit_first_reads = [1]
        partial_block = None
        partial_key = None
        partial_count = 0
        phase_first_read = None
        phase_end = None
        phase_outputs = None
        phase_lses = None

        def next_phase_end(residual_idx):
            for block_end in block_ends:
                if block_end > residual_idx:
                    return block_end
            return residual_idx + 1

        def invalidate_phase():
            nonlocal phase_first_read, phase_end, phase_outputs, phase_lses
            phase_first_read = None
            phase_end = None
            phase_outputs = None
            phase_lses = None

        def ensure_phase(residual_idx):
            nonlocal phase_first_read, phase_end, phase_outputs, phase_lses
            if residual_idx < 1:
                return
            if phase_outputs is not None and phase_first_read <= residual_idx < phase_end:
                return
            phase_first_read = residual_idx
            phase_end = next_phase_end(residual_idx)
            if len(completed_values) == 1:
                logit_offset = phase_first_read - completed_logit_first_reads[0]
                phase_lse = completed_logits[0][logit_offset:phase_end - completed_logit_first_reads[0]].float()
                phase_outputs = completed_values[0].unsqueeze(0).expand(phase_lse.size(0), -1, -1, -1).unbind(0)
                phase_lses = phase_lse.unbind(0)
                return
            phase_logits = torch.stack(
                [
                    source_logits[phase_first_read - first_read:phase_end - first_read]
                    for source_logits, first_read in zip(completed_logits, completed_logit_first_reads)
                ],
                dim=1,
            ).contiguous()
            phase_output, phase_lse = attention_residual_phase1_from_logits(
                completed_values,
                phase_logits,
            )
            phase_outputs = phase_output.unbind(0)
            phase_lses = phase_lse.unbind(0)

        def current_partial_source(summary_idx):
            return self._lrid_block_source(partial_block, partial_key, None, partial_count, summary_idx)

        def read_residual(residual_idx):
            if residual_idx == 0:
                output = completed_values[0]
            else:
                ensure_phase(residual_idx)
                phase_idx = residual_idx - phase_first_read
                interblock_output = phase_outputs[phase_idx]
                if partial_block is None:
                    output = interblock_output
                else:
                    partial_source = current_partial_source(residual_idx)
                    query = query_bank[self._attnres_query_idx(residual_idx)].to(partial_source[1].dtype)
                    output = self._apply_lrid_training_phase2(
                        partial_source[0],
                        partial_source[1].reshape(partial_source[1].size(0), partial_source[1].size(1), key_dim),
                        query,
                        interblock_output,
                        phase_lses[phase_idx],
                        normalize_output=fused_read_norm,
                        partial_count=partial_count,
                        read_idx=residual_idx,
                        source_summary_idx=residual_idx,
                    )
            return norm(output) if fused_read_norm and (residual_idx == 0 or partial_block is None) else output

        def append_partial_if_block_end(residual_idx):
            nonlocal completed_logit_first_reads, partial_block, partial_key, partial_count
            if residual_idx not in block_end_set:
                return
            completed_source = current_partial_source(residual_idx)
            completed_values.append(completed_source[0])
            completed_logits.append(make_source_logits(completed_source, residual_idx, partial_count, residual_idx))
            completed_logit_first_reads.append(residual_idx)
            partial_block = None
            partial_key = None
            partial_count = 0
            invalidate_phase()

        def add_partial(layer_output, lrid_key):
            nonlocal partial_block, partial_key, partial_count
            if partial_block is None:
                partial_block = layer_output
                partial_key = lrid_key
                partial_count = 1
            else:
                partial_block = partial_block + layer_output
                partial_key = partial_key + lrid_key
                partial_count += 1

        for layer_idx, block in enumerate(self.transformer.layers):
            x = read_residual(2 * layer_idx)
            attn_out, lrid_key = block.forward_attention(
                x,
                past_kv=None,
                use_cache=False,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                x_is_normalized=fused_read_norm,
                emit_lrid_key=True,
            )
            add_partial(attn_out, lrid_key)
            append_partial_if_block_end(2 * layer_idx + 1)

            x = read_residual(2 * layer_idx + 1)
            mlp_out, lrid_key = block.forward_mlp(
                x,
                x_is_normalized=fused_read_norm,
                emit_lrid_key=True,
            )
            add_partial(mlp_out, lrid_key)
            append_partial_if_block_end(2 * layer_idx + 2)

        x = read_residual(total_sublayers)
        if not fused_read_norm:
            x = norm(x)

        if return_hidden:
            return x

        if self.config.weight_tying:
            return F.linear(x, self.transformer.wte.weight, None)
        return self.lm_head(x)

    def _forward_block_lrid_attnres_split_fused_training(
        self,
        x,
        past_kv=None,
        use_cache=False,
        cu_doc_len=None,
        max_doc_len=None,
        return_hidden=False,
        fused_read_norm=False,
    ):
        if past_kv is not None or use_cache:
            raise NotImplementedError("KV-cache generation is not supported with attention residuals yet.")
        if self.config.lrid_input_dependent_query:
            raise NotImplementedError("Cached split LRID training does not support input-dependent queries.")

        total_sublayers = 2 * self.config.n_layer
        block_ends = sorted(self.attnres_block_ends)
        block_end_set = self.attnres_block_ends
        key_dim = self.config.lrid_rank
        query_bank = torch.stack(
            [
                self.transformer.lrid_queries[self._attnres_query_idx(read_idx)].reshape(key_dim)
                for read_idx in range(1, total_sublayers + 1)
            ],
            dim=0,
        )
        if self.config.attn_res_query_norm:
            query_bank = norm(query_bank.float()).to(x.dtype)
        else:
            query_bank = query_bank.to(x.dtype)
        query_bank = query_bank.contiguous()

        def source_key(source):
            key = source[1].reshape(source[1].size(0), source[1].size(1), key_dim)
            if self.config.attnres_key_norm:
                key = norm(key.float()).to(source[0].dtype)
            return key

        def add_logit_bias(logit, count, read_idx, source_summary_idx):
            bias = self._attnres_block_count_logit_bias(
                count,
                read_idx=read_idx,
                source_summary_idx=source_summary_idx,
            )
            if torch.is_tensor(bias):
                return logit + bias.to(device=logit.device, dtype=torch.float32).reshape(())
            if float(bias) != 0.0:
                return logit + float(bias)
            return logit

        def make_source_logits(source, first_read, source_count=1, source_summary_idx=None):
            key = source_key(source)
            queries = query_bank[first_read - 1:].to(key.dtype)
            logits = (
                F.linear(key, queries)
                .permute(2, 0, 1)
                .contiguous()
                * self.config.lrid_logit_scale
            )
            logit_biases = self._attnres_block_read_logit_biases(
                source_count,
                first_read,
                source_summary_idx,
                logits.device,
            )
            if logit_biases is not None:
                logits = logits + logit_biases.view(-1, 1, 1)
            return logits

        embedding_source = self._embedding_lrid_source(x)
        completed_values = [embedding_source[0]]
        completed_logits = [make_source_logits(embedding_source, 1, 1)]
        completed_logit_first_reads = [1]
        partial_values = [None, None]
        partial_keys = [None, None]
        partial_queries = [None, None]
        partial_counts = [0, 0]
        partial_summary_idxs = [None, None]
        phase_first_read = None
        phase_end = None
        phase_outputs = None
        phase_lses = None

        def next_phase_end(residual_idx):
            for block_end in block_ends:
                if block_end > residual_idx:
                    return block_end
            return residual_idx + 1

        def invalidate_phase():
            nonlocal phase_first_read, phase_end, phase_outputs, phase_lses
            phase_first_read = None
            phase_end = None
            phase_outputs = None
            phase_lses = None

        def ensure_phase(residual_idx):
            nonlocal phase_first_read, phase_end, phase_outputs, phase_lses
            if residual_idx < 1:
                return
            if phase_outputs is not None and phase_first_read <= residual_idx < phase_end:
                return
            phase_first_read = residual_idx
            phase_end = next_phase_end(residual_idx)
            if len(completed_values) == 1:
                logit_offset = phase_first_read - completed_logit_first_reads[0]
                phase_lse = completed_logits[0][logit_offset:phase_end - completed_logit_first_reads[0]].float()
                phase_outputs = completed_values[0].unsqueeze(0).expand(phase_lse.size(0), -1, -1, -1).unbind(0)
                phase_lses = phase_lse.unbind(0)
                return
            phase_logits = torch.stack(
                [
                    source_logits[phase_first_read - first_read:phase_end - first_read]
                    for source_logits, first_read in zip(completed_logits, completed_logit_first_reads)
                ],
                dim=1,
            ).contiguous()
            phase_output, phase_lse = attention_residual_phase1_from_logits(
                completed_values,
                phase_logits,
            )
            phase_outputs = phase_output.unbind(0)
            phase_lses = phase_lse.unbind(0)

        def current_partial_sources():
            sources = []
            for value, key, query, count, summary_idx in zip(
                partial_values,
                partial_keys,
                partial_queries,
                partial_counts,
                partial_summary_idxs,
            ):
                if value is None:
                    continue
                source = self._lrid_block_source(
                    value,
                    key,
                    query if self.config.lrid_input_dependent_query else None,
                    count,
                    summary_idx,
                )
                sources.append((source, count, summary_idx))
            return sources

        def merge_partial_sources(partial_sources, interblock_output, interblock_lse, residual_idx):
            query = query_bank[self._attnres_query_idx(residual_idx)]
            partial_logits = []
            partial_values_for_output = []
            for source, count, summary_idx in partial_sources:
                key = source_key(source)
                logit = (
                    torch.sum(key * query.to(key.dtype).view(1, 1, key_dim), dim=-1)
                    * self.config.lrid_logit_scale
                )
                logit = add_logit_bias(logit.float(), count, residual_idx, summary_idx)
                partial_logits.append(logit)
                partial_values_for_output.append(source[0])

            partial_logits = torch.stack(partial_logits, dim=0)
            total_lse = torch.logsumexp(
                torch.cat([interblock_lse.float().unsqueeze(0), partial_logits], dim=0),
                dim=0,
            )
            completed_weight = torch.exp(interblock_lse.float() - total_lse).to(interblock_output.dtype)
            output = interblock_output * completed_weight.unsqueeze(-1)
            partial_weights = torch.exp(partial_logits - total_lse.unsqueeze(0)).to(interblock_output.dtype)
            for weight, value in zip(partial_weights, partial_values_for_output):
                output = output + weight.unsqueeze(-1) * value
            return norm(output) if fused_read_norm else output

        def read_residual(residual_idx):
            if residual_idx == 0:
                output = completed_values[0]
                return norm(output) if fused_read_norm else output

            ensure_phase(residual_idx)
            phase_idx = residual_idx - phase_first_read
            interblock_output = phase_outputs[phase_idx]
            partial_sources = current_partial_sources()
            if len(partial_sources) == 0:
                output = interblock_output
                return norm(output) if fused_read_norm else output
            if len(partial_sources) == 1:
                source, count, summary_idx = partial_sources[0]
                query = query_bank[self._attnres_query_idx(residual_idx)].to(source[1].dtype)
                return self._apply_lrid_training_phase2(
                    source[0],
                    source[1].reshape(source[1].size(0), source[1].size(1), key_dim),
                    query,
                    interblock_output,
                    phase_lses[phase_idx],
                    normalize_output=fused_read_norm,
                    partial_count=count,
                    read_idx=residual_idx,
                    source_summary_idx=summary_idx,
                )
            return merge_partial_sources(
                partial_sources,
                interblock_output,
                phase_lses[phase_idx],
                residual_idx,
            )

        def append_partials_if_block_end(residual_idx):
            if residual_idx not in block_end_set:
                return
            appended = False
            for slot, (value, key, query, count, summary_idx) in enumerate(
                zip(
                    partial_values,
                    partial_keys,
                    partial_queries,
                    partial_counts,
                    partial_summary_idxs,
                )
            ):
                if value is None:
                    continue
                completed_source = self._lrid_block_source(
                    value,
                    key,
                    query if self.config.lrid_input_dependent_query else None,
                    count,
                    summary_idx,
                )
                completed_values.append(completed_source[0])
                completed_logits.append(
                    make_source_logits(
                        completed_source,
                        residual_idx,
                        count,
                        summary_idx,
                    )
                )
                completed_logit_first_reads.append(residual_idx)
                partial_values[slot] = None
                partial_keys[slot] = None
                partial_queries[slot] = None
                partial_counts[slot] = 0
                partial_summary_idxs[slot] = None
                appended = True
            if appended:
                invalidate_phase()

        def add_partial(slot, value, lrid_key, lrid_query, summary_idx, key_projector, query_projector=None):
            self._lrid_split_add_partial(
                partial_values,
                partial_keys,
                partial_queries,
                partial_counts,
                partial_summary_idxs,
                slot,
                value,
                lrid_key,
                lrid_query,
                summary_idx,
                key_projector=key_projector,
                query_projector=query_projector,
            )

        for layer_idx, block in enumerate(self.transformer.layers):
            x = read_residual(2 * layer_idx)
            attn_out, lrid_key = block.forward_attention(
                x,
                past_kv=None,
                use_cache=False,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                x_is_normalized=fused_read_norm,
                emit_lrid_key=True,
            )
            add_partial(
                0,
                attn_out,
                lrid_key,
                None,
                2 * layer_idx + 1,
                block.attn.c_proj.project_key_from_value,
                block.attn.c_proj.project_query_from_value,
            )
            append_partials_if_block_end(2 * layer_idx + 1)

            x = read_residual(2 * layer_idx + 1)
            mlp_out, lrid_key = block.forward_mlp(
                x,
                x_is_normalized=fused_read_norm,
                emit_lrid_key=True,
            )
            add_partial(
                1,
                mlp_out,
                lrid_key,
                None,
                2 * layer_idx + 2,
                block.mlp.fc2.project_key_from_value,
                block.mlp.fc2.project_query_from_value,
            )
            append_partials_if_block_end(2 * layer_idx + 2)

        x = read_residual(total_sublayers)
        if not fused_read_norm:
            x = norm(x)

        if return_hidden:
            return x

        if self.config.weight_tying:
            return F.linear(x, self.transformer.wte.weight, None)
        return self.lm_head(x)

    def _use_block_attnres_fused_training_path(self, past_kv, use_cache):
        return (
            self.use_attnres
            and self.config.attnres_backend != "fast"
            and self.config.use_fused_attnres
            and self.config.attnres_training_cache_phase1
            and not self.use_lrid
            and self.attnres_type == "block"
            and past_kv is None
            and not use_cache
        )

    def _use_block_lrid_attnres_fused_training_path(self, past_kv, use_cache):
        return (
            self.use_attnres
            and self.config.attnres_backend != "fast"
            and self.config.use_fused_attnres
            and self.config.attnres_training_cache_phase1
            and self.use_lrid
            and self.attnres_type == "block"
            and past_kv is None
            and not use_cache
            and self.config.lrid_num_heads == 1
            and not self.config.lrid_input_dependent_query
            and not self.config.lrid_key_from_value
            and not self.config.lrid_key_from_value_shared
            and not self.config.lrid_query_from_value
            and not self.config.lrid_query_from_value_shared
        )

    def _use_block_lrid_attnres_split_fused_training_path(self, past_kv, use_cache):
        return (
            self.config.attnres_block_split_sublayers
            and self._use_block_lrid_attnres_fused_training_path(past_kv, use_cache)
        )

    def _project_lrid_source_key(self, value):
        if self.config.lrid_key_from_output_tail:
            return self._lrid_output_tail_key(value)
        return self.transformer.lrid_embedding_key(value)

    def _project_lrid_source_query(self, value):
        return self.transformer.lrid_embedding_key.project_query_from_value(value)

    def _static_lrid_embedding_key(self, embedding):
        return self.transformer.lrid_static_embedding_key(embedding)

    def _add_static_lrid_embedding_key(self, key):
        if self.config.lrid_add_static_embedding_key:
            key = key + self.transformer.lrid_static_embedding_key(key)
        return key

    def _add_static_lrid_source_key(self, key):
        if self.config.lrid_add_static_source_key:
            key = key + self.transformer.lrid_static_source_key(key)
        return key

    def _lrid_output_tail_key(self, value):
        return value[..., -self.config.lrid_rank:].contiguous()

    def _lrid_source(self, value, key=None, query=None, key_value=None, query_value=None, add_static_key=False):
        if key is None:
            if self.config.lrid_key_from_output_tail:
                key = self._lrid_output_tail_key(value)
            elif self.config.lrid_key_from_value_shared:
                key = self._project_lrid_source_key(value if key_value is None else key_value)
            else:
                raise RuntimeError("LR AttnRes source key is missing")
        if add_static_key:
            key = self._add_static_lrid_source_key(key)
        if self.config.lrid_input_dependent_query:
            if self.config.lrid_query_from_value_shared:
                query = self._project_lrid_source_query(value if query_value is None else query_value)
            return value, key, query
        return value, key

    def _embedding_lrid_source(self, embedding):
        if self.config.lrid_key_from_output_tail:
            key = self._lrid_output_tail_key(embedding)
        elif self.config.lrid_static_embedding_key:
            key = self._static_lrid_embedding_key(embedding)
        else:
            key = self._project_lrid_source_key(embedding)
            key = self._add_static_lrid_embedding_key(key)
        return self._lrid_source(embedding, key)

    def _apply_lrid_attnres(
        self,
        residual_idx,
        sources,
        query_override=None,
        normalize_output=False,
        source_counts=None,
        source_logit_biases=None,
    ):
        if len(sources) == 1:
            output = sources[0][0]
            return norm(output) if normalize_output else output

        value_sources = [source[0] for source in sources]
        key_sources = [source[1] for source in sources]
        num_heads = self.config.lrid_num_heads
        key_head_dim = self.config.lrid_rank // num_heads
        value_head_dim = self.config.n_embd // num_heads

        query_idx = self._attnres_query_idx(residual_idx)
        static_query = self.transformer.lrid_queries[query_idx]
        if self.config.lrid_input_dependent_query:
            dynamic_query = query_override if query_override is not None else sources[-1][2]
            if dynamic_query is None:
                raise RuntimeError("Input-dependent LR AttnRes query is missing for the active source")
            dynamic_query = dynamic_query.reshape(*dynamic_query.shape[:-1], num_heads, key_head_dim)
            gate = self.transformer.lrid_query_gates[query_idx].view(1, 1, num_heads, 1)
            query = static_query.unsqueeze(0).unsqueeze(0) + gate * dynamic_query
        else:
            query = static_query

        if self.config.attn_res_query_norm:
            query = norm(query.float())

        if self.config.attnres_backend == "fast":
            self._assert_fast_runtime_read(
                value_sources,
                source_counts=source_counts,
                source_logit_biases=source_logit_biases,
            )
            if not self.config.attnres_key_norm:
                raise RuntimeError("Fast-AttnRes residual reads require key normalization")
            if self.config.lrid_logit_scale != 1.0:
                raise RuntimeError(
                    "Fast-AttnRes LR reads require --no-lrid_logit_scale "
                    "(the qualified full-graph path uses neutral scale 1.0)"
                )
            fast_query = query.reshape(-1).to(value_sources[0].dtype)
            fast_values = value_sources
            if self._fast_lrid_requires_key_payload():
                fast_values = [
                    torch.cat((value, key.reshape(*value.shape[:-1], -1)), dim=-1)
                    for value, key in zip(value_sources, key_sources)
                ]
            output = self._ensure_fast_attnres_op()(
                fast_values,
                fast_query,
                eps=self._fast_attnres_eps(value_sources[0].dtype, lrid=True),
                # Keep this literal. Passing ModelConfig's scalar into the
                # v2.0.1 public validator becomes symbolic under full-graph
                # torch.compile and fails at math.isfinite(scale).
                scale=1.0,
            )
            output = output[..., :self.config.n_embd]
            return norm(output) if normalize_output else output

        if self.config.use_fused_attnres:
            return lrid_attention_residual_read(
                value_sources,
                key_sources,
                query,
                num_heads,
                self.config.lrid_logit_scale,
                self.config.attnres_key_norm,
                normalize_output=normalize_output,
                source_counts=source_counts,
                source_logit_biases=source_logit_biases,
            )

        values = torch.stack(value_sources, dim=0)
        keys = torch.stack(key_sources, dim=0)
        keys = keys.reshape(*keys.shape[:-1], num_heads, key_head_dim)
        values = values.reshape(*values.shape[:-1], num_heads, value_head_dim)
        if self.config.attnres_key_norm:
            keys = norm(keys.float()).to(values.dtype)

        query = query.to(keys.dtype)
        if self.config.lrid_input_dependent_query:
            logits = torch.einsum("sbthr,bthr->sbth", keys, query) * self.config.lrid_logit_scale
        else:
            logits = torch.einsum("sbthr,hr->sbth", keys, query) * self.config.lrid_logit_scale
        if source_counts is not None and source_logit_biases is not None:
            raise RuntimeError("source_counts and source_logit_biases are mutually exclusive")
        if source_counts is not None:
            log_counts = torch.as_tensor(source_counts, device=logits.device, dtype=torch.float32).log()
            logits = logits + log_counts.view(-1, 1, 1, 1)
        if source_logit_biases is not None:
            bias_values = [
                bias.to(device=logits.device, dtype=torch.float32).reshape(())
                if torch.is_tensor(bias)
                else torch.tensor(float(bias), device=logits.device, dtype=torch.float32)
                for bias in source_logit_biases
            ]
            logit_bias = torch.stack(bias_values)
            logits = logits + logit_bias.view(-1, 1, 1, 1)
        weights = F.softmax(logits.float(), dim=0).to(values.dtype)
        output = torch.einsum("sbth,sbthd->bthd", weights, values)
        output = output.reshape(output.size(0), output.size(1), self.config.n_embd)
        return norm(output) if normalize_output else output
    
    def _init_attnres_query(self, query, std=0.02, init_cutoff_factor=None):
        if self.config.attn_res_query_init == "zero":
            nn.init.zeros_(query)
        elif self.config.attn_res_query_init == "normal":
            nn.init.normal_(query, mean=0.0, std=std)
        elif self.config.attn_res_query_init == "trunc_normal":
            cutoff = (init_cutoff_factor if init_cutoff_factor is not None else 3.0) * std
            nn.init.trunc_normal_(query, mean=0.0, std=std, a=-cutoff, b=cutoff)

    def _init_lrid_dynamic_query_projection(self, module, std=0.02, init_cutoff_factor=None):
        if not module.use_fused_query:
            return
        query_weight = module.proj.weight[module.query_offset:]
        if init_cutoff_factor is not None:
            cutoff = init_cutoff_factor * std
            nn.init.trunc_normal_(query_weight, mean=0.0, std=std, a=-cutoff, b=cutoff)
        else:
            nn.init.normal_(query_weight, mean=0.0, std=std)

    def _init_weights(self, module, std=0.02, init_cutoff_factor=None):
        if isinstance(module, nn.Linear):
            if init_cutoff_factor is not None:
                cutoff = init_cutoff_factor * std
                nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff, b=cutoff)
            else:
                nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, (nn.Embedding, TokenEmbedding)):
            if init_cutoff_factor is not None:
                cutoff = init_cutoff_factor * std
                nn.init.trunc_normal_(module.weight, mean=0.0, std=std, a=-cutoff, b=cutoff)
            else:
                nn.init.normal_(module.weight, mean=0.0, std=std)
        elif isinstance(module, LRIDStaticKey):
            module.reset_parameters(std, init_cutoff_factor)
        elif isinstance(module, AttentionResidual):
            self._init_attnres_query(module.query, std, init_cutoff_factor)

    def _sample_next_token(self, logits, temperature=1.0, top_k=None):
        if temperature < 0.0:
            raise ValueError("temperature must be non-negative")
        if top_k is not None and top_k < 1:
            raise ValueError("top_k must be >= 1 when set")

        if temperature == 0.0:
            return torch.argmax(logits, dim=-1, keepdim=True)

        logits = logits / temperature

        if top_k is not None:
            top_k = min(top_k, logits.size(-1))
            v, _ = torch.topk(logits, top_k)
            logits = logits.masked_fill(logits < v[:, [-1]], float("-inf"))

        probs = F.softmax(logits, dim=-1)
        return torch.multinomial(probs, num_samples=1)
    
    def get_lm_head_weight(self):
        if self.config.weight_tying:
            return self.transformer.wte.weight
        return self.lm_head.weight

    def get_lm_head_bias(self):
        return None

    def forward(
        self,
        idx,
        past_kv=None,
        use_cache=False,
        cu_doc_len=None,
        max_doc_len=None,
        return_hidden=False,
    ):
        _, T = idx.size()
        assert T <= self.config.block_size, f"Token length {T} exceeds max sequence length {self.config.block_size}"
        
        x = self.transformer.wte(idx)
        fused_read_norm = (
            self.use_attnres
            and self.config.use_fused_attnres
            and self.config.norm_pos == "before"
            and self.config.attnres_fuse_read_norm
            and (not self.use_lrid or self.config.lrid_num_heads == 1)
        )
        if (
            self.use_attnres
            and self.attnres_type == "block"
            and self.config.attnres_block_split_sublayers
            and past_kv is None
            and not use_cache
        ):
            if self.use_lrid:
                if self._use_block_lrid_attnres_split_fused_training_path(past_kv, use_cache):
                    return self._forward_block_lrid_attnres_split_fused_training(
                        x,
                        past_kv=past_kv,
                        use_cache=use_cache,
                        cu_doc_len=cu_doc_len,
                        max_doc_len=max_doc_len,
                        return_hidden=return_hidden,
                        fused_read_norm=fused_read_norm,
                    )
                return self._forward_block_lrid_attnres_split_sublayers(
                    x,
                    past_kv=past_kv,
                    use_cache=use_cache,
                    cu_doc_len=cu_doc_len,
                    max_doc_len=max_doc_len,
                    return_hidden=return_hidden,
                    fused_read_norm=fused_read_norm,
                )
            return self._forward_block_attnres_split_sublayers(
                x,
                past_kv=past_kv,
                use_cache=use_cache,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                return_hidden=return_hidden,
                fused_read_norm=fused_read_norm,
            )
        if self._use_block_attnres_fused_training_path(past_kv, use_cache):
            return self._forward_block_attnres_fused_training(
                x,
                past_kv=past_kv,
                use_cache=use_cache,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                return_hidden=return_hidden,
                fused_read_norm=fused_read_norm,
            )
        if self._use_block_lrid_attnres_fused_training_path(past_kv, use_cache):
            return self._forward_block_lrid_attnres_fused_training(
                x,
                past_kv=past_kv,
                use_cache=use_cache,
                cu_doc_len=cu_doc_len,
                max_doc_len=max_doc_len,
                return_hidden=return_hidden,
                fused_read_norm=fused_read_norm,
            )
        if self.use_attnres:
            if past_kv is not None or use_cache:
                raise NotImplementedError("KV-cache generation is not supported with attention residuals yet.")
            embedding = x
            if self.use_lrid:
                embedding_source = self._embedding_lrid_source(embedding)
            if self.attnres_type == "full":
                residual_sources = [embedding_source] if self.use_lrid else [embedding]
            else:
                block_ends = self.attnres_block_ends
                completed_blocks = [embedding_source] if self.use_lrid else [embedding]
                completed_block_counts = [1]
                completed_block_summary_idxs = [None]
                partial_block = None
                partial_count = 0
                if self.use_lrid:
                    partial_key = None
                    if self.config.lrid_input_dependent_query:
                        partial_query = None
        
        if past_kv is None:
            past_kv = [None] * len(self.transformer.layers)
        new_kv = [] if use_cache else None
        
        for layer_idx, block in enumerate(self.transformer.layers):
            if self.use_attnres:
                if self.use_lrid:
                    if self.attnres_type == "full":
                        x = self._apply_lrid_attnres(
                            2 * layer_idx,
                            residual_sources,
                            normalize_output=fused_read_norm,
                        )
                    else:
                        attn_res_idx = 2 * layer_idx
                        sources = completed_blocks if partial_block is None else completed_blocks + [
                            self._lrid_block_source(partial_block, partial_key, partial_query if self.config.lrid_input_dependent_query else None, partial_count, attn_res_idx)
                        ]
                        source_logit_biases = self._attnres_block_source_logit_biases(
                            attn_res_idx,
                            completed_block_counts,
                            completed_block_summary_idxs,
                            None if partial_block is None else partial_count,
                            None if partial_block is None else attn_res_idx,
                        )
                        x = self._apply_lrid_attnres(
                            attn_res_idx,
                            sources,
                            normalize_output=fused_read_norm,
                            source_logit_biases=source_logit_biases,
                        )

                    attn_out = block.forward_attention(
                        x,
                        past_kv=past_kv[layer_idx],
                        use_cache=False,
                        cu_doc_len=cu_doc_len,
                        max_doc_len=max_doc_len,
                        x_is_normalized=fused_read_norm,
                        emit_lrid_key=True,
                    )
                    if self.config.lrid_input_dependent_query:
                        attn_out, lrid_key, lrid_query = attn_out
                    else:
                        attn_out, lrid_key = attn_out
                    layer_output = attn_out

                    if self.attnres_type == "full":
                        residual_sources.append(
                            self._lrid_source(
                                layer_output,
                                lrid_key,
                                lrid_query if self.config.lrid_input_dependent_query else None,
                                add_static_key=True,
                            )
                        )
                        x = self._apply_lrid_attnres(
                            2 * layer_idx + 1,
                            residual_sources,
                            normalize_output=fused_read_norm,
                        )
                    else:
                        after_attn_idx = 2 * layer_idx + 1
                        if partial_block is None:
                            partial_block = layer_output
                            partial_count = 1
                        else:
                            partial_block = partial_block + layer_output
                            partial_count += 1
                        block_source_value = self._attnres_block_summary(partial_block, partial_count, after_attn_idx)
                        if self.config.lrid_key_from_value_shared:
                            partial_key = None
                        elif self.config.lrid_key_from_value:
                            partial_key = block.attn.c_proj.project_key_from_value(block_source_value)
                        else:
                            partial_key = lrid_key if partial_count == 1 else partial_key + lrid_key
                        if self.config.lrid_input_dependent_query:
                            if self.config.lrid_query_from_value_shared:
                                partial_query = None
                            elif self.config.lrid_query_from_value:
                                partial_query = block.attn.c_proj.project_query_from_value(block_source_value)
                            else:
                                partial_query = lrid_query
                        is_block_end = after_attn_idx in block_ends
                        if is_block_end:
                            completed_blocks.append(
                                self._lrid_block_source(partial_block, partial_key, partial_query if self.config.lrid_input_dependent_query else None, partial_count, after_attn_idx)
                            )
                            completed_block_counts.append(partial_count)
                            completed_block_summary_idxs.append(after_attn_idx)
                            partial_block = None
                            partial_key = None
                            partial_count = 0
                            if self.config.lrid_input_dependent_query:
                                partial_query = None
                        sources = completed_blocks if partial_block is None else completed_blocks + [
                            self._lrid_block_source(partial_block, partial_key, partial_query if self.config.lrid_input_dependent_query else None, partial_count, after_attn_idx)
                        ]
                        source_logit_biases = self._attnres_block_source_logit_biases(
                            after_attn_idx,
                            completed_block_counts,
                            completed_block_summary_idxs,
                            None if partial_block is None else partial_count,
                            None if partial_block is None else after_attn_idx,
                        )
                        x = self._apply_lrid_attnres(
                            after_attn_idx,
                            sources,
                            query_override=partial_query if self.config.lrid_input_dependent_query else None,
                            normalize_output=fused_read_norm,
                            source_logit_biases=source_logit_biases,
                        )

                    mlp_out = block.forward_mlp(
                        x,
                        x_is_normalized=fused_read_norm,
                        emit_lrid_key=True,
                    )
                    if self.config.lrid_input_dependent_query:
                        mlp_out, lrid_key, lrid_query = mlp_out
                    else:
                        mlp_out, lrid_key = mlp_out
                    layer_output = mlp_out
                    x = mlp_out

                    if self.attnres_type == "full":
                        residual_sources.append(
                            self._lrid_source(
                                layer_output,
                                lrid_key,
                                lrid_query if self.config.lrid_input_dependent_query else None,
                                add_static_key=True,
                            )
                        )
                    else:
                        after_mlp_idx = 2 * layer_idx + 2
                        if partial_block is None:
                            partial_block = layer_output
                            partial_count = 1
                        else:
                            partial_block = partial_block + layer_output
                            partial_count += 1
                        block_source_value = self._attnres_block_summary(partial_block, partial_count, after_mlp_idx)
                        if self.config.lrid_key_from_value_shared:
                            partial_key = None
                        elif self.config.lrid_key_from_value:
                            partial_key = block.mlp.fc2.project_key_from_value(block_source_value)
                        else:
                            partial_key = lrid_key if partial_count == 1 else partial_key + lrid_key
                        if self.config.lrid_input_dependent_query:
                            if self.config.lrid_query_from_value_shared:
                                partial_query = None
                            elif self.config.lrid_query_from_value:
                                partial_query = block.mlp.fc2.project_query_from_value(block_source_value)
                            else:
                                partial_query = lrid_query
                        is_block_end = after_mlp_idx in block_ends
                        if is_block_end:
                            completed_blocks.append(
                                self._lrid_block_source(partial_block, partial_key, partial_query if self.config.lrid_input_dependent_query else None, partial_count, after_mlp_idx)
                            )
                            completed_block_counts.append(partial_count)
                            completed_block_summary_idxs.append(after_mlp_idx)
                            partial_block = None
                            partial_key = None
                            partial_count = 0
                            if self.config.lrid_input_dependent_query:
                                partial_query = None
                    continue
                elif self.attnres_type == "full":
                    x = self._apply_attnres(
                        2 * layer_idx,
                        residual_sources,
                        normalize_output=fused_read_norm,
                    )
                else:
                    attn_res_idx = 2 * layer_idx
                    sources = completed_blocks if partial_block is None else completed_blocks + [self._attnres_block_summary(partial_block, partial_count, attn_res_idx)]
                    source_logit_biases = self._attnres_block_source_logit_biases(
                        attn_res_idx,
                        completed_block_counts,
                        completed_block_summary_idxs,
                        None if partial_block is None else partial_count,
                        None if partial_block is None else attn_res_idx,
                    )
                    x = self._apply_attnres(
                        2 * layer_idx,
                        sources,
                        normalize_output=fused_read_norm,
                        source_logit_biases=source_logit_biases,
                    )

                attn_out = block.forward_attention(
                    x,
                    past_kv=past_kv[layer_idx],
                    use_cache=False,
                    cu_doc_len=cu_doc_len,
                    max_doc_len=max_doc_len,
                    x_is_normalized=fused_read_norm,
                )
                layer_output = attn_out

                if self.attnres_type == "full":
                    residual_sources.append(layer_output)
                    x = self._apply_attnres(
                        2 * layer_idx + 1,
                        residual_sources,
                        normalize_output=fused_read_norm,
                    )
                else:
                    after_attn_idx = 2 * layer_idx + 1
                    if partial_block is None:
                        partial_block = layer_output
                        partial_count = 1
                    else:
                        partial_block = partial_block + layer_output
                        partial_count += 1
                    is_block_end = after_attn_idx in block_ends
                    if is_block_end:
                        completed_blocks.append(self._attnres_block_summary(partial_block, partial_count, after_attn_idx))
                        completed_block_counts.append(partial_count)
                        completed_block_summary_idxs.append(after_attn_idx)
                        partial_block = None
                        partial_count = 0
                    sources = completed_blocks if partial_block is None else completed_blocks + [self._attnres_block_summary(partial_block, partial_count, after_attn_idx)]
                    source_logit_biases = self._attnres_block_source_logit_biases(
                        after_attn_idx,
                        completed_block_counts,
                        completed_block_summary_idxs,
                        None if partial_block is None else partial_count,
                        None if partial_block is None else after_attn_idx,
                    )
                    x = self._apply_attnres(
                        after_attn_idx,
                        sources,
                        normalize_output=fused_read_norm,
                        source_logit_biases=source_logit_biases,
                    )

                mlp_out = block.forward_mlp(x, x_is_normalized=fused_read_norm)
                layer_output = mlp_out
                x = mlp_out

                if self.attnres_type == "full":
                    residual_sources.append(layer_output)
                else:
                    after_mlp_idx = 2 * layer_idx + 2
                    if partial_block is None:
                        partial_block = layer_output
                        partial_count = 1
                    else:
                        partial_block = partial_block + layer_output
                        partial_count += 1
                    is_block_end = after_mlp_idx in block_ends
                    if is_block_end:
                        completed_blocks.append(self._attnres_block_summary(partial_block, partial_count, after_mlp_idx))
                        completed_block_counts.append(partial_count)
                        completed_block_summary_idxs.append(after_mlp_idx)
                        partial_block = None
                        partial_count = 0
            else:
                block_out = block(
                    x,
                    past_kv=past_kv[layer_idx],
                    use_cache=use_cache,
                    cu_doc_len=cu_doc_len,
                    max_doc_len=max_doc_len,
                )

                if use_cache:
                    x, present_kv = block_out
                    new_kv.append(present_kv)
                else:
                    x = block_out

        if self.use_attnres:
            if self.use_lrid:
                if self.attnres_type == "full":
                    x = self._apply_lrid_attnres(
                        2 * self.config.n_layer,
                        residual_sources,
                        normalize_output=fused_read_norm,
                    )
                else:
                    sources = completed_blocks if partial_block is None else completed_blocks + [
                        self._lrid_block_source(partial_block, partial_key, partial_query if self.config.lrid_input_dependent_query else None, partial_count, 2 * self.config.n_layer)
                    ]
                    source_logit_biases = self._attnres_block_source_logit_biases(
                        2 * self.config.n_layer,
                        completed_block_counts,
                        completed_block_summary_idxs,
                        None if partial_block is None else partial_count,
                        None if partial_block is None else 2 * self.config.n_layer,
                    )
                    x = self._apply_lrid_attnres(
                        2 * self.config.n_layer,
                        sources,
                        normalize_output=fused_read_norm,
                        source_logit_biases=source_logit_biases,
                    )
            elif self.attnres_type == "full":
                x = self._apply_attnres(
                    2 * self.config.n_layer,
                    residual_sources,
                    normalize_output=fused_read_norm,
                )
            else:
                sources = completed_blocks if partial_block is None else completed_blocks + [self._attnres_block_summary(partial_block, partial_count, 2 * self.config.n_layer)]
                source_logit_biases = self._attnres_block_source_logit_biases(
                    2 * self.config.n_layer,
                    completed_block_counts,
                    completed_block_summary_idxs,
                    None if partial_block is None else partial_count,
                    None if partial_block is None else 2 * self.config.n_layer,
                )
                x = self._apply_attnres(
                    2 * self.config.n_layer,
                    sources,
                    normalize_output=fused_read_norm,
                    source_logit_biases=source_logit_biases,
                )
        
        if not fused_read_norm:
            x = norm(x)

        if return_hidden:
            if use_cache:
                return x, new_kv
            return x
        
        if self.config.weight_tying:
            logits = F.linear(x, self.transformer.wte.weight, None)
        else:
            logits = self.lm_head(x)
        
        if use_cache:
            return logits, new_kv
        return logits
    
    @torch.no_grad()
    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None, max_context=None):
        self.eval()
        device = next(self.parameters()).device
        idx = idx.to(device)
        _, T = idx.size()

        if max_context is None:
            max_context = self.config.block_size

        if max_new_tokens < 0:
            raise ValueError("max_new_tokens must be non-negative")
        if max_context < 1:
            raise ValueError("max_context must be >= 1")

        if T > max_context:
            idx = idx[:, -max_context:]
            T = idx.size(1)

        if self.use_attnres or idx.size(1) + max_new_tokens > max_context:
            generated = idx
            for _ in range(max_new_tokens):
                idx_cond = generated[:, -max_context:]
                logits = self(idx_cond)
                logits = logits[:, -1, :]
                next_token = self._sample_next_token(logits, temperature=temperature, top_k=top_k)
                generated = torch.cat((generated, next_token), dim=1)

            return generated

        past_kv = None

        if T > 0:
            start = 0
            while start < T:
                end = min(start + self.config.block_size, T)
                idx_cond = idx[:, start:end]
                logits, past_kv = self(idx_cond, past_kv=past_kv, use_cache=True)
                start = end

        for generation_step in range(max_new_tokens):
            if generation_step == 0 and T > 0:
                # Prefill logits at the final prompt position already predict
                # the first continuation token. Feeding the final prompt token
                # into the populated cache again duplicates it and changes the
                # prediction.
                next_token_logits = logits[:, -1, :]
            else:
                idx_cond = idx[:, -1:] if idx.size(1) > 0 else idx
                logits, past_kv = self(idx_cond, past_kv=past_kv, use_cache=True)
                next_token_logits = logits[:, -1, :]
            next_token = self._sample_next_token(
                next_token_logits,
                temperature=temperature,
                top_k=top_k,
            )
            idx = torch.cat((idx, next_token), dim=1)

        return idx