Upload aether_pkg/v2_attentions/nsa.py with huggingface_hub

aether_pkg/v2_attentions/nsa.py

@@ -0,0 +1,254 @@
+"""
+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
+
+5/7 fixes applied to placeholder bugs:
+  - 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) -> 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)
+
+        # 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, device=q.device).unsqueeze(1)  # [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) -> 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)
+
+        # 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)
+        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=True)
+        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) -> 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)
+
+        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
+        idx = torch.arange(S, device=q.device)
+        # mask: True = blocked
+        # idx[None,:] is k_pos (column), idx[:,None] is q_pos (row)
+        # block if k_pos > q_pos (future) or k_pos < q_pos - window + 1 (too far past)
+        mask = (idx[None, :] > idx[:, None]) | (idx[None, :] < idx[:, 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]]:
+        out_c = self.compressed(hidden_states)
+        out_s = self.selected(hidden_states)
+        out_w = self.sliding(hidden_states)
+
+        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"]