checkpoint-134000/modeling_convgpt_v2.py

# coding=utf-8
"""ConvGPT-v2: configurable hybrid causal 1D/2D ConvLM with tiny retrieval routers."""

from __future__ import annotations

import math
from functools import lru_cache
from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GenerationMixin
from transformers.activations import ACT2FN
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .configuration_convgpt_v2 import ConvGPTV2Config

logger = logging.get_logger(__name__)


@dataclass
class ConvGPTV2PrefixCache:
    """Generation cache for ConvGPT-v2.

    Stores the accepted prefix plus per-layer activations needed to decode the
    next token without re-running every layer over the whole prefix. If an old
    two-tensor legacy cache is provided, the model falls back to the exact
    full-prefix path until a rich cache is produced again.
    """

    input_ids: torch.LongTensor
    attention_mask: Optional[torch.Tensor] = None
    layer_inputs: Optional[Tuple[torch.Tensor, ...]] = None
    layer_post_conv: Optional[Tuple[torch.Tensor, ...]] = None

    @property
    def has_incremental_state(self) -> bool:
        return self.layer_inputs is not None and self.layer_post_conv is not None

    @property
    def device(self) -> torch.device:
        return self.input_ids.device

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return int(self.input_ids.shape[1])

    def get_max_cache_shape(self) -> Optional[int]:
        return None

    def reorder_cache(self, beam_idx: torch.LongTensor) -> "ConvGPTV2PrefixCache":
        self.input_ids = self.input_ids.index_select(0, beam_idx.to(self.input_ids.device))
        if self.attention_mask is not None:
            self.attention_mask = self.attention_mask.index_select(0, beam_idx.to(self.attention_mask.device))
        if self.layer_inputs is not None:
            self.layer_inputs = tuple(
                layer_state.index_select(0, beam_idx.to(layer_state.device))
                for layer_state in self.layer_inputs
            )
        if self.layer_post_conv is not None:
            self.layer_post_conv = tuple(
                layer_state.index_select(0, beam_idx.to(layer_state.device))
                for layer_state in self.layer_post_conv
            )
        return self

    def to_legacy_cache(self):
        return (self.input_ids, self.attention_mask)

    @classmethod
    def from_legacy_cache(cls, cache) -> "ConvGPTV2PrefixCache":
        if isinstance(cache, cls):
            return cache
        if isinstance(cache, tuple) and len(cache) == 2 and torch.is_tensor(cache[0]):
            return cls(input_ids=cache[0], attention_mask=cache[1])
        raise TypeError(f"Unsupported ConvGPT-v2 cache type: {type(cache)!r}")

try:
    import triton
    import triton.language as tl

    _TRITON_AVAILABLE = True
except Exception:
    triton = None
    tl = None
    _TRITON_AVAILABLE = False


