functional mask fix

aoti_attention.py

@@ -99,13 +99,16 @@ def sparse_attention_functional(
     k_pool = key_tiled.view(batch, heads, n_tiles, 64, dim).sum(dim=3, dtype=torch.float32) / tile_divisor
     scores = torch.matmul(q_pool, k_pool.transpose(-2, -1)) / math.sqrt(dim)
 
-    # Prefix-exempt top-k mask, assembled from fresh tensors (functional, export-traceable).
-    video_scores = scores[..., num_prefix_tiles:, num_prefix_tiles:]
-    indices = video_scores.topk(topk, dim=-1, sorted=False).indices
-    mask_video = torch.zeros_like(video_scores, dtype=torch.bool).scatter(-1, indices, True)
-    prefix_cols = torch.zeros_like(video_scores, dtype=torch.bool)
-    prefix_cols[..., :num_prefix_tiles] = True
-    mask_video = mask_video | prefix_cols
+    # Prefix-exempt top-k mask, assembled from fresh tensors (functional, export-traceable):
+    # video query rows keep the top-k VIDEO kv tiles, prefix kv tiles are always selected,
+    # prefix query rows are dense over everything.
+    video_kv_scores = scores[..., num_prefix_tiles:, num_prefix_tiles:]          # [B, H, nq_video, kv_video]
+    indices = video_kv_scores.topk(topk, dim=-1, sorted=False).indices
+    mask_video_kv = torch.zeros_like(video_kv_scores, dtype=torch.bool).scatter(-1, indices, True)
+    prefix_cols = torch.ones(
+        (batch, heads, num_q_video, num_prefix_tiles), dtype=torch.bool, device=query.device
+    )
+    mask_video = torch.cat([prefix_cols, mask_video_kv], dim=-1)                 # [B, H, nq_video, n_tiles]
     mask = torch.cat(
         [
             torch.ones((batch, heads, num_prefix_tiles, n_tiles), dtype=torch.bool, device=query.device),