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