[Attention] support softmax lse

[Valentine233]

Support softmax lse for prefill/decode/chunk_prefill.
Fix speculative decoding issue with topk>1 https://jira.devtools.intel.com/browse/SGLANGT-1430.
The wheel size increases for 1M and the compiling time hardly changes ~20min.

[Copilot]

Pull request overview

This PR adjusts the FlashAttention SYCL forward path for chunk prefill (two-launch dispatcher) so that softmax_lse is returned with the correct non-empty shape even though the value is not stitched/computed on that path, addressing a speculative decoding issue when topk > 1.

Changes:

• Allocate and return softmax_lse in the chunkprefill path with shape (nheads, total_q) instead of an empty {0} placeholder.
• Derive total_q and num_heads from q to construct the correct softmax_lse shape.
• Add TODO notes indicating prefill softmax_lse is currently unsupported.

[Valentine233]

Figure out that softmax_lse would indeed affect the calculation.
The previous fixing only makes sure the right shape with empty values.
As the next step, I will implement the complete softmax_lse calculation for prefill and chunk_prefill.

[ckvermaAI]

repro_softmax_lse.py
@Valentine233 Please check the attached unit test. Currently it fails with this PR, please address the failing test cases.
(These test cases needs to be passed to make ring attention functional in sglang)

[Valentine233]

repro_softmax_lse.py @Valentine233 Please check the attached unit test. Currently it fails with this PR, please address the failing test cases. (These test cases needs to be passed to make ring attention functional in sglang)

Hi @ckvermaAI, there are several checks in your reproducer. Could you elaborate what the exact failure is, core dump or accuracy issue?

[mingfeima]

need rebase.

[Valentine233]

There are some hang and accuracy issues after rebase. The fix is WIP.

[Valentine233]

Issue fixed.