decoder.py

from dataclasses import dataclass
import os
from typing import Optional, Tuple

import torch
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline
from learnorm import LeaRNorm
from layout import ContinuousSnakeLayout


_CCE_EXT = None


def _bf16_prefers_native_mm(hidden: torch.Tensor) -> bool:
    if hidden.dtype != torch.bfloat16 or not hidden.is_cuda:
        return False
    if not hasattr(torch.cuda, "is_bf16_supported"):
        return False
    try:
        if not torch.cuda.is_bf16_supported():
            return False
        major, _minor = torch.cuda.get_device_capability(hidden.device)
    except Exception:
        return False
    return major >= 8


def _get_cce_ext():
    global _CCE_EXT
    if _CCE_EXT is not None:
        return _CCE_EXT

    if not torch.cuda.is_available():
        return None

    # Set the arch list for the current GPU.
    if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
        try:
            caps = sorted({torch.cuda.get_device_capability(i)
                           for i in range(torch.cuda.device_count())})
            archs = [f"{ma}.{mi}" for ma, mi in caps]
            archs[-1] = archs[-1] + "+PTX"
            os.environ["TORCH_CUDA_ARCH_LIST"] = ";".join(archs)
        except Exception:
            pass

    cpp_src = r"""
#include <torch/extension.h>
torch::Tensor fused_row_max_sumexp(torch::Tensor logits, int64_t n, int64_t c);
torch::Tensor fused_exp_sub(torch::Tensor logits, torch::Tensor bias);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("fused_row_max_sumexp", &fused_row_max_sumexp);
  m.def("fused_exp_sub", &fused_exp_sub);
}
"""
    cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cfloat>

// One block per row. Each thread handles ceil(C/blockDim.x) elements.
// Output: [N, 2] with col0=max and col1=sum_exp. Single pass, online.
__global__ void fused_row_max_sumexp_kernel(
    const float* __restrict__ logits,
    float* __restrict__ out,
    int64_t N, int64_t C) {
  int64_t row = blockIdx.x;
  if (row >= N) return;
  const float* row_ptr = logits + row * C;

  // Online max + sum_exp per thread.
  float local_max = -FLT_MAX;
  float local_sum = 0.0f;
  for (int64_t j = threadIdx.x; j < C; j += blockDim.x) {
    float val = row_ptr[j];
    if (val > local_max) {
      local_sum *= expf(local_max - val);
      local_max = val;
    }
    local_sum += expf(val - local_max);
  }

  // Warp-shuffle reduction with online merge.
  for (int off = 16; off > 0; off >>= 1) {
    float o_max = __shfl_down_sync(0xFFFFFFFF, local_max, off);
    float o_sum = __shfl_down_sync(0xFFFFFFFF, local_sum, off);
    float new_max = fmaxf(local_max, o_max);
    local_sum = local_sum * expf(local_max - new_max)
              + o_sum * expf(o_max - new_max);
    local_max = new_max;
  }

  // Cross-warp reduction via shared memory.
  __shared__ float s_max[32], s_sum[32];
  int lane = threadIdx.x & 31;
  int wid = threadIdx.x >> 5;
  if (lane == 0) { s_max[wid] = local_max; s_sum[wid] = local_sum; }
  __syncthreads();

  int nwarps = (blockDim.x + 31) >> 5;
  if (wid == 0) {
    local_max = (lane < nwarps) ? s_max[lane] : -FLT_MAX;
    local_sum = (lane < nwarps) ? s_sum[lane] : 0.0f;
    for (int off = 16; off > 0; off >>= 1) {
      float o_max = __shfl_down_sync(0xFFFFFFFF, local_max, off);
      float o_sum = __shfl_down_sync(0xFFFFFFFF, local_sum, off);
      float new_max = fmaxf(local_max, o_max);
      local_sum = local_sum * expf(local_max - new_max)
                + o_sum * expf(o_max - new_max);
      local_max = new_max;
    }
    if (lane == 0) {
      out[row * 2]     = local_max;
      out[row * 2 + 1] = local_sum;
    }
  }
}

