rxnn 0.2.52__py3-none-any.whl → 0.2.53__py3-none-any.whl

rxnn/experimental/attention.py

@@ -318,76 +318,102 @@ class SparseQueryAttention(MultiHeadAttention):
 
 
 # Others
-
 class FlexAttention(MultiHeadAttention):
     def __init__(
-            self,
-            embed_dim: int,
-            num_heads: int,
-            num_global_tokens: int = 16,
-            window_size: int = 128,
-            **kwargs
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        rope: RotaryPositionalEmbedding = None,
+        rope_only_for_query: bool = False,
+        rope_only_for_keys: bool = False,
+        use_relative_embeddings: bool = False,
+        max_seq_len: int = 1024,
+        use_flash_attention: bool = True,
+        is_causal: bool = False,
+        use_bias: bool = False,
+        num_global_tokens: int = 16,
+        window_size: int = 128,
     ):
-        super().__init__(embed_dim, num_heads, **kwargs)
+        super(FlexAttention, self).__init__(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            rope_only_for_query=rope_only_for_query,
+            rope_only_for_keys=rope_only_for_keys,
+            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,
+        )
+        self.head_dim = embed_dim // num_heads
         self.num_global_tokens = num_global_tokens
         self.window_size = window_size
-        self.global_tokens = nn.Parameter(torch.zeros(1, num_global_tokens, embed_dim))
 
-    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask=None):
+        # Learnable global tokens
+        self.global_tokens = nn.Parameter(torch.randn(1, num_global_tokens, embed_dim))
+
+
+    def _split_heads(self, x: torch.Tensor) -> torch.Tensor:
+        b, t, d = x.size()
+        return x.view(b, t, self.num_heads, self.head_dim).transpose(1, 2)
+
+    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+        b, h, t, d = x.size()
+        return self._transpose_output(x, b, t, h * d)
+
+    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
         b, t, d = query.size()
-        head_dim = d // self.num_heads
 
-        # Split into global and local
-        x = torch.cat([self.global_tokens.expand(b, -1, -1), query], dim=1)
-        seq_len = x.size(1)
-        num_windows = (seq_len - self.num_global_tokens + self.window_size - 1) // self.window_size
+        # Prepend global tokens to the input query
+        global_tokens = self.global_tokens.expand(b, -1, -1)
+        x = torch.cat([global_tokens, query], dim=1)
 
         # Project Q, K, V
-        q, k, v = self._forward_qkv(x, key, value, b, seq_len, d)
-
-        # Process Global-to-Global Attention
-        global_q = q[:, :, :self.num_global_tokens]  # [B, H, G, head_dim]
-        global_k = k[:, :, :self.num_global_tokens]
-        global_v = v[:, :, :self.num_global_tokens]
-        global_attn = self._calculate_attn_weights(global_q, global_k, d) @ global_v
-
-        # Process Global-to-Local Attention
-        local_k = k[:, :, self.num_global_tokens:]  # [B, H, (num_windows * window_size), head_dim]
-        local_v = v[:, :, self.num_global_tokens:]
-        # Apply RoPE to local_k if needed
+        q = self._split_heads(self.q_proj(x))
+        k = self._split_heads(self.k_proj(key))
+        v = self._split_heads(self.v_proj(value))
+
+        # Apply RoPE
         if self.rope:
