Fix mask support and fallback context for GPU tests

[MagellaX]

Summary by cubic

Fixes attention mask shape handling and stabilizes GPU tests by forcing the safe math SDPA kernel on CUDA. Updates both forward and backward paths in fused_online_attention.

• Bug Fixes
  • Support attention masks with shapes: 2D (B, K), 3D (B, Q, K), and 4D (B, ..., K); previously only 2D (B, K) was accepted.
  • On CUDA, use torch.backends.cuda.sdp_kernel(enable_math=True, enable_flash=False, enable_mem_efficient=False) with a safe nullcontext fallback in forward and backward to prevent environment-specific failures.

[cursor]

This PR is being reviewed by Cursor Bugbot

Details

You are on the Bugbot Free tier. On this plan, Bugbot will review limited PRs each billing cycle.

To receive Bugbot reviews on all of your PRs, visit the Cursor dashboard to activate Pro and start your 14-day free trial.

[cursor]

Bug: Incomplete Mask Validation Causes Runtime Failures

The 4D mask validation is incomplete. It only checks shape[0] (batch size) and shape[-1] (seq_len_k), missing compatibility checks for shape[1] (num_heads) and shape[2] (seq_len_q). This allows incompatible masks to pass validation, leading to runtime failures during mask expansion instead of falling back to PyTorch SDPA.

 

[cubic-dev-ai]

1 issue found across 1 file

Prompt for AI agents (all 1 issues)

Understand the root cause of the following 1 issues and fix them.


<file name="stream_attention/core/fused_online_attention.py">

<violation number="1" location="stream_attention/core/fused_online_attention.py:911">
4D attention mask validation only checks batch and K lengths, allowing masks with mismatched head or Q dimensions to slip through and fail during mask expansion. Tighten the 4D check to also validate shape[1] (allow 1 or num_heads) and shape[2] (seq_len_q).</violation>
</file>

_{React with 👍 or 👎 to teach cubic. Mention @cubic-dev-ai to give feedback, ask questions, or re-run the review.}

[cubic-dev-ai]

4D attention mask validation only checks batch and K lengths, allowing masks with mismatched head or Q dimensions to slip through and fail during mask expansion. Tighten the 4D check to also validate shape[1] (allow 1 or num_heads) and shape[2] (seq_len_q).

Prompt for AI agents

Address the following comment on stream_attention/core/fused_online_attention.py at line 911:

<comment>4D attention mask validation only checks batch and K lengths, allowing masks with mismatched head or Q dimensions to slip through and fail during mask expansion. Tighten the 4D check to also validate shape[1] (allow 1 or num_heads) and shape[2] (seq_len_q).</comment>

<file context>
@@ -896,11 +896,21 @@ def forward(
+                    and attention_mask.shape[2] == seq_len_k
+                )
+            elif dims == 4:
+                mask_supported = attention_mask.shape[0] == batch_size and attention_mask.shape[-1] == seq_len_k
+            else:
+                mask_supported = False
</file context>

Suggested change

	mask_supported = attention_mask.shape[0] == batch_size and attention_mask.shape[-1] == seq_len_k
	mask_supported = (attention_mask.shape[0] == batch_size and attention_mask.shape[1] in (1, num_heads) and attention_mask.shape[2] == seq_len_q and attention_mask.shape[3] == seq_len_k)