rxnn 0.2.53__tar.gz → 0.2.54__tar.gz

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

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

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

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

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

@@ -367,7 +367,7 @@ class FlexAttention(MultiHeadAttention):
     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)
 
@@ -380,39 +380,42 @@ 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 = F.scaled_dot_product_attention(global_q, k, v, attn_mask=mask, dropout_p=self.dropout.p if self.training else 0, is_causal=self.is_causal)
+
+        # 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 = F.scaled_dot_product_attention(window_q, window_k, window_v, attn_mask=None, dropout_p=self.dropout.p if self.training else 0, is_causal=self.is_causal)
+
+            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:, :]