rxnn 0.2.53__tar.gz → 0.2.55__tar.gz

{rxnn-0.2.53 → rxnn-0.2.55}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.53
+Version: 0.2.55
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.2.53 → rxnn-0.2.55}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.2.53"
+version = "0.2.55"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.2.53 → rxnn-0.2.55}/src/rxnn/experimental/attention.py

@@ -319,6 +319,8 @@ class SparseQueryAttention(MultiHeadAttention):
 
 # Others
 class FlexAttention(MultiHeadAttention):
+    """Flex attention layer, with RoPE support"""
+
     def __init__(
         self,
         embed_dim: int,
@@ -355,8 +357,7 @@ class FlexAttention(MultiHeadAttention):
         # Learnable global tokens
         self.global_tokens = nn.Parameter(torch.randn(1, num_global_tokens, embed_dim))
 
-
-    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
+    def _split_heads(self, x: torch.Tensor, is_query: bool = False) -> torch.Tensor:
         b, t, d = x.size()
         return x.view(b, t, self.num_heads, self.head_dim).transpose(1, 2)
 
@@ -364,15 +365,19 @@ class FlexAttention(MultiHeadAttention):
         b, h, t, d = x.size()
         return self._transpose_output(x, b, t, h * d)
 
+    def _sdpa(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout.p if self.training else 0,
+                                       is_causal=self.is_causal)
+
     def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
         b, t, d = query.size()
 
-        # Prepend global tokens to the input query
+        # Prepend global tokens
         global_tokens = self.global_tokens.expand(b, -1, -1)
         x = torch.cat([global_tokens, query], dim=1)
 
         # Project Q, K, V
-        q = self._split_heads(self.q_proj(x))
+        q = self._split_heads(self.q_proj(x), is_query=True)
         k = self._split_heads(self.k_proj(key))
         v = self._split_heads(self.v_proj(value))
 
@@ -380,42 +385,112 @@ class FlexAttention(MultiHeadAttention):
         if self.rope:
             q, k = self._apply_rope(q, k, separate=True)
 
-        # Split Q into global and local parts
-        global_q = q[:, :, :self.num_global_tokens]  # (B, H, G, head_dim)
-        local_q = q[:, :, self.num_global_tokens:]   # (B, H, L, head_dim)
+        # Split Q into global and local
+        global_q = q[:, :, :self.num_global_tokens]  # (B, H, G, D)
+        local_q = q[:, :, self.num_global_tokens:]   # (B, H, L, D)
+
+        # Global attention
+        global_attn = self._sdpa(global_q, k, v, mask=mask)
+
+        # Local attention with windowed slicing (no large masks)
         L = local_q.size(2)
         S = k.size(2)
+        windowed_attn = []
 
-        # Global attention: global_q attends to all K/V
-        global_attn = F.scaled_dot_product_attention(
-            global_q, k, v, attn_mask=mask if not self.is_causal else None,
-            dropout_p=self.dropout.p if self.training else 0, is_causal=self.is_causal
-        )
+        for i in range(0, L, self.window_size):
+            start = i
+            end = min(i + self.window_size, L)
+            window_q = local_q[:, :, start:end]  # (B, H, W, D)
 
-        # Local attention: local_q attends to windowed K/V
-        # Vectorized window mask
-        indices = torch.arange(S, device=local_q.device)
-        local_pos = torch.arange(L, device=local_q.device)
-        local_window = (local_pos // self.window_size).unsqueeze(-1)  # (L, 1)
-        key_window = (indices // self.window_size).expand(L, -1)      # (L, S)
-        window_mask = (local_window == key_window).to(device=local_q.device)
-
-        local_attn = F.scaled_dot_product_attention(
-            local_q, k, v, attn_mask=window_mask if not self.is_causal else None,
-            dropout_p=self.dropout.p if self.training else 0, is_causal=self.is_causal
-        )
+            # Use only relevant keys/values (same window)
+            k_window_start = max(0, start - self.window_size)
+            k_window_end = min(S, end + self.window_size)
+            window_k = k[:, :, k_window_start:k_window_end]
+            window_v = v[:, :, k_window_start:k_window_end]
+
+            window_attn = self._sdpa(window_q, window_k, window_v, mask=None)
+
+            windowed_attn.append(window_attn)
 
-        # Combine global and local attention outputs
-        attn = torch.cat([global_attn, local_attn], dim=2)  # (B, H, G+L, head_dim)
+        # Concat local attention
+        local_attn = torch.cat(windowed_attn, dim=2)
 
-        # Merge heads and project back
-        output = self._merge_heads(attn)  # (B, G+L, D)
-        output = self.out_proj(output)    # (B, G+L, D)
+        # Combine global and local
+        attn = torch.cat([global_attn, local_attn], dim=2)
+        output = self._merge_heads(attn)
+        output = self.out_proj(output)
 
-        # Return only the local tokens (original query tokens)
         return output[:, self.num_global_tokens:, :]
 
 
+class FlexSparseQueryAttention(FlexAttention):
+    """Combined Flex and Sparse Query attention layer, with RoPE support"""
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        num_groups: int,
+        num_query_groups: int,
+        dropout: float = 0.0,
+        rope: RotaryPositionalEmbedding = None,
+        rope_only_for_query: bool = False,
+        rope_only_for_keys: bool = False,
+        use_relative_embeddings: bool = False,
+        max_seq_len: int = 1024,
+        use_flash_attention: bool = True,
+        is_causal: bool = False,
+        use_bias: bool = False,
+        num_global_tokens: int = 16,
+        window_size: int = 128,
+    ):
+        self.num_groups = num_groups
+        self.num_query_groups = num_query_groups
+        super(FlexSparseQueryAttention, self).__init__(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            rope_only_for_query=rope_only_for_query,
+            rope_only_for_keys=rope_only_for_keys,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            num_global_tokens=num_global_tokens,
+            window_size=window_size,
+        )
+        assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"
+
+    def _init_kv(self, embed_dim: int):
+        self.k_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_groups), bias=self.use_bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_groups), bias=self.use_bias)
+
+    def _init_q(self, embed_dim: int):
+        self.q_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_query_groups), bias=self.use_bias)
+
+    def _init_out(self, embed_dim: int):
+        """Initialize output projection"""
+        self.out_proj = nn.Linear(embed_dim // (self.num_heads // self.num_query_groups), embed_dim)
+
+    def _split_heads(self, x: torch.Tensor, is_query: bool = False) -> torch.Tensor:
+        b, t, d = x.size()
+        return x.view(b, t, self.num_query_groups if is_query else self.num_groups, self.head_dim).transpose(1, 2)
+
+    def _transpose_output(self, attn_output: torch.Tensor, b: int, t: int, d: int):
+        """Transpose attention output back to (B, T, D) shape"""
+        return attn_output.transpose(1, 2).contiguous().view(b, t, d // (self.num_heads // self.num_query_groups))
+
+    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+        b, h, t, d = x.size()
+        return self._transpose_output(x, b, t, self.embed_dim)
+
+    def _sdpa(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
+        is_gqa = self.num_query_groups != self.num_groups
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=self.dropout.p if self.training else 0,
+                                       is_causal=self.is_causal, enable_gqa=is_gqa)
+
+
 class InfiniteAttention(MultiHeadAttention):
     def __init__(
             self,