[XPU] Support softmax_lse in sgl_kernel::fwd API

[Valentine233]

Motivation

sgl-project/sgl-kernel-xpu#330 supports softmax_lse in XPU attention now, so sgl_kernel::fwd needs to know if softmax_lse is required.

Modifications

• Write in place out and softmax_lse. softmax_lse is None if it isn't required.
• Make the return to be None.

Accuracy Tests

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #32465198483
Latest PR Test (Extra): ❌ Run #32465198031
Latest PR Test (AMD ROCm 7.2): ❌ Run #32465198154

[Valentine233]

@sunjiweiswift @mingfeima Please help review, thanks.

[rahulvijayaraghavan]

@Valentine233
Any change required in xpu_backend.py where flash_attn_with_kvcache / flash_attn_varlenc kernels are called?

SGLang framework expects 3D shape (batch_size, nheads, seqlen)

sglang/python/sglang/kernels/aot/python/sgl_kernel/flash_attn.py

Line 158 in 256981c

softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The

But XPU kernel is returning 2D shape. So some alignment is required in xpu_backend.py

[Valentine233]

@Valentine233 Any change required in xpu_backend.py where flash_attn_with_kvcache / flash_attn_varlenc kernels are called?

SGLang framework expects 3D shape (batch_size, nheads, seqlen)

sglang/python/sglang/kernels/aot/python/sgl_kernel/flash_attn.py

Line 158 in 256981c

softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The

But XPU kernel is returning 2D shape. So some alignment is required in xpu_backend.py

@rahulvijayaraghavan Thanks to point out this. The LSE shape is not explicitly assumed in xpu_backend.py. Instead, it should be changed in sgl-kernel-xpu from 2D (num_heads, total_q) to 3D (batch_size, nheads, seqlen) shape. We can merge batch_size and seqlen for 2D shape, but there is a problem for 3D shape: If each batch has different seqlen_q, how to deal with LSE's seqlen_q? From this point of view, 2D is a better choice than 3D. In addition, merge_state_v2 only accepts 2D LSE.

[rahulvijayaraghavan]

@Valentine233 @ckvermaAI tested the kernel PR for Wan model in SP mode which uses this feature. He added a reshape for for a specific varlenc call. Similar alignment is required for all calls in xpu_backend.py if kernel returns 2D shape.

@ckvermaAI
Can you share the reshape adjustment made in xpu_backend.py

[ckvermaAI]

diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py
index 613f2df661..0b92ef5893 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/xpu_backend.py
@@ -116,6 +125,12 @@ class XPUAttentionImpl(AttentionImpl):
         if return_softmax_lse:
             out_tensor, softmax_lse = out[:2]
             result = out_tensor.reshape(bsz, seqlen_q, nheads_q, d)
+            # The varlen kernel returns the LSE packed as [num_heads, total_q];
+            # total_q is the row-major flattening of (batch, seq) built by
+            # _get_cu_seqlens. Consumers (ring attention's per-hop merge,
+            # _merge_attention_partials) expect the dense-FA layout [B, H, S],
+            # which is what FlashAttentionImpl produces on its non-varlen path.
+            softmax_lse = softmax_lse.reshape(nheads_q, bsz, seqlen_q).transpose(0, 1)
             return result, softmax_lse
 
         result = out.reshape(bsz, seqlen_q, nheads_q, d)

[Valentine233]

@rahulvijayaraghavan @ckvermaAI According to your input, 3D LSE is required for multimodal gen models, but it's not necessary to modify all the calls in xpu_backend.py. Please feel free to submit PR for LSE reshape as needed and it is essential to add a check before reshape.

[mingfeima]

@Valentine233 this change should be only related to api change from sgl-kernel-xpu? not related to torch version?

[Valentine233]

@Valentine233 this change should be only related to api change from sgl-kernel-xpu? not related to torch version?

The change is xpu only, and the API sgl_kernel::fwd is defined in sgl-kernel-xpu, so I don't think it's related to torch version.

[sunjiweiswift]

How to determine whether LSE-Softmax is needed?

[Valentine233]

In the definition of mha_fwd https://github.com/sgl-project/sgl-kernel-xpu/pull/330/changes#diff-0c0fb34a127ec671cce962e4f85e9d124264131f65016cc0299651b724ccefd8L49-R81, the input softmax_lse is an optional tensor, so we can tell from softmax_lse.has_value().

[rahulvijayaraghavan]

According to your input, 3D LSE is required for multimodal gen models, but it's not necessary to modify all the calls in xpu_backend.py

No. It is not specific to multimodal. We expect same shape in all cases if return_softmax_lse=True.

We have shape differences with CUDA flash_attn API definition which needs to be handled in xpu_backend.py

CUDA (returns 3D softmax_lse)

sglang/python/sglang/kernels/aot/python/sgl_kernel/flash_attn.py

Lines 156 to 160 in 256981c

	Return:
	out: (batch_size, seqlen, nheads, headdim).
	softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
	logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
	normalization factor).

XPU returns 2D softmax_lse
https://github.com/sgl-project/sgl-kernel-xpu/blob/e44f649ca6383c7d7bae23f7cbe218b1f978ff4c/python/sgl_kernel/flash_attn.py#L223-L229

[Valentine233]

@rahulvijayaraghavan The comment for CUDA is not complete here:

sglang/python/sglang/kernels/aot/python/sgl_kernel/flash_attn.py

Lines 156 to 160 in 256981c

	Return:
	out: (batch_size, seqlen, nheads, headdim).
	softmax_lse [optional, if return_softmax_lse=True]: (batch_size, nheads, seqlen). The
	logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
	normalization factor).

In the CUDA implementation, varlen path uses 2D LSE and dense path uses 3D LSE. Please check this https://github.com/sgl-project/sgl-flash-attn/blob/main/hopper/flash_api.cpp#L887-L892.