rxnn 0.1.19__tar.gz → 0.1.20__tar.gz

{rxnn-0.1.19 → rxnn-0.1.20}/PKG-INFO

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

{rxnn-0.1.19 → rxnn-0.1.20}/pyproject.toml

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

{rxnn-0.1.19 → rxnn-0.1.20}/src/rxnn/experimental/attention.py

@@ -65,9 +65,9 @@ class GroupedMoeAttention(GroupedQueryAttention):
         self.router = MoeRouter(embed_dim, self.num_experts, top_k=self.num_groups)
 
         hidden_dim = embed_dim // self.num_heads
-        self.wk = nn.Parameter(torch.empty(self.num_experts, embed_dim, hidden_dim))
+        self.wk = nn.Parameter(torch.empty(self.num_experts, hidden_dim, embed_dim))
         self.bk = nn.Parameter(torch.zeros(self.num_experts, hidden_dim)) if self.use_bias else None
-        self.wv = nn.Parameter(torch.empty(self.num_experts, embed_dim, hidden_dim))
+        self.wv = nn.Parameter(torch.empty(self.num_experts, hidden_dim, embed_dim))
         self.bv = nn.Parameter(torch.zeros(self.num_experts, hidden_dim)) if self.use_bias else None
         self._init_experts()
 
@@ -80,34 +80,34 @@ class GroupedMoeAttention(GroupedQueryAttention):
 
     def _process_grouped_experts(self, x: torch.Tensor, w: torch.Tensor, b: torch.Tensor, weights: torch.Tensor, indices: torch.Tensor):
         B, S, G = indices.shape
-        x_flat = x.view(-1, x.size(-1))
+        x_flat = x.view(-1, x.size(-1))  # [B*S, D]
 
-        # Flatten batch and sequence dimensions
-        indices_flat = indices.view(-1, G)
-        weights_flat = weights.view(-1, G, 1)
+        indices_flat = indices.view(-1, G)  # [B*S, G]
+        weights_flat = weights.view(-1, G)  # [B*S, G]
 
-        # Create expanded indices for expert processing
-        mask = torch.zeros(B * S, self.num_experts, device=x.device, dtype=torch.bool)
-        for g in range(G):
-            mask.scatter_(1, indices_flat[:, g].unsqueeze(1), True)
-
-        output = torch.zeros(B * S, G, w.size(2), device=x.device, dtype=x.dtype)
+        output = torch.zeros(B * S, G, w.size(1), device=x.device, dtype=x.dtype)  # [B*S, G, hidden_dim]
 
         for e in range(self.num_experts):
-            token_mask = mask[:, e]
-            if not token_mask.any():
+            # 1. Find tokens where expert `e` is used in ANY group
+            expert_mask = (indices_flat == e).any(dim=1)  # [B*S]
+            if not expert_mask.any():
                 continue
 
-            # Get positions where expert e is used in any group
-            x_slice = x_flat[token_mask]
-            proj = F.linear(x_slice, w[e].t(), b[e] if b is not None else None)
+            # 2. Project tokens using expert `e`
+            x_slice = x_flat[expert_mask]  # [num_selected, D]
+            proj = F.linear(x_slice, w[e], b[e] if b is not None else None)  # [num_selected, hidden_dim]
 
-            # Find which groups use this expert for selected tokens
-            group_mask = (indices_flat[token_mask] == e)
+            # 3. Scatter projections into correct groups
+            for g in range(G):
+                group_mask = indices_flat[expert_mask, g] == e  # [num_selected]
+                if not group_mask.any():
+                    continue
 
-            # Accumulate projections for relevant groups
-            weighted_proj = proj.unsqueeze(1) * weights_flat[token_mask] * group_mask.unsqueeze(-1).float()
-            output[token_mask] += weighted_proj.sum(dim=1)
+                # Get tokens in this group using expert `e`
+                group_tokens = expert_mask.nonzero()[group_mask].squeeze(1)
+                # Weight and scatter
+                weighted_proj = proj[group_mask] * weights_flat[group_tokens, g].unsqueeze(-1)
+                output[group_tokens, g] += weighted_proj
 
         return output.view(B, S, G, -1)
 
@@ -118,7 +118,6 @@ class GroupedMoeAttention(GroupedQueryAttention):
         # Key/Value processing
         B, S, D = key.shape
         key_flat = key.view(-1, D)
-        print('key_flat: ', key_flat.shape)
         weights_k_flat, indices_k_flat = self.router(key_flat)
         # Reshape back to original dimensions
         weights_k = weights_k_flat.view(B, S, -1)
@@ -199,7 +198,7 @@ class DeepMoeAttention(GroupedMoeAttention):
         self.query_router = MoeRouter(embed_dim, self.num_query_experts, top_k=self.num_query_groups)
 
         hidden_dim = embed_dim // self.num_heads
-        self.wq = nn.Parameter(torch.empty(self.num_query_experts, embed_dim, hidden_dim))
+        self.wq = nn.Parameter(torch.empty(self.num_query_experts, hidden_dim, embed_dim))
         self.bq = nn.Parameter(torch.zeros(self.num_query_experts, hidden_dim)) if self.use_bias else None
         self._init_query_experts()
 
@@ -217,7 +216,7 @@ class DeepMoeAttention(GroupedMoeAttention):
         # Query processing
         B, T, D = query.shape
         # Flatten for query routing
-        query_flat = query.view(B * T, D)
+        query_flat = query.view(-1, D)
         weights_q_flat, indices_q_flat = self.query_router(query_flat)
         # Reshape back
         weights_q = weights_q_flat.view(B, T, -1)

{rxnn-0.1.19 → rxnn-0.1.20}/src/rxnn/transformers/moe.py

@@ -16,26 +16,20 @@ class MoeRouter(nn.Module):
 
     def calculate_aux_loss(self, top_k_indices: torch.Tensor, probs: torch.Tensor) -> torch.Tensor:
         expert_mask = F.one_hot(top_k_indices, self.num_experts).float()
-        print('expert mask: ', expert_mask.shape)
         expert_usage = expert_mask.sum(dim=0).mean(dim=0)
-        print('expert usage: ', expert_usage.shape)
         mean_probs = probs.mean(dim=0)
-        print('mean probs: ', mean_probs.shape)
         return (expert_usage * mean_probs).sum() * self.num_experts
 
 
     def forward(self, x: torch.Tensor):
         # Input shape: [batch*seq_len, embed_dim]
         logits = self.gate(x)
-        print('router logits: ', logits.shape)
         probs = F.softmax(logits, dim=-1)
-        print('router probs: ', probs.shape)
         # Get top-k experts for each token
         top_k_weights, top_k_indices = probs.topk(self.top_k, dim=-1)
 
         # Normalize weights (sum to 1 for each token)
         top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-9)
-        print('top k: ', top_k_weights.shape, top_k_indices.shape)
         # Load Balance Loss
         self.aux_loss = self.calculate_aux_loss(top_k_indices, probs)