rxnn 0.1.26__py3-none-any.whl → 0.1.28__py3-none-any.whl

rxnn/experimental/attention.py

@@ -238,6 +238,231 @@ class DeepMoeAttention(GroupedMoeAttention):
         # Key/Value processing
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
 
+# Vectorized
+
+class GroupedMoeAttentionVectorized(GroupedQueryAttention):
+    """
+    Vectorized implementation calculates all expert heads for each token and selecting active tokens later. Linear layers
+    for Attention are rather small, compared to MoE Feed Forward layers, so it's possible that it will be faster than filtering
+    experts - it has to be tested.
+
+    Grouped MoE Attention (GMA) - GQA extended with Mixture-of-Experts (MoE) routing.
+
+    Instead of mapping keys/values to static head groups, it dynamically selects head expert groups. It has the same
+    number of total keys/values heads as query heads, but uses only a selected group for attention calculation.
+    - with num_groups set to 1, it will be MoE MultiQueryAttention
+
+    Compared to traditional GQA/MQA, it should provide better performance, because lot less data could be lost using
+    this approach - we are training the full number of keys/values heads, while using only a group.
+
+    In case of efficiency, it should be close to GQA/MQA linear performance, but with a small MoE routing overhead.
+
+    Optionally, it could use even more expert heads than attention heads - in example:
+    - 512 dim divided into 16 heads with 32 dim, using 4 head groups - may use i.e., 24 total expert heads - still only
+    4 will be used for attention calculation, while 16 is used to split dimensions (in that case it will have 16 query heads)
+
+    © 2025 Adam Filipek
+    """
+
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            num_groups: int,
+            dropout: float = 0.0,
+            rope: RotaryPositionalEmbedding = None,
+            rope_only_for_query: bool = False,
+            use_relative_embeddings: bool = False,
+            max_seq_len: int = 1024,
+            use_flash_attention: bool = False,
+            is_causal: bool = False,
+            use_bias: bool = False,
+            num_experts: int = None,
+            *args,
+            **kwargs,
+    ):
+        self.num_experts = num_experts if num_experts is not None else num_heads
+        super(GroupedMoeAttentionVectorized, self).__init__(
+            embed_dim,
+            num_heads,
+            num_groups=num_groups,
+            dropout=dropout,
+            rope=rope,
+            rope_only_for_query=rope_only_for_query,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            *args,
+            **kwargs,
+        )
+
+    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()
+
+    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)
+
+    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)
+
+        # Key/Value MoE routing
+        B, S, D = key.shape
+        key_flat = key.reshape(B * S, 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)
+
+        # Compute all experts' projections
+        # Shape: (B*S, num_experts, head_dim)
+        k_all = torch.einsum('be,ehd->beh', key_flat, self.wk)  # [B*S, num_experts, head_dim]
+        v_all = torch.einsum('be,ehd->beh', value.view(B*S, D), self.wv)
+
+        # Reshape to [B, S, num_experts, head_dim]
+        k_all = k_all.view(B, S, self.num_experts, -1)
+        v_all = v_all.view(B, S, self.num_experts, -1)
+
+        # Gather top-k experts and weights
+        # Expand indices to [B, S, num_groups, head_dim]
+        expanded_indices = indices.unsqueeze(-1).expand(-1, -1, -1, k_all.size(-1))
+        selected_k = torch.gather(k_all, 2, expanded_indices)  # [B, S, num_groups, head_dim]
+        selected_v = torch.gather(v_all, 2, expanded_indices)
+
+        # Weighted sum
+        weighted_k = (selected_k * weights).sum(dim=2)  # [B, S, head_dim]
+        weighted_v = (selected_v * weights).sum(dim=2)
+
+        # Reshape to GQA format
+        k = weighted_k.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3)  # [B, G, S, head_dim]
+        v = weighted_v.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3)
+
+        if not self.use_flash_attention:
+            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)
+            v = v.unsqueeze(2).expand(-1, -1, group_heads, -1, -1)  # (B, G, group_heads, S, head_dim)
+
+            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)
+
+        return q, k, v
+
+
+class DeepMoeAttentionVectorized(GroupedMoeAttentionVectorized):
+    """
+    Vectorized implementation calculates all expert heads for each token and selecting active tokens later. Linear layers
+    for Attention are rather small, compared to MoE Feed Forward layers, so it's possible that it will be faster than filtering
+    experts - it has to be tested.
+
+    Deep MoE Attention (SMA) - Grouped MoE Attention extended even more for sublinear computational efficiency.
+
+    In addition to using Mixture-of-Experts (MoE) for key/value head groups, SMA is also using dynamically selected
+    query heads - with that approach, each token could attend to every other token, but only partially - only some part of
+    information from each token is used to identify related information parts from other tokens. So, DMA is not spatially
+    sparse (has access to all tokens), but rather structurally sparse (has access only to the part of token's information).
+
+    This solution could reduce the computational complexity of attention operation to sublinear level (<O(N)) and provide
+    a viable and efficient alternative to spatial sparse attention mechanisms like Flex Attention.
+
+    © 2025 Adam Filipek
+    """
+
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            num_groups: int,
+            dropout: float = 0.0,
+            rope: RotaryPositionalEmbedding = None,
+            rope_only_for_query: bool = False,
+            use_relative_embeddings: bool = False,
+            max_seq_len: int = 1024,
+            use_flash_attention: bool = False,
+            is_causal: bool = False,
+            use_bias: bool = False,
+            num_experts: int = None,
+            num_query_experts: int = None,
+            num_query_groups: int = None,
+            *args,
+            **kwargs,
+    ):
+        self.num_query_experts = num_query_experts if num_query_experts is not None else num_heads
+        self.num_query_groups = num_query_groups if num_query_groups is not None else num_groups
+        super(DeepMoeAttentionVectorized, self).__init__(
+            embed_dim,
+            num_heads,
+            num_groups=num_groups,
+            dropout=dropout,
+            rope=rope,
+            rope_only_for_query=rope_only_for_query,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            num_experts=num_experts,
+            *args,
+            **kwargs,
+        )
+
+    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()
+
+    def _init_query_experts(self):
+        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"""
+        out_hidden_dim = embed_dim // self.num_heads * self.num_query_groups
+        self.out_proj = nn.Linear(out_hidden_dim, embed_dim)
+
+    def _transpose_output(self, attn_output: torch.Tensor, b: int, t: int, d: int):
+        """Transpose attention output back to (B, T, D) shape"""
+        out_hidden_dim = d // self.num_heads * self.num_query_groups
+        return attn_output.transpose(1, 2).contiguous().view(b, t, out_hidden_dim)
+
+    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(B * T, 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)
+
+        # Compute all query experts
+        q_all = torch.einsum('be,ehd->beh', query_flat, self.wq)  # [B*T, num_query_experts, head_dim]
+        q_all = q_all.view(B, T, self.num_query_experts, -1)
+
+        # Gather top-k experts
+        expanded_indices = indices_q.unsqueeze(-1).expand(-1, -1, -1, q_all.size(-1))
+        selected_q = torch.gather(q_all, 2, expanded_indices)  # [B, T, num_query_groups, head_dim]
+
+        # Weighted sum
+        q = (selected_q * weights_q).sum(dim=2)  # [B, T, head_dim]
+        q = q.view(B, T, self.num_query_groups, -1).permute(0, 2, 1, 3)  # [B, H_q, T, head_dim]
+
+        return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
+
 
 # Others
 
@@ -389,7 +614,7 @@ def init_moe_attention(
         num_query_experts: int = None,
         num_query_groups: int = None,
 ) -> GroupedQueryAttention:
-    assert attention_type in ['gma', 'dma'], "Error, attention type should be one of: 'gma', 'dma', 'gma_v', 'dma_v'"
+    assert attention_type in ['gma', 'dma', 'gma_v', 'dma_v'], "Error, attention type should be one of: 'gma', 'dma', 'gma_v', 'dma_v'"
 
     if attention_type == "gma":
         return GroupedMoeAttention(
@@ -406,7 +631,7 @@ def init_moe_attention(
             use_bias=use_bias,
             num_experts=num_experts,
         )
-    else:
+    elif attention_type == "dma":
         return DeepMoeAttention(
             embed_dim,
             num_heads,
@@ -423,3 +648,35 @@ def init_moe_attention(
             num_query_experts=num_query_experts,
             num_query_groups=num_query_groups,
         )
+    elif attention_type == "gma_v":
+        return GroupedMoeAttentionVectorized(
+            embed_dim,
+            num_heads,
+            gqa_groups,
+            dropout=dropout,
+            rope=rope,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            rope_only_for_query=rope_only_for_query,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            num_experts=num_experts,
+        )
+    else:
+        return DeepMoeAttentionVectorized(
+            embed_dim,
+            num_heads,
+            gqa_groups,
+            dropout=dropout,
+            rope=rope,
+            use_relative_embeddings=use_relative_embeddings,
+            max_seq_len=max_seq_len,
+            rope_only_for_query=rope_only_for_query,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            num_experts=num_experts,
+            num_query_experts=num_query_experts,
+            num_query_groups=num_query_groups,
+        )

rxnn/experimental/models.py

@@ -65,7 +65,7 @@ class MoeAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="tex
         assert ff_activation in ['relu', 'gelu',
                                  'swish', 'silu', 'linear',
                                  'sigmoid'], 'Feed-forward activation could be "relu", "gelu", "swish", "silu", "linear", "sigmoid".'
-        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma"'
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'gma_v', 'dma_v'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "gma_v", "dma_v".'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)

{rxnn-0.1.26.dist-info → rxnn-0.1.28.dist-info}/METADATA

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

{rxnn-0.1.26.dist-info → rxnn-0.1.28.dist-info}/RECORD

@@ -1,7 +1,7 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=rJSQjA7_9YqcM4Y8SyJSuZQjyz8j4XPhC5jcrcrRK2M,17891
-rxnn/experimental/models.py,sha256=8KAo7BtRkke9qRlzGRtQa9-EZ34roGWrn0N_T6L-6Wc,4561
+rxnn/experimental/attention.py,sha256=GnK_J7o_4fJ5O50ETx4oG-p7dOCsPRMwVGv3BIbUIbg,29439
+rxnn/experimental/models.py,sha256=QEuFBB9iEg5AbKQLwGJkAwPjMfaVeTqazhKDWPRkm7o,4598
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
@@ -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.26.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.26.dist-info/METADATA,sha256=qKx2OFl-ca7OCZoSf99l5lOmg2czSCN8UN_Esqmfco0,16627
-rxnn-0.1.26.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.26.dist-info/RECORD,,
+rxnn-0.1.28.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.28.dist-info/METADATA,sha256=zpVetjl-0pFz7Z4e4GUlybS-rBHKFk3AYIS6fM46diU,16627
+rxnn-0.1.28.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.28.dist-info/RECORD,,