rxnn 0.1.54__py3-none-any.whl → 0.1.55__py3-none-any.whl

rxnn/experimental/attention.py

@@ -226,6 +226,98 @@ class DeepMoeAttention(GroupedMoeAttention):
 
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
 
+class SparseQueryAttention(MultiHeadAttention):
+    """Sparse Grouped Query attention layer, with RoPE support"""
+
+    def __init__(
+            self,
+            embed_dim: int,
+            num_heads: int,
+            num_groups: int,
+            num_query_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,
+            *args,
+            **kwargs,
+    ):
+        self.num_groups = num_groups
+        self.num_query_groups = num_query_groups
+        super(SparseQueryAttention, self).__init__(
+            embed_dim,
+            num_heads,
+            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,
+        )
+        assert num_heads % num_groups == 0, "num_heads must be divisible by num_groups"
+
+    def _init_kv(self, embed_dim: int):
+        self.k_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_groups), bias=self.use_bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_groups), bias=self.use_bias)
+
+    def _init_q(self, embed_dim: int):
+        self.q_proj = nn.Linear(embed_dim, embed_dim // (self.num_heads // self.num_query_groups))
+
+    def _init_out(self, embed_dim: int):
+        """Initialize output projection"""
+        self.out_proj = nn.Linear(embed_dim // (self.num_heads // self.num_query_groups), 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"""
+        return attn_output.transpose(1, 2).contiguous().view(b, t, d // (self.num_heads // self.num_query_groups))
+
+    def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
+        """Override query, key, and value projections for GQA case - split data into heads and groups"""
+        head_dim = d // self.num_heads
+        if not self.rel_embed:
+            q = self.q_proj(query).view(b, t, self.num_query_heads, head_dim).transpose(1, 2)
+            k = self.k_proj(key).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
+            v = self.v_proj(value).view(b, -1, self.num_groups, head_dim).transpose(1, 2)
+        else:
+            group_heads = self.num_heads // self.num_groups
+            query_heads = self.num_query_heads // self.num_query_groups
+            # Process Q
+            q = self.q_proj(query).view(b, -1, self.num_query_groups, head_dim).transpose(1, 2)  # (B, Q_G, T, head_dim)
+
+            # Process K and V
+            k = self.k_proj(key).view(b, -1, self.num_groups, head_dim).transpose(1, 2)  # (B, G, S, head_dim)
+            v = self.v_proj(value).view(b, -1, self.num_groups, head_dim).transpose(1, 2)  # (B, G, S, head_dim)
+
+            # Expand and flatten to 4D tensors
+            q = q.unsqueeze(2).expand(-1, -1, query_heads, -1, -1)  # (B, Q_G, query_heads, T, head_dim)
+            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)
+
+            q = q.flatten(start_dim=1, end_dim=2)  # (B, Q, T, 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
+
+    def _calculate_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, b: int, t: int, d: int, mask: torch.Tensor = None):
+        is_gqa = self.num_query_groups != self.num_groups
+        if self.use_flash_attention:
+            # Compute attention with FlashAttention
+            return self._flash_attention(q.contiguous(), k.contiguous(), v.contiguous(), b, t, d, mask=mask, enable_gqa=is_gqa)
+        else:
+            # Compute attention using optimized PyTorch implementation
+            return self._torch_attention(q.contiguous(), k.contiguous(), v.contiguous(), b, t, d, mask=mask, enable_gqa=is_gqa)
+
+
+
+
 class GroupedMoeAttentionSimplified(GroupedQueryAttention):
     """
     Grouped MoE Attention (GMA) - GQA extended with Mixture-of-Experts (MoE) routing.
@@ -607,8 +699,8 @@ def init_moe_attention(
         num_experts: int = None,
         num_query_experts: int = None,
         num_query_groups: int = None,
-) -> GroupedQueryAttention:
-    assert attention_type in ['gma', 'dma', 'gma_s', 'dma_s'], "Error, attention type should be one of: 'gma', 'dma', 'gma_s', 'dma_s'"
+) -> MultiHeadAttention:
+    assert attention_type in ['gma', 'dma', 'gma_s', 'dma_s', 'sqa'], "Error, attention type should be one of: 'gma', 'dma', 'gma_s', 'dma_s', 'sqa'"
 
     if attention_type == "gma":
         return GroupedMoeAttention(
@@ -642,6 +734,21 @@ def init_moe_attention(
             num_query_experts=num_query_experts,
             num_query_groups=num_query_groups,
         )
+    elif attention_type == 'sqa':
+        return SparseQueryAttention(
+            embed_dim,
+            num_heads,
+            gqa_groups,
+            num_query_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,
+        )
     elif attention_type == "gma_s":
         return GroupedMoeAttentionSimplified(
             embed_dim,

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', 'gma_s', 'dma_s'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "gma_s", "dma_s".'
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)

{rxnn-0.1.54.dist-info → rxnn-0.1.55.dist-info}/METADATA

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

{rxnn-0.1.54.dist-info → rxnn-0.1.55.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=ZYdRxz4ik7knk3VS_9Opzy6ZqVF98FIhSNjsmIUhGfk,29532
-rxnn/experimental/models.py,sha256=-BQn7gWlSHLpkAQdthPW5L9ZNzIBqSJS9tkm2N88jgw,4711
+rxnn/experimental/attention.py,sha256=oPknT_PVcNwZvDwpZM7gmP4M_md_FW8oYwJDdQk1avM,34544
+rxnn/experimental/models.py,sha256=iprFSQDPK75zebDJBJ1i-mnNS9jlGf9RAIk-S0E9D-Q,4689
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 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=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 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.54.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.54.dist-info/METADATA,sha256=FF9XlvOeROGLpVR5pHuuceoeXTzbMNJhEusmQdfPTD0,16627
-rxnn-0.1.54.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.54.dist-info/RECORD,,
+rxnn-0.1.55.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.55.dist-info/METADATA,sha256=qiMp63aMlBdbvVvSJDL2bfW5XoR0PzNwN6pWdkfCuOM,16627
+rxnn-0.1.55.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.55.dist-info/RECORD,,