[FMHA] Add optional ALiBi slopes to the paged FMHA kernels by desmond-intel · Pull Request #573 · sgl-project/sgl-kernel-xpu

desmond-intel · 2026-10-23T06:35:42Z

Motivation

ALiBi models have no RoPE, so attention needs an ALiBi bias to see any position information. One example is Baichuan2-13B. On the intel_xpu backend the bias is dropped, and the model degenerates: Baichuan2-13B-Chat gsm8k is 0.00 on B60 (vLLM RTX 6000 Blackwell: 53.5).

Modifications

Add an optional fp32 alibi_slopes [num_heads] argument. The kernel computes the bias -slope * |i + s_k - s_q - j| in place, on the existing relative-bias instantiations, so no new templates are compiled.

xe20/collective/xe_fmha_fwd_mainloop.hpp: ptr_alibi_slopes in the prefill and decode mainloop args. When it is set, apply_alibi_bias computes the bias from the slopes and the identity-tensor (row, col) coordinates; otherwise apply_relative_bias does the rel_bias lookup. Both add it through add_scaled_bias. Packed decode rows are clamped to the real query-head count.
xe20/kernel/xe_fmha_fwd_kernel.hpp: in split decode, the slope index is based on the kv head (head * head_group_q).
xe20/xe_fmha_fwd_{prefill,decode,split_decode}_kernel.cpp.in, runners: route alibi_slopes_ptr to the HasRelBias kernels, the same way as rel_bias_ptr; the fp8 and non-paged runners reject it. The new runner field is the last member of Arguments, so the existing layout is unchanged for cached JIT modules.
flash_attentionXe20.cpp: thread alibi_slopes through prefill / decode / top-level mha_fwd. For max_seqlen_q > 1 it takes the existing rel_bias route (prefill for all rows). It checks fp32 [H], a paged KV cache, head_dim 128, no fp8 KV, and no combination with softcap / LSE / sinks / rel_bias / bidirectional_block_ids / sliding window. The op schema and sgl_flash_kernel_ops.h are updated.
flash_attentionXe35.cpp: takes alibi_slopes and rejects it, like rel_bias.
python/sgl_kernel/flash_attn.py: alibi_slopes kwarg on flash_attn_with_kvcache. The page_size == 1 helper rejects it.
tests/test_flash_attention.py: test_alibi_attention checks the kernel against an fp32 reference across prefill, mixed, decode, forced split decode, GQA/MHA, bf16/fp16 and page 64/128.

The bias is computed rather than materialized as a rel_bias surface, so it costs O(H) memory and stays in fp32.

The change is fully gated: with alibi_slopes=None, every existing path is unchanged.

Test

Intel Arc Pro B60:

test_alibi_attention 44/44.
Existing test_relative_attention* 129 pass.
Full test_flash_attention.py + test_flash_attn_bidirectional.py on top of Upstream flash attn xe3 #574: 11,013 pass.

End to end, Baichuan2-13B-Chat tp2 --attention-backend intel_xpu, gsm8k 5-shot: 0.00 → 0.52 on 50 docs; full 1319 with the fix 0.5474 flexible / 0.5459 strict (vLLM RTX 6000 Blackwell 53.5).

While testing I noticed that split-decode rel_bias with num_heads_k > 1 differs from num_splits=1 on kv heads > 0 (up to 0.2 abs). Could head_q_start there be set only under LSE? The existing rel_bias decode tests use num_heads_k=1.

The companion sglang change is sgl-project/sglang#45999.

Copilot

🔵 Needs a closer look

Cached JIT module compatibility remains unresolved, and the shared numerical kernel changes warrant hardware-backed human review.

1 open issue

Preserve runner argument layouts or invalidate affected JIT caches · New

What changed in this PR

Adds optional ALiBi positional bias to Intel XPU paged attention, enabling models without RoPE to use the existing FMHA paths without materializing a bias tensor.

Changes:

Threads fp32 per-head slopes through the Python API, Torch registration, and kernel runners.
Computes ALiBi bias in prefill, decode, and split decode, with validation for supported configurations.
Adds reference-based tests covering attention modes, head grouping, dtypes, and page sizes.

