The fused SwiGLU kernel stored swiglu_limit but never applied it.
Paper §4.2.3: gate capped at swiglu_limit, linear clamped to [-limit, +limit].
Non-fused reference path already applies clamping correctly.
Fix: add fmin/fmax clamping in FP32 before BF16 conversion.
Reverted tP to p_tmem_s.outer (needed for make_fragment_A profile).
P store now writes BF16 to TMEM using tOrP0's layout, matching PV A-fragment reads.
This fixes the layout mismatch at hd>64 where QK C-fragment composition
writes to different TMEM columns than PV A-fragment reads.
Key insight: tP (PV A-fragment base) used p_tmem_s.outer layout,
but P store used QK C-fragment composition layout. These diverge at hd>64.
Fix: tP now uses the same QK C-fragment composition layout (tStP_layout)
as the P store. PV A-fragment is derived from tP, so it automatically
uses the same layout. No double-offset since tP includes P offset.
- Changed P store from FP32 QK C-fragment layout to BF16 PV A-fragment layout
- rP_bf16_reg stores directly to TMEM using tOrP0 layout
- Ensures softmax writes P to same TMEM columns that PV GEMM reads
CRITICAL BUG: P was stored using QK C-fragment composition layout, but PV
A-fragment reads using p_tmem_s.outer (PV A-operand layout). These layouts
match at hd=64 (cos 0.999998) but diverge at hd>64 (cos 0.784 at hd=128).
The fix: tStP0 and tScP now use p_tmem_s.outer instead of
composition(tStS.layout, (128, p_cols_fp32)). This ensures the softmax
writes P in the same layout that the PV GEMM expects.
- Promote HEAD_DIM from module constant to constructor parameter
- FmhaKernel(head_dim=64, s_k=128, ...) — default 64 for regression
- All references to HEAD_DIM replaced with self.head_dim
- PV MMA tiler, V layout, softmax corr_tiles all parameterized
- TMEM budget warning when num_tmem_alloc_cols > 512
- New test: test_fmha_v3_stage_d1.py tests hd=64 (regression) and hd=512
- Stage C test preserved as-is for reference