torch::Tensor fused_row_max_sumexp(
    torch::Tensor logits, int64_t n, int64_t c) {
  TORCH_CHECK(logits.is_cuda() && logits.is_contiguous()
              && logits.scalar_type() == torch::kFloat32);
  const at::cuda::CUDAGuard device_guard(logits.device());
  auto out = torch::empty({n, 2}, logits.options());
  int threads = 256;
  if (c <= 128) threads = 128;
  else if (c <= 512) threads = 256;
  else threads = 512;
  auto stream = at::cuda::getDefaultCUDAStream();
  fused_row_max_sumexp_kernel<<<(int)n, threads, 0, stream.stream()>>>(
      logits.data_ptr<float>(), out.data_ptr<float>(), n, c);
  return out;
}

// Fused exp(logits - bias) where bias is [N,1] broadcast over C columns.
// In-place write to the logits buffer, which saves one allocation.
__global__ void fused_exp_sub_kernel(
    float* __restrict__ logits,
    const float* __restrict__ bias,
    int64_t N, int64_t C) {
  int64_t idx = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
  int64_t total = N * C;
  if (idx >= total) return;
  int64_t row = idx / C;
  logits[idx] = expf(logits[idx] - bias[row]);
}

torch::Tensor fused_exp_sub(
    torch::Tensor logits, torch::Tensor bias) {
  TORCH_CHECK(logits.is_cuda() && logits.is_contiguous()
              && logits.scalar_type() == torch::kFloat32);
  TORCH_CHECK(bias.is_cuda() && bias.is_contiguous()
              && bias.scalar_type() == torch::kFloat32);
  const at::cuda::CUDAGuard device_guard(logits.device());
  int64_t N = logits.size(0);
  int64_t C = logits.size(1);
  int threads = 256;
  int blocks = (int)((N * C + threads - 1) / threads);
  auto stream = at::cuda::getDefaultCUDAStream();
  fused_exp_sub_kernel<<<blocks, threads, 0, stream.stream()>>>(
      logits.data_ptr<float>(), bias.data_ptr<float>(), N, C);
  return logits;
}
"""
    _CCE_EXT = load_inline(
        name="cce_fused_v1",
        cpp_sources=[cpp_src],
        cuda_sources=[cuda_src],
        extra_cuda_cflags=["-O3", "--use_fast_math", "-lineinfo"],
        extra_cflags=["-O3"],
        with_cuda=True,
        verbose=os.getenv("CCE_BUILD_VERBOSE", "0") == "1",
    )
    return _CCE_EXT


class _ChunkedCrossEntropyFn(torch.autograd.Function):
    """Exact log-sum-exp cross-entropy without materializing [N, V] logits.

    Only `lse = log sum_v exp(h @ E_v)` is kept on the autograd tape, one
    [N] tensor. Backward recovers the per-chunk softmax via
    `probs = exp(logits - lse)`, identical to the legacy
    (m, s_acc) split since `exp(logits - m) / s_acc == exp(logits - lse)`.

    Pascal note: bf16 inputs to `aten::mm` dispatch
    `magma_sgemmEx_kernel<float, __nv_bfloat16, ...>` which does a per-load
    upcast wrapper around fp32 multiply-add. That wrapper is ~25% slower
    than plain cuBLAS `sgemm` (FFMA-bound) on sm_61 because the bf16->fp32
    convert kernel runs serially with the GEMM. We force the GEMM inputs
    to fp32 here so PyTorch dispatches real FFMA `sgemm`. The model storage
    dtype is unchanged — the cast happens at the matmul boundary only.
    """

    @staticmethod
    def forward(ctx, hidden, target, emb_w, chunk_size: int):
        n, d = hidden.shape
        v = emb_w.shape[0]
        native_bf16 = _bf16_prefers_native_mm(hidden)
        # All matmul / reduce work runs in fp32 to land on cuBLAS sgemm,
        # FFMA, instead of magma_sgemmEx<bf16>. The fp32 cast of
        # `hidden` is a TRANSIENT working buffer: we save the original
        # bf16 tensor on the tape so VRAM stays bf16-sized. emb_w is
        # already fp32, the master param. Output `loss` is an fp32
        # scalar; autograd casts at the backward boundary if needed.
        h32 = hidden.float() if (hidden.dtype != torch.float32 and not native_bf16) else hidden
        ew32 = emb_w if emb_w.dtype == torch.float32 else emb_w.float()
        if native_bf16:
            true_w = emb_w.index_select(0, target).to(hidden.dtype)
            true_logits = (hidden * true_w).sum(dim=1).float()
        else:
            true_w = ew32.index_select(0, target)
            true_logits = (h32 * true_w).sum(dim=1)

        ext = _get_cce_ext() if hidden.is_cuda else None
        m = None
        s_acc = None
        for s in range(0, v, chunk_size):
            e = min(s + chunk_size, v)
            if native_bf16:
                w_native = emb_w[s:e].to(hidden.dtype)
                logits = hidden @ w_native.t()
                logits = logits.float() if logits.dtype != torch.float32 else logits
            else:
                w = ew32[s:e]
                device_type = 'cuda' if hidden.is_cuda else 'cpu'
                with torch.amp.autocast(device_type, enabled=False):
                    logits = h32 @ w.t()  # fp32 sgemm = FFMA on Pascal
            c = e - s
            if ext is not None:
                if logits.dtype != torch.float32:
                    logits = logits.float()
                if not logits.is_contiguous():
                    logits = logits.contiguous()
                ms = ext.fused_row_max_sumexp(logits, n, c)
                m_chunk = ms[:, 0]
                s_chunk = ms[:, 1]
            else:
                m_chunk = logits.max(dim=1).values
                s_chunk = torch.exp(logits - m_chunk.unsqueeze(1)).sum(dim=1)
            if m is None:
                m = m_chunk
                s_acc = s_chunk
            else:
                m_new = torch.maximum(m, m_chunk)
                s_acc = s_acc * torch.exp(m - m_new) + s_chunk * torch.exp(m_chunk - m_new)
                m = m_new
        lse = m + torch.log(s_acc)
        loss = (lse - true_logits).mean()
        ctx.chunk_size = int(chunk_size)
        # Save the ORIGINAL hidden, possibly bf16, on the tape;
        # backward will re-upcast for its three GEMMs. Same for emb_w.
        ctx.save_for_backward(hidden, target, emb_w, lse)
        return loss

    @staticmethod
    def backward(ctx, grad_out):
        hidden, target, emb_w, lse = ctx.saved_tensors
        n, d = hidden.shape
        v = emb_w.shape[0]
        cs = ctx.chunk_size
        scale = grad_out / float(n)
        native_bf16 = _bf16_prefers_native_mm(hidden)

        # Transient fp32 upcast for FFMA. Saved tensors stay bf16-sized.
        h32 = hidden.float() if (hidden.dtype != torch.float32 and not native_bf16) else hidden
        ew32 = emb_w if emb_w.dtype == torch.float32 else emb_w.float()

        ext = _get_cce_ext() if hidden.is_cuda else None
        grad_hidden32 = torch.zeros_like(h32)
        grad_emb_w32 = torch.zeros_like(ew32)

        for s in range(0, v, cs):
            e = min(s + cs, v)
            if native_bf16:
                w_native = emb_w[s:e].to(hidden.dtype)
                logits = hidden @ w_native.t()
                logits = logits.float() if logits.dtype != torch.float32 else logits
                h_for_gw = hidden
                w_for_gh = w_native
            else:
                w = ew32[s:e]
                logits = h32 @ w.t()                          # sgemm FFMA
                h_for_gw = h32
                w_for_gh = w
            if ext is not None:
                probs = ext.fused_exp_sub(logits, lse)        # in-place, 1 kernel
            else:
                probs = torch.exp(logits - lse.unsqueeze(1))
            if native_bf16:
                probs_gemm = probs.to(hidden.dtype)
            else:
                probs_gemm = probs
            grad_hidden32 = grad_hidden32 + (probs_gemm @ w_for_gh).float()
            grad_emb_w32[s:e] = grad_emb_w32[s:e] + (probs_gemm.t() @ h_for_gw).float()

        if native_bf16:
            true_w = emb_w.index_select(0, target).to(hidden.dtype)
            grad_hidden32 = grad_hidden32 - true_w.float()
            grad_emb_w32.index_add_(0, target, -hidden.float())
        else:
            true_w = ew32.index_select(0, target)
            grad_hidden32 = grad_hidden32 - true_w
            grad_emb_w32.index_add_(0, target, -h32)

        grad_hidden32 = grad_hidden32 * scale
        grad_emb_w32 = grad_emb_w32 * scale

        # Cast grads back to producer dtypes.
        grad_hidden = grad_hidden32 if hidden.dtype == torch.float32 else grad_hidden32.to(hidden.dtype)
        grad_emb_w = grad_emb_w32 if emb_w.dtype == torch.float32 else grad_emb_w32.to(emb_w.dtype)
        return grad_hidden, None, grad_emb_w, None


class _ChunkedCeCosineVecFn(torch.autograd.Function):
    """Joint chunked CE + embedding-space cosine loss in one vocab sweep.

    Forward keeps the same O(N) tape shape as chunked CE plus one [N, Demb]
    expected-embedding accumulator. No [N, V] logits/probs tensor is kept
    across chunks; each vocab block is materialized once, used for both CE and
    expected-embedding accumulation, then released.
    """

    @staticmethod
    def forward(ctx, hidden, target, emb_w, chunk_size: int):
        n, _d = hidden.shape
        v, de = emb_w.shape
        native_bf16 = _bf16_prefers_native_mm(hidden)
        h32 = hidden.float() if (hidden.dtype != torch.float32 and not native_bf16) else hidden
        ew32 = emb_w if emb_w.dtype == torch.float32 else emb_w.float()
        true_w = ew32.index_select(0, target)
        if native_bf16:
            true_logits = (hidden * true_w.to(hidden.dtype)).sum(dim=1).float()
        else:
            true_logits = (h32 * true_w).sum(dim=1)

        m = None
        s_acc = None
        num_acc = torch.zeros((n, de), device=hidden.device, dtype=torch.float32)
        for s in range(0, v, chunk_size):
            e = min(s + chunk_size, v)
            w32 = ew32[s:e]
            if native_bf16:
                logits = hidden @ emb_w[s:e].to(hidden.dtype).t()
                logits = logits.float() if logits.dtype != torch.float32 else logits
            else:
                logits = h32 @ w32.t()
            m_chunk = logits.max(dim=1).values
            unnorm_chunk = torch.exp(logits - m_chunk.unsqueeze(1))
            s_chunk = unnorm_chunk.sum(dim=1)
            num_chunk = unnorm_chunk @ w32
            if m is None:
                m = m_chunk
                s_acc = s_chunk
                num_acc = num_chunk
            else:
                m_new = torch.maximum(m, m_chunk)
                old_scale = torch.exp(m - m_new).unsqueeze(1)
                new_scale = torch.exp(m_chunk - m_new).unsqueeze(1)
                s_acc = s_acc * old_scale.squeeze(1) + s_chunk * new_scale.squeeze(1)
                num_acc = num_acc * old_scale + num_chunk * new_scale
                m = m_new

        lse = m + torch.log(s_acc)
        pred_emb = num_acc / s_acc.unsqueeze(1)
        ce_loss = (lse - true_logits).mean()

        pred_norm = pred_emb.norm(dim=1).clamp_min(1e-8)
        true_norm = true_w.norm(dim=1).clamp_min(1e-8)
        cosine = (pred_emb * true_w).sum(dim=1) / (pred_norm * true_norm)
        vec_loss = (1.0 - cosine).mean()

        ctx.chunk_size = int(chunk_size)
        ctx.save_for_backward(hidden, target, emb_w, lse, pred_emb, true_w)
        return ce_loss, vec_loss

    @staticmethod
    def backward(ctx, grad_ce_out, grad_vec_out):
        hidden, target, emb_w, lse, pred_emb, true_w = ctx.saved_tensors
        n, _d = hidden.shape
        v = emb_w.shape[0]
        cs = ctx.chunk_size
        native_bf16 = _bf16_prefers_native_mm(hidden)
        h32 = hidden.float() if (hidden.dtype != torch.float32 and not native_bf16) else hidden
        ew32 = emb_w if emb_w.dtype == torch.float32 else emb_w.float()

        ce_scale = grad_ce_out / float(n)

        pred_norm = pred_emb.norm(dim=1, keepdim=True).clamp_min(1e-8)
        true_norm = true_w.norm(dim=1, keepdim=True).clamp_min(1e-8)
        cosine = (pred_emb * true_w).sum(dim=1, keepdim=True) / (pred_norm * true_norm)
        grad_pred_emb = -(
            true_w / (pred_norm * true_norm)
            - cosine * pred_emb / pred_norm.pow(2)
        )
        grad_pred_emb = grad_pred_emb * (grad_vec_out / float(n))
        gy_dot_pred = (grad_pred_emb * pred_emb).sum(dim=1, keepdim=True)

        grad_hidden32 = torch.zeros_like(h32, dtype=torch.float32)
        grad_emb_w32 = torch.zeros_like(ew32, dtype=torch.float32)

        for s in range(0, v, cs):
            e = min(s + cs, v)
            w32 = ew32[s:e]
            if native_bf16:
                w_native = emb_w[s:e].to(hidden.dtype)
                logits = hidden @ w_native.t()
                logits = logits.float() if logits.dtype != torch.float32 else logits
                h_for_gw = hidden.float()
            else:
                logits = h32 @ w32.t()
                h_for_gw = h32.float()
            probs = torch.exp(logits - lse.unsqueeze(1))

            grad_hidden32 = grad_hidden32 + (probs @ w32).float() * ce_scale
            grad_emb_w32[s:e] = grad_emb_w32[s:e] + (probs.t() @ h_for_gw).float() * ce_scale

            gy_dot_w = grad_pred_emb @ w32.t()
            grad_logits_vec = probs * (gy_dot_w - gy_dot_pred)
            grad_hidden32 = grad_hidden32 + (grad_logits_vec @ w32).float()
            grad_emb_w32[s:e] = grad_emb_w32[s:e] + (probs.t() @ grad_pred_emb).float()
            grad_emb_w32[s:e] = grad_emb_w32[s:e] + (grad_logits_vec.t() @ h_for_gw).float()

        true_w_h = true_w if hidden.dtype == torch.float32 else true_w.to(hidden.dtype)
        grad_hidden32 = grad_hidden32 - true_w_h.float() * ce_scale
        grad_emb_w32.index_add_(0, target, -h32.float() * ce_scale)

        grad_hidden = grad_hidden32 if hidden.dtype == torch.float32 else grad_hidden32.to(hidden.dtype)
        grad_emb_w = grad_emb_w32 if emb_w.dtype == torch.float32 else grad_emb_w32.to(emb_w.dtype)
        return grad_hidden, None, grad_emb_w, None


@dataclass
class DecoderConfig:
    vocab_size: int = 248320
    hidden_dim: int = 32
    grid_g: int = 2
    grid_h: int = 8
    grid_w: int = 16
    bos_token_id: int = 1
    eos_token_id: int = 2
    pad_token_id: int = 0
    dtype: torch.dtype = torch.float32
    vocab_chunk_size: int = 8192
    use_chunked_loss: bool = True
    vector_error_lambda: float = 0.0
    # ALBERT-style factorized tied embedding. If > 0, the token table lives
    # in r=factor_dim space and a shared projection lifts it to hidden_dim
    # for both encode and the tied LM head. 0 keeps the legacy full-width
    # embedding.
    factor_dim: int = 0

    @property
    def seq_len(self) -> int:
        return self.grid_g * self.grid_h * self.grid_w


class ViperTiedDecoder(torch.nn.Module):
    """Trainable token embed + tied LM head for volumetric core outputs.

    Contracts:
    - encode: [B, L] -> ([B, D, G, H, W], valid_mask)
    - vol_to_seq / seq_to_vol: reshape between volumetric and seq views
    - decode: [B, L, D] -> logits [B, L, V]
    - tied weights: the LM head uses embedding.weight
    """

    def __init__(self, cfg: DecoderConfig):
        super().__init__()
        self.cfg = cfg
        self._emb_width = cfg.factor_dim if cfg.factor_dim > 0 else cfg.hidden_dim
        self.embedding = torch.nn.Embedding(cfg.vocab_size, self._emb_width)
        try:
            emb_gain = float(os.getenv("VIPER_EMB_GAIN", os.getenv("VIPER_INIT_GAIN", "0.6")))
        except Exception:
            emb_gain = 0.6
        try:
            emb_proj_gain = float(os.getenv("VIPER_EMB_PROJ_GAIN", os.getenv("VIPER_INIT_GAIN", "0.6")))
        except Exception:
            emb_proj_gain = 0.6
        # Tied LM head requires unit-variance logits at init:
        #   legacy:    logits = learnorm(h) @ E.T, E in [V, d]  -> std(E)=1/sqrt(d)
        #   factored:  logits = (learnorm(h) @ P) @ E.T,
        #              where P in [d, r] and E in [V, r]. Std of the
        #              r-space projected hidden is std(h)*std(P)*sqrt(d);
        #              with P ~ N(0, 1/sqrt(d)) it stays ~1. Then logits
        #              std = std(E)*sqrt(r), so std(E)=1/sqrt(r).
        torch.nn.init.normal_(self.embedding.weight, mean=0.0, std=emb_gain * (self._emb_width ** -0.5))
        if cfg.factor_dim > 0:
            self.emb_proj = torch.nn.Linear(cfg.factor_dim, cfg.hidden_dim, bias=False)
            torch.nn.init.normal_(self.emb_proj.weight, mean=0.0, std=emb_proj_gain * (cfg.hidden_dim ** -0.5))
        else:
            self.emb_proj = None
        self.norm = LeaRNorm(cfg.hidden_dim)
        # Embed lookup output magnitude. The tied-LM-head init makes
        # E ~ N(0, 1/sqrt(emb_width)) — correct for unit-variance LOGITS
        # at output side (norm(h)@P then @E.T). On the INPUT side it
        # produces std(seq) ≈ 1/sqrt(emb_width) ≈ 0.04, which is too weak
        # to drive the core through LeaRNorm (LeaRNorm only scales by
        # ≤sqrt(2), it does NOT restore unit RMS like RMSNorm). Multiply
        # the embed output by sqrt(hidden_dim) so the sequence entering
        # core_input_stem has std ~1, matching the magnitude assumed by
        # downstream LeaRNorm/siloid math without altering the LM head.
        self._embed_scale = float(cfg.hidden_dim) ** 0.5
        self.layout = ContinuousSnakeLayout((cfg.grid_g, cfg.grid_h, cfg.grid_w))

    def _embed(self, tokens: torch.Tensor) -> torch.Tensor:
        # [B, L] -> [B, L, hidden_dim], lifting through factor projection if any,
        # then scaled to unit-RMS for the core forward path. The LM-head reuse
        # path (decode/ce_blocks) never calls _embed; the head reads
        # self.embedding.weight directly so the unit-variance-logits init
        # contract there is unaffected.
        seq = self.embedding(tokens)
        if self.emb_proj is not None:
            seq = self.emb_proj(seq)
        return seq * self._embed_scale

    def _to_vocab(self, hidden: torch.Tensor) -> torch.Tensor:
        # Reverse of emb_proj: [..., hidden_dim] -> [..., emb_width].
        # For nn.Linear(r, d) with weight W shape (d, r), forward y = x @ W.T,
        # the reverse x = y @ W takes hidden back into r-space without any
        # extra parameters. When no factorization, this is a no-op.
        if self.emb_proj is None:
            return hidden
        return torch.matmul(hidden, self.emb_proj.weight)

    def _check_tokens(self, tokens: torch.Tensor) -> None:
        if tokens.shape[1] != self.cfg.seq_len:
            raise ValueError(f"seq_len mismatch: got {tokens.shape[1]} expected {self.cfg.seq_len}")
        if tokens.dtype != torch.long:
            raise ValueError(f"tokens must be torch.long, got {tokens.dtype}")
        max_id = int(tokens.max().item())
        if max_id >= self.cfg.vocab_size:
            raise ValueError(f"token id {max_id} is out of embedding range [0, {self.cfg.vocab_size - 1}]")

    def token_lengths(self, tokens: torch.Tensor) -> torch.Tensor:
        bsz, seqlen = tokens.shape
        idx = torch.arange(seqlen, device=tokens.device).unsqueeze(0).expand(bsz, -1)
        eos_mask = tokens.eq(self.cfg.eos_token_id)
        eos_pos = torch.where(eos_mask, idx, torch.full_like(idx, seqlen))
        first_eos = eos_pos.min(dim=1).values
        has_eos = eos_mask.any(dim=1)
        non_pad = tokens.ne(self.cfg.pad_token_id)
        last_non_pad = torch.where(non_pad, idx + 1, torch.zeros_like(idx)).max(dim=1).values
        return torch.where(has_eos, first_eos + 1, last_non_pad).clamp_min(1).clamp_max(seqlen)

    def lengths_from_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
        """Backward-compat alias kept for older training/tests."""
        return self.token_lengths(tokens)

    def make_valid_mask(self, lengths: torch.Tensor) -> torch.Tensor:
        idx = torch.arange(self.cfg.seq_len, device=lengths.device).unsqueeze(0)
        return idx < lengths.unsqueeze(1)

    def _sync_layout(self) -> None:
        grid = (self.cfg.grid_g, self.cfg.grid_h, self.cfg.grid_w)
        if self.layout.grid != grid:
            self.layout.set_grid(grid)

    def seq_to_vol(self, seq: torch.Tensor) -> torch.Tensor:
        # [B, L, D] -> [B, D, G, H, W] through the sequence layout
        # (Faced Continuous Snake Layout).
        b, l, d = seq.shape
        if l != self.cfg.seq_len or d != self.cfg.hidden_dim:
            raise ValueError(f"pack shape mismatch: got {tuple(seq.shape)}")
        self._sync_layout()
        return self.layout.seq_to_vol(seq)

    def vol_to_seq(self, volume: torch.Tensor) -> torch.Tensor:
        # [B, D, G, H, W] -> [B, L, D] through the same sequence layout.
        b, d, g, h, w = volume.shape
        if (d, g, h, w) != (self.cfg.hidden_dim, self.cfg.grid_g, self.cfg.grid_h, self.cfg.grid_w):
            raise ValueError(f"unpack shape mismatch: got {tuple(volume.shape)}")
        self._sync_layout()
        return self.layout.vol_to_seq(volume)

    def _vocab_blocks(self, chunk_size: int):
        if chunk_size <= 0:
            raise ValueError(f"chunk_size must be > 0, got {chunk_size}")
        v = self.cfg.vocab_size
        for s in range(0, v, chunk_size):
            e = min(s + chunk_size, v)
            yield s, e, self.embedding.weight[s:e]

    def _norm(self, seq: torch.Tensor) -> torch.Tensor:
        if seq.shape[-1] != self.cfg.hidden_dim:
            raise ValueError(f"seq hidden dim mismatch: {seq.shape[-1]} vs {self.cfg.hidden_dim}")
        return self.norm(seq)

    def _norm_loss(self, hidden: torch.Tensor) -> torch.Tensor:
        """Small auxiliary regularizer on normalized hidden states.

        Uses only magnitude constraints (RMS band + absmax cap), so it does not
        impose a directional target that could fight token CE optimization.
        """
        eps = 1e-6
        rms = torch.sqrt(hidden.pow(2).mean(dim=-1) + eps)
        low = F.relu(0.85 - rms)
        high = F.relu(rms - 1.15)
        rms_pen = (low.pow(2) + high.pow(2)).mean()
        absmax = hidden.abs().amax(dim=-1)
        absmax_pen = F.relu(absmax - 4.0).pow(2).mean()
        return rms_pen + 0.05 * absmax_pen

    def encode(self, tokens: torch.Tensor, valid_mask: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
        self._check_tokens(tokens)
        if valid_mask is None:
            lengths = self.token_lengths(tokens)
            valid_mask = self.make_valid_mask(lengths)
        seq = self._embed(tokens)
        seq = seq * valid_mask.unsqueeze(-1).to(seq.dtype)
        return self.seq_to_vol(seq), valid_mask

    def decode(self, seq: torch.Tensor) -> torch.Tensor:
        # [B, L, D] -> [B, L, V]
        hidden = self._to_vocab(self._norm(seq))
        return F.linear(hidden, self.embedding.weight)

    def decode_blocks(self, seq: torch.Tensor, chunk_size: Optional[int] = None) -> torch.Tensor:
        # [B, L, D] -> [B, L, V], projected chunk-wise to reduce peak VRAM in projection kernels.
        hidden = self._to_vocab(self._norm(seq))
        cs = self.cfg.vocab_chunk_size if chunk_size is None else chunk_size
        v = self.cfg.vocab_size
        emb_w = self.embedding.weight
        logits_parts = []
        for s in range(0, v, cs):
            e = min(s + cs, v)
            logits_parts.append(F.linear(hidden, emb_w[s:e]))
        return torch.cat(logits_parts, dim=-1)

    def topk_blocks(self, seq: torch.Tensor, k: int = 16, chunk_size: Optional[int] = None):
        # Returns top-k scores/ids without building full [B,L,V].
        hidden = self._to_vocab(self._norm(seq))
        b, l, d = hidden.shape
        flat = hidden.reshape(b * l, d)
        cs = self.cfg.vocab_chunk_size if chunk_size is None else chunk_size
        v = self.cfg.vocab_size
        emb_w = self.embedding.weight

        best_scores = torch.full((flat.shape[0], k), float("-inf"), device=flat.device, dtype=flat.dtype)
        best_ids = torch.zeros((flat.shape[0], k), device=flat.device, dtype=torch.long)
        for s in range(0, v, cs):
            e = min(s + cs, v)
            w = emb_w[s:e]
            logits_c = flat @ w.t()  # [N, C]
            kk = min(k, e - s)
            vals_c, ids_c = torch.topk(logits_c, k=kk, dim=1)
            ids_c = ids_c + s

            cand_scores = torch.cat([best_scores, vals_c], dim=1)
            cand_ids = torch.cat([best_ids, ids_c], dim=1)
            pick = torch.topk(cand_scores, k=k, dim=1).indices
            best_scores = cand_scores.gather(1, pick)
            best_ids = cand_ids.gather(1, pick)
        return best_scores.view(b, l, k), best_ids.view(b, l, k)

    def ce_loss(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        # shifted causal LM loss
        if labels.shape != logits.shape[:2]:
            raise ValueError(f"labels shape {tuple(labels.shape)} != logits[:2] {tuple(logits.shape[:2])}")
        pred = logits[:, :-1, :].contiguous()
        target = labels[:, 1:].contiguous()
        return F.cross_entropy(
            pred.reshape(-1, pred.shape[-1]),
            target.reshape(-1),
            ignore_index=self.cfg.pad_token_id,
        )

    def ce_blocks(
        self,
        seq: torch.Tensor,
        labels: torch.Tensor,
        chunk_size: Optional[int] = None,
        return_parts: bool = False,
        vector_error_lambda: float = 0.0,
    ):
        # Exact CE loss without full logits materialization. Under factorized
        # tied embedding all vocab dot-products happen in the r-dim space,
        # so the vocab GEMM stays [N, r] @ [r, C] regardless of hidden_dim.
        hidden_norm = self._norm(seq)
        hidden = hidden_norm[:, :-1, :].contiguous().view(-1, self.cfg.hidden_dim)
        target = labels[:, 1:].contiguous().view(-1)
        keep = target.ne(self.cfg.pad_token_id)
        if not keep.any():
            zero = hidden.sum() * 0.0
            if return_parts:
                return zero, zero, zero
            return zero

        hidden = hidden[keep]
        target = target[keep]
        norm_loss = self._norm_loss(hidden)
        hidden = self._to_vocab(hidden)
        emb_w = self.embedding.weight
        # NOTE: _ChunkedCrossEntropyFn recomputes the per-target dot product
        # internally (via grad_emb_w.index_add_ on backward and the lse-true
        # term on forward). Computing true_w / true_logits here would only
        # waste an [N, r] tensor and a reduction.

        cs = self.cfg.vocab_chunk_size if chunk_size is None else chunk_size
        if float(vector_error_lambda) > 0.0:
            ce_loss, vec_loss = _ChunkedCeCosineVecFn.apply(hidden, target, emb_w, cs)
        else:
            ce_loss = _ChunkedCrossEntropyFn.apply(hidden, target, emb_w, cs)
            vec_loss = ce_loss * 0.0
        if return_parts:
            return ce_loss, norm_loss, vec_loss
        return ce_loss

    def forward(
        self,
        volume: Optional[torch.Tensor] = None,
        seq: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        return_logits: Optional[bool] = None,
        chunk_size: Optional[int] = None,
    ):
        if (volume is None) == (seq is None):
            raise ValueError("provide exactly one of volume or seq")
        if seq is None:
            seq = self.vol_to_seq(volume)
        if return_logits is None:
            return_logits = labels is None

        out = {}
        logits = None
        if return_logits:
            logits = self.decode_blocks(seq, chunk_size=chunk_size)
            out["logits"] = logits
        if labels is not None:
            if self.cfg.use_chunked_loss:
                out["loss"] = self.ce_blocks(seq, labels, chunk_size=chunk_size)
            else:
                if logits is None:
                    logits = self.decode_blocks(seq, chunk_size=chunk_size)
                out["loss"] = self.ce_loss(logits, labels)
        return out