[Bug]: ARK return_lse gives the same LSE for every 16 query rows (sdpa, sagev1, sageattn) on XPU

Start by tracing the LSE output from ark.sdpa, ark.sagev1, and ark.sageattn(kernel="v1_pvhalf") through the FMHA epilogue described in #1991; the issue does not name source files or tests. Use the supplied XPU reproduction to compare each query row against its reference, including the smooth_k=True case. Done means returning one correct log2 LSE per row for the original unsmoothed logits across the reported APIs.

Problem Description

With return_lse=True, ark.sdpa, ark.sagev1 and ark.sageattn(kernel="v1_pvhalf") return an LSE tensor with the right shape, [B, H, S_q], but every block of 16 consecutive query rows holds the same value. For sdpa that value is exactly the LSE of the last row in the block: rows 0-15 all get row 15's LSE, rows 16-31 get row 31's, and so on. The attention output itself is fine. Only the LSE is wrong.

I see this with sagev1 and sageattn in both HND and NHD, bf16 and fp16, head dim 64 and 128, and with GQA. The script below shows it for sdpa too.

It's easy to miss with plain torch.randn inputs. In that case every row's LSE is about log2(kv_len), so the error stays under 0.5. Once the rows really differ, the error is around 9 (log2 units) in the repro. I think that's why test_lse.py passes. It uses randn q/k/v and checks that a split-KV merge matches the full run, but it never compares the LSE against a reference.

There's also a second, smaller issue with smooth_k=True, the default for sagev1 and sageattn. The LSE comes from the smoothed logits, so it's missing the q . mean(k) * scale term that smoothing subtracts. With 16 equal rows per block (so the first bug stays out of the way) and K offset by 0.5, the error is 3.7. Adding that term back brings it down to 0.03. For ring attention or split-KV this matters, because each KV chunk has its own mean(k), so the partial LSEs aren't on the same scale.

Expected: one LSE per query row, in log2 units as described in #1991, computed for the original unsmoothed logits, i.e. log2(sum_j exp(scale * q . k_j)).

Reproduction Steps

1. pip install auto-round-lib==0.16.0 on an XPU machine
2. Run this script:

import math

import auto_round_kernel as ark
import torch

torch.manual_seed(0)
B, H, SQ, SK, D = 1, 4, 64, 1000, 128
scale = D**-0.5
LOG2 = math.log(2.0)

q = torch.randn(B, H, SQ, D, device="xpu").bfloat16()
k = torch.randn(B, H, SK, D, device="xpu").bfloat16()
v = torch.randn(B, H, SK, D, device="xpu").bfloat16()

# give each query row a different magnitude, so each row has a different LSE
row_scale = torch.linspace(0.25, 4.0, SQ, device="xpu").view(1, 1, SQ, 1)
q_rows = (q.float() * row_scale).bfloat16()
# same, but the 16 rows of each block are identical
q_blocks = q_rows[:, :, ::16].repeat_interleave(16, dim=2)


def ref_lse_log2(q, k):
    # log2-sum-exp of the original (unsmoothed) scaled logits, one value per query row
    return torch.logsumexp((q.float() @ k.float().transpose(-1, -2)) * scale, dim=-1) / LOG2


apis = {
    "sdpa": lambda q, k, v: ark.sdpa(q, k, v, scale=scale, tensor_layout="HND", return_lse=True),
    "sagev1": lambda q, k, v: ark.sagev1(q, k, v, scale=scale, tensor_layout="HND", return_lse=True),
    "sagev1 smooth_k=False": lambda q, k, v: ark.sagev1(
        q, k, v, scale=scale, tensor_layout="HND", return_lse=True, smooth_k=False
    ),
    "sageattn v1_pvhalf": lambda q, k, v: ark.sageattn(
        q, k, v, tensor_layout="HND", sm_scale=scale, return_lse=True, kernel="v1_pvhalf"
    ),
}

print("1) every row different: only one LSE value per 16 rows")
_, lse = apis["sdpa"](q_rows, k, v)
ref = ref_lse_log2(q_rows, k)
print("   row   ark sdpa lse   reference (log2)")
for r in range(12, 20):
    print(f"   {r:3d}   {lse[0, 0, r].item():12.3f}   {ref[0, 0, r].item():16.3f}")

print("\n2) max |lse - reference| per API")
print(f"   {'api':24s} {'rows differ':>12s} {'16 equal rows':>14s} {'16 equal rows, K+0.5':>21s}")
for name, fn in apis.items():
    errs = []
    for qq, kk in ((q_rows, k), (q_blocks, k), (q_blocks, (k.float() + 0.5).bfloat16())):
        _, lse = fn(qq, kk, v)
        errs.append((lse.float() - ref_lse_log2(qq, kk)).abs().max().item())
    print(f"   {name:24s} {errs[0]:12.3f} {errs[1]:14.3f} {errs[2]:21.3f}")

print("\n3) smooth_k=True: adding back q . mean(k) * scale fixes the K+0.5 case")
k_off = (k.float() + 0.5).bfloat16()
_, lse = ark.sagev1(q_blocks, k_off, v, scale=scale, tensor_layout="HND", return_lse=True)
shift = (q_blocks.float() * k_off.float().mean(dim=2, keepdim=True)).sum(-1) * scale / LOG2
ref = ref_lse_log2(q_blocks, k_off)
print(f"   max |lse - reference|         = {(lse - ref).abs().max().item():.3f}")
print(f"   max |lse + shift - reference| = {(lse + shift - ref).abs().max().item():.3f}")

Environment Information

• OS: Ubuntu 25.04, kernel 6.14.0-37-generic
• Python: 3.13.18
• auto-round-lib: 0.16.0 (PyPI)
• torch: 2.14.0+xpu
• Hardware: Intel Arc Pro B60 (Battlemage), driver 1.15.39122+11

Error Logs

No error, just wrong values. Script output:

1) every row different: only one LSE value per 16 rows
   row   ark sdpa lse   reference (log2)
    12         10.858             10.833
    13         10.858             10.610
    14         10.858             10.971
    15         10.858             10.858
    16         12.986             10.881
    17         12.986             10.886
    18         12.986             11.056
    19         12.986             11.521

2) max |lse - reference| per API
   api                       rows differ  16 equal rows  16 equal rows, K+0.5
   sdpa                            9.093          0.000                 0.000
   sagev1                          8.994          0.175                 3.727
   sagev1 smooth_k=False           9.095          0.022                 0.051
   sageattn v1_pvhalf              8.994          0.175                 3.727

3) smooth_k=True: adding back q . mean(k) * scale fixes the K+0.5 case
   max |lse - reference|         = 3.727
   max |lse + shift - reference| = 0.032

Additional Context

We hit this while adding ARK to the sage_attn diffusion backend in SGLang on XPU. For now we only use ARK when no LSE is needed, and ring attention falls back to our own kernel.

Since sdpa returns exactly the last row's value for every 16-row block, my guess is that the LSE store added to the FMHA epilogue in #1991 writes all rows of a 16-row group to the same slots, and the last one wins. That's only a guess from the outside; I haven't read the kernel code.

One small doc thing: the sagev1 docstring says LSE[b, h, q] = log(sum_j exp(score)), but the values are log2, as #1991 describes. It might be worth saying log2 there.

Happy to test a fix.

JH...!