# 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", ]