nsa.py

"""
NSA: Native Sparse Attention (PRODUCTION - causal-safe)
========================================================

DeepSeek 2502.11089 (ACL 2025 Best Paper)
Hardware-Aligned and Natively Trainable Sparse Attention.

3 branches with learnable gating:
  1. Compressed: block-mean K/V + causal block mask (FIX 5/7)
  2. Selected: top-k blocks (fallback to full causal)
  3. Sliding: local window with causal mask

Implementation notes:
  - NSACompressBranch: added causal block mask (q at pos t can only see blocks fully <= t-1)
  - Class renamed NSAttention -> NSAAttention (modeling compat)
  - forward signature: layer_idx kwarg + (Tensor, past_kv) tuple return

NSAAttention is the public name imported by modeling_aether_v2_7way.py.
"""
from __future__ import annotations
import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


class NSACompressBranch(nn.Module):
    """Block-compressed long-range attention WITH causal block mask."""

    def __init__(self, config):
        super().__init__()
        self.h = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        # accept either nsa_compressed_block_size or nsa_compress_block_size
        self.block_size = getattr(config, "nsa_compressed_block_size",
                                  getattr(config, "nsa_compress_block_size", 32))

        self.q_proj = nn.Linear(self.h, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.h, bias=False)

        self.block_compress_k = nn.Linear(self.head_dim, self.head_dim, bias=False)
        self.block_compress_v = nn.Linear(self.head_dim, self.head_dim, bias=False)

    def compress_blocks(self, x: torch.Tensor) -> torch.Tensor:
        """x: [B, S, H, D] -> [B, num_blocks, H, D] via block-mean (truncate non-aligned)."""
        B, S, H, D = x.shape
        num_blocks = S // self.block_size
        if num_blocks == 0:
            return x[:, :0]  # empty
        x = x[:, : num_blocks * self.block_size]
        x = x.view(B, num_blocks, self.block_size, H, D)
        return x.mean(dim=2)

    def forward(self, hidden_states: torch.Tensor, past_key_value=None, cache_idx: int = 0) -> torch.Tensor:
        B, S, _ = hidden_states.shape

        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)

        # AetherCache: cache raw K,V (projection becomes O(1)); block-compress is a cheap reduce
        if past_key_value is not None:
            kc, vc = past_key_value.update(k.transpose(1, 2), v.transpose(1, 2), cache_idx)
            k = kc.transpose(1, 2)
            v = vc.transpose(1, 2)
        S_kv = k.shape[1]

        # Compress K, V into blocks
        k_compressed = self.compress_blocks(k)  # [B, num_blocks, kv_h, D]
        v_compressed = self.compress_blocks(v)
        num_blocks = k_compressed.shape[1]

        if num_blocks == 0:
            # Sequence shorter than block_size — fallback to identity (no compress contribution)
            out = torch.zeros(B, S, self.h, device=hidden_states.device, dtype=hidden_states.dtype)
            return out

        k_compressed = self.block_compress_k(k_compressed)
        v_compressed = self.block_compress_v(v_compressed)

        # GQA repeat
        repeat = self.num_heads // self.num_kv_heads
        k_compressed = k_compressed.repeat_interleave(repeat, dim=2)
        v_compressed = v_compressed.repeat_interleave(repeat, dim=2)

        # [B, H, S, D]
        q = q.transpose(1, 2)
        k_c = k_compressed.transpose(1, 2)  # [B, H, num_blocks, D]
        v_c = v_compressed.transpose(1, 2)

        scores = torch.matmul(q, k_c.transpose(-2, -1)) / math.sqrt(self.head_dim)
        # ★ FIX 5/7: causal block mask
        # query at position t can attend to block i iff (i+1)*bs <= t  (block fully completed by t-1)
        # equivalently: block_end_inclusive = (i+1)*bs - 1 < t  →  i*bs + bs - 1 < t  →  i < (t - bs + 1) / bs
        # simpler: mask True where block start > t (block has not started)... but we need stricter:
        # block i is "in the past" iff its END (i+1)*bs - 1 < t, i.e. (i+1)*bs <= t.
        q_pos = torch.arange(S_kv - S, S_kv, device=q.device).unsqueeze(1)  # AetherCache: absolute [S, 1]
        block_end = (torch.arange(num_blocks, device=q.device) + 1) * self.block_size  # [num_blocks]
        block_end = block_end.unsqueeze(0)  # [1, num_blocks]
        # mask: True = blocked (future block)
        mask = block_end > q_pos  # [S, num_blocks], True if block_end > q_pos (block not fully past)
        mask = mask.unsqueeze(0).unsqueeze(0)  # [1, 1, S, num_blocks]
        scores = scores.masked_fill(mask, float("-inf"))

        # If a query has no past block, all -inf -> nan softmax. Set first block prob to 0 by giving it 0 logit.
        # Cheap guard: rows that are all -inf get a 0-tensor output.
        all_neg_inf = mask.all(dim=-1, keepdim=True)  # [1,1,S,1]
        # set those rows to a small finite value so softmax produces uniform; we'll zero output below
        scores = torch.where(all_neg_inf.expand_as(scores), torch.zeros_like(scores), scores)

        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v_c)  # [B, H, S, D]
        # Zero out positions that have no past block (would be invalid)
        out = out.masked_fill(all_neg_inf, 0.0)
        out = out.transpose(1, 2).reshape(B, S, -1)
        return self.o_proj(out)


