fix(ttx): support SDPA without attention mask

[shengw-bd]

No description provided.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[github-actions]

Claude Code Review

Verdict: Comment -- Adds optional attention mask support to SDPA kernels via a HAS_MASK constexpr specialization; logic looks correct but a couple of concerns worth addressing.

Summary

Makes the mask argument optional in both the ILU and NPU A2 SDPA kernels by adding a HAS_MASK constexpr specialization and passing a dummy pointer when no mask is supplied. Adds a corresponding no-mask accuracy test.

Must fix

None.

Suggestions

Suggestions (3)

• [MAJOR] NPU no-mask path drops sequence-length guarding -- mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:65-68 -- Unlike the ILU kernel which still applies mask_valid (bounds against SEQ) when HAS_MASK=False, the NPU inner loop now performs no masking at all. If SEQ is not a multiple of BLOCK_N/BLOCK_M, out-of-range positions contribute to softmax. Confirm SEQ is always block-aligned here, or add a bounds-based tl.where.
• [MAJOR] Dummy mask pointer dtype/stride mismatch -- mojo_opset/backends/ttx/kernels/ilu/sdpa.py:189 and mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:825 -- Passing q (bf16/fp16, different shape/strides) as the mask pointer relies on the compiler fully DCE-ing all mask uses under HAS_MASK=False. If any stride/dtype metadata is still consumed by the kernel signature, this will silently misbehave. A zero-length/zero-stride tensor of the expected dtype would be safer.
• [MINOR] Test coverage is thin -- mojo_opset/tests/accuracy/operators/test_attention.py:924-948 -- Only one shape and only the default backend is exercised; the NPU A2 change has no direct no-mask test in this diff.

Nits

Nits (1)

• [NIT] Stale commented-out lines retained -- mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:67-68,826-827 -- Consider removing the dead qk += (1 - mask.to(...)) / mask = 1 - mask.to(...) comments while touching this code.

Notes

• [CHECK] libentry() / autotune caching: adding HAS_MASK to key= should force separate cache entries, but verify the autotuner actually re-specializes rather than reusing a compiled kernel across the two boolean values.
• [CHECK] For the ILU path, mask_c = q_c (bf16/fp16) is later read via stride_m0/stride_m1 computed from a bool tensor when has_mask=True; ensure the stride_m* arguments passed when has_mask=False are also from q_c (not shown in the diff) so pointer arithmetic remains consistent.

[Neuromancer42]

LGTM