Commit Graph

3 Commits

Author SHA1 Message Date
24993428a2 fix: D4 test reference computation only applies causal mask when is_causal=True 2026-05-26 10:56:04 +00:00
e3e01071f4 fix: swa_len as Int32 scalar instead of CuTe tensor
CuTeDSL @cute.kernel cannot handle dynamic-shape tensors as parameters.
Pass swa_len as Int32 scalar instead of a 1D tensor.
This works for batch_size=1 (current config).
Updated D3 and D4 tests to pass swa_len as int.
2026-05-26 10:54:41 +00:00
841a3e87b2 D4: Causal mask on SWA branch
- Add is_causal flag to FmhaKernel constructor
- Mask positions where k_coord > m_coord to -inf (causal attention)
- Combined with D3 SWA mask: both conditions use OR logic
- Same tTMEM_LOADcS coordinate mapping as D3
- const_expr guarded: zero overhead when is_causal=False
- New test: test_d4_causal_mask.py with causal + combined masking
2026-05-26 10:52:30 +00:00