rxnn 0.1.31__py3-none-any.whl → 0.1.33__py3-none-any.whl

rxnn/experimental/attention.py

@@ -323,7 +323,7 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
 
         # Key/Value MoE routing
         B, S, D = key.shape
-        key_flat = key.reshape(B * S, D)
+        key_flat = key.reshape(-1, D)
         weights, indices = self.router(key_flat)  # (B*S, num_groups), (B*S, num_groups)
         weights = weights.view(B, S, self.num_groups, 1)
         indices = indices.view(B, S, self.num_groups)
@@ -331,35 +331,28 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
         # Compute all experts' projections
         # Shape: (B*S, num_experts, head_dim)
         k_all = torch.einsum('bd,edh->beh', key_flat, self.wk)  # [B*S, num_experts, head_dim]
-        v_all = torch.einsum('bd,edh->beh', value.view(B*S, D), self.wv)
+        v_all = torch.einsum('bd,edh->beh', value.view(-1, D), self.wv)
 
-        # Reshape to [B, S, num_experts, head_dim]
-        k_all = k_all.view(B, S, self.num_experts, -1)
-        v_all = v_all.view(B, S, self.num_experts, -1)
-
-        print('k_all', k_all.size())
-        print('v_all', v_all.size())
-
-
-        # Gather top-k experts and weights
-        # Expand indices to [B, S, num_groups, head_dim]
-        expanded_indices = indices.unsqueeze(-1).expand(-1, -1, -1, k_all.size(-1))
-        selected_k = torch.gather(k_all, 2, expanded_indices)  # [B, S, num_groups, head_dim]
-        selected_v = torch.gather(v_all, 2, expanded_indices)
+        if self.use_bias:
+            k_all += self.bk
+            v_all += self.bv
 
-        print('selected_k', selected_k.size())
-        print('selected_v', selected_v.size())
+        # Get results for all heads
+        k_all = k_all.view(B, S, self.num_experts, -1) # [B, S, num_experts, head_dim]
+        v_all = v_all.view(B, S, self.num_experts, -1) # [B, S, num_experts, head_dim]
 
-        # Weighted sum
-        weighted_k = (selected_k * weights).sum(dim=2)  # [B, S, head_dim]
-        weighted_v = (selected_v * weights).sum(dim=2)
+        # Gather top-k experts using expanded indices
+        expanded_indices = indices.unsqueeze(-1).expand(-1, -1, -1, k_all.size(-1)) # [B, S, num_groups, head_dim]
+        selected_k = torch.gather(k_all, 2, expanded_indices) # [B, S, num_groups, head_dim]
+        selected_v = torch.gather(v_all, 2, expanded_indices) # [B, S, num_groups, head_dim]
 
-        print('weighted_k', weighted_k.size())
-        print('weighted_v', weighted_v.size())
+        # Weighted
+        weighted_k = selected_k * weights # [B, S, num_groups, head_dim]
+        weighted_v = selected_v * weights # [B, S, num_groups, head_dim]
 
         # Reshape to GQA format
-        k = weighted_k.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3)  # [B, G, S, head_dim]
-        v = weighted_v.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3)
+        k = weighted_k.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3) # [B, num_groups, S, head_dim]
+        v = weighted_v.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3) # [B, num_groups, S, head_dim]
 
         if not self.use_flash_attention:
             group_heads = self.num_heads // self.num_groups
@@ -370,10 +363,6 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
             k = k.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
             v = v.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
 
-        print('q', q.size())
-        print('k', k.size())
-        print('v', v.size())
-
         return q, k, v
 
 
@@ -458,13 +447,16 @@ class DeepMoeAttentionVectorized(GroupedMoeAttentionVectorized):
 
     def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
         B, T, D = query.shape
-        query_flat = query.reshape(B * T, D)
+        query_flat = query.reshape(-1, D)
         weights_q, indices_q = self.query_router(query_flat)
         weights_q = weights_q.view(B, T, self.num_query_groups, 1)
         indices_q = indices_q.view(B, T, self.num_query_groups)
 
         # Compute all query experts
         q_all = torch.einsum('bd,edh->beh', query_flat, self.wq)  # [B*T, num_query_experts, head_dim]
+        if self.use_bias:
+            q_all += self.bq
+
         q_all = q_all.view(B, T, self.num_query_experts, -1)
 
         # Gather top-k experts
