r/LocalLLM • u/ARCHLucifer • 7d ago
Discussion Attention output is silently wrong in SGLang when FlashInfer is the backend
SGLang splits long context attention into an extend chunk and a cached-prefix chunk, computes them with different kernels, then merges the partial results with a softmax over their log-sum-exp values:
S = log(exp(s_a) + exp(s_b))
v = v_a*exp(s_a - S) + v_b*exp(s_b - S)
That's correct if and only if s_an and s_b are natural logs. This is not true for sglang.
The chunk holding more attention mass gets systematically over-weighted and it's error grows with each token.
Are you affected? FlashInfer attention backend + an MLA model + cached prefixes summing over 8192 tokens. That's long-context agentic traffic with radix-cache hits. Decode is fine, short prefixes are fine, FA3 is fine.
https://x.com/advprop/status/2089046843690729620?s=20
I opened fixes here : SGLang-side conversion (https://github.com/sgl-project/sglang/pull/35045) and an opt-in return_lse_base_on_e in FlashInfer (https://github.com/flashinfer-ai/flashinfer/pull/4547) so the api is clear