Add optional forward LSE output

[AakarshAMD]

Summary

• Add optional forward per-row log-sum-exp (LSE) output to FlyDSL flash attention through return_lse=True.
• Support the generic kernel, gfx950 dualwave software-pipelined kernel, and split-K combine kernel.
• Compute LSE directly from the online-softmax state already maintained by the kernels:
  • Generic: sm_scale * m + ln(l)
  • gfx950 dualwave/split-K: convert the internal log2-domain maximum to natural-log LSE.
• Preserve the existing API and compiled kernel path when return_lse=False.
• Enable backward implementations to recompute softmax probabilities from (Q, K, LSE) without storing the full attention matrix.

Details

Kernel implementation

Files:

• kernels/attention/flash_attn_utils.py
• kernels/attention/flash_attn_generic.py
• kernels/attention/flash_attn_gfx950.py

The implementation adds a RETURN_LSE compile-time trait to the generic and gfx950 dualwave kernel configurations.

Because it is a compile-time trait:

• LSE instructions are emitted only for return_lse=True.
• The normal kernel does not execute a runtime condition.
• LSE and non-LSE variants use separate JIT cache entries.
• Existing callers retain the original kernel path.

_make_lse_rsrc creates a bounded buffer-resource descriptor for each batch's [H, Sq] slice of the float32 LSE tensor.

Three LSE store paths are implemented:

• GenericStoreHelper.store_lse

  • Stores LSE from the generic kernel.
  • Runs after normal online-softmax accumulation.
  • Stores -inf for fully masked rows.

• DualwaveStoreHelper.store_lse_from_ml

  • Stores LSE from the gfx950 dualwave kernel.
  • Converts the kernel's log2-domain running maximum: m_row * ln(2) + ln(l_row).

• DualwaveSplitKCombineHelper.store_lse

  • Stores LSE from the split-K combine kernel.
  • Computes the result from the combined maximum and denominator after all splits have been reduced.

All paths use a single-writer-per-row scheme. Invalid rows and non-writer lanes are redirected to an out-of-bounds sentinel index, which is safely discarded by the bounded buffer resource.

Upstream integration

The change was rebased onto the latest origin/main, including the new gfx942 flash-attention optimizations.

The conflict resolution preserves:

• skip_kv_pad_mask
• gfx942 DMA and global-prefetch paths
• vectorized K handling
• the upstream cache traits
• M128/M256 automatic dispatch
• padding-mask dispatch

return_lse is threaded through every new recursive builder and dispatch path, so the upstream implementation remains the structural baseline with LSE added as an independent option.

The rebased LSE code follows the updated typed-value style and does not reintroduce the deprecated ArithValue wrapper.

Low-level safety guards

The generic and gfx950 launch/compile entry points raise an error when a kernel compiled with RETURN_LSE=True is invoked without an LSE output tensor:

if traits.RETURN_LSE and lse is None:
    raise ValueError("return_lse=True requires an lse output tensor")

This protects direct low-level builder users from accidentally writing LSE into a placeholder tensor.

Public API

File:

• kernels/attention/flash_attn_interface.py

The public function gains:

return_lse: bool = False

Default behavior remains unchanged:

out = flydsl_flash_attn_func(q, k, v)

With LSE enabled:

out, lse = flydsl_flash_attn_func(q, k, v, return_lse=True)

The returned LSE tensor is:

• float32
• shape [B, H, Sq]
• natural-log domain
• includes the softmax scale

For packed variable-length attention, its shape is [B, H, max_seqlen_q]; padded rows are undefined.

Unsupported combinations fail explicitly:

• fp8 with return_lse=True
• paged KV with return_lse=True

The public API allocates the tensor only when requested:

lse = torch.empty((B, H, Sq), dtype=torch.float32, device=q.device)

It returns (out, lse) when enabled and the original out tensor otherwise.

Compatibility

• return_lse defaults to False.
• Existing callers do not need changes.
• Existing return types remain unchanged unless the new option is explicitly enabled.
• fp8 and paged-KV behavior is unchanged.
• The upstream gfx942 optimization paths are preserved.
• The CUDA/CK/AIter implementations are unaffected.

Diff

kernels/attention/flash_attn_generic.py   | 41
kernels/attention/flash_attn_gfx950.py    | 32
kernels/attention/flash_attn_interface.py | 43
kernels/attention/flash_attn_utils.py     | 80

4 files changed, 179 insertions(+), 17 deletions(-)

Test Plan

Tested on AMD Instinct MI350X (gfx950) using ROCm.

Correctness

External LSE correctness suite:

15 passed

The suite covers:

• generic attention
• gfx950 dualwave attention
• split-K combine
• MHA and GQA
• causal and non-causal modes
• supported head dimensions
• fully masked rows
• LSE numerical agreement with a reference implementation

All official performance-test rows passed with:

MaxErr = 3.91e-03

Performance methodology

• GPU: AMD Instinct MI350X (gfx950)
• dtype: bf16
• mode: causal
• warmup: 50 iterations
• measured iterations: 100
• timing: official GPU profiler timing
• FLYDSL_PERFTEST_USE_EVENTS unset
• baseline: latest origin/main
• comparison: rebased branch with return_lse=True
• CK and ASM measured through the official --compare harness
• runs executed sequentially on one GPU

Performance

Values are latency / TFLOPS.

Shape (B,S,H,Hkv,D)	FlyDSL baseline	FlyDSL +LSE	Change	CK +LSE	ASM +LSE
(2,2048,16,16,128) MHA	51.4 µs / 668.2	51.3 µs / 670.4	-0.2%	82.2 µs / 418.3	60.9 µs / 564.5
(2,4096,16,16,128) MHA	156.0 µs / 881.5	154.7 µs / 888.6	-0.8%	175.9 µs / 781.5	124.0 µs / 1108.7
(1,8192,16,16,128) MHA	263.0 µs / 1045.2	262.8 µs / 1046.0	-0.1%	342.3 µs / 803.1	230.1 µs / 1194.7
(2,2048,16,2,128) GQA	51.3 µs / 669.9	50.4 µs / 682.3	-1.8%	83.4 µs / 412.3	62.0 µs / 554.5
(4,2048,8,8,64) D64	40.5 µs / 424.5	38.7 µs / 444.3	-4.4%	61.6 µs / 278.9	N/A
(8,1024,16,16,128) MHA	57.4 µs / 599.1	58.5 µs / 587.8	+1.9%	81.2 µs / 423.6	49.6 µs / 693.1
(2,4096,16,16,128) repeat	156.1 µs / 880.4	155.2 µs / 886.0	-0.6%	178.6 µs / 769.5	124.1 µs / 1107.8

Performance conclusion

• No broad kernel-performance regression from optional LSE output.
• Five of six distinct shapes improved or remained effectively unchanged.
• The only slower result was 1024 MHA at +1.9%.
• D64 improved by 4.4% in this run.
• FlyDSL with LSE remained faster than CK on every tested shape.
• CK and ASM already return LSE in both runs, so their measurements are comparison baselines rather than before/after LSE measurements.

Known follow-up

The public API currently allocates a fresh LSE tensor for every call with return_lse=True.

A future optional lse= output parameter, matching the existing out= convention, would allow latency-sensitive callers to reuse a preallocated buffer and remove that API-level allocation cost. This is not part of this change.