feat: add return_lse support for SDPA and SAGEV1 kernel — enable sequence-parallel (Ring Attention) merge

[luoyu-intel]

Summary

Add optional return_lse (Log-Sum-Exp) support to the ARK FMHA and SAGEV1 kernels, enabling correct sequence-parallel / Ring Attention merge via LSE rescaling. When return_lse=False (default), behavior is identical to the previous code — zero performance regression verified by benchmark.

Changes

1. Kernel layer (sycl-tla / stla)

• xe_fmha_fwd_epilogue.hpp: Epilogue now writes LSE = max + log₂(sum) to an optional float32 output buffer. The kernel already computes running max and sum internally for online softmax — this change only stores already-available values.
• xe_fmha_fwd_kernel.hpp: KernelArguments gains float* Lse field; to_underlying_arguments forwards LSE pointer to epilogue Arguments. Epilogue operator() receives head_q and idx_b for correct LSE output layout indexing.

2. Wrapper layer (ARK)

• sycl_tla_sdpa.hpp: Options struct and both KernelRunner::run instantiations thread lse pointer through.
• sdpa.cpp: All dispatchers (sdpa_impl, sdpa_impl_qks8_pvhalf, sdpa_impl_qks8_pvi8, sage_prefill, sage_prefill_varlen, sdpa_varlen_impl, flash_attn_prefill, flash_attn_decode) accept optional float* lse = nullptr.
• xpu_wrapper.hpp: sagev1_impl, sagev1, sagev1_pvi8, sagev1_varlen_impl thread LSE pointer through.
• sycl_tla_common.hpp: Forward declarations updated.
• ark.cpp: All pybind wrappers accept torch_ptr lse = 0.

3. Python API

• sdpa(), sdpa_varlen(), sagev1(), sagev1_pvi8(), sageattn(), sageattn_varlen() — all gain return_lse: bool = False parameter.
• When True, returns (O, LSE) tuple with LSE shape [batch, num_heads, seq_len] (float32).
• When False, returns O only — backward compatible, zero overhead.

4. Environment variable

• Added IGC_RemoveUnusedIdImplicitLocalIDs=0 to __init__.py to prevent driver-upgrade-induced performance regressions. Reference: IGC Register-Allocation Regression on Xe2 (BMG) intel-graphics-compiler#412

Config	Bench	IGC=1 (us)	IGC=0 (us)	Diff	TFLOPS(1)	TFLOPS(0)
96_H128_K8192	SDPA	20281.0	20226.9	-0.27%	81.3	81.5
	SAGEV1	23709.6	17168.3	-27.59% 🚨	69.6	96.1
96_H128_K16384	SDPA	40346.7	40401.1	+0.13%	81.8	81.6
	SAGEV1	48073.8	33648.4	-30.01% 🚨	68.6	98.0
96_H128_K32768	SDPA	81660.6	81155.0	-0.62%	80.8	81.3
	SAGEV1	97502.4	66668.0	-31.62% 🚨	67.7	99.0

LSE Formula

The attention kernel uses exp2-based softmax internally:
lse[q] = max_j(score_{q,j}) + log₂(sum_j exp₂(score_{q,j} - max_j))
LSE merge for split-KV (Ring Attention):
lse_max = max(lse₁, lse₂)
w_k = 2^(lse_k - lse_max)
O_merged = (O₁·w₁ + O₂·w₂) / (w₁ + w₂)

Verification

Benchmark (no performance regression with return_lse=False)

Config	baseline (us)	with LSE (us)	diff
q96_hd128_k8192	17152.8	17152.4	-0.00%
q96_hd128_k16384	33621.1	33620.8	-0.00%
q96_hd128_k32768	66593.5	66582.6	-0.02%
q128_hd64_k8192	14046.7	14063.1	+0.12%
q16_hd128_k8192	3229.3	3232.9	+0.11%
avg			+0.03%

Max diff +0.12% — well within measurement noise.

Benchmark (no performance regression with return_lse=False|True)

Config	Bench	LSE=False (us)	LSE=True (us)	Diff	TFLOPS (F)	TFLOPS (T)
96_H128_K8192	SDPA	20109.9	20657.8	+2.72%	82.0	79.8
	SAGEV1	17167.9	17836.0	+3.89%	96.1	92.5
