biondizzle b6b581777a D3: In-kernel SWA sequence length 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

2026-05-26 10:51:23 +00:00

..

__init__.py

Restructure: cutedsl/ -> dsv4/ with proper layering

2026-05-21 17:30:44 +00:00

cudagraph_test.py

Restructure: cutedsl/ -> dsv4/ with proper layering

2026-05-21 17:30:44 +00:00

layertest.py

Restructure: cutedsl/ -> dsv4/ with proper layering

2026-05-21 17:30:44 +00:00

test_cotiled_diag.py

fix: proper v_major from tensor

2026-05-24 01:55:37 +00:00

test_cutedsl.py

Restructure: cutedsl/ -> dsv4/ with proper layering

2026-05-21 17:30:44 +00:00

test_d1_3_cotiled.py

D1.3: Fix while loop in cotiled diag - precompute num_tmem_alloc_cols

2026-05-24 00:07:22 +00:00

test_d1_3_layout_diag.py

D1.3: Enhanced diagnostic - test QK C-fragment as source for make_tiled_copy_C

2026-05-23 22:24:15 +00:00

test_d1_3_smem_diag.py

D1.3: Add SMEM-P coordinate diagnostic test

2026-05-23 23:23:05 +00:00

test_d1_3_smem_direct.py

D1.5: Replace TMEM round-trip normalize with correction epilog (one-way: TMEM→reg→SMEM→GMEM)

2026-05-24 00:24:24 +00:00

test_d1_3_smem_vs_tmem.py

D1.3: Add SMEM-P vs TMEM-P comparison test

2026-05-24 00:10:18 +00:00

test_d1_3_unnorm_debug.py

D1.3: Add unnormalized debug test to isolate SMEM-P vs O round-trip error

2026-05-24 00:15:41 +00:00

test_d1_3_write_read.py

D1.3: Add SMEM-P write/read diagnostic

2026-05-24 00:13:28 +00:00

test_d1_debug.py

debug: hd=64 with CUDA_LAUNCH_BLOCKING

2026-05-23 03:42:53 +00:00

test_d1_diag2.py

D1: Add diagnostic test (TMEM-P vs SMEM-P at various hd)

2026-05-24 03:22:23 +00:00

test_d1_diag.py

fix: use mV.iterator

2026-05-23 03:25:29 +00:00

test_d1_hd512_merge.py

D1.4: Fix merge test - use use_smem_p=False for hd=256 kernel (SMEM budget)

2026-05-24 16:36:48 +00:00

test_d1_hd512_only.py

D1.4: Remove --opt-level 0 from hd512 test (use default opt level)

2026-05-24 16:42:01 +00:00

test_d1_hd512.py

d1: add hd=512 test

2026-05-23 03:20:46 +00:00

test_d1_kv_merge_v2.py

D1: fix per-row LSE output + add KV merge test v2 with per-row LSE

2026-05-24 22:21:51 +00:00

test_d1_kv_merge_v3.py

D1: corrected KV merge test with proper normalized output formula

2026-05-24 22:24:27 +00:00

test_d1_kv_merge.py

D1: add KV merge test using log-sum-exp (avoids TMEM round-trip)

2026-05-24 22:17:24 +00:00

test_d1_lse_verify.py

fix lse verify

2026-05-24 22:23:08 +00:00

test_d1_lse.py

D1: LSE diagnostic at various hd

2026-05-24 03:23:16 +00:00

test_d1_multi_kv.py

D1: add multi-KV-tile O rescale test (s_k=256,384,512)

2026-05-24 22:00:42 +00:00

test_d1_qk512.py

Fix: add cutlass import to test_d1_qk512

2026-05-24 14:20:32 +00:00

test_d1_raw.py

D1: test raw unnormalized output via epilogue_tma_store

2026-05-23 03:33:59 +00:00

test_d1_regression.py

D1.4: Fix regression test for un-normalized O output (D5a)

2026-05-24 15:13:16 +00:00

test_d1_rescale_debug.py

fix debug test

2026-05-24 22:04:51 +00:00

test_d1_rescale_diag.py

D1: add rescale diagnostic

2026-05-24 22:18:12 +00:00

test_d1_rescale_min.py

D1: add KV merge test using log-sum-exp (avoids TMEM round-trip)

