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

rxnn/experimental/attention.py

@@ -68,18 +68,22 @@ class GroupedMoeAttention(GroupedQueryAttention):
     def _init_kv(self, embed_dim: int):
         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.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.bv = nn.Parameter(torch.zeros(self.num_experts, hidden_dim)) if self.use_bias else None
-        self._init_experts()
+        moe_dim = hidden_dim * self.num_experts
+        self.k_proj = nn.Linear(embed_dim, moe_dim, bias=self.use_bias)
+        self.v_proj = nn.Linear(embed_dim, moe_dim, bias=self.use_bias)
+        # self.wk = nn.Parameter(torch.empty(self.num_experts, embed_dim, hidden_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.bv = nn.Parameter(torch.zeros(self.num_experts, hidden_dim)) if self.use_bias else None
+        # self._init_experts()
 
     def _init_experts(self):
-        torch.nn.init.xavier_uniform_(self.wk)
-        torch.nn.init.xavier_uniform_(self.wv)
-        if self.use_bias:
-            torch.nn.init.zeros_(self.bk)
-            torch.nn.init.zeros_(self.bv)
+        pass
+        # torch.nn.init.xavier_uniform_(self.wk)
+        # torch.nn.init.xavier_uniform_(self.wv)
+        # if self.use_bias:
+        #     torch.nn.init.zeros_(self.bk)
+        #     torch.nn.init.zeros_(self.bv)
 
     def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int,
                      skip_query_processing: bool = False):
@@ -91,25 +95,15 @@ class GroupedMoeAttention(GroupedQueryAttention):
         key_flat = key.reshape(-1, D)
         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)
+        indices = indices.view(B, S, self.num_groups).unsqueeze(-1).expand(-1, -1, S, -1)
 
         # 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]
-        v_all = torch.einsum('bd,edh->beh', value.view(-1, D), self.wv)
-
-        if self.use_bias:
-            k_all += self.bk
-            v_all += self.bv
-
-        # 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]
+        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]
 
         # 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]
+        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]
 
         # Weighted
         weighted_k = (selected_k * weights).to(selected_k.device, dtype=selected_k.dtype) # [B, S, num_groups, head_dim]
@@ -187,14 +181,17 @@ class DeepMoeAttention(GroupedMoeAttention):
     def _init_q(self, embed_dim: int):
         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.bq = nn.Parameter(torch.zeros(self.num_query_experts, hidden_dim)) if self.use_bias else None
-        self._init_query_experts()
+        moe_dim = hidden_dim * self.num_query_experts
+        self.q_proj = nn.Linear(embed_dim, moe_dim)
+        # self.wq = nn.Parameter(torch.empty(self.num_query_experts, embed_dim, hidden_dim))
+        # self.bq = nn.Parameter(torch.zeros(self.num_query_experts, hidden_dim)) if self.use_bias else None
+        # self._init_query_experts()
 
     def _init_query_experts(self):
-        torch.nn.init.xavier_uniform_(self.wq)
-        if self.use_bias:
-            torch.nn.init.zeros_(self.bq)
+        pass
+        # torch.nn.init.xavier_uniform_(self.wq)
+        # if self.use_bias:
+        #     torch.nn.init.zeros_(self.bq)
 
     def _init_out(self, embed_dim: int):
         """Initialize output projection"""
@@ -209,23 +206,17 @@ 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_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)
+        weights, indices = self.query_router(query_flat)
+        weights = weights.view(B, T, self.num_query_groups, 1)
+        indices = indices.view(B, T, self.num_query_groups).unsqueeze(-1).expand(-1, -1, T, -1) # [B, num_query_groups, T, head_dim]
 
-        # Compute all query experts
-        q_all = torch.einsum('bd,edh->beh', query_flat, self.wq)  # [B*T, num_query_experts, head_dim]
-        if self.use_bias:
-            q_all += self.bq
-
-        q_all = q_all.view(B, T, self.num_query_experts, -1) # [B, T, num_query_experts, head_dim]
+        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]
 
-        # 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]
+        # Gather top-k experts using expanded indices
+        selected_q = torch.gather(q_all, 1, indices) # [B, num_query_groups, T, head_dim]
 
         # Weighted sum
-        q = (selected_q * weights_q).to(selected_q.device, dtype=selected_q.dtype)  # [B, T, num_query_groups, head_dim]
+        q = (selected_q * weights).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-0.1.43.dist-info → rxnn-0.1.45.dist-info}/METADATA

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

{rxnn-0.1.43.dist-info → rxnn-0.1.45.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=tyeZ2xkZfFrlTWzqjKS13NnIRSLGe1-oR7J20hN6BgI,29602
+rxnn/experimental/attention.py,sha256=gj1Zi7GGxn2DYJiBjlbMaNWwfmuUxVxi7j94AaJlm44,29387
 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
@@ -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.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,,
+rxnn-0.1.45.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.45.dist-info/METADATA,sha256=E-oJ9C52lOKVC4ZCnYQRSfqh9zhGFjLfc8OF_BWlUfc,16627
+rxnn-0.1.45.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.45.dist-info/RECORD,,