rxnn 0.1.41__py3-none-any.whl → 0.1.43__py3-none-any.whl

rxnn/experimental/attention.py

@@ -83,20 +83,16 @@ class GroupedMoeAttention(GroupedQueryAttention):
 
     def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int,
                      skip_query_processing: bool = False):
-        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) if not skip_query_processing else query
 
         # Key/Value MoE routing
         B, S, D = key.shape
-        print('key/value type', key.dtype, value.dtype)
         key_flat = key.reshape(-1, D)
-        print('key flat type', key_flat.dtype)
         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)
-        print('weights/indices type', weights.dtype, indices.dtype)
+
         # 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]
@@ -106,43 +102,23 @@ class GroupedMoeAttention(GroupedQueryAttention):
             k_all += self.bk
             v_all += self.bv
 
-        print('k all/v all before get all')
-        print(k_all.size(), k_all.dtype)
-        print(v_all.size(), v_all.dtype)
-
         # 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]
 
-        print('k all/v all get all')
-        print(k_all.size(), k_all.dtype)
-        print(v_all.size(), v_all.dtype)
-
         # 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('selected k/selected v')
-        print(selected_k.size(), selected_k.dtype)
-        print(selected_v.size(), selected_v.dtype)
-
         # Weighted
-        weighted_k = selected_k * weights # [B, S, num_groups, head_dim]
-        weighted_v = selected_v * weights # [B, S, num_groups, head_dim]
-
-        print('weighted')
-        print(weighted_k.size(), weighted_k.dtype)
-        print(weighted_v.size(), weighted_v.dtype)
+        weighted_k = (selected_k * weights).to(selected_k.device, dtype=selected_k.dtype) # [B, S, num_groups, head_dim]
+        weighted_v = (selected_v * weights).to(selected_k.device, dtype=selected_k.dtype)  # [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, 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]
 
-        print('out 1')
-        print(k.size(), k.dtype)
-        print(v.size(), v.dtype)
-
         if self.rel_embed:
             group_heads = self.num_heads // self.num_groups
 
@@ -152,10 +128,6 @@ class GroupedMoeAttention(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('out 2')
-        print(k.size(), k.dtype)
-        print(v.size(), v.dtype)
-
         return q, k, v
 
 
@@ -253,7 +225,7 @@ class DeepMoeAttention(GroupedMoeAttention):
         selected_q = torch.gather(q_all, 2, expanded_indices)  # [B, T, num_query_groups, head_dim]
 
         # Weighted sum
-        q = selected_q * weights_q  # [B, T, num_query_groups, head_dim]
+        q = (selected_q * weights_q).to(selected_q.device, dtype=selected_q.dtype)  # [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)

rxnn/transformers/attention.py

@@ -106,9 +106,6 @@ class MultiHeadAttention(nn.Module):
         # with sdpa_kernel(backends=[SDPBackend.FLASH_ATTENTION]):
         #     return self._torch_attention(q, k, v, b, t, d, mask=mask, enable_gqa=enable_gqa)
         from flash_attn import flash_attn_func
-        print(q.size(), q.dtype)
-        print(k.size(), k.dtype)
-        print(v.size(), v.dtype)
         attn_output = flash_attn_func(q, k, v, dropout_p=self.dropout.p if self.training else 0.0, causal=self.is_causal)
         return self._transpose_output(attn_output, b, t, d)
 
@@ -137,10 +134,6 @@ class MultiHeadAttention(nn.Module):
         return self._calculate_output(attn_weights, v, b, t, d)
 
     def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
-        print('MHA forward')
-        print(query.size(), query.dtype)
-        print(key.size(), key.dtype)
-        print(value.size(), value.dtype)
         b, t, d = query.size()
         q, k, v = self._forward_qkv(query, key, value, b, t, d)
         if not self.rel_embed:

{rxnn-0.1.41.dist-info → rxnn-0.1.43.dist-info}/METADATA

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

{rxnn-0.1.41.dist-info → rxnn-0.1.43.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=tvyjyBWLh9Bcu3MXlo3yx1hA48La8Zks4wWYtcnOYc8,30365
+rxnn/experimental/attention.py,sha256=tyeZ2xkZfFrlTWzqjKS13NnIRSLGe1-oR7J20hN6BgI,29602
 rxnn/experimental/models.py,sha256=IzUVc5s-cA__8jsG2mVvzUDmzPRcfBcI5btaOjnPYhA,4598
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 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=ZDdu1plCLhh-c5x-ZSMD1Avl36LMuTJSq617hD_hLCg,15942
+rxnn/transformers/attention.py,sha256=dC0UmC-_kjX8US6Sf0Fi5zw5kJ-P6orH3JDHeBB5gI8,15695
 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.41.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.41.dist-info/METADATA,sha256=B9AeelUbhWbH1uYB17WIEz4gjo_1c0V5jTM5S2QgulE,16627
-rxnn-0.1.41.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.41.dist-info/RECORD,,
+rxnn-0.1.43.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.43.dist-info/METADATA,sha256=7iCT0KM5FyGKZwe9ZbxnULu9z-rwyko4Xbklok1ZsOs,16627
+rxnn-0.1.43.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.43.dist-info/RECORD,,