if _TRITON_AVAILABLE:
    @triton.jit
    def _depthwise_neighbor_gather_kernel(
        x_ptr,
        idx_ptr,
        mask_ptr,
        w_ptr,
        out_ptr,
        c_stride_x,
        p_stride_x,
        p_stride_idx,
        k_stride_idx,
        p_stride_mask,
        k_stride_mask,
        c_stride_w,
        k_stride_w,
        c_stride_out,
        p_stride_out,
        P,
        K,
        N_P_BLOCKS: tl.constexpr,
        BLOCK_P: tl.constexpr,
        BLOCK_K: tl.constexpr,
    ):
        linear_pid = tl.program_id(0)
        pid = linear_pid % N_P_BLOCKS
        bc = linear_pid // N_P_BLOCKS

        offs_p = pid * BLOCK_P + tl.arange(0, BLOCK_P)
        offs_k = tl.arange(0, BLOCK_K)
        p_mask = offs_p < P
        k_mask = offs_k < K
        mask_pk = p_mask[:, None] & k_mask[None, :]

        idx = tl.load(idx_ptr + offs_p[:, None] * p_stride_idx + offs_k[None, :] * k_stride_idx, mask=mask_pk, other=0)
        valid = tl.load(mask_ptr + offs_p[:, None] * p_stride_mask + offs_k[None, :] * k_stride_mask, mask=mask_pk, other=0).to(tl.int1)

        x = tl.load(x_ptr + bc * c_stride_x + idx * p_stride_x, mask=mask_pk & valid, other=0.0)
        w = tl.load(w_ptr + bc * c_stride_w + offs_k * k_stride_w, mask=k_mask, other=0.0)
        acc = tl.sum(x * w[None, :], axis=1)
        tl.store(out_ptr + bc * c_stride_out + offs_p * p_stride_out, acc, mask=p_mask)


    @triton.jit
    def _chunk_token_score_forward_kernel(
        q_ptr,
        k_ptr,
        pos_ptr,
        valid_ptr,
        out_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        C: tl.constexpr,
        BLOCK_C: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        offs_c = tl.arange(0, BLOCK_C)
        c_mask = offs_c < C
        b = row // T
        t = row - b * T
        q_base = (b * T + t) * D
        cand_pos = tl.load(pos_ptr + row * C + offs_c, mask=c_mask, other=0)
        cand_valid = tl.load(valid_ptr + row * C + offs_c, mask=c_mask, other=0).to(tl.int1)
        acc = tl.zeros((BLOCK_C,), tl.float32)
        for d0 in range(0, D, BLOCK_D):
            offs_d = d0 + tl.arange(0, BLOCK_D)
            d_mask = offs_d < D
            q = tl.load(q_ptr + q_base + offs_d, mask=d_mask, other=0.0).to(tl.float32)
            k = tl.load(
                k_ptr + (b * T + cand_pos[:, None]) * D + offs_d[None, :],
                mask=c_mask[:, None] & cand_valid[:, None] & d_mask[None, :],
                other=0.0,
            ).to(tl.float32)
            acc += tl.sum(k * q[None, :], axis=1)
        score = acc * tl.rsqrt(tl.full((), D, tl.float32))
        neg_inf = tl.full((BLOCK_C,), -3.4028234663852886e38, tl.float32)
        score = tl.where(c_mask & cand_valid, score, neg_inf)
        tl.store(out_ptr + row * C + offs_c, score, mask=c_mask)

    @triton.jit
    def _chunk_token_score_grad_q_kernel(
        grad_scores_ptr,
        k_ptr,
        pos_ptr,
        valid_ptr,
        grad_q_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        C: tl.constexpr,
        BLOCK_C: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        d_block = tl.program_id(1)
        offs_c = tl.arange(0, BLOCK_C)
        offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D)
        c_mask = offs_c < C
        d_mask = offs_d < D
        b = row // T
        cand_pos = tl.load(pos_ptr + row * C + offs_c, mask=c_mask, other=0)
        cand_valid = tl.load(valid_ptr + row * C + offs_c, mask=c_mask, other=0).to(tl.int1)
        grad_scores = tl.load(grad_scores_ptr + row * C + offs_c, mask=c_mask & cand_valid, other=0.0).to(tl.float32)
        k = tl.load(
            k_ptr + (b * T + cand_pos[:, None]) * D + offs_d[None, :],
            mask=c_mask[:, None] & cand_valid[:, None] & d_mask[None, :],
            other=0.0,
        ).to(tl.float32)
        grad_q = tl.sum(k * grad_scores[:, None], axis=0) * tl.rsqrt(tl.full((), D, tl.float32))
        tl.store(grad_q_ptr + row * D + offs_d, grad_q, mask=d_mask)

    @triton.jit
    def _chunk_token_score_grad_k_kernel(
        grad_scores_ptr,
        q_ptr,
        pos_ptr,
        valid_ptr,
        grad_k_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        C: tl.constexpr,
        BLOCK_C: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        c_block = tl.program_id(1)
        d_block = tl.program_id(2)
        offs_c = c_block * BLOCK_C + tl.arange(0, BLOCK_C)
        offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D)
        c_mask = offs_c < C
        d_mask = offs_d < D
        b = row // T
        cand_pos = tl.load(pos_ptr + row * C + offs_c, mask=c_mask, other=0)
        cand_valid = tl.load(valid_ptr + row * C + offs_c, mask=c_mask, other=0).to(tl.int1)
        grad_scores = tl.load(grad_scores_ptr + row * C + offs_c, mask=c_mask & cand_valid, other=0.0).to(tl.float32)
        q = tl.load(q_ptr + row * D + offs_d, mask=d_mask, other=0.0).to(tl.float32)
        grad = grad_scores[:, None] * q[None, :] * tl.rsqrt(tl.full((), D, tl.float32))
        tl.atomic_add(
            grad_k_ptr + (b * T + cand_pos[:, None]) * D + offs_d[None, :],
            grad,
            sem="relaxed",
            mask=c_mask[:, None] & cand_valid[:, None] & d_mask[None, :],
        )

    @triton.jit
    def _weighted_value_forward_kernel(
        weights_ptr,
        v_ptr,
        pos_ptr,
        valid_ptr,
        out_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        K: tl.constexpr,
        BLOCK_K: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        d_block = tl.program_id(1)
        offs_k = tl.arange(0, BLOCK_K)
        offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D)
        k_mask = offs_k < K
        d_mask = offs_d < D
        b = row // T
        pos = tl.load(pos_ptr + row * K + offs_k, mask=k_mask, other=0)
        valid = tl.load(valid_ptr + row * K + offs_k, mask=k_mask, other=0).to(tl.int1)
        weights = tl.load(weights_ptr + row * K + offs_k, mask=k_mask & valid, other=0.0).to(tl.float32)
        v = tl.load(
            v_ptr + (b * T + pos[:, None]) * D + offs_d[None, :],
            mask=k_mask[:, None] & valid[:, None] & d_mask[None, :],
            other=0.0,
        ).to(tl.float32)
        out = tl.sum(v * weights[:, None], axis=0)
        tl.store(out_ptr + row * D + offs_d, out, mask=d_mask)

    @triton.jit
    def _weighted_value_grad_v_kernel(
        grad_out_ptr,
        weights_ptr,
        pos_ptr,
        valid_ptr,
        grad_v_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        K: tl.constexpr,
        BLOCK_K: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        k_block = tl.program_id(1)
        d_block = tl.program_id(2)
        offs_k = k_block * BLOCK_K + tl.arange(0, BLOCK_K)
        offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D)
        k_mask = offs_k < K
        d_mask = offs_d < D
        b = row // T
        pos = tl.load(pos_ptr + row * K + offs_k, mask=k_mask, other=0)
        valid = tl.load(valid_ptr + row * K + offs_k, mask=k_mask, other=0).to(tl.int1)
        weights = tl.load(weights_ptr + row * K + offs_k, mask=k_mask & valid, other=0.0).to(tl.float32)
        grad_out = tl.load(grad_out_ptr + row * D + offs_d, mask=d_mask, other=0.0).to(tl.float32)
        grad = weights[:, None] * grad_out[None, :]
        tl.atomic_add(
            grad_v_ptr + (b * T + pos[:, None]) * D + offs_d[None, :],
            grad,
            sem="relaxed",
            mask=k_mask[:, None] & valid[:, None] & d_mask[None, :],
        )

    @triton.jit
    def _weighted_value_grad_weights_kernel(
        grad_out_ptr,
        v_ptr,
        pos_ptr,
        valid_ptr,
        grad_weights_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        K: tl.constexpr,
        BLOCK_K: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        offs_k = tl.arange(0, BLOCK_K)
        k_mask = offs_k < K
        b = row // T
        pos = tl.load(pos_ptr + row * K + offs_k, mask=k_mask, other=0)
        valid = tl.load(valid_ptr + row * K + offs_k, mask=k_mask, other=0).to(tl.int1)
        acc = tl.zeros((BLOCK_K,), tl.float32)
        for d0 in range(0, D, BLOCK_D):
            offs_d = d0 + tl.arange(0, BLOCK_D)
            d_mask = offs_d < D
            grad_out = tl.load(grad_out_ptr + row * D + offs_d, mask=d_mask, other=0.0).to(tl.float32)
            v = tl.load(
                v_ptr + (b * T + pos[:, None]) * D + offs_d[None, :],
                mask=k_mask[:, None] & valid[:, None] & d_mask[None, :],
                other=0.0,
            ).to(tl.float32)
            acc += tl.sum(v * grad_out[None, :], axis=1)
        tl.store(grad_weights_ptr + row * K + offs_k, acc, mask=k_mask)



    # Experimental fused chunk+token selector retained for future tuning.
    # Disabled in router dispatch because the first benchmark was not faster
    # than chunk PyTorch top-k + fused token top-k.
    @triton.jit
    def _chunk2_token_select_forward_kernel(
        q_ptr,
        chunk_k_ptr,
        token_k_ptr,
        top_scores_ptr,
        top_pos_ptr,
        top_valid_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        NUM_CHUNKS: tl.constexpr,
        CHUNK_SIZE: tl.constexpr,
        KTOP: tl.constexpr,
        INCLUDE_CURRENT: tl.constexpr,
        BLOCK_CHUNKS: tl.constexpr,
        BLOCK_CAND: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        b = row // T
        t = row - b * T
        q_base = (b * T + t) * D

        offs_chunks = tl.arange(0, BLOCK_CHUNKS)
        chunk_mask = offs_chunks < NUM_CHUNKS
        token_chunk = t // CHUNK_SIZE
        if INCLUDE_CURRENT:
            allowed = offs_chunks <= token_chunk
        else:
            allowed = offs_chunks < token_chunk
        chunk_mask = chunk_mask & allowed & (offs_chunks * CHUNK_SIZE < T)

        chunk_scores = tl.zeros((BLOCK_CHUNKS,), tl.float32)
        for d0 in range(0, D, BLOCK_D):
            offs_d = d0 + tl.arange(0, BLOCK_D)
            d_mask = offs_d < D
            q = tl.load(q_ptr + q_base + offs_d, mask=d_mask, other=0.0).to(tl.float32)
            ck = tl.load(
                chunk_k_ptr + (b * NUM_CHUNKS + offs_chunks[:, None]) * D + offs_d[None, :],
                mask=chunk_mask[:, None] & d_mask[None, :],
                other=0.0,
            ).to(tl.float32)
            chunk_scores += tl.sum(ck * q[None, :], axis=1)
        neg_inf_chunks = tl.full((BLOCK_CHUNKS,), -3.4028234663852886e38, tl.float32)
        chunk_scores = tl.where(chunk_mask, chunk_scores * tl.rsqrt(tl.full((), D, tl.float32)), neg_inf_chunks)
        chunk0 = tl.argmax(chunk_scores, axis=0, tie_break_left=True)
        score0 = tl.max(chunk_scores, axis=0)
        chunk_scores = tl.where(offs_chunks == chunk0, neg_inf_chunks, chunk_scores)
        chunk1 = tl.argmax(chunk_scores, axis=0, tie_break_left=True)
        score1 = tl.max(chunk_scores, axis=0)
        valid0 = score0 > -3.0e38
        valid1 = score1 > -3.0e38

        offs_c = tl.arange(0, BLOCK_CAND)
        c_mask = offs_c < (CHUNK_SIZE * 2)
        in_first = offs_c < CHUNK_SIZE
        cand_chunk = tl.where(in_first, chunk0, chunk1)
        cand_valid_chunk = tl.where(in_first, valid0, valid1)
        cand_offset = tl.where(in_first, offs_c, offs_c - CHUNK_SIZE)
        cand_pos = cand_chunk * CHUNK_SIZE + cand_offset
        cand_valid = c_mask & cand_valid_chunk & (cand_pos < T)
        if not INCLUDE_CURRENT:
            cand_valid = cand_valid & (cand_pos < t)

        token_scores = tl.zeros((BLOCK_CAND,), tl.float32)
        for d0 in range(0, D, BLOCK_D):
            offs_d = d0 + tl.arange(0, BLOCK_D)
            d_mask = offs_d < D
            q = tl.load(q_ptr + q_base + offs_d, mask=d_mask, other=0.0).to(tl.float32)
            tk = tl.load(
                token_k_ptr + (b * T + cand_pos[:, None]) * D + offs_d[None, :],
                mask=cand_valid[:, None] & d_mask[None, :],
                other=0.0,
            ).to(tl.float32)
            token_scores += tl.sum(tk * q[None, :], axis=1)
        neg_inf = tl.full((BLOCK_CAND,), -3.4028234663852886e38, tl.float32)
        token_scores = tl.where(cand_valid, token_scores * tl.rsqrt(tl.full((), D, tl.float32)), neg_inf)

        for i in tl.static_range(0, KTOP):
            max_val = tl.max(token_scores, axis=0)
            max_idx = tl.argmax(token_scores, axis=0, tie_break_left=True)
            is_valid = max_val > -3.0e38
            selected_pos = tl.sum(tl.where(offs_c == max_idx, cand_pos, 0), axis=0)
            selected_pos = tl.where(is_valid, selected_pos, 0)
            tl.store(top_scores_ptr + row * KTOP + i, max_val, mask=True)
            tl.store(top_pos_ptr + row * KTOP + i, selected_pos, mask=True)
            tl.store(top_valid_ptr + row * KTOP + i, is_valid, mask=True)
            token_scores = tl.where(offs_c == max_idx, neg_inf, token_scores)

    @triton.jit
    def _chunk_token_select_forward_kernel(
        q_ptr,
        k_ptr,
        pos_ptr,
        valid_ptr,
        top_scores_ptr,
        top_pos_ptr,
        top_valid_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        C: tl.constexpr,
        KTOP: tl.constexpr,
        BLOCK_C: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        offs_c = tl.arange(0, BLOCK_C)
        c_mask = offs_c < C
        b = row // T
        t = row - b * T
        q_base = (b * T + t) * D
        cand_pos = tl.load(pos_ptr + row * C + offs_c, mask=c_mask, other=0)
        cand_valid = tl.load(valid_ptr + row * C + offs_c, mask=c_mask, other=0).to(tl.int1)
        scores = tl.zeros((BLOCK_C,), tl.float32)
        for d0 in range(0, D, BLOCK_D):
            offs_d = d0 + tl.arange(0, BLOCK_D)
            d_mask = offs_d < D
            q = tl.load(q_ptr + q_base + offs_d, mask=d_mask, other=0.0).to(tl.float32)
            k = tl.load(
                k_ptr + (b * T + cand_pos[:, None]) * D + offs_d[None, :],
                mask=c_mask[:, None] & cand_valid[:, None] & d_mask[None, :],
                other=0.0,
            ).to(tl.float32)
            scores += tl.sum(k * q[None, :], axis=1)
        neg_inf = tl.full((BLOCK_C,), -3.4028234663852886e38, tl.float32)
        scores = tl.where(c_mask & cand_valid, scores * tl.rsqrt(tl.full((), D, tl.float32)), neg_inf)

        for i in tl.static_range(0, KTOP):
            max_val = tl.max(scores, axis=0)
            max_idx = tl.argmax(scores, axis=0, tie_break_left=True)
            is_valid = max_val > -3.0e38
            selected_pos = tl.load(pos_ptr + row * C + max_idx, mask=is_valid, other=0)
            tl.store(top_scores_ptr + row * KTOP + i, max_val, mask=True)
            tl.store(top_pos_ptr + row * KTOP + i, selected_pos, mask=True)
            tl.store(top_valid_ptr + row * KTOP + i, is_valid, mask=True)
            scores = tl.where(offs_c == max_idx, neg_inf, scores)

    @triton.jit
    def _chunk_token_select_grad_kernel(
        grad_scores_ptr,
        q_ptr,
        k_ptr,
        top_pos_ptr,
        top_valid_ptr,
        grad_q_ptr,
        grad_k_ptr,
        T: tl.constexpr,
        D: tl.constexpr,
        KTOP: tl.constexpr,
        BLOCK_K: tl.constexpr,
        BLOCK_D: tl.constexpr,
    ):
        row = tl.program_id(0)
        d_block = tl.program_id(1)
        offs_k = tl.arange(0, BLOCK_K)
        offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D)
        k_mask = offs_k < KTOP
        d_mask = offs_d < D
        b = row // T
        pos = tl.load(top_pos_ptr + row * KTOP + offs_k, mask=k_mask, other=0)
        valid = tl.load(top_valid_ptr + row * KTOP + offs_k, mask=k_mask, other=0).to(tl.int1)
        grad_scores = tl.load(grad_scores_ptr + row * KTOP + offs_k, mask=k_mask & valid, other=0.0).to(tl.float32)
        q = tl.load(q_ptr + row * D + offs_d, mask=d_mask, other=0.0).to(tl.float32)
        k = tl.load(
            k_ptr + (b * T + pos[:, None]) * D + offs_d[None, :],
            mask=k_mask[:, None] & valid[:, None] & d_mask[None, :],
            other=0.0,
        ).to(tl.float32)
        scale = tl.rsqrt(tl.full((), D, tl.float32))
        grad_q = tl.sum(k * grad_scores[:, None], axis=0) * scale
        tl.store(grad_q_ptr + row * D + offs_d, grad_q, mask=d_mask)
        grad_k = grad_scores[:, None] * q[None, :] * scale
        tl.atomic_add(
            grad_k_ptr + (b * T + pos[:, None]) * D + offs_d[None, :],
            grad_k,
            sem="relaxed",
            mask=k_mask[:, None] & valid[:, None] & d_mask[None, :],
        )


def _ceil_pow2(n: int) -> int:
    return 1 << (n - 1).bit_length()


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


class ConvGPTV2RotaryEmbedding(nn.Module):
    """Simple 1D RoPE cache for router query/key projections."""

    def __init__(self, dim: int, theta: float = 10000.0):
        super().__init__()
        if dim % 2 != 0:
            dim -= 1
        self.dim = max(0, int(dim))
        self.theta = float(theta)
        if self.dim > 0:
            inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim))
        else:
            inv_freq = torch.empty(0, dtype=torch.float32)
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @torch.no_grad()
    def forward(self, position_ids: torch.Tensor, dtype: torch.dtype, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
        if self.dim == 0:
            empty = torch.empty((*position_ids.shape, 0), device=device, dtype=dtype)
            return empty, empty
        inv_freq = self.inv_freq.to(device=device)
        pos = position_ids.to(device=device, dtype=torch.float32)
        with torch.autocast(device_type=device.type if device.type != "mps" else "cpu", enabled=False):
            freqs = torch.einsum("bt,d->btd", pos, inv_freq.float())
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos()
            sin = emb.sin()
        return cos.to(dtype=dtype), sin.to(dtype=dtype)


def _apply_rope_to_btd(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    rope_dim = cos.shape[-1]
    if rope_dim == 0:
        return x
    x_rope = x[..., :rope_dim]
    x_pass = x[..., rope_dim:]
    x_rot = (x_rope * cos) + (_rotate_half(x_rope) * sin)
    return torch.cat((x_rot, x_pass), dim=-1) if x_pass.shape[-1] else x_rot


def _token_position_ids(
    position_ids: Optional[torch.LongTensor],
    bsz: int,
    seq_len: int,
    device: torch.device,
) -> torch.LongTensor:
    if position_ids is not None:
        return position_ids.to(device=device, dtype=torch.long)
    return torch.arange(seq_len, device=device, dtype=torch.long).view(1, seq_len).expand(bsz, -1)


def _chunk_start_position_ids(token_pos: torch.LongTensor, chunk_size: int, num_chunks: int) -> torch.LongTensor:
    seq_len = token_pos.shape[1]
    starts = (torch.arange(num_chunks, device=token_pos.device, dtype=torch.long) * chunk_size).clamp(max=max(seq_len - 1, 0))
    return token_pos.index_select(1, starts)


class _ChunkTokenScoreFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, token_k: torch.Tensor, flat_pos: torch.Tensor, flat_valid: torch.Tensor) -> torch.Tensor:
        bsz, seq_len, dim = q.shape
        cand = flat_pos.shape[-1]
        q_c = q.contiguous()
        k_c = token_k.contiguous()
        pos_c = flat_pos.contiguous()
        valid_c = flat_valid.contiguous()
        out = torch.empty((bsz, seq_len, cand), device=q.device, dtype=torch.float32)
        block_c = _ceil_pow2(cand)
        block_d = 64
        _chunk_token_score_forward_kernel[(bsz * seq_len,)](
            q_c, k_c, pos_c, valid_c, out,
            seq_len, dim, cand,
            BLOCK_C=block_c, BLOCK_D=block_d,
            num_warps=4,
        )
        ctx.save_for_backward(q_c, k_c, pos_c, valid_c)
        return out

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        q, token_k, flat_pos, flat_valid = ctx.saved_tensors
        bsz, seq_len, dim = q.shape
        cand = flat_pos.shape[-1]
        grad_scores = grad_out.contiguous()
        grad_q = torch.empty(q.shape, device=q.device, dtype=torch.float32)
        grad_k = torch.zeros(token_k.shape, device=token_k.device, dtype=torch.float32)
        block_c_q = _ceil_pow2(cand)
        block_c_k = min(32, _ceil_pow2(cand))
        block_d = 64
        _chunk_token_score_grad_q_kernel[(bsz * seq_len, triton.cdiv(dim, block_d))](
            grad_scores, token_k, flat_pos, flat_valid, grad_q,
            seq_len, dim, cand,
            BLOCK_C=block_c_q, BLOCK_D=block_d,
            num_warps=4,
        )
        _chunk_token_score_grad_k_kernel[(bsz * seq_len, triton.cdiv(cand, block_c_k), triton.cdiv(dim, block_d))](
            grad_scores, q, flat_pos, flat_valid, grad_k,
            seq_len, dim, cand,
            BLOCK_C=block_c_k, BLOCK_D=block_d,
            num_warps=4,
        )
        return grad_q, grad_k, None, None






class _Chunk2TokenSelectFunction(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        q: torch.Tensor,
        chunk_k: torch.Tensor,
        token_k: torch.Tensor,
        k_top: int,
        chunk_size: int,
        include_current_chunk: bool,
    ):
        bsz, seq_len, dim = q.shape
        num_chunks = chunk_k.shape[1]
        q_c = q.contiguous()
        chunk_k_c = chunk_k.contiguous()
        token_k_c = token_k.contiguous()
        top_scores = torch.empty((bsz, seq_len, k_top), device=q.device, dtype=torch.float32)
        top_pos = torch.empty((bsz, seq_len, k_top), device=q.device, dtype=torch.long)
        top_valid = torch.empty((bsz, seq_len, k_top), device=q.device, dtype=torch.bool)
        _chunk2_token_select_forward_kernel[(bsz * seq_len,)](
            q_c, chunk_k_c, token_k_c, top_scores, top_pos, top_valid,
            seq_len, dim, num_chunks, chunk_size, int(k_top), bool(include_current_chunk),
            BLOCK_CHUNKS=_ceil_pow2(num_chunks), BLOCK_CAND=_ceil_pow2(chunk_size * 2), BLOCK_D=64,
            num_warps=4,
        )
        ctx.save_for_backward(q_c, token_k_c, top_pos, top_valid)
        ctx.k_top = int(k_top)
        return top_scores, top_pos, top_valid

    @staticmethod
    def backward(ctx, grad_top_scores: torch.Tensor, grad_top_pos=None, grad_top_valid=None):
        q, token_k, top_pos, top_valid = ctx.saved_tensors
        bsz, seq_len, dim = q.shape
        k_top = ctx.k_top
        grad_scores = grad_top_scores.contiguous()
        grad_q = torch.empty(q.shape, device=q.device, dtype=torch.float32)
        grad_k = torch.zeros(token_k.shape, device=token_k.device, dtype=torch.float32)
        _chunk_token_select_grad_kernel[(bsz * seq_len, triton.cdiv(dim, 64))](
            grad_scores, q, token_k, top_pos, top_valid, grad_q, grad_k,
            seq_len, dim, k_top,
            BLOCK_K=_ceil_pow2(k_top), BLOCK_D=64,
            num_warps=4,
        )
        return grad_q, None, grad_k, None, None, None

class _ChunkTokenSelectFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, q: torch.Tensor, token_k: torch.Tensor, flat_pos: torch.Tensor, flat_valid: torch.Tensor, k_top: int):
        bsz, seq_len, dim = q.shape
        cand = flat_pos.shape[-1]
        q_c = q.contiguous()
        k_c = token_k.contiguous()
        pos_c = flat_pos.contiguous()
        valid_c = flat_valid.contiguous()
        top_scores = torch.empty((bsz, seq_len, k_top), device=q.device, dtype=torch.float32)
        top_pos = torch.empty((bsz, seq_len, k_top), device=q.device, dtype=torch.long)
        top_valid = torch.empty((bsz, seq_len, k_top), device=q.device, dtype=torch.bool)
        _chunk_token_select_forward_kernel[(bsz * seq_len,)](
            q_c, k_c, pos_c, valid_c, top_scores, top_pos, top_valid,
            seq_len, dim, cand, k_top,
            BLOCK_C=_ceil_pow2(cand), BLOCK_D=64,
            num_warps=4,
        )
        ctx.save_for_backward(q_c, k_c, top_pos, top_valid)
        ctx.k_top = k_top
        return top_scores, top_pos, top_valid

    @staticmethod
    def backward(ctx, grad_top_scores: torch.Tensor, grad_top_pos=None, grad_top_valid=None):
        q, token_k, top_pos, top_valid = ctx.saved_tensors
        bsz, seq_len, dim = q.shape
        k_top = ctx.k_top
        grad_scores = grad_top_scores.contiguous()
        grad_q = torch.empty(q.shape, device=q.device, dtype=torch.float32)
        grad_k = torch.zeros(token_k.shape, device=token_k.device, dtype=torch.float32)
        _chunk_token_select_grad_kernel[(bsz * seq_len, triton.cdiv(dim, 64))](
            grad_scores, q, token_k, top_pos, top_valid, grad_q, grad_k,
            seq_len, dim, k_top,
            BLOCK_K=_ceil_pow2(k_top), BLOCK_D=64,
            num_warps=4,
        )
        return grad_q, grad_k, None, None, None

class _WeightedValueFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, weights: torch.Tensor, token_v: torch.Tensor, top_pos: torch.Tensor, top_valid: torch.Tensor) -> torch.Tensor:
        bsz, seq_len, dim = token_v.shape
        k_top = top_pos.shape[-1]
        weights_c = weights.contiguous()
        v_c = token_v.contiguous()
        pos_c = top_pos.contiguous()
        valid_c = top_valid.contiguous()
        out = torch.empty((bsz, seq_len, dim), device=token_v.device, dtype=token_v.dtype)
        block_k = _ceil_pow2(k_top)
        block_d = 64
        _weighted_value_forward_kernel[(bsz * seq_len, triton.cdiv(dim, block_d))](
            weights_c, v_c, pos_c, valid_c, out,
            seq_len, dim, k_top,
            BLOCK_K=block_k, BLOCK_D=block_d,
            num_warps=4,
        )
        ctx.save_for_backward(weights_c, v_c, pos_c, valid_c)
        return out

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor):
        weights, token_v, top_pos, top_valid = ctx.saved_tensors
        bsz, seq_len, dim = token_v.shape
        k_top = top_pos.shape[-1]
        grad_out_c = grad_out.contiguous()
        grad_weights = torch.empty(weights.shape, device=weights.device, dtype=torch.float32)
        grad_v = torch.zeros(token_v.shape, device=token_v.device, dtype=torch.float32)
        block_k_full = _ceil_pow2(k_top)
        block_k_v = min(32, _ceil_pow2(k_top))
        block_d = 64
        _weighted_value_grad_weights_kernel[(bsz * seq_len,)](
            grad_out_c, token_v, top_pos, top_valid, grad_weights,
            seq_len, dim, k_top,
            BLOCK_K=block_k_full, BLOCK_D=block_d,
            num_warps=4,
        )
        _weighted_value_grad_v_kernel[(bsz * seq_len, triton.cdiv(k_top, block_k_v), triton.cdiv(dim, block_d))](
            grad_out_c, weights, top_pos, top_valid, grad_v,
            seq_len, dim, k_top,
            BLOCK_K=block_k_v, BLOCK_D=block_d,
            num_warps=4,
        )
        return grad_weights, grad_v, None, None