-            # Compute frequencies for entire local sequence
-            local_k = self.rope.forward_one(local_k)
-
-        global_local_attn = self._calculate_attn_weights(global_q, local_k, d) @ local_v
-
-        # Process Local-to-Local Attention (per window)
-        local_q = q[:, :, self.num_global_tokens:]  # [B, H, (num_windows * window_size), head_dim]
-        local_q = local_q.view(b, self.num_heads, num_windows, self.window_size, head_dim)
-        local_k = local_k.view(b, self.num_heads, num_windows, self.window_size, head_dim)
-        local_v = local_v.view(b, self.num_heads, num_windows, self.window_size, head_dim)
-
-        local_attn = []
-        for i in range(num_windows):
-            window_q = local_q[:, :, i]  # [B, H, window_size, head_dim]
-            window_k = local_k[:, :, i]
-            window_v = local_v[:, :, i]
-
-            # Apply RoPE to window_q and window_k
-            if self.rope:
-                # Compute frequencies for this window
-                window_q, window_k = self.rope(window_q, window_k)
-
-            # Calculate attention for this window
-            attn = self._calculate_attn_weights(window_q, window_k, d)
-            attn_i = torch.einsum('bhij, bhjd -> bhid', attn, window_v)
-            local_attn.append(attn_i)
-        local_attn = torch.cat(local_attn, dim=2).view(b, self.num_heads, -1, head_dim)
-
-        # Combine all attention outputs
-        combined_attn = torch.cat([global_attn, global_local_attn, local_attn], dim=2)
-        output = self._calculate_output(combined_attn, v, b, t, d)
-        return self.out_proj(output)
+            q, k = self._apply_rope(q, k, separate=True)
+
+        # Split Q into global and local parts
+        global_q = q[:, :, :self.num_global_tokens]  # (B, H, G, head_dim)
+        local_q = q[:, :, self.num_global_tokens:]   # (B, H, L, head_dim)
+        L = local_q.size(2)
+        S = k.size(2)
+
+        # Global attention: global_q attends to all K/V
+        global_attn = F.scaled_dot_product_attention(
+            global_q, k, v, attn_mask=mask if not self.is_causal else None,
+            dropout_p=self.dropout.p if self.training else 0, is_causal=self.is_causal
+        )
+
+        # Local attention: local_q attends to windowed K/V
+        # Vectorized window mask
+        indices = torch.arange(S, device=local_q.device)
+        local_pos = torch.arange(L, device=local_q.device)
+        local_window = (local_pos // self.window_size).unsqueeze(-1)  # (L, 1)
+        key_window = (indices // self.window_size).expand(L, -1)      # (L, S)
+        window_mask = (local_window == key_window).to(device=local_q.device)
+
+        local_attn = F.scaled_dot_product_attention(
+            local_q, k, v, attn_mask=window_mask if not self.is_causal else None,
+            dropout_p=self.dropout.p if self.training else 0, is_causal=self.is_causal
+        )
+
+        # Combine global and local attention outputs
+        attn = torch.cat([global_attn, local_attn], dim=2)  # (B, H, G+L, head_dim)
+
+        # Merge heads and project back
+        output = self._merge_heads(attn)  # (B, G+L, D)
+        output = self.out_proj(output)    # (B, G+L, D)
+
+        # Return only the local tokens (original query tokens)
+        return output[:, self.num_global_tokens:, :]
 
 
 class InfiniteAttention(MultiHeadAttention):
@@ -467,8 +493,10 @@ def init_experimental_attention(
         num_experts: int = None,
         num_query_experts: int = None,
         num_query_groups: int = None,
+        num_global_tokens: int = 16,
+        window_size: int = 128,
 ) -> MultiHeadAttention:
-    assert attention_type in ['gma', 'dma', 'sqa'], "Error, attention type should be one of: 'gma', 'dma', 'sqa'"
+    assert attention_type in ['gma', 'dma', 'sqa', 'flex'], "Error, attention type should be one of: 'gma', 'dma', 'sqa', 'flex'"
 
     if attention_type == "gma":
         return GroupedMoeAttention(
@@ -504,6 +532,21 @@ def init_experimental_attention(
             num_query_experts=num_query_experts,
             num_query_groups=num_query_groups,
         )
+    elif attention_type == "flex":
+        return FlexAttention(
+            embed_dim,
+            num_heads,
+            dropout=dropout,
+            rope=rope,
+            max_seq_len=max_seq_len,
+            rope_only_for_query=rope_only_for_query,
+            rope_only_for_keys=rope_only_for_keys,
+            use_flash_attention=use_flash_attention,
+            is_causal=is_causal,
+            use_bias=use_bias,
+            num_global_tokens=num_global_tokens,
+            window_size=window_size,
+        )
     else:
         return SparseQueryAttention(
             embed_dim,

rxnn/experimental/models.py

@@ -32,6 +32,8 @@ class ExperimentalAttentionTransformerConfig(TypedDict):
     att_num_experts: int
     att_num_query_experts: int
     att_num_query_groups: int
+    att_num_global_tokens: int
+    att_window_size: int
 
 
 class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline_tag="text-generation", license="apache-2.0"):
@@ -63,13 +65,15 @@ class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline
             att_num_experts: int = None,
             att_num_query_experts: int = None,
             att_num_query_groups: int = None,
+            att_num_global_tokens: int = 16,
+            att_window_size: int = 128,
             **kwargs
     ):
         super(ExperimentalAttentionTransformer, self).__init__(**kwargs)
         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', 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa', 'flex'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa", "flex".'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
@@ -84,8 +88,8 @@ class ExperimentalAttentionTransformer(nn.Module, PyTorchModelHubMixin, pipeline
             att_init = lambda: init_experimental_attention(embed_dim, att_heads, att_type, att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
                                                            max_seq_len=seq_len, is_causal=True, num_experts=att_num_experts,
-                                                           num_query_experts=att_num_query_experts,
-                                                           num_query_groups=att_num_query_groups)
+                                                           num_query_experts=att_num_query_experts, num_query_groups=att_num_query_groups,
+                                                           num_global_tokens=att_num_global_tokens, window_size=att_window_size)
 
         use_moe_att = att_type in ['gma', 'dma']
 

rxnn/transformers/attention.py

@@ -69,12 +69,14 @@ class MultiHeadAttention(nn.Module):
         v = self.v_proj(value).view(b, -1, self.num_heads, d // self.num_heads).transpose(1, 2)
         return q, k, v
 
-    def _apply_rope(self, q: torch.Tensor, k: torch.Tensor):
+    def _apply_rope(self, q: torch.Tensor, k: torch.Tensor, separate: bool = False):
         if self.rope is not None:
             if self.rope_only_for_query:
                 q = self.rope.forward_one(q)
             elif self.rope_only_for_keys:
                 k = self.rope.forward_one(k)
+            elif separate:
+                q, k = self.rope.forward_one(q), self.rope.forward_one(k)
             else:
                 q, k = self.rope(q, k)
         return q, k

{rxnn-0.2.52.dist-info → rxnn-0.2.53.dist-info}/METADATA

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

{rxnn-0.2.52.dist-info → rxnn-0.2.53.dist-info}/RECORD

@@ -1,8 +1,8 @@
 rxnn/.DS_Store,sha256=BxZLo9tFs48JMq6jhumiCnCPLTeCwl619CFSg4ClRAY,6148
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=46qwZLJuZMpIBrZ-r9DaQEPPmmZkO464C3Tkm_Mq-cs,23445
-rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4888
+rxnn/experimental/attention.py,sha256=z9zahN1KTDcVabUB0RwWsT6oK8MTuz65haVwYRHAzy4,24689
+rxnn/experimental/models.py,sha256=HPOIRpnX_oiI10wsVC4J6rzo3T6dj10aNWGYpa9S1UU,5115
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/memory/attention.py,sha256=wnYjd3UnmzAA79-7QxpMoEk3O1qRy8LmW1JcE_Fotck,3094
@@ -24,7 +24,7 @@ rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,14
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/training/utils.py,sha256=C0OS2RAGQ3L7D_G3CWupu_BpAFhkovMByBKm355Ibfc,6087
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
+rxnn/transformers/attention.py,sha256=KRnKT6XUqAXElxV9y72mSpdTeiMgCKCCLqqxCFNTHmA,16372
 rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
 rxnn/transformers/layers.py,sha256=wG8C9doafpLUsGtUTg-xdrHt7EQMEdB10vcSD-O1nVg,7999
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
@@ -33,7 +33,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.52.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.52.dist-info/METADATA,sha256=_0O3SB2pBSkic0UMXV4_nlYJjFw9bqUFroOs3A4FuW0,25997
-rxnn-0.2.52.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.52.dist-info/RECORD,,
+rxnn-0.2.53.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.53.dist-info/METADATA,sha256=CqOV8zz8qcGYyZXbeRv7jLMd40AeF8iU7vyfRsMuslE,25997
+rxnn-0.2.53.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.53.dist-info/RECORD,,