@@ -472,12 +464,11 @@ class DeepMoeAttentionVectorized(GroupedMoeAttentionVectorized):
         selected_q = torch.gather(q_all, 2, expanded_indices)  # [B, T, num_query_groups, head_dim]
 
         # Weighted sum
-        q = (selected_q * weights_q).sum(dim=2)  # [B, T, head_dim]
-        q = q.view(B, T, self.num_query_groups, -1).permute(0, 2, 1, 3)  # [B, H_q, T, head_dim]
+        q = selected_q * weights_q  # [B, T, num_query_groups, head_dim]
+        q = q.view(B, T, self.num_query_groups, -1).permute(0, 2, 1, 3)  # [B, num_query_groups, T, head_dim]
 
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
 
-
 # Others
 
 class FlexAttention(MultiHeadAttention):

rxnn/transformers/attention.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.backends.cuda import sdp_kernel, SDPBackend
 import math
 from .positional import RotaryPositionalEmbedding, RelativePositionalEmbedding
 
@@ -17,7 +18,7 @@ class MultiHeadAttention(nn.Module):
             rope_only_for_query: bool = False,
             use_relative_embeddings: bool = False,
             max_seq_len: int = 1024,
-            use_flash_attention: bool = False,
+            use_flash_attention: bool = True,
             is_causal: bool = False,
             use_bias: bool = False,
             *args,
@@ -101,13 +102,14 @@ class MultiHeadAttention(nn.Module):
 
     def _flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,
                          mask: torch.Tensor = None, enable_gqa: bool = False):
-        attn_output = F.scaled_dot_product_attention(
-            q, k, v,
-            attn_mask=mask if not self.is_causal else None,
-            dropout_p=self.dropout.p if self.training else 0.0,
-            is_causal=self.is_causal,
-            enable_gqa=enable_gqa,
-        )
+        with sdp_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
+            attn_output = F.scaled_dot_product_attention(
+                q, k, v,
+                attn_mask=mask if not self.is_causal else None,
+                dropout_p=self.dropout.p if self.training else 0.0,
+                is_causal=self.is_causal,
+                enable_gqa=enable_gqa,
+            )
         return self._transpose_output(attn_output, b, t, d)
 
     def _calculate_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int,

{rxnn-0.1.31.dist-info → rxnn-0.1.33.dist-info}/METADATA

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

{rxnn-0.1.31.dist-info → rxnn-0.1.33.dist-info}/RECORD

@@ -1,6 +1,6 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=lqRIR2V7pA5LZQ1x_nhUvQBeJ83DLD6tXdoee-W6ZdU,29833
+rxnn/experimental/attention.py,sha256=h9pEv_70NKpD5KOQOFP3h-IJKzh7Wbnaxka4Bd3rdd8,29745
 rxnn/experimental/models.py,sha256=QEuFBB9iEg5AbKQLwGJkAwPjMfaVeTqazhKDWPRkm7o,4598
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,7 +16,7 @@ rxnn/training/dataset.py,sha256=vQ5mDF3bA0HXya474n4D4iL8Mn3AEpJukgzFNVkxjGU,5106
 rxnn/training/scheduler.py,sha256=ow6oALzWjWQmHSpcJEjv6tg4g4CDMvr73TypxfcefMc,712
 rxnn/training/tokenizer.py,sha256=4Y41f07uo2KPA_7bp3FCcwGKbXoS2hsckOoXUsXfQxY,8052
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/transformers/attention.py,sha256=FHATZVf_kt3OHnG02zEeG9QdUXLncKDjrhyT28Pk0E4,14185
+rxnn/transformers/attention.py,sha256=Nox986BH9qq4rDYLiYmfj1DeMeULF3akexIl99MPccM,14331
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
 rxnn/transformers/layers.py,sha256=n_jZTqEF_vLkF31AkB5XGErfm2sQFd9CRqJUHKRFkKI,6956
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
@@ -25,7 +25,7 @@ rxnn/transformers/moe.py,sha256=6Cffyo0QjmEWc4rK1ncOmLRCQbY0OpQJ4D7xH_4nTN4,4738
 rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.31.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.31.dist-info/METADATA,sha256=jy1wfyPS-nD34GAMFxcWR_dFn6ZJogJ2-IDxIVmLdmI,16627
-rxnn-0.1.31.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.31.dist-info/RECORD,,
+rxnn-0.1.33.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.33.dist-info/METADATA,sha256=m3DWDnTu7Lx1kHYPIAQCdKU8t4QZBdqG0QcSIFvB924,16627
+rxnn-0.1.33.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.33.dist-info/RECORD,,