def _can_use_chunk_token_triton(q: torch.Tensor, token_k: torch.Tensor, token_v: torch.Tensor, flat_pos: torch.Tensor) -> bool:
    if not _TRITON_AVAILABLE or q.device.type != "cuda":
        return False
    if q.dtype not in {torch.float16, torch.bfloat16, torch.float32}:
        return False
    if token_k.dtype != q.dtype or token_v.dtype != q.dtype:
        return False
    dim = q.shape[-1]
    cand = flat_pos.shape[-1]
    return dim <= 1024 and cand <= 256


def _chunk_token_scores(q: torch.Tensor, token_k: torch.Tensor, flat_pos: torch.Tensor, flat_valid: torch.Tensor) -> torch.Tensor:
    if _can_use_chunk_token_triton(q, token_k, token_k, flat_pos):
        return _ChunkTokenScoreFunction.apply(q, token_k, flat_pos, flat_valid)
    dim = q.shape[-1]
    bsz, seq_len, _ = q.shape
    cand = flat_pos.shape[-1]
    gather_idx = flat_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
    cand_k = token_k.gather(1, gather_idx).view(bsz, seq_len, cand, dim)
    token_scores = (q.unsqueeze(-2) * cand_k).sum(dim=-1) / math.sqrt(dim)
    return token_scores.masked_fill(~flat_valid, torch.finfo(token_scores.dtype).min)






