feat: improve operator accuracy coverage and support mask-free TTX SDPA

[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 -- Optional-mask support looks correct, but the ILU no-mask path silently uses q_c as a dummy pointer which relies on Triton's specialization actually eliding the load.

Summary

Adds optional attention mask support to the ILU and NPU A2 SDPA kernels by introducing a HAS_MASK constexpr and gating the mask load/apply behind it. Also expands test coverage for SDPA (no-mask case), gemm dtypes, position embedding, and adds a tiny-values dynamic-quant test.

Must fix

None.

Suggestions

Suggestions (3)

• [MAJOR] Dummy mask pointer aliasing Q buffer -- mojo_opset/backends/ttx/kernels/ilu/sdpa.py:189 and mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:825 -- Passing q as the mask argument when HAS_MASK=False works only if the specialization fully elides the load; if the autotuner or a future refactor ever evaluates the mask expressions under HAS_MASK=False, this will read garbage from Q with a bool-dtype pointer. Safer to pass a small persistent zero tensor or a null-ish pointer alias with matching dtype/stride, or at least assert HAS_MASK gates every use in the kernel.
• [MAJOR] Autotune key includes HAS_MASK but cache key must actually change -- mojo_opset/backends/ttx/kernels/ilu/sdpa.py:39 -- HAS_MASK is passed as a kwarg constexpr; confirm the @libentry/autotune wrapper actually rekeys on it (Triton keys off runtime args, not constexprs, unless explicitly declared). Otherwise the first-seen config will be reused across mask/no-mask calls.
• [MINOR] Stale commented-out code retained -- mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:65-66,820-821 -- The # qk += (1 - mask.to(tl.float32)) * (-1e6) etc. blocks were already dead before this PR; while touching these lines, consider removing them to reduce noise.

Nits

Nits (1)

• [NIT] test_sdpa no-mask branch uses hand-tuned atol=6e-2, rtol=8e-2 -- mojo_opset/tests/accuracy/operators/test_attention.py:929 -- fairly loose; a comment justifying the tolerance (bf16 accumulation over seq=512) would help future readers.

Notes

• [CHECK] mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:822 -- Reusing q (bfloat16/fp16) as a stand-in for a bool mask pointer: verify the compiled NPU kernel truly never issues the tl.load(mask_ptr) under HAS_MASK=False, since dtype mismatch between the declared pointer type and the actual buffer could still cause issues under some Triton frontends.
• [CHECK] mojo_opset/tests/accuracy/operators/test_attention.py:936 -- Swapping which side calls forward_diff_with (now sdpa calls it against sdpa_ref instead of the reverse) changes which implementation is "under test" in the assertion direction; confirm this matches the framework's convention.

[github-actions]

Claude Code Review

Verdict: Comment -- Optional-mask specialization for SDPA and small-scale guard for dynamic quant look reasonable, but a couple of behavioral changes deserve a second look.

Summary

Adds a HAS_MASK constexpr specialization to the ILU and NPU SDPA kernels so callers can pass mask=None, and changes dynamic quant to fall back to a unit scale when abs_max is near zero instead of emitting a zero scale. Tests are extended to cover the no-mask SDPA path, tiny-value quant, and additional gemm/rope shapes.

Must fix

None.

Suggestions

Suggestions (3)

• [MAJOR] Dynamic quant scale semantics change -- mojo_opset/backends/ttx/kernels/ilu/quant.py:180-181,224-225 -- Previously qscale was stored as-is (could be 0) with inv_qscale=0; now qscale itself is clamped to 1.0 when < 1e-6. Downstream consumers that read qscale (e.g. dequant paths) will see 1.0 instead of a tiny/zero value, which changes the reconstructed magnitude. Confirm all readers of the stored scale expect this new convention, and consider clamping only inv_qscale while keeping the true qscale written to memory.
• [MAJOR] 1e-6 threshold vs fp32 abs_max/127 -- mojo_opset/backends/ttx/kernels/ilu/quant.py:180 -- The threshold is applied after dividing by 127, so it triggers for abs_max < 1.27e-4. For bf16/fp16 inputs whose legitimate abs_max can fall in that range (e.g. heavily pre-scaled activations), quantization silently degrades to all-zero output with scale=1. Consider a much smaller epsilon (e.g. tied to dtype eps) or gating on exact zero.
• [MINOR] Dummy mask pointer aliases Q -- mojo_opset/backends/ttx/kernels/ilu/sdpa.py:188, mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:825 -- Passing q as the mask pointer relies on Triton fully eliminating the load under HAS_MASK=False. Safer to pass a tiny scratch tensor (or 0) to avoid surprises if a future edit reintroduces an unconditional load or bounds check.

Nits

Nits (2)

• [NIT] Leftover commented-out qk mask expressions -- mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:66-67 -- Two dead comment lines carried through the refactor; drop them.
• [NIT] Test parameter list construction -- mojo_opset/tests/accuracy/operators/test_gemm.py:155-165 -- The nested comprehension + append reads awkwardly; a small helper or separate parametrize stack would be clearer.

Notes

• [CHECK] mojo_opset/backends/ttx/kernels/ilu/sdpa.py:108 -- With HAS_MASK=False, attn_mask = mask_valid now applies -inf to padded positions where before an all-true mask tensor could have suppressed that; verify the no-mask numerics match the torch reference for non-power-of-two SEQ.
• [CHECK] mojo_opset/backends/ttx/kernels/npu/a2/sdpa.py:65-69 -- Under HAS_MASK=False there is no sequence-bounds masking of qk at all (unlike the ILU variant). If SEQ isn't a multiple of BLOCK_N, out-of-range K/V lanes could contribute to softmax; confirm the tiling guarantees alignment or add a validity mask.