rxnn 0.1.39__tar.gz → 0.1.41__tar.gz

{rxnn-0.1.39 → rxnn-0.1.41}/PKG-INFO

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

{rxnn-0.1.39 → rxnn-0.1.41}/pyproject.toml

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

{rxnn-0.1.39 → rxnn-0.1.41}/src/rxnn/experimental/attention.py

@@ -90,11 +90,13 @@ class GroupedMoeAttention(GroupedQueryAttention):
 
         # 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]
@@ -104,24 +106,44 @@ 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)
+
         # 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]
 
-        if not self.rel_embed:
+        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
 
             k = k.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
@@ -130,6 +152,10 @@ 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
 
 
@@ -220,10 +246,10 @@ class DeepMoeAttention(GroupedMoeAttention):
         if self.use_bias:
             q_all += self.bq
 
-        q_all = q_all.view(B, T, self.num_query_experts, -1)
+        q_all = q_all.view(B, T, self.num_query_experts, -1) # [B, T, num_query_experts, head_dim]
 
         # Gather top-k experts
-        expanded_indices = indices_q.unsqueeze(-1).expand(-1, -1, -1, q_all.size(-1))
+        expanded_indices = indices_q.unsqueeze(-1).expand(-1, -1, -1, q_all.size(-1)) # [B, T, num_query_groups, head_dim]
         selected_q = torch.gather(q_all, 2, expanded_indices)  # [B, T, num_query_groups, head_dim]
 
         # Weighted sum

{rxnn-0.1.39 → rxnn-0.1.41}/src/rxnn/transformers/attention.py

@@ -137,6 +137,10 @@ 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: