refactor: simplify attention mask handling

astrai/inference/core/executor.py

@@ -43,15 +43,20 @@ class Executor:
         )
 
         task_ids = [t.task_id for t in tasks]
+        position_ids = (
+            torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device)
+            .unsqueeze(0)
+            .expand(batch_sz, -1)
+        )
+        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
+            prompt_len, device=self.device
+        )
 
         with torch.inference_mode():
             self.model(
                 input_ids,
-                position_ids=torch.arange(
-                    start_pos, prompt_len, dtype=torch.long, device=self.device
-                )
-                .unsqueeze(0)
-                .expand(batch_sz, -1),
+                input_mask=input_mask,
+                position_ids=position_ids,
                 paged_cache=self.kv_cache.bind_tasks(task_ids, prompt_len, self.device),
             )
 
@@ -84,7 +89,10 @@ class Executor:
         position_ids = torch.tensor(
             [t.next_pos for t in tasks], dtype=torch.long, device=self.device
         )
-        total_len = position_ids.max().item() + 1
+        total_len = max(t.next_pos for t in tasks) + 1
+        input_mask = position_ids[:, None, None] >= torch.arange(
+            total_len, device=self.device
+        )
 
         task_ids = [t.task_id for t in tasks]
 
@@ -122,6 +130,7 @@ class Executor:
         with torch.inference_mode():
             outputs = self.model(
                 input_ids.unsqueeze(1),
+                input_mask=input_mask,
                 paged_cache=self.kv_cache.bind_tasks(
                     task_ids,
                     total_len,

astrai/model/components/attention.py

@@ -76,9 +76,8 @@ class GQA(nn.Module):
         rotary_emb: Tensor,
         attn_mask: Tensor = None,
         paged_cache: Optional[CacheView] = None,
+        is_causal: bool = False,
     ) -> Tensor:
-        is_causal = attn_mask is None
-
         q = self._split_heads(self.q_proj(x), self.n_heads)
         k = self._split_heads(self.k_proj(x), self.n_kv_heads)
         v = self._split_heads(self.v_proj(x), self.n_kv_heads)
@@ -163,9 +162,9 @@ class MLA(nn.Module):
         rotary_emb: Tensor,
         attn_mask: Tensor = None,
         paged_cache: Optional[CacheView] = None,
+        is_causal: bool = False,
     ) -> Tensor:
         bsz, seq_len, _ = x.size()
-        is_causal = attn_mask is None
 
         q = self.q_proj(x)
         q = q.view(bsz, seq_len, self.n_heads, self.head_dim)

astrai/model/components/decoder_block.py

@@ -26,12 +26,14 @@ class DecoderBlock(nn.Module):
         rotary_emb: Tensor,
         attention_mask: Optional[Tensor] = None,
         paged_cache: Optional[CacheView] = None,
+        is_causal: bool = False,
     ) -> Tensor:
         attn_output = self.attention(
             self.input_norm(x),
             rotary_emb,
             attention_mask,
             paged_cache,
+            is_causal,
         )
         x = attn_output + x
         x = self.mlp(self.post_attention_norm(x)) + x

astrai/model/encoder.py

@@ -59,10 +59,10 @@ class EmbeddingEncoder(AutoModel):
         x = self.embed_tokens(input_ids)
 
         rotary_emb = self.rotary_embedding(x, position_ids)
-        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=False)
+        attn_mask = process_attention_mask(input_mask)
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache=None)
+            x = layer(x, rotary_emb, attn_mask)
 
         hidden_states = self.norm(x)
 

astrai/model/transformer.py

@@ -15,32 +15,15 @@ from astrai.model.components.rope import RotaryEmbedding
 
 
 def process_attention_mask(
-    input_tensor: Tensor,
-    position_ids: Optional[Tensor],
-    input_mask: Optional[Tensor] = None,
-    is_causal: bool = False,
+    input_mask: Optional[Tensor],
 ) -> Optional[Tensor]:
-    if position_ids is None:
-        return None
-    if input_mask is not None and input_mask.dim() > 2:
-        return input_mask
-
-    device = input_tensor.device
-    B = input_tensor.size(0)
-    T = position_ids.max().item() + 1
-
     if input_mask is None:
-        if position_ids.min().item() == 0 and is_causal:
         return None
-        attend = torch.ones(B, 1, T, dtype=torch.bool, device=device)
-    else:
-        attend = input_mask[:, :T].to(device=device, dtype=torch.bool).unsqueeze(1)
-
-    if is_causal:
-        causal = position_ids.unsqueeze(-1) >= torch.arange(T, device=device)
-        attend = attend & causal
-
-    return attend.unsqueeze(1)
+    if input_mask.dim() == 2:
+        return input_mask[:, None, None, :]
+    if input_mask.dim() == 3:
+        return input_mask[:, None, :, :]
+    return input_mask
 
 
 @AutoModel.register("autoregressive_lm")
@@ -119,10 +102,11 @@ class AutoRegressiveLM(AutoModel):
 
         x = self.embed_tokens(input_ids)
         rotary_emb = self.rotary_embedding(x, position_ids)
-        attn_mask = process_attention_mask(x, position_ids, input_mask, is_causal=True)
+        attn_mask = process_attention_mask(input_mask)
+        use_sdpa_causal_mask = attn_mask is None
 
         for layer in self.layers:
-            x = layer(x, rotary_emb, attn_mask, paged_cache)
+            x = layer(x, rotary_emb, attn_mask, paged_cache, use_sdpa_causal_mask)
 
         hidden_states = self.norm(x)
         logits = self.lm_head(hidden_states)

astrai/trainer/strategy.py

@@ -395,7 +395,9 @@ class GRPOStrategy(BaseStrategy):
         # response tokens.  get_logprobs shifts the mask by one position, so
         # the first response token's logprob (predicted from the last prompt
         # token) is correctly included.
-        full_masks = torch.cat([torch.zeros_like(prompt_expanded), masks_flat], dim=-1)
+        full_masks = torch.cat(
+            [torch.zeros_like(prompt_expanded, dtype=torch.bool), masks_flat], dim=-1
+        )
 
         # get_logprobs returns [B*G, S-1] (S = prompt_len + response_len).
         # Response token logprobs occupy the last ``response_len`` positions