96_H128_K16384	SDPA	40131.7	40550.1	+1.04%	82.2	81.3
	SAGEV1	33640.0	34247.2	+1.81%	98.1	96.3
96_H128_K32768	SDPA	80966.3	81522.0	+0.69%	81.5	80.9
	SAGEV1	66663.8	67225.5	+0.84%	99.0	98.1

Unit tests (21/21 passed)

• API contract: return_lse=False → Tensor, return_lse=True → (O, LSE), O from both paths identical (atol=0)
• LSE merge correctness: Split KV into 2 or 3 chunks, compute attention + LSE per chunk, merge via rescaling, compare against full attention — O and LSE match within tolerance
• Shapes: SDPA (64/96/128/192 head dim), SAGEV1 (64/128), causal/non-causal, power-of-2 and non-power-of-2 KV lengths

[Copilot]

Pull request overview

Adds optional Log-Sum-Exp (LSE) output plumbing for XPU attention kernels (SDPA + SAGEV1) to support sequence-parallel / Ring Attention merges, plus Python API exposure and unit tests. Also sets an Intel IGC environment variable to stabilize performance across driver/compiler changes.

Changes:

• Thread an optional float* lse output pointer through SYCL-TLA wrapper options, kernel arguments, and epilogues.
• Extend the Python-facing API with return_lse to optionally return (O, LSE) for SDPA/SAGEV1 paths and add unit tests validating merge correctness.
• Set IGC_RemoveUnusedIdImplicitLocalIDs=0 at import time to avoid performance regressions tied to IGC behavior.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
auto_round_extension/ark/test/test_lse.py	New XPU unit tests covering return_lse contract and split-KV merge via LSE rescaling.
auto_round_extension/ark/auto_round_kernel/wrapper/include/xpu_wrapper.hpp	Threads optional lse pointer through XPU wrapper entry points into ARK SDPA/SAGE dispatchers.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_sdpa.hpp	Adds Options::lse and forwards it into SDPA/SAGE kernel runner arguments.
auto_round_extension/ark/auto_round_kernel/wrapper/include/sycl_tla_common.hpp	Updates forward declarations to accept an optional float* lse.
auto_round_extension/ark/auto_round_kernel/wrapper/include/stla/xe_sage_fwd_kernel.hpp	Forwards LSE metadata into the epilogue and passes (head_q, idx_b) for correct LSE layout indexing.
auto_round_extension/ark/auto_round_kernel/sdpa.cpp	Propagates lse into option structs and dispatcher calls for prefill/decode and varlen paths.
auto_round_extension/ark/auto_round_kernel/ark.cpp	Extends pybind bridge function signatures to accept an optional LSE pointer.
auto_round_extension/ark/auto_round_kernel/init.py	Python API: adds return_lse, allocates LSE buffers, passes pointers into native bindings, and sets IGC env var.

Comments suppressed due to low confidence (2)

auto_round_extension/ark/auto_round_kernel/init.py:635

• sdpa_varlen() can return either O or (O, LSE) but the return type annotation and docstring still describe only O. This makes the public Python API contract unclear/inaccurate, especially for the LSE shape in varlen mode.

    return_lse: bool = False,
) -> torch.Tensor:
    """Scaled dot-product attention (SDPA) prefill+decode with variable-length sequences.

    Q, K, V use a flat 3-D layout where tokens from all sequences in the batch
    are concatenated along dim-0 (the *total* dimension).  Per-sequence

auto_round_extension/ark/auto_round_kernel/init.py:1235

• sageattn_varlen() can return either O or (O, LSE), but the return type annotation/docstring still indicate only O. This makes it easy to misuse the returned LSE (shape/domain) in sequence-parallel merges.

    return_lse: bool = False,
    **kwargs,
) -> torch.Tensor:
    """SAGE attention with variable-length sequences (no padding).

    Q/K/V are flat 3-D tensors: [total_tokens, num_heads, head_dim].
    Per-sequence boundaries given by ``cu_seqlens_q`` / ``cu_seqlens_k``.

[a32543254]

LGTM

[chensuyue]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.

[AutoRoundBot]

/azp run Unit-Test-CUDA-AutoRound

[azure-pipelines]

Azure Pipelines could not run because the pipeline triggers exclude this branch/path.