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

rxnn/experimental/attention.py

@@ -319,7 +319,7 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
         head_dim = d // self.num_heads
 
         # Process Query as in GQA
-        q = self.q_proj(query).view(b, t, self.num_heads, -1).transpose(1, 2)
+        q = self.q_proj(query).view(b, t, self.num_heads, -1).transpose(1, 2) if not skip_query_processing else query
 
         # Key/Value MoE routing
         B, S, D = key.shape
@@ -337,16 +337,26 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
         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)
 
+        print('selected_k', selected_k.size())
+        print('selected_v', selected_v.size())
+
         # Weighted sum
         weighted_k = (selected_k * weights).sum(dim=2)  # [B, S, head_dim]
         weighted_v = (selected_v * weights).sum(dim=2)
 
+        print('weighted_k', weighted_k.size())
+        print('weighted_v', weighted_v.size())
+
         # 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)
@@ -360,6 +370,10 @@ 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
 
 

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

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

{rxnn-0.1.29.dist-info → rxnn-0.1.31.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=FLIyfRWZC1Ppfx9u8qUcNcfMSscuibhKEQ8zCJHXcfk,29439
+rxnn/experimental/attention.py,sha256=lqRIR2V7pA5LZQ1x_nhUvQBeJ83DLD6tXdoee-W6ZdU,29833
 rxnn/experimental/models.py,sha256=QEuFBB9iEg5AbKQLwGJkAwPjMfaVeTqazhKDWPRkm7o,4598
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -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.29.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.29.dist-info/METADATA,sha256=u3tlE_n8tZsqgW44vPssCmmJ03wHKM57tGS1xXI3HcQ,16627
-rxnn-0.1.29.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.29.dist-info/RECORD,,
+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,,