5 Commits

Author	SHA1	Message	Date
biondizzle	5b55cf0767	fix: k_seg is already 3D from slicing, don't add extra unsqueeze(-1)	2026-05-26 11:02:44 +00:00
biondizzle	2252d7c865	fix: make K/V segments contiguous before passing to kernel (TMA needs contiguous tensors)	2026-05-26 11:00:36 +00:00
biondizzle	5407dc768a	test: minor comment fix in D5b test	2026-05-26 10:59:51 +00:00
biondizzle	bb7ec341cb	fix: D5b test uses reference attn_sum for normalization, correct D5 merge formula ... - exp(LSE) != row_sum (it's row_sum * exp(max(S*scale))) - Normalize using reference attn_sum (same as other tests) - D5 merge uses normalized O + LSE: O = sum(exp(lse)*O_norm)/sum(exp(lse)) - Added 4-tile KV merge test (s_k=512)	2026-05-26 10:59:04 +00:00
biondizzle	6c73069cb9	D5b: Per-row LSE output + Python KV merge test ... - 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	2026-05-26 10:57:54 +00:00