refactor: unify attention mask to single attn_mask tensor

- CudaBackend.fwd_decode passes attn_mask directly instead of kv_cache.decode_mask
- TorchNativeBackend derives pos_mask from attn_mask[:,0,0] on decode
- Drop decode_mask and page_table fields from KVCache and bind_tasks

astrai/extension/attention_backend.py

@@ -274,12 +274,12 @@ class TorchNativeBackend(AttentionBackend):
             kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
 
             max_len = kv_cache.max_len
-            if kv_cache.page_table is not None:
-                indices = kv_cache.page_table
-            else:
             indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-            if kv_cache.decode_mask is not None:
+            # Zero out padding positions so gather never touches invalid slots.
-                pos_mask = kv_cache.decode_mask
+            # Decode: attn_mask[:,0,0] is exactly the per-position validity
+            # mask ([B, max_len], True=keep).  Prefill: fall back to seq_lens.
+            if q.size(1) == 1 and attn_mask is not None and attn_mask.dim() == 4:
+                pos_mask = attn_mask[:, 0, 0]
             else:
                 pos_mask = (
                     torch.arange(max_len, device=q.device)[None, :]
@@ -344,10 +344,6 @@ class CudaBackend(AttentionBackend):
         kv_indptr = torch.zeros(b + 1, dtype=torch.int32, device=q.device)
         kv_indptr[1:] = kv_cache.seq_lens.cumsum(0).to(torch.int32)
 
-        mask = None
-        if b > 1 and kv_cache.decode_mask is not None:
-            mask = kv_cache.decode_mask
-
         out = attn_paged_decode(
             q_3d,
             kv_cache.k_buffer[layer_id],
@@ -356,7 +352,7 @@ class CudaBackend(AttentionBackend):
             kv_cache.req_pool_indices,
             kv_indptr,
             kv_cache.max_len,
-            mask=mask,
+            mask=attn_mask,
             is_causal=is_causal,
         )
         return out.unsqueeze(1).flatten(2)

astrai/inference/core/cache.py

@@ -203,10 +203,6 @@ class KVCache:
         seq_lens: [batch_size] — per-request total sequence lengths
         out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
         max_len: max(seq_lens) as Python int — avoids GPU sync in decode
-        page_table: [batch, max_len] — precomputed gather indices for decode;
-            None for prefill or when not yet computed.
-        decode_mask: [batch, max_len] bool — precomputed position validity
-            mask for decode; None for prefill or single-batch decode.
     """
 
     k_buffer: Tensor
@@ -216,8 +212,6 @@ class KVCache:
     seq_lens: Tensor
     out_cache_loc: Tensor
     max_len: int = 0
-    page_table: Optional[Tensor] = None
-    decode_mask: Optional[Tensor] = None
 
 
 class PagePool:
@@ -434,21 +428,11 @@ class PagePool:
             out_cache_loc = self._req_pool.req_to_token[
                 req_pool_indices, start_pos:seq_len
             ]
-            page_table = None
-            decode_mask = None
         else:
             write_pos = seq_lens_t - 1
             out_cache_loc = self._req_pool.req_to_token[
                 req_pool_indices, write_pos
             ].unsqueeze(-1)
-            ml = max(seq_lens)
-            page_table = self._req_pool.req_to_token[req_pool_indices, :ml]
-            if len(task_ids) > 1:
-                decode_mask = (
-                    torch.arange(ml, device=device)[None, :] < seq_lens_t[:, None]
-                )
-            else:
-                decode_mask = None
 
         return KVCache(
             k_buffer=self._storage.k_buffer,
@@ -458,8 +442,6 @@ class PagePool:
             seq_lens=seq_lens_t,
             out_cache_loc=out_cache_loc,
             max_len=max(seq_lens),
-            page_table=page_table,
-            decode_mask=decode_mask,
         )
 
     # ---- internals ----