rxnn 0.2.53__py3-none-any.whl → 0.2.54__py3-none-any.whl

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:, :]
 
 

{rxnn-0.2.53.dist-info → rxnn-0.2.54.dist-info}/METADATA

@@ -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.dist-info → rxnn-0.2.54.dist-info}/RECORD

@@ -1,7 +1,7 @@
 rxnn/.DS_Store,sha256=BxZLo9tFs48JMq6jhumiCnCPLTeCwl619CFSg4ClRAY,6148
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=z9zahN1KTDcVabUB0RwWsT6oK8MTuz65haVwYRHAzy4,24689
+rxnn/experimental/attention.py,sha256=zQM6Og62IZVGogsDBReYrHSiRZmDaebl1FcH2e6sHyY,24589
 rxnn/experimental/models.py,sha256=HPOIRpnX_oiI10wsVC4J6rzo3T6dj10aNWGYpa9S1UU,5115
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -33,7 +33,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.53.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.53.dist-info/METADATA,sha256=CqOV8zz8qcGYyZXbeRv7jLMd40AeF8iU7vyfRsMuslE,25997
-rxnn-0.2.53.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.53.dist-info/RECORD,,
+rxnn-0.2.54.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.54.dist-info/METADATA,sha256=t6l1VezLNpdpgaXaqB-YhrfAhEUlWZm9-wwzBZ_Xk34,25997
+rxnn-0.2.54.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.54.dist-info/RECORD,,