- Fix LSE output: all 128 rows now write (mLSE[sfw_idx, 0, 0])
instead of only row 0 (mLSE[0])
- Each softmax thread (sfw_idx 0..127) independently writes its LSE
- This enables accurate Python-side KV merge for multi-KV-tile
- New test: test_d5b_perrow_lse.py with LSE verification + KV merge
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.
- 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
- Add apply_swa_mask flag to FmhaKernel constructor
- After TMEM load of S, use tTMEM_LOADcS coordinates to map register
fragment positions to (row, col) in QK matrix
- Mask positions >= swa_lens[batch_idx] to -inf before softmax
- Supports multi-KV-tile (kt*128 + k_coord for absolute position)
- swa_lens parameter passed as CuTe tensor, indexed by block_idx_z
- Dummy tensor (max int) when swa_lens=None (no masking)
- New test: test_d3_inkernel_mask.py with proper in-kernel masking
- Replaces pre-masking approach (BF16 min on K) which can't produce -inf
- Head-packed approach: Q is (n_h*T, hd, 1), kernel treats each row independently
- Grid: (1, 1, batch) — M dimension handled by head packing
- n_h=128, T=1 → M=128, one MMA tile, all heads in single CTA
- Tested: cos 0.999995 for both n_h=1 and n_h=128
The warp shuffle approach failed because __shfl_down_sync with 16 threads
has undefined behavior for the odd nibble. Use the same pattern as the
working deinterleave_quantize.cu: 1 CTA per 16-element block, 16 threads
per CTA, each thread reads all 16 elements sequentially and computes
amax + quantize + pack.
