Upload models/linear.py

models/linear.py

@@ -0,0 +1,610 @@
+"""Linear (Kimi Delta Attention) decoder-only transformer.
+
+Pure-PyTorch chunkwise-parallel KDA scan.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Dict, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from einops import rearrange
+
+from models.lm_loss import (
+    lm_cross_entropy_from_logits,
+    token_superposition_attention_mask,
+    token_superposition_embeddings,
+)
+from models.baseline import (
+    BaselineConfig,
+    RMSNorm,
+    SwiGLU,
+    MoELayer,
+    combine_lm_and_aux_loss,
+    init_moe_router_weights,
+    _validate_moe_config,
+    count_parameters,
+    model_summary,
+)
+
+
+class _ShortConvolution(nn.Module):
+    """Causal depthwise 1-D conv with optional cached state for O(1) decode."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        kernel_size: int,
+        activation: str = "silu",
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.conv = nn.Conv1d(
+            hidden_size,
+            hidden_size,
+            kernel_size=kernel_size,
+            groups=hidden_size,
+            padding=kernel_size - 1,
+            bias=bias,
+        )
+        self.activation = activation
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cache: Optional[torch.Tensor] = None,
+        return_cache: bool = False,
+    ):
+        T = x.size(1)
+        K = self.kernel_size
+
+        if cache is None:
+            y = self.conv(x.transpose(1, 2))[..., :T].transpose(1, 2)
+        else:
+            x_full = torch.cat([cache, x], dim=1)
+            y = F.conv1d(
+                x_full.transpose(1, 2),
+                self.conv.weight,
+                self.conv.bias,
+                stride=1,
+                padding=0,
+                groups=self.conv.groups,
+            ).transpose(1, 2)
+
+        if self.activation == "silu":
+            y = F.silu(y)
+
+        if not return_cache:
+            return y
+
+        if K <= 1:
+            new_cache = x.new_zeros(x.size(0), 0, x.size(-1))
+        else:
+            combined = torch.cat([cache, x], dim=1) if cache is not None else x
+            if combined.size(1) >= K - 1:
+                new_cache = combined[:, -(K - 1):].contiguous()
+            else:
+                pad = combined.new_zeros(
+                    combined.size(0), (K - 1) - combined.size(1), combined.size(-1)
+                )
+                new_cache = torch.cat([pad, combined], dim=1)
+        return y, new_cache
+
+
+class _RMSNormGatedSigmoid(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
+        dtype = x.dtype
+        x_f = x.float()
+        rms_inv = x_f.pow(2).mean(dim=-1, keepdim=True).add_(self.eps).rsqrt()
+        y = (x_f * rms_inv).to(dtype) * self.weight
+        return y * torch.sigmoid(gate.to(dtype))
+
+
+def _kda_gate(
+    g: torch.Tensor,
+    A_log: torch.Tensor,
+    dt_bias: torch.Tensor,
+) -> torch.Tensor:
+    """Log-space decay gate: ``-exp(A_log) * softplus(g + dt_bias)``."""
+    H, K = g.shape[-2], g.shape[-1]
+    g = g.float() + dt_bias.float().view(H, K)
+    dt = F.softplus(g)
+    A = A_log.float().view(1, 1, H, 1)
+    return -A.exp() * dt
+
+
+def _kda_chunk_scan(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_g: torch.Tensor,
+    beta: torch.Tensor,
+    chunk_size: int = 64,
+    use_qk_l2norm: bool = True,
+    initial_state: Optional[torch.Tensor] = None,
+    output_final_state: bool = False,
+):
+    """Chunkwise-parallel KDA scan in pure PyTorch.
+
+    Recurrence: ``S_i = (I - beta_i k_i k_i^T) D_i S_{i-1} + beta_i k_i v_i^T``,
+    ``o_i = q_i @ S_i``, with ``D_i = diag(exp(log_g_i))``. Chunk-level
+    parallelism comes from the similarity transform ``~S_i = W_i^{-1} S_i``
+    plus a single triangular solve per chunk.
+    """
+    B, T, H, K = q.shape
+    V = v.shape[-1]
+    orig_dtype = v.dtype
+    device = q.device
+
+    # The body runs in fp32: per-channel decays accumulate aggressively and
+    # CUDA's triangular_solve has no bf16/fp16 kernel.
+    with torch.autocast(device_type=device.type, enabled=False):
+        if use_qk_l2norm:
+            q = F.normalize(q, dim=-1)
+            k = F.normalize(k, dim=-1)
+        scale = K ** -0.5
+
+        q = q.float() * scale
+        k = k.float()
+        v = v.float()
+        log_g = log_g.float()
+        beta = beta.float()
+
+        pad = (chunk_size - T % chunk_size) % chunk_size
+        if pad > 0:
+            q = F.pad(q, (0, 0, 0, 0, 0, pad))
+            k = F.pad(k, (0, 0, 0, 0, 0, pad))
+            v = F.pad(v, (0, 0, 0, 0, 0, pad))
+            log_g = F.pad(log_g, (0, 0, 0, 0, 0, pad))
+            beta = F.pad(beta, (0, 0, 0, pad))
+        Nc = (T + pad) // chunk_size
+        C = chunk_size
+
+        q = rearrange(q, "b (n c) h k -> b h n c k", c=C)
+        k = rearrange(k, "b (n c) h k -> b h n c k", c=C)
+        v = rearrange(v, "b (n c) h v -> b h n c v", c=C)
+        log_g = rearrange(log_g, "b (n c) h k -> b h n c k", c=C)
+        beta = rearrange(beta, "b (n c) h -> b h n c", c=C)
+
+        # Clamp the cumulative log-decay to [-15, 0]: at default A/dt_bias
+        # ranges a 64-token cumsum can drop below -80, and exp(-cum) then
+        # overflows fp32 and NaNs the triangular solve.
+        cum_log_g = log_g.cumsum(dim=-2).clamp(min=-15.0)
+        W = cum_log_g.exp()
+        W_inv = (-cum_log_g).exp()
+
+        u_mat = k * W_inv
+        w_mat = k * W
+        q_tilde = q * W
+
+        beta_e = beta.unsqueeze(-1)
+        beta_w = beta_e * w_mat
+        beta_v = beta_e * v
+
+        L = torch.einsum("bhnik,bhnjk->bhnij", beta_w, u_mat)
+        upper_incl_diag = torch.triu(
+            torch.ones(C, C, dtype=torch.bool, device=device), diagonal=0
+        )
+        L = L.masked_fill(upper_incl_diag, 0)
+
+        I_plus_L = L + torch.eye(C, dtype=L.dtype, device=device)
+        effective_v = torch.linalg.solve_triangular(
+            I_plus_L, beta_v, upper=False, unitriangular=True
+        )
+        effective_w = torch.linalg.solve_triangular(
+            I_plus_L, beta_w, upper=False, unitriangular=True
+        )
+
+        intra_attn = torch.einsum("bhnik,bhnjk->bhnij", q_tilde, u_mat)
+        strict_upper = torch.triu(
+            torch.ones(C, C, dtype=torch.bool, device=device), diagonal=1
+        )
+        intra_attn = intra_attn.masked_fill(strict_upper, 0)
+
+        if initial_state is not None:
+            S = initial_state.to(dtype=q.dtype, device=q.device)
+        else:
+            S = q.new_zeros(B, H, K, V)
+        outputs: List[torch.Tensor] = []
+        for n in range(Nc):
+            delta = effective_v[:, :, n] - effective_w[:, :, n] @ S
+            o_inter = q_tilde[:, :, n] @ S
+            o_chunk = o_inter + intra_attn[:, :, n] @ delta
+            outputs.append(o_chunk)
+
+            state_update = torch.einsum(
+                "bhck,bhcv->bhkv", u_mat[:, :, n], delta
+            )
+            S = W[:, :, n, -1].unsqueeze(-1) * (S + state_update)
+
+        out = torch.stack(outputs, dim=2)
+        out = rearrange(out, "b h n c v -> b (n c) h v")
+        if pad > 0:
+            out = out[:, :T]
+        out = out.to(orig_dtype)
+
+    if output_final_state:
+        # State stays fp32 so cached decode preserves precision.
+        return out, S
+    return out
+
+
+def _kda_recurrent_step(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    log_g: torch.Tensor,
+    beta: torch.Tensor,
+    state: torch.Tensor,
+    use_qk_l2norm: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Single-token KDA step matching ``_kda_chunk_scan`` for ``T == 1``."""
+    assert q.size(1) == 1 and k.size(1) == 1 and v.size(1) == 1
+    orig_dtype = v.dtype
+    K = q.size(-1)
+    device = q.device
+
+    with torch.autocast(device_type=device.type, enabled=False):
+        if use_qk_l2norm:
+            q = F.normalize(q, dim=-1)
+            k = F.normalize(k, dim=-1)
+        scale = K ** -0.5
+
+        q_t = (q[:, 0].float()) * scale
+        k_t = k[:, 0].float()
+        v_t = v[:, 0].float()
+        g_t = log_g[:, 0].float()
+        b_t = beta[:, 0].float()
+
+        S = state.to(torch.float32)
+        S = S * g_t.exp().unsqueeze(-1)
+        kS = torch.einsum("bhk,bhkv->bhv", k_t, S)
+        update = torch.einsum(
+            "bhk,bhv->bhkv", (b_t.unsqueeze(-1) * k_t), (v_t - kS)
+        )
+        S = S + update
+        o = torch.einsum("bhk,bhkv->bhv", q_t, S).unsqueeze(1)
+    return o.to(orig_dtype), S
+
+
+@dataclass
+class LinearConfig(BaselineConfig):
+    head_dim: int = 64
+    conv_size: int = 4
+    chunk_size: int = 64
+    A_init_range: Tuple[float, float] = (1, 16)
+
+    expand: int = 2
+    rope_base: float = 10000.0
+
+    def __post_init__(self):
+        if self.d_model % self.num_heads != 0:
+            raise ValueError("d_model must be divisible by num_heads")
+        if self.partial_rope_dim is not None:
+            if self.partial_rope_dim % 2 != 0:
+                raise ValueError(
+                    f"partial_rope_dim ({self.partial_rope_dim}) must be even"
+                )
+        if self.head_dim < 1:
+            raise ValueError("head_dim must be >= 1")
+        if self.chunk_size < 1:
+            raise ValueError("chunk_size must be >= 1")
+        if self.conv_size < 1:
+            raise ValueError("conv_size must be >= 1")
+        _validate_moe_config(self)
+
+
+class KimiDeltaAttention(nn.Module):
+    def __init__(self, config: LinearConfig):
+        super().__init__()
+        self.hidden_size = config.d_model
+        self.num_heads = config.num_heads
+        self.head_dim = config.head_dim
+        self.head_k_dim = self.head_dim
+        self.conv_size = config.conv_size
+        self.chunk_size = config.chunk_size
+
+        projection_size = self.num_heads * self.head_dim
+
+        self.q_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
+
+        self.q_conv1d = _ShortConvolution(projection_size, self.conv_size, "silu")
+        self.k_conv1d = _ShortConvolution(projection_size, self.conv_size, "silu")
+        self.v_conv1d = _ShortConvolution(projection_size, self.conv_size, "silu")
+
+        A = torch.empty(self.num_heads, dtype=torch.float32).uniform_(
+            *config.A_init_range
+        )
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+
+        self.dt_bias = nn.Parameter(torch.empty(projection_size, dtype=torch.float32))
+        self.dt_bias._no_weight_decay = True
+
+        self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
+
+        self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
+
+        self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+        self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=True)
+
+        self.o_norm = _RMSNormGatedSigmoid(self.head_dim, eps=config.norm_eps)
+        self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        # Inverse-softplus init (Mamba-2 / KDA scheme).
+        dt = torch.exp(
+            torch.rand(self.num_heads * self.head_dim)
+            * (math.log(0.1) - math.log(0.001))
+            + math.log(0.001)
+        )
+        dt = torch.clamp(dt, min=1e-4)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))
+        with torch.no_grad():
+            self.dt_bias.copy_(inv_dt)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        cache: Optional[Dict[str, Optional[torch.Tensor]]] = None,
+    ) -> torch.Tensor:
+        use_cache = cache is not None
+
+        q_in = self.q_proj(x)
+        k_in = self.k_proj(x)
+        v_in = self.v_proj(x)
+        if use_cache:
+            q, cache["conv_state_q"] = self.q_conv1d(
+                q_in, cache=cache.get("conv_state_q"), return_cache=True
+            )
+            k, cache["conv_state_k"] = self.k_conv1d(
+                k_in, cache=cache.get("conv_state_k"), return_cache=True
+            )
+            v, cache["conv_state_v"] = self.v_conv1d(
+                v_in, cache=cache.get("conv_state_v"), return_cache=True
+            )
+        else:
+            q = self.q_conv1d(q_in)
+            k = self.k_conv1d(k_in)
+            v = self.v_conv1d(v_in)
+
+        q = rearrange(q, "... (h d) -> ... h d", d=self.head_dim)
+        k = rearrange(k, "... (h d) -> ... h d", d=self.head_dim)
+        v = rearrange(v, "... (h d) -> ... h d", d=self.head_dim)
+
+        g_raw = self.f_b_proj(self.f_a_proj(x))
+        g_raw = rearrange(g_raw, "... (h d) -> ... h d", d=self.head_dim)
+        log_g = _kda_gate(g_raw, self.A_log, self.dt_bias)
+
+        beta = self.b_proj(x).float().sigmoid()
+
+        # Zero q/k/v, log_g, beta at padded positions so they contribute no
+        # content and no decay to the recurrent state.
+        if attention_mask is not None:
+            mask_4d = attention_mask.unsqueeze(-1).unsqueeze(-1)
+            q = q * mask_4d.to(q.dtype)
+            k = k * mask_4d.to(k.dtype)
+            v = v * mask_4d.to(v.dtype)
+            log_g = log_g * mask_4d.to(log_g.dtype)
+            beta = beta * attention_mask.unsqueeze(-1).to(beta.dtype)
+
+        if use_cache:
+            prev_state = cache.get("recurrent_state")
+            if prev_state is not None and x.size(1) == 1:
+                o, new_state = _kda_recurrent_step(
+                    q, k, v, log_g, beta, prev_state, use_qk_l2norm=True
+                )
+            else:
+                o, new_state = _kda_chunk_scan(
+                    q=q, k=k, v=v, log_g=log_g, beta=beta,
+                    chunk_size=self.chunk_size,
+                    use_qk_l2norm=True,
+                    initial_state=prev_state,
+                    output_final_state=True,
+                )
+            cache["recurrent_state"] = new_state
+        else:
+            o = _kda_chunk_scan(
+                q=q, k=k, v=v, log_g=log_g, beta=beta,
+                chunk_size=self.chunk_size,
+                use_qk_l2norm=True,
+            )
+
+        gate = self.g_b_proj(self.g_a_proj(x))
+        gate = rearrange(gate, "... (h d) -> ... h d", d=self.head_dim)
+        o = self.o_norm(o, gate)
+
+        o = rearrange(o, "b t h d -> b t (h d)")
+        return self.o_proj(o)
+
+
+class LinearTransformerBlock(nn.Module):
+    def __init__(self, config: LinearConfig):
+        super().__init__()
+        self.use_moe = config.use_moe
+
+        self.kda_norm = RMSNorm(config.d_model, eps=config.norm_eps)
+        self.kda = KimiDeltaAttention(config)
+
+        self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
+        if config.use_moe:
+            self.ffn = MoELayer(config)
+        else:
+            self.ffn = SwiGLU(config.d_model, config.d_ff)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        is_causal: bool = True,
+        cache: Optional[Dict[str, Optional[torch.Tensor]]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+        x = x + self.kda(self.kda_norm(x), attention_mask=attention_mask, cache=cache)
+
+        if self.use_moe:
+            ffn_out, aux_loss, topk_indices = self.ffn(self.ffn_norm(x))
+            x = x + ffn_out
+            return x, aux_loss, topk_indices
+        else:
+            x = x + self.ffn(self.ffn_norm(x))
+            return x, torch.zeros((), device=x.device, dtype=x.dtype), None
+
+
+class LinearTransformer(nn.Module):
+    def __init__(self, config: LinearConfig):
+        super().__init__()
+        self.config = config
+
+        self.token_emb = nn.Embedding(config.vocab_size, config.d_model)
+
+        self.layers = nn.ModuleList([
+            LinearTransformerBlock(config) for _ in range(config.num_layers)
+        ])
+
+        self.final_norm = RMSNorm(config.d_model, eps=config.norm_eps)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.token_emb.weight
+
+        self._init_weights()
+
+    def _init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    torch.nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        init_moe_router_weights(self, self.config.router_init_std)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        is_causal: bool = True,
+        caches: Optional[List[Dict[str, Optional[torch.Tensor]]]] = None,
+        token_superposition_bag_size: int = 1,
+    ) -> Dict[str, Any]:
+        x = token_superposition_embeddings(
+            self.token_emb, input_ids, token_superposition_bag_size,
+        )
+        attention_mask = token_superposition_attention_mask(
+            attention_mask, token_superposition_bag_size,
+        )
+
+        aux_loss = torch.zeros((), device=input_ids.device, dtype=x.dtype)
+        topk_indices_list: List[Optional[torch.Tensor]] = []
+        for i, layer in enumerate(self.layers):
+            layer_cache = caches[i] if caches is not None else None
+            x, layer_aux, layer_topk = layer(
+                x, attention_mask=attention_mask,
+                is_causal=is_causal, cache=layer_cache,
+            )
+            aux_loss = aux_loss + layer_aux
+            topk_indices_list.append(layer_topk)
+
+        x = self.final_norm(x)
+        logits = self.lm_head(x)
+
+        lm_loss: Optional[torch.Tensor] = None
+        if labels is not None:
+            lm_loss = lm_cross_entropy_from_logits(
+                logits,
+                labels,
+                token_superposition_bag_size=token_superposition_bag_size,
+                ignore_index=-100,
+            )
+        loss = combine_lm_and_aux_loss(
+            lm_loss,
+            aux_loss if self.config.use_moe else None,
+            self.training,
+        )
+
+        return {
+            "logits": logits,
+            "loss": loss,
+            "lm_loss": lm_loss,
+            "aux_loss": aux_loss if self.config.use_moe else None,
+            "topk_indices": topk_indices_list if self.config.use_moe else None,
+        }
+
+    def update_router_biases(self, topk_indices_list: List[Optional[torch.Tensor]]) -> None:
+        if not self.config.use_moe:
+            return
+        for layer, topk_indices in zip(self.layers, topk_indices_list):
+            if topk_indices is not None and isinstance(layer.ffn, MoELayer):
+                layer.ffn.update_bias(topk_indices)
+
+    @torch.no_grad()
+    def get_balance_stats(self) -> Dict[str, float]:
+        if not self.config.use_moe:
+            return {}
+        stats = {}
+        for idx, layer in enumerate(self.layers):
+            if hasattr(layer.ffn, "bias"):
+                bias = layer.ffn.bias
+                stats[f"layer{idx}_bias_mean"] = bias.abs().mean().item()
+                stats[f"layer{idx}_bias_max"] = bias.abs().max().item()
+        return stats
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        max_new_tokens: int = 100,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        eos_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        self.train(False)
+
+        caches: List[Dict[str, Optional[torch.Tensor]]] = [
+            {
+                "recurrent_state": None,
+                "conv_state_q": None,
+                "conv_state_k": None,
+                "conv_state_v": None,
+            }
+            for _ in self.layers
+        ]
+
+        def _sample(logits: torch.Tensor) -> torch.Tensor:
+            logits = logits / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits = logits.masked_fill(logits < v[:, [-1]], float("-inf"))
+            probs = F.softmax(logits, dim=-1)
+            return torch.multinomial(probs, num_samples=1)
+
+        outputs = self.forward(input_ids, is_causal=True, caches=caches)
+        next_token = _sample(outputs["logits"][:, -1, :])
+        input_ids = torch.cat([input_ids, next_token], dim=-1)
+
+        if eos_token_id is not None and (next_token == eos_token_id).all():
+            return input_ids
+
+        for _ in range(max_new_tokens - 1):
+            outputs = self.forward(next_token, is_causal=True, caches=caches)
+            next_token = _sample(outputs["logits"][:, -1, :])
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            if eos_token_id is not None and (next_token == eos_token_id).all():
+                break
+
+        return input_ids