class NSASelectBranch(nn.Module):
    """Top-k block selection. Currently falls back to full causal attention."""

    def __init__(self, config):
        super().__init__()
        self.h = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.block_size = getattr(config, "nsa_compressed_block_size",
                                  getattr(config, "nsa_compress_block_size", 32))
        self.top_k = getattr(config, "nsa_selected_top_k",
                             getattr(config, "nsa_select_top_k", 16))

        self.q_proj = nn.Linear(self.h, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.h, bias=False)

        self.block_score = nn.Linear(self.head_dim, 1, bias=False)

    def forward(self, hidden_states: torch.Tensor, past_key_value=None, cache_idx: int = 0) -> torch.Tensor:
        B, S, _ = hidden_states.shape

        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)

        # AetherCache: standard KV cache
        if past_key_value is not None:
            kc, vc = past_key_value.update(k.transpose(1, 2), v.transpose(1, 2), cache_idx)
            k = kc.transpose(1, 2)
            v = vc.transpose(1, 2)

        # GQA repeat
        repeat = self.num_heads // self.num_kv_heads
        k_full = k.repeat_interleave(repeat, dim=2)
        v_full = v.repeat_interleave(repeat, dim=2)

        # Use SDPA causal (FIX 5/7: strict causal) — causal only when prefill
        q = q.transpose(1, 2)
        k_full = k_full.transpose(1, 2)
        v_full = v_full.transpose(1, 2)
        out = F.scaled_dot_product_attention(q, k_full, v_full, dropout_p=0.0, is_causal=(S > 1))
        out = out.transpose(1, 2).reshape(B, S, -1)
        return self.o_proj(out)


class NSASlidingBranch(nn.Module):
    """Local sliding window with causal mask."""

    def __init__(self, config):
        super().__init__()
        self.h = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.window_size = getattr(config, "nsa_sliding_size",
                                   getattr(config, "nsa_sliding_window", 512))

        self.q_proj = nn.Linear(self.h, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.h, self.num_kv_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.h, bias=False)

    def forward(self, hidden_states: torch.Tensor, past_key_value=None, cache_idx: int = 0) -> torch.Tensor:
        B, S, _ = hidden_states.shape

        q = self.q_proj(hidden_states).view(B, S, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(B, S, self.num_kv_heads, self.head_dim)

        # AetherCache: standard KV cache
        if past_key_value is not None:
            kc, vc = past_key_value.update(k.transpose(1, 2), v.transpose(1, 2), cache_idx)
            k = kc.transpose(1, 2)
            v = vc.transpose(1, 2)
        S_kv = k.shape[1]

        repeat = self.num_heads // self.num_kv_heads
        k = k.repeat_interleave(repeat, dim=2)
        v = v.repeat_interleave(repeat, dim=2)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Sliding window + causal (AetherCache: absolute positions)
        q_abs = torch.arange(S_kv - S, S_kv, device=q.device)
        k_abs = torch.arange(S_kv, device=q.device)
        mask = (k_abs[None, :] > q_abs[:, None]) | (k_abs[None, :] < q_abs[:, None] - self.window_size + 1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(mask, float("-inf"))
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).reshape(B, S, -1)
        return self.o_proj(out)


class NSAAttention(nn.Module):
    """
    Native Sparse Attention (DeepSeek 2502.11089) — modeling-compat wrapper.

    Public name: NSAAttention (matches modeling_aether_v2_7way.py import).
    Accepts layer_idx kwarg (stored, currently unused).
    Returns (Tensor, past_key_value=None) tuple to match other attention modules.
    """

    def __init__(self, config, layer_idx: int = 0):
        super().__init__()
        self.layer_idx = layer_idx
        self.compressed = NSACompressBranch(config)
        self.selected = NSASelectBranch(config)
        self.sliding = NSASlidingBranch(config)

        # Learnable gate (sigmoid normalized)
        self.gate_c = nn.Parameter(torch.zeros(1))
        self.gate_s = nn.Parameter(torch.zeros(1))
        self.gate_w = nn.Parameter(torch.zeros(1))

        self.norm = nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value=None,
        use_cache: bool = False,
        **kwargs,
    ) -> Tuple[torch.Tensor, Optional[object]]:
        # AetherCache: branches are stateless fns of the full sequence -> cache the layer input
        # and recompute over (history + new), returning only the new positions. Prefill unchanged.
        cidx = int(getattr(self, "cache_idx", self.layer_idx))
        # compress uses the layer's own slot (keeps get_seq_length() valid); others get high slots
        out_c = self.compressed(hidden_states, past_key_value, cidx)
        out_s = self.selected(hidden_states, past_key_value, 200 + 2 * cidx)
        out_w = self.sliding(hidden_states, past_key_value, 200 + 2 * cidx + 1)

        g_c = torch.sigmoid(self.gate_c)
        g_s = torch.sigmoid(self.gate_s)
        g_w = torch.sigmoid(self.gate_w)

        out = g_c * out_c + g_s * out_s + g_w * out_w
        out = self.norm(out)
        return out, past_key_value


# Backward compat alias (in case of NSAttention import)
NSAttention = NSAAttention

__all__ = ["NSAAttention", "NSAttention", "NSACompressBranch", "NSASelectBranch", "NSASlidingBranch"]