rxnn 0.1.48__tar.gz → 0.1.49__tar.gz

{rxnn-0.1.48 → rxnn-0.1.49}/PKG-INFO

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

{rxnn-0.1.48 → rxnn-0.1.49}/pyproject.toml

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

{rxnn-0.1.48 → rxnn-0.1.49}/src/rxnn/experimental/attention.py

@@ -94,16 +94,17 @@ class GroupedMoeAttention(GroupedQueryAttention):
         B, S, D = key.shape
         key_flat = key.reshape(-1, D)
         weights, indices = self.router(key_flat)  # (B*S, num_groups), (B*S, num_groups)
-        weights = weights.view(B, self.num_groups, S, 1)
-        indices = indices.view(B, self.num_groups, S).unsqueeze(-1).expand(-1, -1, S, -1)
+        weights = weights.view(B, S, self.num_groups, 1)
+        indices = indices.view(B, S, self.num_groups)
 
         # Compute all experts' projections
-        k_all = self.k_proj(key_flat).view(B, S, self.num_experts, -1).permute(0, 2, 1, 3) # [B, num_experts, S, head_dim]
-        v_all = self.v_proj(value).view(B, S, self.num_experts, -1).permute(0, 2, 1, 3) # [B, num_experts, S, head_dim]
+        k_all = self.k_proj(key_flat).view(B, S, self.num_experts, -1) # [B, S, num_experts, head_dim]
+        v_all = self.v_proj(value).view(B, S, self.num_experts, -1) # [B, S, num_experts, head_dim]
 
         # Gather top-k experts using expanded indices
-        selected_k = torch.gather(k_all, 1, indices) # [B, num_groups, S, head_dim]
-        selected_v = torch.gather(v_all, 1, indices) # [B, num_groups, S, head_dim]
+        expanded_indices = indices.unsqueeze(-1).expand(-1, -1, -1, k_all.size(-1)) # [B, num_groups, S, head_dim]
+        selected_k = torch.gather(k_all, 2, expanded_indices) # [B, num_groups, S, head_dim]
+        selected_v = torch.gather(v_all, 2, expanded_indices) # [B, num_groups, S, head_dim]
 
         # Weighted
         weighted_k = (selected_k * weights).to(selected_k.device, dtype=selected_k.dtype) # [B, S, num_groups, head_dim]
@@ -122,8 +123,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(q.size(), k.size(), v.size())
-
         return q, k, v
 
 
@@ -208,17 +207,18 @@ class DeepMoeAttention(GroupedMoeAttention):
     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(-1, D)
-        weights, indices = self.query_router(query_flat)
-        weights = weights.view(B, self.num_query_groups, T, 1)
-        indices = indices.view(B, self.num_query_groups, T).unsqueeze(-1).expand(-1, -1, T, -1) # [B, num_query_groups, T, head_dim]
+        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)
 
-        q_all = self.q_proj(query_flat).view(B, T, self.num_query_experts, -1).permute(0, 2, 1, 3) # [B, num_query_experts, T, head_dim]
+        q_all = self.q_proj(query_flat).view(B, T, self.num_query_experts, -1) # [B, num_groups, S, head_dim]
 
-        # Gather top-k experts using expanded indices
-        selected_q = torch.gather(q_all, 1, indices) # [B, num_query_groups, T, head_dim]
+        # Gather top-k experts
+        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
-        q = (selected_q * weights).to(selected_q.device, dtype=selected_q.dtype)  # [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)