2026-05-24 22:17:24 +00:00

test_d1_smem128.py

D1: SMEM-P test at hd=128

2026-05-24 03:48:37 +00:00

test_d1_sweep.py

D1: paired atoms epilogue (no TMEM round-trip)

2026-05-23 03:29:51 +00:00

test_d1_tmem_only.py

D1: Fix SMEM-P (coordinate store), LSE (FP32), add TMEM-P-only test

2026-05-24 03:27:14 +00:00

test_d1_tmem_trip.py

D1: add KV merge test using log-sum-exp (avoids TMEM round-trip)

2026-05-24 22:17:24 +00:00

test_d2_headpacked.py

D2: comprehensive head-packed test (n_h=1, 64, 128, hd=64, 128)

2026-05-25 17:16:05 +00:00

test_d2_multicta.py

D2: add num_query_heads/batch_size params + head-packed test

2026-05-25 16:50:49 +00:00

test_d2_multihead.py

D1: revert per-row LSE to sfw_idx=0 for now (debugging D2 regression)

2026-05-24 22:28:11 +00:00

test_d2_perhead.py

D2: add per-head launch test

2026-05-24 22:48:22 +00:00

test_d2_regression.py

fix: use reference attn_sum for normalization (kernel LSE per-row may be wrong)

2026-05-25 17:13:34 +00:00

test_d2_scale.py

D2: add scale test (more heads, larger hd)

2026-05-24 22:49:44 +00:00

test_d3_inkernel_mask.py

D3: In-kernel SWA sequence length masking

2026-05-26 10:51:23 +00:00

test_d3_swa_mask.py

fix typo: from_dlset → from_dlpack

2026-05-25 17:28:43 +00:00

test_d15_multi_kv.py

D1.5: add multi-KV-tile attention test with Python KV merge

2026-05-25 17:18:50 +00:00

test_fmha_v3_stage_c.py

fix: revert to composition layout for hand-constructed atoms (matching CUTLASS)

2026-05-23 02:54:54 +00:00

test_fmha_v3_stage_d1.py

D1: Full test with TMEM-P at hd=64,128,256,512

2026-05-24 04:07:40 +00:00

test_fmha_v3_stage_d5b.py

D5b: Fix reference computation - use logsumexp for stable LSE, fix o_unnorm definition

2026-05-23 21:43:04 +00:00

test_fmha_v3.py

FIX: (None,0,None,0) for ALL tma_partition outputs — verified shapes on B200

2026-05-22 23:35:55 +00:00

test_fp4_roundtrip.py

Restructure: cutedsl/ -> dsv4/ with proper layering

2026-05-21 17:30:44 +00:00

test_nvfp4_diag.py

NVFP4-0.2-0.4: add FP4 primitives diagnostic test

2026-05-25 03:07:53 +00:00

test_nvfp4_gpu_quantize.py

fix test 4: use silu(gate)+swiglu interleaved (matching fused kernel output)

2026-05-25 16:24:04 +00:00

test_nvfp4_primitives.py

diag: add 2-CTA check + fix LayoutEnum in MMA test

2026-05-23 08:45:26 +00:00

test_nvfp4_quant_kernel.py

NVFP4-1.1 Step 3: post-SWiGLU quantization test suite (all PASS)

2026-05-25 09:08:01 +00:00

test_nvfp4_quantize_kernel.py

fix: quantize_activation_nvfp4 returns 2 values, not 3

2026-05-25 03:17:13 +00:00

test_paired_epilog.py

test: paired atoms epilog from old commit 6ee28d8

2026-05-23 03:32:53 +00:00

test_smem_budget.py

D1.4: Reduce pv_n_tile to 128 for hd=512 to fit SMEM budget (192KB)

2026-05-24 08:07:32 +00:00

test_smem_p_coord.py

test: add try/except for SMEM-P coord test

2026-05-24 02:15:07 +00:00

test_smem_p_diag.py

shit left dangling

2026-05-23 23:58:57 +00:00

test_smem_p_write.py

test: SMEM-P coordinate verification test

2026-05-24 01:58:32 +00:00

test_tmem_budget.py

D1.2: fix probe for hd=512 (MMA max N=256, use pv_n_tile)

2026-05-23 06:41:42 +00:00