Add forward LSE output to FlyDSL flash attention kernels

[amd-wsung102]

Motivation

This PR is motivated by the task described here https://amd-hub.atlassian.net/jira/software/c/projects/AIMODELS/boards/2992?assignee=712020%3Acedccbdf-bb1e-4e04-8c3b-edcdd7744eda&selectedIssue=AIMODELS-983, which prepares for a MSLK integration.

Summary

Extends the FlyDSL flash attention forward kernels to optionally emit the per-row log-sum-exp (LSE) alongside the attention output. LSE is the normalization statistic the backward pass needs to recompute the softmax probabilities P without re-running the online softmax, so this is the enablement step for a pure-FlyDSL forward+backward attention path.

Opt-in via return_lse=True; default behavior is unchanged.

Changes

• flash_attn_interface.py: add return_lse to flydsl_flash_attn_func; thread it through every build-cache helper; allocate fp32 [B, num_heads, Sq] and return (out, lse) when enabled. Cache is keyed on return_lse so the existing no-LSE kernels are never perturbed.
• flash_attn_generic.py: return_lse build flag; LSE tensor in the kernel/launch/compile signatures; store sm_scale*m_raw + ln(l) on both the normal store and the zeroed (fully-masked) tile.
• flash_attn_gfx950.py: same for the dualwave kernel (dense store m_row*ln2 + ln(l_row)), plus the split-K combine kernel finalizing m_max*ln2 + ln(den) from the per-split (m, l). One writer per row via half-wave + OOB-sentinel guards.
• tests/kernels/test_flash_attn_lse.py: numerical validation vs a float32 PyTorch reference.
• docs/flash_attn_lse_contract.md: the A1↔A3 interface contract.

Output tensor

• dtype float32, shape [B, num_heads, Sq] (varlen: [B, num_heads, max_seqlen_q], padded rows undefined).

Supported / unsupported paths

Path	LSE
dense (gfx950 dualwave / generic)	yes
dense split-K (num_kv_splits > 1)	yes
varlen (packed cu_seqlens)	yes
GQA / MQA, causal, cross-attn (Sq≠Skv)	yes
fp8	no — raises NotImplementedError
paged KV (block_table / seqlen_k)	no — raises NotImplementedError

fp8 and paged KV are decode-oriented and out of scope for the training backward path; they fail loudly rather than silently returning wrong LSE.

Validation

tests/kernels/test_flash_attn_lse.py — 33 cases pass on gfx950, covering dense MHA/GQA/MQA (D=64/128, bf16/f16), split-K, varlen, cross-attention, and fully-masked (-inf) rows. LSE matches the fp32 reference within ~8e-3 (bf16) / ~4e-3 (f16); fp32-accumulating paths (varlen light, D=64) match to ~1e-6.

Performance

The LSE store is one fp32 write per query row (single lane per row). Measured on gfx950 (bf16, H=16, D=128, S ∈ {1024, 2048, 4096}), pure kernel overhead is within run-to-run noise (≈0–1.3%). The full-API delta (~3%) is dominated by the extra LSE tensor allocation, not the GPU store.

Compatibility

Fully backwards compatible — return_lse defaults to False and returns a bare output tensor; LSE build variants are cached separately.

Test plan

• pytest tests/kernels/test_flash_attn_lse.py on gfx950 (33 passed)
• CI on target architectures