def _select_chunk2_token_scores(
    q: torch.Tensor,
    chunk_k: torch.Tensor,
    token_k: torch.Tensor,
    k_top: int,
    chunk_size: int,
    include_current_chunk: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if (
        _TRITON_AVAILABLE
        and q.device.type == "cuda"
        and q.dtype in {torch.float16, torch.bfloat16, torch.float32}
        and chunk_k.dtype == q.dtype
        and token_k.dtype == q.dtype
        and q.shape[-1] <= 1024
        and chunk_size * 2 <= 256
    ):
        return _Chunk2TokenSelectFunction.apply(q, chunk_k, token_k, int(k_top), int(chunk_size), bool(include_current_chunk))
    raise RuntimeError("chunk2 token Triton selector called outside supported envelope")

def _select_chunk_token_scores(
    q: torch.Tensor,
    token_k: torch.Tensor,
    flat_pos: torch.Tensor,
    flat_valid: torch.Tensor,
    k_top: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    if _can_use_chunk_token_triton(q, token_k, token_k, flat_pos):
        return _ChunkTokenSelectFunction.apply(q, token_k, flat_pos, flat_valid, int(k_top))
    token_scores = _chunk_token_scores(q, token_k, flat_pos, flat_valid)
    safe_token_scores = torch.where(flat_valid.any(dim=-1, keepdim=True), token_scores, torch.zeros_like(token_scores))
    top_scores, top_idx = safe_token_scores.topk(k_top, dim=-1)
    top_valid = flat_valid.gather(-1, top_idx)
    top_pos = flat_pos.gather(-1, top_idx)
    top_scores = top_scores.masked_fill(~top_valid, torch.finfo(top_scores.dtype).min)
    return top_scores, top_pos, top_valid

def _weighted_token_values(weights: torch.Tensor, token_v: torch.Tensor, top_pos: torch.Tensor, top_valid: torch.Tensor) -> torch.Tensor:
    if _can_use_chunk_token_triton(token_v, token_v, token_v, top_pos):
        return _WeightedValueFunction.apply(weights, token_v, top_pos, top_valid)
    dim = token_v.shape[-1]
    bsz, seq_len, _ = token_v.shape
    k_top = top_pos.shape[-1]
    gather_idx = top_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
    cand_v = token_v.gather(1, gather_idx).view(bsz, seq_len, k_top, dim)
    return (weights.unsqueeze(-1) * cand_v).sum(dim=-2)


def _part1by1(n: int) -> int:
    n &= 0x0000FFFF
    n = (n | (n << 8)) & 0x00FF00FF
    n = (n | (n << 4)) & 0x0F0F0F0F
    n = (n | (n << 2)) & 0x33333333
    n = (n | (n << 1)) & 0x55555555
    return n


def _morton_code(row: int, col: int) -> int:
    # Interleave col as x bits and row as y bits.
    return _part1by1(col) | (_part1by1(row) << 1)


def _hilbert_d2xy(order: int, d: int) -> Tuple[int, int]:
    """Convert Hilbert distance to x/y for a power-of-two square."""
    x = y = 0
    t = d
    s = 1
    while s < order:
        rx = 1 & (t // 2)
        ry = 1 & (t ^ rx)
        if ry == 0:
            if rx == 1:
                x = s - 1 - x
                y = s - 1 - y
            x, y = y, x
        x += s * rx
        y += s * ry
        t //= 4
        s *= 2
    return x, y


@lru_cache(maxsize=64)
def _build_curve_indices_cached(grid_size: int, packing: str, pack_order: str = "sequence_to_curve") -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Return (seq_to_grid, grid_to_seq) flattened indices for a square grid.

    seq_to_grid[pos] is the flat grid cell where token position `pos` should be
    placed. grid_to_seq[cell] is the source token position for that flat grid
    cell. These are inverse permutations.
    """
    packing = packing.lower().replace("-", "_")
    if packing in {"z_order", "zorder"}:
        packing = "morton"

    n = grid_size
    total = n * n

    if packing == "row_major":
        cells = list(range(total))
    elif packing == "snake":
        cells = []
        for r in range(n):
            cols = range(n) if r % 2 == 0 else range(n - 1, -1, -1)
            cells.extend(r * n + c for c in cols)
    elif packing == "morton":
        cells = sorted(range(total), key=lambda idx: _morton_code(idx // n, idx % n))
    elif packing == "hilbert":
        order = _ceil_pow2(n)
        cells = []
        # For non-power-of-two grid sizes, walk a power-of-two Hilbert curve and
        # keep only coordinates inside the requested square.
        for d in range(order * order):
            x, y = _hilbert_d2xy(order, d)
            if x < n and y < n:
                cells.append(y * n + x)
                if len(cells) == total:
                    break
    else:
        raise ValueError(f"Unsupported packing: {packing}")

    if pack_order == "curve_to_sequence":
        # Interpret the curve as an ordering over grid cells, but fill normal
        # row-major cells with sequence positions from that curve. This is useful
        # as an ablation for whether curve locality or physical placement matters.
        grid_to_seq = torch.tensor(cells, dtype=torch.long)
        seq_to_grid = torch.empty_like(grid_to_seq)
        seq_to_grid[grid_to_seq] = torch.arange(total, dtype=torch.long)
    else:
        seq_to_grid = torch.tensor(cells, dtype=torch.long)
        grid_to_seq = torch.empty_like(seq_to_grid)
        grid_to_seq[seq_to_grid] = torch.arange(total, dtype=torch.long)
    return seq_to_grid, grid_to_seq


def build_curve_indices(grid_size: int, packing: str, pack_order: str = "sequence_to_curve") -> Tuple[torch.Tensor, torch.Tensor]:
    """Return cloned curve-index tensors, backed by a process-local cache.

    A 32-layer 256x256 Hilbert ConvGPT-v2 previously rebuilt the same Hilbert
    permutation dozens of times during model construction. Cache the immutable
    CPU tensors and hand modules clones so buffer registration stays isolated.
    """
    seq_to_grid, grid_to_seq = _build_curve_indices_cached(grid_size, packing, pack_order)
    return seq_to_grid.clone(), grid_to_seq.clone()


@lru_cache(maxsize=128)
def _build_conv2d_causal_tables_cached(
    grid_size: int,
    kernel_size: int,
    dilation: int,
    packing: str,
    pack_order: str,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Vectorized cached causal 2D lookup tables.

    Model construction used to rebuild these tables with Python loops once per
    ConvGPT-v2 layer. A 256x256, 32-layer checkpoint therefore spent ~2 minutes
    before generation even started. The tables only depend on geometry, so build
    them once per `(grid, kernel, dilation, packing)` and clone on registration.
    """
    n = int(grid_size)
    k = int(kernel_size)
    d = int(dilation)
    p = n * n
    center = k // 2
    _seq_to_grid, grid_to_seq = _build_curve_indices_cached(n, packing, pack_order)

    cells = torch.arange(p, dtype=torch.long)
    rows = torch.div(cells, n, rounding_mode="floor").view(p, 1)
    cols = (cells % n).view(p, 1)

    offsets = []
    for kr in range(k):
        for kc in range(k):
            offsets.append(((kr - center) * d, (kc - center) * d))
    delta_r = torch.tensor([x[0] for x in offsets], dtype=torch.long).view(1, k * k)
    delta_c = torch.tensor([x[1] for x in offsets], dtype=torch.long).view(1, k * k)

    nr = rows + delta_r
    nc = cols + delta_c
    in_bounds = (nr >= 0) & (nr < n) & (nc >= 0) & (nc < n)
    neighbor_cells = (nr.clamp(0, n - 1) * n + nc.clamp(0, n - 1)).to(torch.long)

    center_seq = grid_to_seq.index_select(0, cells).view(p, 1)
    neighbor_seq = grid_to_seq.index_select(0, neighbor_cells.reshape(-1)).view(p, k * k)
    neighbor_mask = in_bounds & (neighbor_seq <= center_seq)
    neighbor_indices = torch.where(neighbor_mask, neighbor_cells, torch.zeros_like(neighbor_cells)).to(torch.long)
    patch_mask = neighbor_mask.transpose(0, 1).contiguous().view(1, 1, k * k, p)
    return patch_mask, neighbor_indices.contiguous(), neighbor_mask.contiguous()


class ConvGPTV2RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.float()
        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
        return self.weight * hidden_states.to(input_dtype)


class ConvGPTV2RMSNorm2d(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        input_dtype = x.dtype
        variance = x.float().pow(2).mean(dim=1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps).to(dtype=x.dtype)
        return self.weight.view(1, -1, 1, 1).to(dtype=input_dtype) * x


class ConvGPTV2MLP(nn.Module):
    def __init__(self, config: ConvGPTV2Config):
        super().__init__()
        self.gate_up_proj = nn.Linear(config.hidden_size, config.intermediate_size * 2, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.act = ACT2FN[config.hidden_act]
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
        return self.dropout(self.down_proj(self.act(gate) * up))


class CausalDepthwiseConv1d(nn.Module):
    def __init__(self, hidden_size: int, kernel_size: int, dilation: int):
        super().__init__()
        self.left_padding = (kernel_size - 1) * dilation
        self.conv = nn.Conv1d(
            hidden_size,
            hidden_size,
            kernel_size=kernel_size,
            dilation=dilation,
            groups=hidden_size,
            bias=False,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # hidden_states: [B, T, D]
        x = hidden_states.transpose(1, 2)
        x = F.pad(x, (self.left_padding, 0))
        x = self.conv(x)
        return x.transpose(1, 2)


class CausalConv1dBranch(nn.Module):
    def __init__(self, config: ConvGPTV2Config, dilation: int):
        super().__init__()
        d = config.hidden_size
        self.norm = ConvGPTV2RMSNorm(d, config.rms_norm_eps)
        self.dwconv = CausalDepthwiseConv1d(d, config.conv1d_kernel_size, dilation)
        self.pw = nn.Linear(d, d * 2, bias=False)
        self.out = nn.Linear(d, d, bias=False)
        self.act = ACT2FN[config.hidden_act]
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        y = self.norm(hidden_states)
        y = self.dwconv(y)
        gate, value = self.pw(y).chunk(2, dim=-1)
        y = self.out(self.act(gate) * value)
        return self.dropout(y)

    def forward_token(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Exact causal 1D branch output for the newest token only."""
        y = self.norm(hidden_states)
        kernel = self.dwconv.conv.kernel_size[0]
        dilation = self.dwconv.conv.dilation[0]
        seq_len = y.shape[1]
        offsets = torch.arange(kernel, device=y.device)
        positions = seq_len - 1 - (kernel - 1 - offsets) * dilation
        valid = positions >= 0
        positions = positions.clamp(min=0).to(torch.long)
        gathered = y.index_select(1, positions) * valid.to(dtype=y.dtype).view(1, -1, 1)
        weight = self.dwconv.conv.weight.view(y.shape[-1], kernel).to(dtype=y.dtype)
        y_token = (gathered * weight.transpose(0, 1).view(1, kernel, y.shape[-1])).sum(dim=1, keepdim=True)
        if self.dwconv.conv.bias is not None:
            y_token = y_token + self.dwconv.conv.bias.view(1, 1, -1)
        gate, value = self.pw(y_token).chunk(2, dim=-1)
        y_token = self.out(self.act(gate) * value)
        return self.dropout(y_token)


class SpaceFillingPacker(nn.Module):
    def __init__(self, config: ConvGPTV2Config):
        super().__init__()
        self.grid_size = config.grid_size
        self.max_grid_tokens = config.grid_size * config.grid_size
        self.hidden_size = config.hidden_size
        seq_to_grid, grid_to_seq = build_curve_indices(
            config.grid_size, config.packing, config.pack_order
        )
        self.register_buffer("seq_to_grid", seq_to_grid, persistent=False)
        self.register_buffer("grid_to_seq", grid_to_seq, persistent=False)

    def seq_to_2d(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, int]:
        bsz, seq_len, dim = hidden_states.shape
        if seq_len > self.max_grid_tokens:
            raise ValueError(
                f"Sequence length {seq_len} exceeds configured grid capacity {self.max_grid_tokens}. "
                "Increase config.grid_size or truncate inputs."
            )
        grid = hidden_states.new_zeros(bsz, self.max_grid_tokens, dim)
        grid_positions = self.seq_to_grid[:seq_len].to(hidden_states.device)
        expanded_positions = grid_positions.view(1, -1, 1).expand(bsz, -1, dim)
        grid.scatter_(1, expanded_positions, hidden_states)
        grid = grid.view(bsz, self.grid_size, self.grid_size, dim).permute(0, 3, 1, 2).contiguous()
        return grid, seq_len

    def d2_to_seq(self, grid: torch.Tensor, seq_len: int) -> torch.Tensor:
        bsz, dim, _, _ = grid.shape
        flat = grid.permute(0, 2, 3, 1).contiguous().view(bsz, self.max_grid_tokens, dim)
        gather_positions = self.seq_to_grid[:seq_len].to(grid.device)
        return flat.index_select(1, gather_positions)


class CausalConv2dBranch(nn.Module):
    def __init__(self, config: ConvGPTV2Config, dilation: int):
        super().__init__()
        d = config.hidden_size
        k = config.conv2d_kernel_size
        self.grid_size = config.grid_size
        self.kernel_size = k
        self.dilation = dilation
        self.padding = dilation * (k // 2)
        self.use_depthwise = config.use_2d_depthwise
        self.backend = config.conv2d_backend
        self.chunk_size = config.conv2d_chunk_size
        self.norm = ConvGPTV2RMSNorm2d(d, config.rms_norm_eps)

        if self.use_depthwise:
            self.spatial = nn.Conv2d(
                d,
                d,
                kernel_size=k,
                dilation=dilation,
                padding=0,
                groups=d,
                bias=False,
            )
            self.pointwise = nn.Conv2d(d, d * 2, kernel_size=1, bias=False)
        else:
            inner = d * config.conv_expand
            self.spatial = nn.Conv2d(
                d,
                inner,
                kernel_size=k,
                dilation=dilation,
                padding=0,
                bias=False,
            )
            self.pointwise = nn.Conv2d(inner, d * 2, kernel_size=1, bias=False)

        patch_mask, neighbor_idx, neighbor_mask = _build_conv2d_causal_tables_cached(
            config.grid_size,
            self.kernel_size,
            self.dilation,
            config.packing,
            config.pack_order,
        )
        self.register_buffer("patch_causal_mask", patch_mask, persistent=False)
        self.register_buffer("neighbor_indices", neighbor_idx, persistent=False)
        self.register_buffer("neighbor_mask", neighbor_mask, persistent=False)
        self.register_buffer("seq_to_grid", self._seq_to_grid(config), persistent=False)
        self.register_buffer("grid_to_seq", self._grid_to_seq(config), persistent=False)
        static_mask = self._build_static_row_major_mask(config)
        self.register_buffer("static_causal_mask", static_mask, persistent=False)
        self.out = nn.Conv2d(d, d, kernel_size=1, bias=False)
        self.act = ACT2FN[config.hidden_act]
        self.dropout = nn.Dropout2d(config.dropout)

    def _grid_to_seq(self, config: ConvGPTV2Config) -> torch.Tensor:
        _, grid_to_seq = build_curve_indices(config.grid_size, config.packing, config.pack_order)
        return grid_to_seq

    def _seq_to_grid(self, config: ConvGPTV2Config) -> torch.Tensor:
        seq_to_grid, _ = build_curve_indices(config.grid_size, config.packing, config.pack_order)
        return seq_to_grid

    def _build_patch_causal_mask(self, config: ConvGPTV2Config) -> torch.Tensor:
        """Per-cell exact causal mask for unfold backend."""
        n = config.grid_size
        k = self.kernel_size
        center = k // 2
        grid_to_seq = self._grid_to_seq(config)
        mask = torch.zeros(k * k, n * n, dtype=torch.bool)
        for r in range(n):
            for c in range(n):
                cell = r * n + c
                center_seq = int(grid_to_seq[cell])
                offset_idx = 0
                for kr in range(k):
                    for kc in range(k):
                        nr = r + (kr - center) * self.dilation
                        nc = c + (kc - center) * self.dilation
                        if 0 <= nr < n and 0 <= nc < n:
                            neighbor_cell = nr * n + nc
                            neighbor_seq = int(grid_to_seq[neighbor_cell])
                            mask[offset_idx, cell] = neighbor_seq <= center_seq
                        offset_idx += 1
        return mask.view(1, 1, k * k, n * n)

    def _build_neighbor_table(self, config: ConvGPTV2Config) -> Tuple[torch.Tensor, torch.Tensor]:
        """Precompute exact causal neighbor indices for chunked gather backend."""
        n = config.grid_size
        k = self.kernel_size
        center = k // 2
        p = n * n
        grid_to_seq = self._grid_to_seq(config)
        indices = torch.zeros(p, k * k, dtype=torch.long)
        mask = torch.zeros(p, k * k, dtype=torch.bool)
        for r in range(n):
            for c in range(n):
                cell = r * n + c
                center_seq = int(grid_to_seq[cell])
                offset_idx = 0
                for kr in range(k):
                    for kc in range(k):
                        nr = r + (kr - center) * self.dilation
                        nc = c + (kc - center) * self.dilation
                        if 0 <= nr < n and 0 <= nc < n:
                            neighbor_cell = nr * n + nc
                            neighbor_seq = int(grid_to_seq[neighbor_cell])
                            if neighbor_seq <= center_seq:
                                indices[cell, offset_idx] = neighbor_cell
                                mask[cell, offset_idx] = True
                        offset_idx += 1
        return indices, mask

    def _build_static_row_major_mask(self, config: ConvGPTV2Config) -> torch.Tensor:
        """Static PixelCNN mask, exact only for row_major packing."""
        k = self.kernel_size
        center = k // 2
        mask = torch.ones(k, k)
        mask[center, center + 1 :] = 0
        mask[center + 1 :, :] = 0
        return mask.view(1, 1, k, k)

    def _depthwise_from_patches(self, patches: torch.Tensor, channels: int) -> torch.Tensor:
        weight = self.spatial.weight.view(channels, self.kernel_size * self.kernel_size)
        return torch.einsum("bcmp,cm->bcp", patches, weight)

    def _dense_from_patches(self, patches: torch.Tensor) -> torch.Tensor:
        weight = self.spatial.weight.view(
            self.spatial.out_channels, patches.shape[1], self.kernel_size * self.kernel_size
        )
        return torch.einsum("bcmp,ocm->bop", patches, weight)

    def _forward_unfold(self, y: torch.Tensor) -> torch.Tensor:
        bsz, channels, height, width = y.shape
        patches = F.unfold(
            y,
            kernel_size=self.kernel_size,
            dilation=self.dilation,
            padding=self.padding,
            stride=1,
        )
        patches = patches.view(bsz, channels, self.kernel_size * self.kernel_size, height * width)
        patches = patches * self.patch_causal_mask.to(device=patches.device, dtype=patches.dtype)
        if self.use_depthwise:
            out = self._depthwise_from_patches(patches, channels)
        else:
            out = self._dense_from_patches(patches)
        return out

    def _forward_chunked_gather(self, y: torch.Tensor) -> torch.Tensor:
        bsz, channels, height, width = y.shape
        p = height * width
        x_flat = y.view(bsz, channels, p)
        out = y.new_empty(bsz, self.spatial.out_channels, p)
        idx_all = self.neighbor_indices.to(y.device)
        mask_all = self.neighbor_mask.to(device=y.device, dtype=y.dtype)
        for p0 in range(0, p, self.chunk_size):
            p1 = min(p0 + self.chunk_size, p)
            idx = idx_all[p0:p1]  # [chunk, K*K]
            # Advanced indexing returns [B, C, chunk, K*K]. This is the same
            # logical tensor as unfold but only materialized for one chunk.
            patches = x_flat[:, :, idx]
            patches = patches * mask_all[p0:p1].view(1, 1, p1 - p0, self.kernel_size * self.kernel_size)
            patches = patches.permute(0, 1, 3, 2).contiguous()  # [B, C, K*K, chunk]
            if self.use_depthwise:
                out[:, :, p0:p1] = self._depthwise_from_patches(patches, channels)
            else:
                out[:, :, p0:p1] = self._dense_from_patches(patches)
        return out

    def _forward_triton_gather(self, y: torch.Tensor) -> torch.Tensor:
        if not _TRITON_AVAILABLE:
            return self._forward_chunked_gather(y)
        if not self.use_depthwise:
            return self._forward_chunked_gather(y)
        if y.device.type != "cuda":
            return self._forward_chunked_gather(y)

        bsz, channels, height, width = y.shape
        p = height * width
        k2 = self.kernel_size * self.kernel_size
        x_flat = y.contiguous().view(bsz * channels, p)
        idx = self.neighbor_indices.to(y.device).contiguous()
        mask = self.neighbor_mask.to(device=y.device, dtype=torch.bool).contiguous()
        weight = self.spatial.weight.view(channels, k2).to(y.dtype).contiguous()
        weight = weight.repeat(bsz, 1)
        out = torch.empty((bsz * channels, p), device=y.device, dtype=y.dtype)

        n_p_blocks = triton.cdiv(p, 128)
        grid = (n_p_blocks * bsz * channels,)
        _depthwise_neighbor_gather_kernel[grid](
            x_flat,
            idx,
            mask,
            weight,
            out,
            x_flat.stride(0),
            x_flat.stride(1),
            idx.stride(0),
            idx.stride(1),
            mask.stride(0),
            mask.stride(1),
            weight.stride(0),
            weight.stride(1),
            out.stride(0),
            out.stride(1),
            p,
            k2,
            N_P_BLOCKS=n_p_blocks,
            BLOCK_P=128,
            BLOCK_K=16,
        )
        return out.view(bsz, channels, p)

    def _forward_triton_sequence_gather(
        self,
        y_seq: torch.Tensor,
        neighbor_seq_indices: torch.Tensor,
        neighbor_seq_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Depthwise causal 2D gather over active sequence cells only.

        y_seq is [B, C, T]. neighbor_seq_indices/mask are [T, K*K] in
        sequence coordinates, not full-grid coordinates. This is the fast path
        for large grids with shorter active sequences, e.g. 256x256 grid and
        seq_len=2048. It avoids doing work over all grid_size^2 cells.
        """
        if not _TRITON_AVAILABLE or y_seq.device.type != "cuda" or not self.use_depthwise:
            raise RuntimeError("active Triton sequence gather requires CUDA Triton depthwise mode")

        bsz, channels, seq_len = y_seq.shape
        k2 = self.kernel_size * self.kernel_size
        x_flat = y_seq.contiguous().view(bsz * channels, seq_len)
        idx = neighbor_seq_indices.to(y_seq.device).contiguous()
        mask = neighbor_seq_mask.to(device=y_seq.device, dtype=torch.bool).contiguous()
        weight = self.spatial.weight.view(channels, k2).to(y_seq.dtype).contiguous().repeat(bsz, 1)
        out = torch.empty((bsz * channels, seq_len), device=y_seq.device, dtype=y_seq.dtype)

        n_p_blocks = triton.cdiv(seq_len, 128)
        grid = (n_p_blocks * bsz * channels,)
        _depthwise_neighbor_gather_kernel[grid](
            x_flat,
            idx,
            mask,
            weight,
            out,
            x_flat.stride(0),
            x_flat.stride(1),
            idx.stride(0),
            idx.stride(1),
            mask.stride(0),
            mask.stride(1),
            weight.stride(0),
            weight.stride(1),
            out.stride(0),
            out.stride(1),
            seq_len,
            k2,
            N_P_BLOCKS=n_p_blocks,
            BLOCK_P=128,
            BLOCK_K=16,
        )
        return out.view(bsz, channels, seq_len)

    def forward_sequence(self, hidden_states: torch.Tensor, active_cells: torch.Tensor) -> Optional[torch.Tensor]:
        """Fast active-cell path for triton_gather.

        Returns [B, T, D] when the active Triton path is available, otherwise
        None so callers can fall back to the full-grid backend. This avoids
        processing grid_size^2 cells when only the first T sequence positions
        are active.
        """
        if self.backend != "triton_gather" or not _TRITON_AVAILABLE or hidden_states.device.type != "cuda" or not self.use_depthwise:
            return None

        bsz, seq_len, channels = hidden_states.shape
        input_dtype = hidden_states.dtype
        y = hidden_states.float()
        variance = y.pow(2).mean(dim=-1, keepdim=True)
        y = hidden_states * torch.rsqrt(variance + self.norm.eps).to(dtype=hidden_states.dtype)
        y = y * self.norm.weight.view(1, 1, -1).to(device=hidden_states.device, dtype=input_dtype)
        y = y.transpose(1, 2).contiguous()  # [B, C, T]

        active_cells = active_cells[:seq_len].to(hidden_states.device, dtype=torch.long)

        # Safe physical indexing for CUDA/Triton. Invalid entries remain masked
        # out, but every tensor passed into index_select is clamped to a valid
        # physical range. This avoids device-side asserts on newer CUDA/Triton
        # stacks while preserving the Triton sequence-gather kernel below.
        max_cells = int(self.neighbor_indices.shape[0])
        active_valid = (active_cells >= 0) & (active_cells < max_cells)
        active_cells_safe = active_cells.clamp(min=0, max=max(max_cells - 1, 0))

        neighbor_indices = self.neighbor_indices.to(hidden_states.device)
        base_neighbor_mask = self.neighbor_mask.to(hidden_states.device)
        neighbor_grid = neighbor_indices.index_select(0, active_cells_safe)
        neighbor_mask = base_neighbor_mask.index_select(0, active_cells_safe)
        neighbor_mask = neighbor_mask & active_valid.view(-1, 1)

        grid_to_seq = self.grid_to_seq.to(hidden_states.device)
        max_grid = int(grid_to_seq.shape[0])
        neighbor_grid_valid = (neighbor_grid >= 0) & (neighbor_grid < max_grid)
        neighbor_grid_safe = neighbor_grid.clamp(min=0, max=max(max_grid - 1, 0))

        neighbor_seq = grid_to_seq.index_select(0, neighbor_grid_safe.reshape(-1)).view_as(neighbor_grid_safe)
        neighbor_mask = neighbor_mask & neighbor_grid_valid & (neighbor_seq >= 0) & (neighbor_seq < seq_len)
        neighbor_seq = neighbor_seq.clamp(min=0, max=max(seq_len - 1, 0)).to(torch.long).contiguous()
        neighbor_mask = neighbor_mask.to(torch.bool).contiguous()

        y = self._forward_triton_sequence_gather(y, neighbor_seq, neighbor_mask)
        if self.spatial.bias is not None:
            y = y + self.spatial.bias.view(1, -1, 1)

        # 1x1 convs over active cells only; equivalent to Conv2d(kernel=1).
        y_tokens = y.transpose(1, 2).contiguous()  # [B, T, C]
        pointwise_weight = self.pointwise.weight.view(self.pointwise.out_channels, self.pointwise.in_channels)
        pointwise_bias = self.pointwise.bias
        gate, value = F.linear(y_tokens, pointwise_weight, pointwise_bias).chunk(2, dim=-1)
        out_weight = self.out.weight.view(self.out.out_channels, self.out.in_channels)
        out_bias = self.out.bias
        y_tokens = F.linear(self.act(gate) * value, out_weight, out_bias)
        return F.dropout(y_tokens, p=self.dropout.p, training=self.training)

    def forward_token(self, hidden_states: torch.Tensor, position_idx: int) -> torch.Tensor:
        """Exact causal 2D branch output for the last token only.

        `hidden_states` contains the complete prefix at this layer. The 2D
        branch is depthwise-spatial + pointwise 1x1, so the newest token's
        output only needs its causal spatial neighbours, not all grid cells.
        """
        if not self.use_depthwise:
            raise NotImplementedError("incremental 2D branch currently requires depthwise spatial conv")

        bsz, seq_len, channels = hidden_states.shape
        if position_idx < 0 or position_idx >= seq_len:
            raise ValueError(f"position_idx {position_idx} is out of range for seq_len {seq_len}")

        input_dtype = hidden_states.dtype
        y = hidden_states.float()
        variance = y.pow(2).mean(dim=-1, keepdim=True)
        y = hidden_states * torch.rsqrt(variance + self.norm.eps).to(dtype=hidden_states.dtype)
        y = y * self.norm.weight.view(1, 1, -1).to(device=hidden_states.device, dtype=input_dtype)

        seq_to_grid = self.seq_to_grid.to(hidden_states.device)
        max_cells = int(self.neighbor_indices.shape[0])
        active_cell = seq_to_grid.narrow(0, position_idx, 1).to(dtype=torch.long)
        active_valid = (active_cell >= 0) & (active_cell < max_cells)
        active_cell_safe = active_cell.clamp(min=0, max=max(max_cells - 1, 0))

        neighbor_grid = self.neighbor_indices.to(hidden_states.device).index_select(0, active_cell_safe).squeeze(0)
        neighbor_mask = self.neighbor_mask.to(hidden_states.device).index_select(0, active_cell_safe).squeeze(0)
        neighbor_mask = neighbor_mask & bool(active_valid.item())

        grid_to_seq = self.grid_to_seq.to(hidden_states.device)
        max_grid = int(grid_to_seq.shape[0])
        neighbor_grid_valid = (neighbor_grid >= 0) & (neighbor_grid < max_grid)
        neighbor_grid_safe = neighbor_grid.clamp(min=0, max=max(max_grid - 1, 0))

        neighbor_seq = grid_to_seq.index_select(0, neighbor_grid_safe)
        neighbor_mask = neighbor_mask & neighbor_grid_valid & (neighbor_seq >= 0) & (neighbor_seq < seq_len)
        neighbor_seq = neighbor_seq.clamp(min=0, max=max(seq_len - 1, 0)).to(torch.long)

        # [B, K*K, C]
        gathered = y.index_select(1, neighbor_seq)
        gathered = gathered * neighbor_mask.to(dtype=y.dtype).view(1, -1, 1)
        weight = self.spatial.weight.view(channels, self.kernel_size * self.kernel_size).to(dtype=y.dtype)
        y_token = (gathered * weight.transpose(0, 1).view(1, -1, channels)).sum(dim=1, keepdim=True)
        if self.spatial.bias is not None:
            y_token = y_token + self.spatial.bias.view(1, 1, -1)

        pointwise_weight = self.pointwise.weight.view(self.pointwise.out_channels, self.pointwise.in_channels)
        gate, value = F.linear(y_token, pointwise_weight, self.pointwise.bias).chunk(2, dim=-1)
        out_weight = self.out.weight.view(self.out.out_channels, self.out.in_channels)
        y_token = F.linear(self.act(gate) * value, out_weight, self.out.bias)
        return F.dropout(y_token, p=self.dropout.p, training=self.training)

    def _forward_masked_conv2d(self, y: torch.Tensor) -> torch.Tensor:
        # Fast exact path for row_major packing only. Config validation prevents
        # using this backend with arbitrary layouts.
        pad = (self.kernel_size - 1) * self.dilation
        y = F.pad(y, (pad // 2, pad - pad // 2, pad, 0))
        weight = self.spatial.weight * self.static_causal_mask.to(
            device=self.spatial.weight.device, dtype=self.spatial.weight.dtype
        )
        return F.conv2d(
            y,
            weight,
            bias=self.spatial.bias,
            stride=self.spatial.stride,
            padding=0,
            dilation=self.spatial.dilation,
            groups=self.spatial.groups,
        ).view(y.shape[0], self.spatial.out_channels, -1)

    def forward(self, grid: torch.Tensor) -> torch.Tensor:
        bsz, channels, height, width = grid.shape
        if height != self.grid_size or width != self.grid_size:
            raise ValueError(
                f"CausalConv2dBranch expected {self.grid_size}x{self.grid_size} grid, got {height}x{width}"
            )
        y = self.norm(grid)
        if self.backend == "unfold":
            y = self._forward_unfold(y)
        elif self.backend == "chunked_gather":
            y = self._forward_chunked_gather(y)
        elif self.backend == "triton_gather":
            y = self._forward_triton_gather(y)
        elif self.backend == "masked_conv2d":
            y = self._forward_masked_conv2d(y)
        else:
            raise ValueError(f"Unknown conv2d_backend: {self.backend}")

        if self.spatial.bias is not None:
            y = y + self.spatial.bias.view(1, -1, 1)
        y = y.view(bsz, self.spatial.out_channels, height, width)
        gate, value = self.pointwise(y).chunk(2, dim=1)
        y = self.out(self.act(gate) * value)
        return self.dropout(y)


class TinyRetrievalRouter(nn.Module):
    """Small top-k learned-memory router for cheap content-based long-range mixing."""

    def __init__(self, config: ConvGPTV2Config):
        super().__init__()
        d = config.hidden_size
        self.num_slots = config.retrieval_num_slots
        self.top_k = max(1, min(config.retrieval_top_k, self.num_slots))
        self.norm = ConvGPTV2RMSNorm(d, config.rms_norm_eps)
        self.query = nn.Linear(d, d, bias=False)
        self.key = nn.Parameter(torch.empty(self.num_slots, d))
        self.value = nn.Parameter(torch.empty(self.num_slots, d))
        self.out = nn.Linear(d, d, bias=False)
        self.dropout = nn.Dropout(config.retrieval_dropout)
        nn.init.normal_(self.key, mean=0.0, std=config.initializer_range)
        nn.init.normal_(self.value, mean=0.0, std=config.initializer_range)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        q = self.query(self.norm(hidden_states))
        scores = torch.matmul(q, self.key.t()) / math.sqrt(q.shape[-1])
        if self.top_k < self.num_slots:
            top_values, top_indices = scores.topk(self.top_k, dim=-1)
            weights = F.softmax(top_values, dim=-1).to(hidden_states.dtype)
            memory = self.value.index_select(0, top_indices.reshape(-1)).view(
                *top_indices.shape, hidden_states.shape[-1]
            )
            routed = (weights.unsqueeze(-1) * memory).sum(dim=-2)
        else:
            weights = F.softmax(scores, dim=-1).to(hidden_states.dtype)
            routed = torch.matmul(weights, self.value.to(hidden_states.dtype))
        return self.dropout(self.out(routed))


class CausalChunkMemoryRouter(nn.Module):
    """Causal top-k retrieval over prefix chunk summaries.

    This is a non-quadratic attention substitute: tokens query pooled chunk
    summaries rather than every previous token. With chunk size C, sequence
    length T, and top-k K, scoring cost is O(T * ceil(T/C)) and retrieval cost
    is O(T * K). The default mask excludes the token's current chunk so future
    tokens in the same chunk cannot leak backward.
    """

    def __init__(self, config: ConvGPTV2Config):
        super().__init__()
        d = config.hidden_size
        self.chunk_size = int(config.chunk_memory_size)
        self.top_k = int(config.chunk_memory_top_k)
        self.include_current_chunk = bool(config.chunk_memory_include_current_chunk)
        self.norm = ConvGPTV2RMSNorm(d, config.rms_norm_eps)
        self.query = nn.Linear(d, d, bias=False)
        self.key = nn.Linear(d, d, bias=False)
        self.value = nn.Linear(d, d, bias=False)
        self.out = nn.Linear(d, d, bias=False)
        self.gate = nn.Parameter(torch.tensor(float(config.chunk_memory_gate_init)))
        rope_dim = int(d * float(getattr(config, "router_rope_fraction", 1.0)))
        rope_dim = min(d, rope_dim - (rope_dim % 2))
        self.rotary_emb = ConvGPTV2RotaryEmbedding(rope_dim, getattr(config, "rope_theta", 10000.0)) if getattr(config, "position_embedding_type", "learned") == "rope_nope" else None
        self.dropout = nn.Dropout(config.retrieval_dropout)

    def _chunk_summaries(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        bsz, seq_len, dim = hidden_states.shape
        num_chunks = math.ceil(seq_len / self.chunk_size)
        pad_len = num_chunks * self.chunk_size - seq_len
        if pad_len:
            padded = F.pad(hidden_states, (0, 0, 0, pad_len))
        else:
            padded = hidden_states
        chunks = padded.view(bsz, num_chunks, self.chunk_size, dim)
        valid = torch.arange(num_chunks * self.chunk_size, device=hidden_states.device).view(1, num_chunks, self.chunk_size) < seq_len
        weights = valid.to(hidden_states.dtype).unsqueeze(-1)
        denom = weights.sum(dim=2).clamp_min(1.0)
        summaries = (chunks * weights).sum(dim=2) / denom
        return summaries, valid.any(dim=2)

    def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        bsz, seq_len, dim = hidden_states.shape
        if seq_len <= self.chunk_size and not self.include_current_chunk:
            return hidden_states.new_zeros(hidden_states.shape)

        normed = self.norm(hidden_states)
        summaries, valid_chunks = self._chunk_summaries(normed)
        num_chunks = summaries.shape[1]
        valid_chunks = valid_chunks.expand(bsz, -1)
        q = self.query(normed)
        k = self.key(summaries)
        if self.rotary_emb is not None:
            token_pos = _token_position_ids(position_ids, bsz, seq_len, hidden_states.device)
            chunk_pos = _chunk_start_position_ids(token_pos, self.chunk_size, num_chunks)
            q_cos, q_sin = self.rotary_emb(token_pos, q.dtype, q.device)
            k_cos, k_sin = self.rotary_emb(chunk_pos, k.dtype, k.device)
            q = _apply_rope_to_btd(q, q_cos, q_sin)
            k = _apply_rope_to_btd(k, k_cos, k_sin)
        v = self.value(summaries)
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(dim)

        token_chunks = torch.div(torch.arange(seq_len, device=hidden_states.device), self.chunk_size, rounding_mode="floor")
        chunk_ids = torch.arange(num_chunks, device=hidden_states.device).view(1, 1, num_chunks)
        max_allowed = token_chunks if self.include_current_chunk else token_chunks - 1
        causal_mask = chunk_ids <= max_allowed.view(1, seq_len, 1)
        causal_mask = causal_mask & valid_chunks.view(bsz, 1, num_chunks)
        scores = scores.masked_fill(~causal_mask, torch.finfo(scores.dtype).min)

        any_allowed = causal_mask.any(dim=-1, keepdim=True)
        safe_scores = torch.where(any_allowed, scores, torch.zeros_like(scores))
        k_top = max(1, min(self.top_k, num_chunks))
        top_scores, top_idx = safe_scores.topk(k_top, dim=-1)
        top_mask = causal_mask.gather(-1, top_idx)
        top_scores = top_scores.masked_fill(~top_mask, torch.finfo(top_scores.dtype).min)
        weights = F.softmax(top_scores, dim=-1).to(hidden_states.dtype)
        weights = torch.where(top_mask.any(dim=-1, keepdim=True), weights, torch.zeros_like(weights))
        gathered = v.gather(1, top_idx.reshape(bsz, -1, 1).expand(-1, -1, dim)).view(bsz, seq_len, k_top, dim)
        routed = (weights.unsqueeze(-1) * gathered).sum(dim=-2)
        return torch.sigmoid(self.gate).to(hidden_states.dtype) * self.dropout(self.out(routed))

    def forward_token(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        """Exact router output for the newest token only."""
        bsz, seq_len, dim = hidden_states.shape
        if seq_len <= self.chunk_size and not self.include_current_chunk:
            return hidden_states.new_zeros((bsz, 1, dim))

        normed = self.norm(hidden_states)
        summaries, valid_chunks = self._chunk_summaries(normed)
        num_chunks = summaries.shape[1]
        q = self.query(normed[:, -1:, :])
        k = self.key(summaries)
        if self.rotary_emb is not None:
            token_pos = _token_position_ids(position_ids, bsz, seq_len, hidden_states.device)
            chunk_pos = _chunk_start_position_ids(token_pos, self.chunk_size, num_chunks)
            q_cos, q_sin = self.rotary_emb(token_pos[:, -1:], q.dtype, q.device)
            k_cos, k_sin = self.rotary_emb(chunk_pos, k.dtype, k.device)
            q = _apply_rope_to_btd(q, q_cos, q_sin)
            k = _apply_rope_to_btd(k, k_cos, k_sin)
        v = self.value(summaries)
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(dim)

        token_chunk = (seq_len - 1) // self.chunk_size
        chunk_ids = torch.arange(num_chunks, device=hidden_states.device).view(1, 1, num_chunks)
        max_allowed = token_chunk if self.include_current_chunk else token_chunk - 1
        causal_mask = chunk_ids <= max_allowed
        causal_mask = causal_mask & valid_chunks.expand(bsz, -1).view(bsz, 1, num_chunks)
        scores = scores.masked_fill(~causal_mask, torch.finfo(scores.dtype).min)

        any_allowed = causal_mask.any(dim=-1, keepdim=True)
        safe_scores = torch.where(any_allowed, scores, torch.zeros_like(scores))
        k_top = max(1, min(self.top_k, num_chunks))
        top_scores, top_idx = safe_scores.topk(k_top, dim=-1)
        top_mask = causal_mask.gather(-1, top_idx)
        top_scores = top_scores.masked_fill(~top_mask, torch.finfo(top_scores.dtype).min)
        weights = F.softmax(top_scores, dim=-1).to(hidden_states.dtype)
        weights = torch.where(top_mask.any(dim=-1, keepdim=True), weights, torch.zeros_like(weights))
        gathered = v.gather(1, top_idx.reshape(bsz, -1, 1).expand(-1, -1, dim)).view(bsz, 1, k_top, dim)
        routed = (weights.unsqueeze(-1) * gathered).sum(dim=-2)
        return torch.sigmoid(self.gate).to(hidden_states.dtype) * self.dropout(self.out(routed))


class CausalChunkTokenMemoryRouter(nn.Module):
    """Two-stage causal sparse token memory.

    Stage 1 selects top-k previous chunks using pooled summaries. Stage 2 lets
    each token attend to the actual tokens inside selected chunks, preserving a
    high-bandwidth copy path without full quadratic token-token attention.
    """

    def __init__(self, config: ConvGPTV2Config):
        super().__init__()
        d = config.hidden_size
        self.chunk_size = int(config.chunk_memory_size)
        self.chunk_top_k = int(config.chunk_memory_top_k)
        self.token_top_k = int(config.chunk_memory_token_top_k)
        self.include_current_chunk = bool(config.chunk_memory_include_current_chunk)
        self.use_triton = bool(getattr(config, "chunk_memory_use_triton", True))
        self.norm = ConvGPTV2RMSNorm(d, config.rms_norm_eps)
        self.query = nn.Linear(d, d, bias=False)
        self.chunk_key = nn.Linear(d, d, bias=False)
        self.token_key = nn.Linear(d, d, bias=False)
        self.token_value = nn.Linear(d, d, bias=False)
        self.out = nn.Linear(d, d, bias=False)
        self.gate = nn.Parameter(torch.tensor(float(config.chunk_memory_gate_init)))
        rope_dim = int(d * float(getattr(config, "router_rope_fraction", 1.0)))
        rope_dim = min(d, rope_dim - (rope_dim % 2))
        self.rotary_emb = ConvGPTV2RotaryEmbedding(rope_dim, getattr(config, "rope_theta", 10000.0)) if getattr(config, "position_embedding_type", "learned") == "rope_nope" else None
        self.dropout = nn.Dropout(config.retrieval_dropout)

    def _chunk_summaries(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        bsz, seq_len, dim = hidden_states.shape
        num_chunks = math.ceil(seq_len / self.chunk_size)
        pad_len = num_chunks * self.chunk_size - seq_len
        padded = F.pad(hidden_states, (0, 0, 0, pad_len)) if pad_len else hidden_states
        chunks = padded.view(bsz, num_chunks, self.chunk_size, dim)
        valid = torch.arange(num_chunks * self.chunk_size, device=hidden_states.device).view(1, num_chunks, self.chunk_size) < seq_len
        weights = valid.to(hidden_states.dtype).unsqueeze(-1)
        summaries = (chunks * weights).sum(dim=2) / weights.sum(dim=2).clamp_min(1.0)
        return summaries, valid.expand(bsz, -1, -1)

    def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        bsz, seq_len, dim = hidden_states.shape
        if seq_len <= self.chunk_size and not self.include_current_chunk:
            return hidden_states.new_zeros(hidden_states.shape)

        normed = self.norm(hidden_states)
        summaries, valid_token_chunks = self._chunk_summaries(normed)
        num_chunks = summaries.shape[1]
        q = self.query(normed)
        chunk_k = self.chunk_key(summaries)
        token_k = self.token_key(normed)
        if self.rotary_emb is not None:
            token_pos = _token_position_ids(position_ids, bsz, seq_len, hidden_states.device)
            chunk_pos = _chunk_start_position_ids(token_pos, self.chunk_size, num_chunks)
            q_cos, q_sin = self.rotary_emb(token_pos, q.dtype, q.device)
            tk_cos, tk_sin = self.rotary_emb(token_pos, token_k.dtype, token_k.device)
            ck_cos, ck_sin = self.rotary_emb(chunk_pos, chunk_k.dtype, chunk_k.device)
            q = _apply_rope_to_btd(q, q_cos, q_sin)
            token_k = _apply_rope_to_btd(token_k, tk_cos, tk_sin)
            chunk_k = _apply_rope_to_btd(chunk_k, ck_cos, ck_sin)
        token_v = self.token_value(normed)

        top_scores = top_pos = top_valid = None
        flat_pos = flat_valid = None
        k_chunks = max(1, min(self.chunk_top_k, num_chunks))

        can_use_chunk2_triton = (
            False  # experimental path currently benchmark-neutral/slower
            and self.use_triton
            and _TRITON_AVAILABLE
            and q.device.type == "cuda"
            and q.dtype in {torch.float16, torch.bfloat16, torch.float32}
            and chunk_k.dtype == q.dtype
            and token_k.dtype == q.dtype
            and token_v.dtype == q.dtype
            and self.chunk_top_k == 2
            and k_chunks == 2
            and self.token_top_k > 0
            and self.token_top_k < self.chunk_top_k * self.chunk_size
            and self.chunk_top_k * self.chunk_size <= 256
            and dim <= 1024
        )
        if can_use_chunk2_triton:
            top_scores, top_pos, top_valid = _select_chunk2_token_scores(
                q,
                chunk_k,
                token_k,
                self.token_top_k,
                self.chunk_size,
                self.include_current_chunk,
            )
        else:
            chunk_scores = torch.matmul(q, chunk_k.transpose(-1, -2)) / math.sqrt(dim)

            token_pos = torch.arange(seq_len, device=hidden_states.device)
            token_chunks = torch.div(token_pos, self.chunk_size, rounding_mode="floor")
            chunk_ids = torch.arange(num_chunks, device=hidden_states.device).view(1, 1, num_chunks)
            max_allowed = token_chunks if self.include_current_chunk else token_chunks - 1
            chunk_mask = chunk_ids <= max_allowed.view(1, seq_len, 1)
            chunk_mask = chunk_mask & valid_token_chunks.any(dim=-1).view(bsz, 1, num_chunks)
            chunk_scores = chunk_scores.masked_fill(~chunk_mask, torch.finfo(chunk_scores.dtype).min)
            any_chunk = chunk_mask.any(dim=-1, keepdim=True)
            safe_chunk_scores = torch.where(any_chunk, chunk_scores, torch.zeros_like(chunk_scores))
            _top_chunk_scores, top_chunk_idx = safe_chunk_scores.topk(k_chunks, dim=-1)
            top_chunk_mask = chunk_mask.gather(-1, top_chunk_idx)

            offsets = torch.arange(self.chunk_size, device=hidden_states.device).view(1, 1, 1, self.chunk_size)
            candidate_pos = top_chunk_idx.unsqueeze(-1) * self.chunk_size + offsets
            candidate_valid = candidate_pos < seq_len
            candidate_valid = candidate_valid & top_chunk_mask.unsqueeze(-1)
            if not self.include_current_chunk:
                candidate_valid = candidate_valid & (candidate_pos < token_pos.view(1, seq_len, 1, 1))
            candidate_pos = candidate_pos.clamp(max=max(seq_len - 1, 0)).to(torch.long)
            flat_pos = candidate_pos.reshape(bsz, seq_len, k_chunks * self.chunk_size)
            flat_valid = candidate_valid.reshape(bsz, seq_len, k_chunks * self.chunk_size)

        if top_scores is None:
            use_triton = self.use_triton and _can_use_chunk_token_triton(q, token_k, token_v, flat_pos)
            any_token = flat_valid.any(dim=-1, keepdim=True)
            if self.token_top_k > 0 and self.token_top_k < flat_pos.shape[-1]:
                k_tokens = self.token_top_k
                if use_triton:
                    top_scores, top_pos, top_valid = _select_chunk_token_scores(q, token_k, flat_pos, flat_valid, k_tokens)
                    top_scores = top_scores.masked_fill(~top_valid, torch.finfo(top_scores.dtype).min)
                else:
                    gather_idx = flat_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
                    cand_k = token_k.gather(1, gather_idx).view(bsz, seq_len, flat_pos.shape[-1], dim)
                    token_scores = (q.unsqueeze(-2) * cand_k).sum(dim=-1) / math.sqrt(dim)
                    token_scores = token_scores.masked_fill(~flat_valid, torch.finfo(token_scores.dtype).min)
                    safe_token_scores = torch.where(any_token, token_scores, torch.zeros_like(token_scores))
                    top_scores, top_idx = safe_token_scores.topk(k_tokens, dim=-1)
                    top_valid = flat_valid.gather(-1, top_idx)
                    top_pos = flat_pos.gather(-1, top_idx)
                    top_scores = top_scores.masked_fill(~top_valid, torch.finfo(top_scores.dtype).min)
                weights = F.softmax(top_scores, dim=-1).to(hidden_states.dtype)
                weights = torch.where(top_valid.any(dim=-1, keepdim=True), weights, torch.zeros_like(weights))
            else:
                if use_triton:
                    token_scores = _chunk_token_scores(q, token_k, flat_pos, flat_valid)
                else:
                    gather_idx = flat_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
                    cand_k = token_k.gather(1, gather_idx).view(bsz, seq_len, flat_pos.shape[-1], dim)
                    token_scores = (q.unsqueeze(-2) * cand_k).sum(dim=-1) / math.sqrt(dim)
                    token_scores = token_scores.masked_fill(~flat_valid, torch.finfo(token_scores.dtype).min)
                top_valid = flat_valid
                top_pos = flat_pos
                weights = F.softmax(token_scores, dim=-1).to(hidden_states.dtype)
                weights = torch.where(any_token, weights, torch.zeros_like(weights))
        else:
            top_scores = top_scores.masked_fill(~top_valid, torch.finfo(top_scores.dtype).min)
            weights = F.softmax(top_scores, dim=-1).to(hidden_states.dtype)
            weights = torch.where(top_valid.any(dim=-1, keepdim=True), weights, torch.zeros_like(weights))

        use_triton_value = self.use_triton and _can_use_chunk_token_triton(token_v, token_v, token_v, top_pos)
        if use_triton_value:
            routed = _weighted_token_values(weights, token_v, top_pos, top_valid)
        else:
            gather_idx = top_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
            cand_v = token_v.gather(1, gather_idx).view(bsz, seq_len, top_pos.shape[-1], dim)
            routed = (weights.unsqueeze(-1) * cand_v).sum(dim=-2)
        return torch.sigmoid(self.gate).to(hidden_states.dtype) * self.dropout(self.out(routed))

    def forward_token(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        """Exact sparse token-memory router output for the newest token only."""
        bsz, seq_len, dim = hidden_states.shape
        if seq_len <= self.chunk_size and not self.include_current_chunk:
            return hidden_states.new_zeros((bsz, 1, dim))

        normed = self.norm(hidden_states)
        summaries, valid_token_chunks = self._chunk_summaries(normed)
        num_chunks = summaries.shape[1]
        q = self.query(normed[:, -1:, :])
        chunk_k = self.chunk_key(summaries)
        token_k = self.token_key(normed)
        if self.rotary_emb is not None:
            token_pos_ids = _token_position_ids(position_ids, bsz, seq_len, hidden_states.device)
            chunk_pos = _chunk_start_position_ids(token_pos_ids, self.chunk_size, num_chunks)
            q_cos, q_sin = self.rotary_emb(token_pos_ids[:, -1:], q.dtype, q.device)
            tk_cos, tk_sin = self.rotary_emb(token_pos_ids, token_k.dtype, token_k.device)
            ck_cos, ck_sin = self.rotary_emb(chunk_pos, chunk_k.dtype, chunk_k.device)
            q = _apply_rope_to_btd(q, q_cos, q_sin)
            token_k = _apply_rope_to_btd(token_k, tk_cos, tk_sin)
            chunk_k = _apply_rope_to_btd(chunk_k, ck_cos, ck_sin)
        token_v = self.token_value(normed)

        k_chunks = max(1, min(self.chunk_top_k, num_chunks))
        chunk_scores = torch.matmul(q, chunk_k.transpose(-1, -2)) / math.sqrt(dim)
        token_idx = seq_len - 1
        token_chunk = token_idx // self.chunk_size
        chunk_ids = torch.arange(num_chunks, device=hidden_states.device).view(1, 1, num_chunks)
        max_allowed = token_chunk if self.include_current_chunk else token_chunk - 1
        chunk_mask = chunk_ids <= max_allowed
        chunk_mask = chunk_mask & valid_token_chunks.any(dim=-1).view(bsz, 1, num_chunks)
        chunk_scores = chunk_scores.masked_fill(~chunk_mask, torch.finfo(chunk_scores.dtype).min)
        any_chunk = chunk_mask.any(dim=-1, keepdim=True)
        safe_chunk_scores = torch.where(any_chunk, chunk_scores, torch.zeros_like(chunk_scores))
        _top_chunk_scores, top_chunk_idx = safe_chunk_scores.topk(k_chunks, dim=-1)
        top_chunk_mask = chunk_mask.gather(-1, top_chunk_idx)

        offsets = torch.arange(self.chunk_size, device=hidden_states.device).view(1, 1, 1, self.chunk_size)
        candidate_pos = top_chunk_idx.unsqueeze(-1) * self.chunk_size + offsets
        candidate_valid = candidate_pos < seq_len
        candidate_valid = candidate_valid & top_chunk_mask.unsqueeze(-1)
        if not self.include_current_chunk:
            candidate_valid = candidate_valid & (candidate_pos < token_idx)
        candidate_pos = candidate_pos.clamp(max=max(seq_len - 1, 0)).to(torch.long)
        flat_pos = candidate_pos.reshape(bsz, 1, k_chunks * self.chunk_size)
        flat_valid = candidate_valid.reshape(bsz, 1, k_chunks * self.chunk_size)
        any_token = flat_valid.any(dim=-1, keepdim=True)

        gather_idx = flat_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
        cand_k = token_k.gather(1, gather_idx).view(bsz, 1, flat_pos.shape[-1], dim)
        token_scores = (q.unsqueeze(-2) * cand_k).sum(dim=-1) / math.sqrt(dim)
        token_scores = token_scores.masked_fill(~flat_valid, torch.finfo(token_scores.dtype).min)

        if self.token_top_k > 0 and self.token_top_k < flat_pos.shape[-1]:
            k_tokens = self.token_top_k
            safe_token_scores = torch.where(any_token, token_scores, torch.zeros_like(token_scores))
            top_scores, top_idx = safe_token_scores.topk(k_tokens, dim=-1)
            top_valid = flat_valid.gather(-1, top_idx)
            top_pos = flat_pos.gather(-1, top_idx)
            top_scores = top_scores.masked_fill(~top_valid, torch.finfo(top_scores.dtype).min)
            weights = F.softmax(top_scores, dim=-1).to(hidden_states.dtype)
            weights = torch.where(top_valid.any(dim=-1, keepdim=True), weights, torch.zeros_like(weights))
        else:
            top_valid = flat_valid
            top_pos = flat_pos
            weights = F.softmax(token_scores, dim=-1).to(hidden_states.dtype)
            weights = torch.where(any_token, weights, torch.zeros_like(weights))

        gather_idx = top_pos.reshape(bsz, -1, 1).expand(-1, -1, dim)
        cand_v = token_v.gather(1, gather_idx).view(bsz, 1, top_pos.shape[-1], dim)
        routed = (weights.unsqueeze(-1) * cand_v).sum(dim=-2)
        return torch.sigmoid(self.gate).to(hidden_states.dtype) * self.dropout(self.out(routed))

class HybridConvGPTV2Block(nn.Module):
    def __init__(self, config: ConvGPTV2Config, layer_idx: int):
        super().__init__()
        self.use_1d_branch = config.use_1d_branch
        self.use_2d_branch = (
            config.use_2d_branch
            and layer_idx >= config.two_d_start_layer
            and ((layer_idx - config.two_d_start_layer) % config.two_d_every == 0)
        )
        if not self.use_1d_branch and not self.use_2d_branch:
            raise ValueError("At least one active branch is required; check use_1d_branch/use_2d_branch/two_d_every")

        d1 = config.conv1d_dilations or [1, 2, 4, 8, 16, 32]
        d2 = config.conv2d_dilations or [1, 2, 4, 8, 16]
        self.branch_1d = CausalConv1dBranch(config, d1[layer_idx % len(d1)]) if self.use_1d_branch else None
        self.branch_2d = CausalConv2dBranch(config, d2[layer_idx % len(d2)]) if self.use_2d_branch else None
        self.packer = SpaceFillingPacker(config) if self.use_2d_branch else None
        self.fusion = config.fusion
        branch_count = int(self.use_1d_branch) + int(self.use_2d_branch)
        if self.fusion == "concat" and branch_count > 1:
            self.fuse_proj = nn.Linear(config.hidden_size * branch_count, config.hidden_size, bias=False)
        elif self.fusion == "gated" and branch_count > 1:
            self.fuse_gate = nn.Linear(config.hidden_size, branch_count, bias=True)
        else:
            self.fuse_proj = None
            self.fuse_gate = None

        self.branch_dropout = nn.Dropout(config.branch_dropout)
        self.gate_1d = nn.Parameter(torch.tensor(float(config.conv1d_residual_gate_init))) if self.use_1d_branch else None
        self.gate_2d = nn.Parameter(torch.tensor(float(config.conv2d_residual_gate_init))) if self.use_2d_branch else None

        self.post_conv_norm = ConvGPTV2RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.mlp = ConvGPTV2MLP(config)
        self.use_retrieval = (
            config.router_type != "none"
            and config.retrieval_every is not None
            and config.retrieval_every > 0
            and (layer_idx + 1) % config.retrieval_every == 0
        )
        if self.use_retrieval and config.router_type == "topk_memory":
            self.router = TinyRetrievalRouter(config)
        elif self.use_retrieval and config.router_type == "chunk_memory":
            self.router = CausalChunkMemoryRouter(config)
        elif self.use_retrieval and config.router_type == "chunk_token_memory":
            self.router = CausalChunkTokenMemoryRouter(config)
        else:
            self.router = None

    def _fuse(self, hidden_states: torch.Tensor, named_branches: list[tuple[str, torch.Tensor]]) -> torch.Tensor:
        if len(named_branches) == 1:
            name, branch = named_branches[0]
            gate = self.gate_1d if name == "1d" else self.gate_2d
            return torch.sigmoid(gate).to(branch.dtype) * branch

        branches = [branch for _, branch in named_branches]
        if self.fusion == "sum":
            mixed = torch.stack(branches, dim=0).sum(dim=0)
        elif self.fusion == "concat":
            mixed = self.fuse_proj(torch.cat(branches, dim=-1))
        else:
            gates = F.softmax(self.fuse_gate(hidden_states), dim=-1).to(hidden_states.dtype)
            mixed = hidden_states.new_zeros(hidden_states.shape)
            for idx, branch in enumerate(branches):
                mixed = mixed + gates[..., idx : idx + 1] * branch

        residual_scale = hidden_states.new_tensor(0.0)
        for name, _branch in named_branches:
            gate = self.gate_1d if name == "1d" else self.gate_2d
            residual_scale = residual_scale + torch.sigmoid(gate).to(hidden_states.dtype)
        residual_scale = residual_scale / len(named_branches)
        return residual_scale * mixed

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: Optional[torch.LongTensor] = None,
        return_post_conv: bool = False,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        branches = []
        if self.branch_1d is not None:
            branches.append(("1d", self.branch_1d(hidden_states)))
        if self.branch_2d is not None:
            active_cells = self.packer.seq_to_grid[: hidden_states.shape[1]]
            seq_delta = self.branch_2d.forward_sequence(hidden_states, active_cells)
            if seq_delta is not None:
                branches.append(("2d", seq_delta))
            else:
                grid, seq_len = self.packer.seq_to_2d(hidden_states)
                grid_delta = self.branch_2d(grid)
                branches.append(("2d", self.packer.d2_to_seq(grid_delta, seq_len)))

        hidden_states = hidden_states + self.branch_dropout(self._fuse(hidden_states, branches))
        post_conv_states = hidden_states
        if self.router is not None:
            hidden_states = hidden_states + self.router(hidden_states, position_ids=position_ids)
        hidden_states = hidden_states + self.mlp(self.post_conv_norm(hidden_states))
        if return_post_conv:
            return hidden_states, post_conv_states
        return hidden_states

    def forward_token(
        self,
        hidden_states: torch.Tensor,
        position_ids: Optional[torch.LongTensor] = None,
        post_conv_prefix: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Exact block output for the newest token only.

        Returns `(new_hidden_state, post_conv_hidden_state)` for the last token.
        The caller owns appending these one-token states to the cache.
        """
        branches = []
        if self.branch_1d is not None:
            branches.append(("1d", self.branch_1d.forward_token(hidden_states)))
        if self.branch_2d is not None:
            branches.append(("2d", self.branch_2d.forward_token(hidden_states, hidden_states.shape[1] - 1)))

        last_hidden = hidden_states[:, -1:, :]
        token_states = last_hidden + self.branch_dropout(self._fuse(last_hidden, branches))
        post_conv_states = token_states
        if self.router is not None:
            router_input = token_states if post_conv_prefix is None else torch.cat(
                [post_conv_prefix.to(token_states.device), token_states], dim=1
            )
            token_states = token_states + self.router.forward_token(router_input, position_ids=position_ids)
        token_states = token_states + self.mlp(self.post_conv_norm(token_states))
        return token_states, post_conv_states


class ConvGPTV2PreTrainedModel(PreTrainedModel):
    config_class = ConvGPTV2Config
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["HybridConvGPTV2Block"]

    def _init_weights(self, module: nn.Module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d)):
            module.weight.data.normal_(mean=0.0, std=std)
            if getattr(module, "bias", None) is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


class ConvGPTV2Model(ConvGPTV2PreTrainedModel):
    def __init__(self, config: ConvGPTV2Config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) if config.position_embedding_type == "learned" else None
        if config.use_row_col_embeddings and config.position_embedding_type == "learned":
            self.row_embeddings = nn.Embedding(config.grid_size, config.hidden_size)
            self.col_embeddings = nn.Embedding(config.grid_size, config.hidden_size)
            seq_to_grid, _ = build_curve_indices(config.grid_size, config.packing, config.pack_order)
            row_ids = torch.div(seq_to_grid[: config.max_position_embeddings], config.grid_size, rounding_mode="floor")
            col_ids = seq_to_grid[: config.max_position_embeddings] % config.grid_size
            self.register_buffer("position_row_ids", row_ids, persistent=False)
            self.register_buffer("position_col_ids", col_ids, persistent=False)
        else:
            self.row_embeddings = None
            self.col_embeddings = None
        self.dropout = nn.Dropout(config.dropout)
        self.layers = nn.ModuleList(
            [HybridConvGPTV2Block(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = ConvGPTV2RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.post_init()

    def _embed_inputs(
        self,
        input_ids: Optional[torch.LongTensor],
        inputs_embeds: Optional[torch.FloatTensor],
        position_ids: Optional[torch.LongTensor],
    ) -> Tuple[torch.Tensor, torch.LongTensor]:
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds")
        if input_ids is None and inputs_embeds is None:
            raise ValueError("You must specify input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        bsz, seq_len, _ = inputs_embeds.shape
        if seq_len > self.config.max_position_embeddings:
            raise ValueError(
                f"Sequence length {seq_len} exceeds max_position_embeddings {self.config.max_position_embeddings}."
            )

        if position_ids is None:
            position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0).expand(bsz, -1)
        hidden_states = inputs_embeds
        if self.position_embeddings is not None:
            hidden_states = hidden_states + self.position_embeddings(position_ids)
        if self.row_embeddings is not None:
            rows = self.position_row_ids.to(inputs_embeds.device).index_select(0, position_ids.reshape(-1)).view_as(position_ids)
            cols = self.position_col_ids.to(inputs_embeds.device).index_select(0, position_ids.reshape(-1)).view_as(position_ids)
            hidden_states = hidden_states + self.row_embeddings(rows) + self.col_embeddings(cols)
        return self.dropout(hidden_states), position_ids

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[torch.Tensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        hidden_states, position_ids = self._embed_inputs(input_ids, inputs_embeds, position_ids)

        all_hidden_states = () if output_hidden_states else None
        layer_inputs = [] if use_cache and not self.training else None
        layer_post_conv = [] if use_cache and not self.training else None
        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)
            if layer_inputs is not None:
                layer_inputs.append(hidden_states.detach())
            if self.gradient_checkpointing and self.training:
                hidden_states = self._gradient_checkpointing_func(layer.__call__, hidden_states)
            else:
                if layer_post_conv is not None:
                    hidden_states, post_conv_states = layer(hidden_states, position_ids=position_ids, return_post_conv=True)
                    layer_post_conv.append(post_conv_states.detach())
                else:
                    hidden_states = layer(hidden_states, position_ids=position_ids)
            if attention_mask is not None:
                hidden_states = hidden_states * attention_mask.to(hidden_states.dtype).unsqueeze(-1)

        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, None, all_hidden_states, None] if v is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=ConvGPTV2PrefixCache(
                input_ids=input_ids.detach(),
                attention_mask=attention_mask.detach() if attention_mask is not None else None,
                layer_inputs=tuple(layer_inputs),
                layer_post_conv=tuple(layer_post_conv),
            ) if layer_inputs is not None and input_ids is not None else None,
            hidden_states=all_hidden_states,
            attentions=None if output_attentions else None,
        )

    @torch.no_grad()
    def incremental_forward(
        self,
        input_ids: torch.LongTensor,
        cache: ConvGPTV2PrefixCache,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        output_hidden_states: bool = False,
        output_attentions: bool = False,
        return_dict: bool = True,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        if input_ids.shape[1] != 1:
            raise ValueError("ConvGPT-v2 incremental_forward expects exactly one new token")
        if not cache.has_incremental_state:
            raise ValueError("ConvGPT-v2 incremental_forward requires per-layer cache state")
        if cache.layer_inputs is None or cache.layer_post_conv is None:
            raise ValueError("ConvGPT-v2 incremental_forward requires complete cache state")
        if len(cache.layer_inputs) != len(self.layers) or len(cache.layer_post_conv) != len(self.layers):
            raise ValueError("ConvGPT-v2 cache layer count does not match model layer count")

        bsz = input_ids.shape[0]
        prefix_len = cache.input_ids.shape[1]
        full_len = prefix_len + 1
        if full_len > self.config.max_position_embeddings:
            raise ValueError(
                f"Sequence length {full_len} exceeds max_position_embeddings {self.config.max_position_embeddings}."
            )
        if position_ids is None:
            position_ids = torch.full((bsz, 1), prefix_len, device=input_ids.device, dtype=torch.long)
        token_states, _ = self._embed_inputs(input_ids, None, position_ids)

        new_layer_inputs = []
        new_layer_post_conv = []
        all_hidden_states = () if output_hidden_states else None
        for layer_idx, layer in enumerate(self.layers):
            layer_input_full = torch.cat([cache.layer_inputs[layer_idx].to(token_states.device), token_states], dim=1)
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (layer_input_full,)
            token_states, post_conv_token = layer.forward_token(
                layer_input_full,
                position_ids=None,
                post_conv_prefix=cache.layer_post_conv[layer_idx],
            )
            if attention_mask is not None:
                token_mask = attention_mask[:, -1:].to(token_states.device)
                token_states = token_states * token_mask.to(token_states.dtype).unsqueeze(-1)
                post_conv_token = post_conv_token * token_mask.to(post_conv_token.dtype).unsqueeze(-1)
            new_layer_inputs.append(layer_input_full.detach())
            new_layer_post_conv.append(torch.cat([cache.layer_post_conv[layer_idx].to(post_conv_token.device), post_conv_token], dim=1).detach())

        hidden_states = self.norm(token_states)
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)
        next_cache = ConvGPTV2PrefixCache(
            input_ids=torch.cat([cache.input_ids.to(input_ids.device), input_ids], dim=1).detach(),
            attention_mask=attention_mask.detach() if attention_mask is not None else None,
            layer_inputs=tuple(new_layer_inputs),
            layer_post_conv=tuple(new_layer_post_conv),
        )

        if not return_dict:
            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, None] if v is not None)
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=None if output_attentions else None,
        )

    @torch.no_grad()
    def build_prefix_cache(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
    ) -> ConvGPTV2PrefixCache:
        hidden_states, position_ids = self._embed_inputs(input_ids, None, position_ids)
        layer_inputs = []
        layer_post_conv = []
        for layer in self.layers:
            layer_inputs.append(hidden_states.detach())
            branches = []
            if layer.branch_1d is not None:
                branches.append(("1d", layer.branch_1d(hidden_states)))
            if layer.branch_2d is not None:
                active_cells = layer.packer.seq_to_grid[: hidden_states.shape[1]]
                seq_delta = layer.branch_2d.forward_sequence(hidden_states, active_cells)
                if seq_delta is not None:
                    branches.append(("2d", seq_delta))
                else:
                    grid, seq_len = layer.packer.seq_to_2d(hidden_states)
                    grid_delta = layer.branch_2d(grid)
                    branches.append(("2d", layer.packer.d2_to_seq(grid_delta, seq_len)))
            hidden_states = hidden_states + layer.branch_dropout(layer._fuse(hidden_states, branches))
            layer_post_conv.append(hidden_states.detach())
            if layer.router is not None:
                hidden_states = hidden_states + layer.router(hidden_states, position_ids=position_ids)
            hidden_states = hidden_states + layer.mlp(layer.post_conv_norm(hidden_states))
            if attention_mask is not None:
                hidden_states = hidden_states * attention_mask.to(hidden_states.dtype).unsqueeze(-1)
        return ConvGPTV2PrefixCache(
            input_ids=input_ids.detach(),
            attention_mask=attention_mask.detach() if attention_mask is not None else None,
            layer_inputs=tuple(layer_inputs),
            layer_post_conv=tuple(layer_post_conv),
        )


class ConvGPTV2ForCausalLM(ConvGPTV2PreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    _supports_cache_class = False

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

    def get_input_embeddings(self):
        return self.model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.lm_head

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

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[torch.Tensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cache_position: Optional[torch.LongTensor] = None,
        logits_to_keep: Union[int, torch.Tensor] = 0,
        **loss_kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        incoming_cache = past_key_values
        ran_incremental = False
        if use_cache and incoming_cache is not None and input_ids is not None:
            if not isinstance(incoming_cache, ConvGPTV2PrefixCache):
                incoming_cache = ConvGPTV2PrefixCache.from_legacy_cache(incoming_cache)
            if attention_mask is None and incoming_cache.attention_mask is not None:
                attention_mask = incoming_cache.attention_mask.to(input_ids.device)
            if input_ids.shape[1] == 1 and incoming_cache.has_incremental_state and inputs_embeds is None and labels is None and not self.training:
                outputs = self.model.incremental_forward(
                    input_ids=input_ids,
                    cache=incoming_cache,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    output_hidden_states=bool(output_hidden_states),
                    output_attentions=bool(output_attentions),
                    return_dict=return_dict,
                )
                ran_incremental = True
            else:
                prefix_ids = incoming_cache.input_ids.to(input_ids.device)
                if input_ids.shape[1] > 0:
                    # Fallback path for legacy/unsupported cache states: rebuild
                    # the exact full prefix.
                    if input_ids.shape[1] == 1 or not torch.equal(input_ids[:, : prefix_ids.shape[1]], prefix_ids):
                        input_ids = torch.cat([prefix_ids, input_ids], dim=1)

        if not ran_incremental:
            outputs = self.model(
                input_ids=input_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_values=past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=use_cache,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
                cache_position=cache_position,
            )
        hidden_states = outputs[0]
        if isinstance(logits_to_keep, int) and logits_to_keep > 0:
            hidden_states_for_logits = hidden_states[:, -logits_to_keep:, :]
        elif torch.is_tensor(logits_to_keep):
            hidden_states_for_logits = hidden_states[:, logits_to_keep, :]
        else:
            hidden_states_for_logits = hidden_states
        logits = self.lm_head(hidden_states_for_logits)

        loss = None
        if labels is not None:
            if logits.shape[1] != labels.shape[1]:
                # logits_to_keep is an inference optimization; for training labels
                # we need full logits to align next-token loss.
                logits = self.lm_head(hidden_states)
            
            loss = self.loss_function(
                logits=logits,
                labels=labels,
                vocab_size=self.config.vocab_size,
                **loss_kwargs,
            )
            num_items = loss_kwargs.get("num_items_in_batch")
            if num_items is not None:
                loss = loss / num_items.to(loss.device)

        next_cache = None
        if use_cache:
            if ran_incremental:
                next_cache = outputs.past_key_values if return_dict else outputs[1]
            elif input_ids is None:
                raise ValueError("ConvGPT-v2 cache generation requires input_ids; inputs_embeds-only caching is unsupported.")
            elif not self.training:
                model_cache = outputs.past_key_values if return_dict and hasattr(outputs, "past_key_values") else None
                if isinstance(model_cache, ConvGPTV2PrefixCache):
                    next_cache = model_cache
                else:
                    next_cache = self.model.build_prefix_cache(input_ids, attention_mask=attention_mask, position_ids=position_ids)
            else:
                next_cache = ConvGPTV2PrefixCache(input_ids=input_ids.detach(), attention_mask=attention_mask.detach() if attention_mask is not None else None)

        if not return_dict:
            output = (logits,) + outputs[1:]
            if use_cache:
                output = output[:1] + (next_cache,) + output[1:]
            return ((loss,) + output) if loss is not None else output
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=next_cache,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    @classmethod
    def _supports_default_dynamic_cache(cls) -> bool:
        return False

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        **kwargs,
    ):
        # Cache stores the accepted prefix plus layer states. After prefill,
        # pass only the newest token; forward uses the incremental layer cache
        # when available and falls back to exact full-prefix recompute otherwise.
        if past_key_values is not None and input_ids is not None:
            model_inputs = {"input_ids": input_ids[:, -1:]}
        else:
            model_inputs = {"input_ids": input_ids}
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache", self.config.use_cache),
                "attention_mask": attention_mask,
            }
        )
        return model_inputs

    def _reorder_cache(self, past_key_values, beam_idx):
        if past_key_values is None:
            return None
        if isinstance(past_key_values, ConvGPTV2PrefixCache):
            return past_key_values.reorder_cache(beam_idx)
        return ConvGPTV2PrefixCache.from_legacy_cache(past_key_values).reorder_cache(beam_idx)


__all__ = [
    "ConvGPTV2Config",
    "ConvGPTV2Model",
    "ConvGPTV2ForCausalLM",
    "ConvGPTV2PreTrainedModel",
    "build_curve_indices",
]