plan update

STAGE_D2.md

@@ -204,10 +204,14 @@ O has shape `(batch, n_h, T, head_dim)`. Each CTA writes its head's output. The
   - hd=128, n_h=8 — PASS
   - hd=256, n_h=2 — PASS
 
+- [x] **D2 Head-Packed:** Q reshaped to (n_h*T, hd, 1), per-row softmax
+  - cos 0.999995 for n_h=1-128 at hd=64
+  - Pro decode (n_h=128, T=1): M=128, one CTA processes all 128 heads
+
 ### 🟡 Blocked (Multi-CTA Grid)
 
 - [ ] **D2.1:** Add `num_query_heads` and `batch_size` to `FmhaKernel.__init__`
-  - Simple to add, but the grid change is blocked (see below)
+  - Added as constructor params, but grid still (1,1,batch) — per-head launch in Python
 
 - [ ] **D2.3–D2.6:** Multi-CTA grid with runtime block coordinates
   - **BLOCKED:** `cute.local_tile` does not support runtime coordinates. Must use `cute.flat_divide` instead.
@@ -216,9 +220,8 @@ O has shape `(batch, n_h, T, head_dim)`. Each CTA writes its head's output. The
   - **CUTLASS reference approach:** Uses `flat_divide` + `tma_partition` inside the TMA warp block, and a custom epilogue that handles the flat_divide coordinate system. Estimated 1-2 day effort.
 
 - [ ] **D2.9:** LSE output for multi-head
-  - Per-row LSE verified correct (max err 0.000001) but CuTe tensor indexing needs work
-  - Currently only row 0 is written (sfw_idx==0 guard)
-  - Full per-row output needed for D5 KV merge
+  - Per-row LSE verified correct (max err 0.000001), all 128 rows now write
+  - row_sums output also working — O_norm = O_unnorm / row_sum
 
 ---