File	Description
`tests/test_flash_attention.py`	Adds ALiBi correctness tests against an fp32 reference.
`src/torch_extension_sycl.cc`	Adds the optional slopes argument to the op schema.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_split_decode_kernel.cpp.in`	Routes ALiBi through bias-enabled split decode.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp`	Passes slopes into the prefill mainloop.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_kernel.cpp.in`	Selects bias-enabled prefill for ALiBi.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_decode_runner.hpp`	Passes slopes into decode and split decode.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_decode_kernel.cpp.in`	Selects bias-enabled decode for ALiBi.
`src/sycl/kernels/flash_attention_v2/kernel/xe_fmha_fwd_kernel.hpp`	Sets the ALiBi head offset for split decode.
`src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp`	Computes positional bias directly from slopes.
`src/sycl/flash_attention.cpp`	Validates slopes and dispatches supported attention paths.
`python/sgl_kernel/flash_attn.py`	Exposes and documents the optional slopes argument.
`include/sgl_flash_kernel_ops.h`	Extends the public C++ declaration.

_{Review effort: Balanced}

ALiBi models (e.g. Baichuan2-13B) skip RoPE, so attention needs the ALiBi bias to see any position information. Add an optional fp32 alibi_slopes [num_heads] argument to flash_attn_with_kvcache. apply_relative_bias computes -slope * |i + s_k - s_q - j| from the identity-tensor coordinates in place of the rel_bias lookup, so it rides the existing HasRelBias instantiations (no new templates) in prefill, decode and split decode. alibi_slopes=None leaves every path unchanged. Tested on Arc Pro B60: test_alibi_attention 40/40, test_relative_attention* 129 passed, test_flash_attention.py 10,436 passed. Baichuan2-13B-Chat gsm8k (intel_xpu, tp2): 0.00 -> 0.52 on 50 docs, 0.5474 on the full 1319.

Move alibi_slopes_ptr to the end of the prefill and decode runner Arguments so the existing member offsets are unchanged for cached JIT modules of unchanged templates. flash_attentionXe35.cpp takes the new alibi_slopes argument and rejects it, like rel_bias. Tested on Arc Pro B60: test_alibi_attention 40/40, test_relative_attention* 129 passed, test_flash_attention.py + test_flash_attn_bidirectional.py 11,005 passed.

mingfeima

generally ok, some minor changes needed.

mingfeima · 2026-10-27T08:22:22Z

    )


+@pytest.mark.skipif(device.type != "xpu", reason="ALiBi requires an XPU device")


you don't have to skip xpu, you need to skip cri.

Done in da4ab30: the test now skips only on Xe35 (is_cri_device()), like the softcap and relative-bias tests. On B60 all 40 cases run and pass.

mingfeima · 2026-10-27T08:26:30Z

ci not triggered due to lack run-ci label.

The test file only runs on XPU, so the device check is not needed. Skip on Xe35 instead, which rejects alibi_slopes, the same way the softcap and relative-bias tests do. On Arc Pro B60 all 40 cases still run and pass.

Valentine233 · 2026-10-28T06:57:56Z

+      for (int i = 0; i < scores.size(); ++i) {
+        int row_idx = get<0>(cS_thread(i));
+        int row_kv_idx = (PackGQA ? 0 : row_idx) + full_tile_offset;
+        ElementS slope = params.ptr_alibi_slopes[bias_head_start + (PackGQA ? row_idx : 0)];


Suggested change

ElementS slope = params.ptr_alibi_slopes[bias_head_start + (PackGQA ? row_idx : 0)];

int slope_idx = bias_head_start + (PackGQA ? sycl::min(row_idx, q_rows - 1) : 0);

ElementS slope = params.ptr_alibi_slopes[slope_idx];

Applied in c517084: apply_alibi_bias now clamps the packed row with sycl::min(row_idx, q_rows - 1).

Valentine233 · 2026-10-28T07:10:55Z

@@ -343,6 +345,21 @@ struct FMHAFwdMainloop<
      int thr_id,
      int full_tile_offset) const {


The non-split kv decode also calls FMHAFwdMainloop. To avoid index out-of-bound access, we should add q_rows like DecodeFwdMainloop.

Suggested change

int full_tile_offset) const {

int full_tile_offset,

int q_rows) const {

Done in c517084: the non-split decode call now passes q_rows = size<0>(Q_2D) (the packed group size), the same way DecodeFwdMainloop does.

Valentine233 · 2026-10-28T07:14:27Z

+        (True, [1, 1], [4097, 2050], 8, 2, 4),
+        (True, [1], [4097], 8, 8, 4),
+    ],
+)


Add test for non-split kv decode, for example num_heads=10, num_heads_k=2 or num_heads=12, num_heads_k=4.

Added num_heads=10, num_heads_k=2 and 12, 4 with num_splits=1 in c517084. The groups of 5 and 3 pad to 8- and 4-row tiles.

Valentine233 · 2026-10-28T07:30:47Z

+    TORCH_CHECK(
+        softcap == 0.f && !return_softmax_lse && !sinks_.has_value() && !rel_bias_.has_value() &&
+            !bidirectional_block_ids_.has_value(),
+        "alibi_slopes does not support softcap, softmax LSE, sinks, rel_bias or bidirectional_block_ids");


If sliding window is also not supported for alibi_slopes, add a check.

Added in c517084: mha_fwd now rejects alibi_slopes with a sliding window.

Valentine233 · 2026-10-28T07:32:06Z

      "sink is only supported for head_size == 64, got @HEAD_DIM@");
 #if @HEAD_DIM@ == 128
-  if (params.rel_bias_ptr != nullptr) {
+  if (params.rel_bias_ptr != nullptr || params.alibi_slopes_ptr != nullptr) {


Add checks for all other .cpp.in files, like fp8, non paged.

Added in c517084: the fp8 and non-paged runners now reject alibi_slopes, the same way they reject rel_bias.

Valentine233 · 2026-10-28T07:41:44Z

      int thr_id,
      int full_tile_offset) const {
    constexpr ElementS kLog2e = ElementS(1.4426950408889634074);
+    if (params.ptr_alibi_slopes != nullptr) {


Instead of putting inside apply_relative_bias, could we create a specific function apply_alibi_bias for alibi_slopes? If there are any common logits, abstract them as a new function.

Split out apply_alibi_bias in both mainloops in c517084. The shared s = qk_scale * s + log2(e) * bias step is now add_scaled_bias, which the rel-bias path uses too.

Valentine233 · 2026-10-28T07:45:25Z

While testing I noticed that split-decode rel_bias with num_heads_k > 1 differs from num_splits=1 on kv heads > 0 (up to 0.2 abs). Could head_q_start there be set only under LSE? The existing rel_bias decode tests use num_heads_k=1.

@yuankuns Could you please have a look at this relative bias bug?

…add checks - Split ALiBi out of apply_relative_bias into apply_alibi_bias in both mainloops. The shared "s = qk_scale * s + log2(e) * bias" step is now add_scaled_bias, used by the relative-bias paths too. - Non-split decode (FMHAFwdMainloop with PackGQA) clamps the packed row to q_rows = size<0>(Q_2D) before reading alibi_slopes, as split decode already does, so padded rows of the last KV head stay in bounds. - mha_fwd rejects alibi_slopes with a sliding window. The fp8 and non-paged runners reject alibi_slopes, like rel_bias. - test_alibi_attention: add non-split decode cases with GQA groups of 5 and 3 (num_heads 10/2 and 12/4). Tested on Arc Pro B60: test_alibi_attention 44/44, test_relative_attention* 129 passed, test_flash_attention.py + test_flash_attn_bidirectional.py 11,013 passed.

Copilot AI balanced review requested due to automatic review settings October 23, 2026 06:35

Copilot started reviewing on behalf of desmond-intel October 23, 2026 06:36

View session

desmond-intel mentioned this pull request Oct 23, 2026

[Fix] [XPU] baichuan: Apply ALiBi slopes on intel_xpu sgl-project/sglang#45999

Open

3 of 5 tasks

Copilot AI reviewed Oct 23, 2026

View reviewed changes

Comment thread src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp Resolved

desmond-intel force-pushed the xpu/alibi-fmha branch from 691af7b to 9e8f4c7 Compare October 26, 2026 17:00

mingfeima requested changes Oct 27, 2026

View reviewed changes

mingfeima requested a review from Valentine233 October 27, 2026 08:25

mingfeima added the run-ci label Oct 27, 2026

[FMHA] Skip the ALiBi test only on Xe35

da4ab30

The test file only runs on XPU, so the device check is not needed. Skip on Xe35 instead, which rejects alibi_slopes, the same way the softcap and relative-bias tests do. On Arc Pro B60 all 40 cases still run and pass.

Valentine233 reviewed Oct 28, 2026

View reviewed changes

desmond-intel requested a review from Valentine233 October 29, 2026 10:11

		)


		@pytest.mark.skipif(device.type != "xpu", reason="ALiBi requires an XPU device")

	ElementS slope = params.ptr_alibi_slopes[bias_head_start + (PackGQA ? row_idx : 0)];
	int slope_idx = bias_head_start + (PackGQA ? sycl::min(row_idx, q_rows - 1) : 0);
	ElementS slope = params.ptr_alibi_slopes[slope_idx];

		@@ -343,6 +345,21 @@ struct FMHAFwdMainloop<
		int thr_id,
		int full_tile_offset) const {

	int full_tile_offset) const {
	int full_tile_offset,
	int q_rows) const {

Conversation

desmond-intel commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Test

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🔵 Needs a closer look

Uh oh!

Uh oh!

mingfeima left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mingfeima commented Oct 27, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Valentine233 Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Valentine233 Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Valentine233 Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Valentine233 commented Oct 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

desmond-intel commented Oct 23, 2026 •

edited

Loading

Valentine233 Oct 28, 2026 •

edited

Loading

Valentine233 Oct 28, 2026 •

edited

Loading

Valentine233 Oct 28, 2026 •

edited

Loading

Valentine233 commented Oct 28, 2026 •

edited

Loading