Improve fp8_paged_mqa_logits by Xia-Weiwen · Pull Request #398 · sgl-project/sgl-kernel-xpu

Xia-Weiwen · 2026-08-20T06:57:36Z

Added a chunking strategy to avoid OOM with large batch.
Improved performance by vectorization, tuning tile/workgroup size, etc.

polisettyvarma

helps in OOM and which model, right ?

Xia-Weiwen · 2026-08-21T10:06:37Z

helps in OOM and which model, right ?

It can probably resolve the OOM issue of DeepSeek. Jianan is going to have a try.

Copilot

Pull request overview

Adds configurable batch chunking to the Xe20 FP8 paged MQA logits GEMM path.

Changes:

Adds a configurable 512 MiB default chunk budget.
Processes batches with chunk-local intermediates.
Adds optional verbose chunk diagnostics.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

airMeng · 2026-08-25T05:31:22Z

+          " chunks)");
+    }
+
+    for (int start = 0; start < B_next; start += chunk_b) {


avoid launching kernels in the loop, fuse into a larger kernel instead

Thanks for the suggestion. My agent says the loop introduces limited kernel launch overhead (~2% in the DeepSeek V4 case). I have added a TODO here to indicate an opportunity for optimizations. I think we can keep the current implementation to resolve the OOM issue which is blocking.

After an offline discussion, we decide to improve the kernel implementation for better performance so I have converted this PR to draft and will work on it later. Thanks.

The chunking strategy is applied on the batch. The gather - GEMM - reduction kernels are launched one by one. And since they have different launch configs, it's hard to fuse them.
The performance of these kernels is also improved.

Copilot

Copilot review overview

🟡 Changes recommended

Fix potentially invalid unaligned vector accesses and add deterministic multi-chunk coverage.

Get a fresh assessment by requesting another Copilot review.

Review effort: Lite
Findings: 2 · 1

Open (3)

Guard Vec4 stores against misaligned destination rows · New
Guard Vec4 loads against misaligned source addresses · New
Add test coverage for multi-chunk batching and partial final chunks · New

Resolved since last review (2)

The configured budget is not actually applied to all chunked temporaries: k_scale_gathered is…
The existing paged-logits tests never exercise this loop with more than one GEMM chunk: the only…

The "Rebuild sgl-kernel from latest main + PR" step recompiles all ~700 FMHA translation units from scratch on every run (the step deliberately `rm -rf build` so Ninja cannot link a PR against the image-baked .o files), so its wall time tracks the size of the AOT instantiation matrix rather than the size of the PR. That matrix has been growing steadily. Measured on bmg-754, which ran every one of these builds, for the same branches over time: 2026-08-11 27.3m baseline 2026-09-16 31.4m organic growth, ~+1m/week 2026-09-18 40.0m logit soft-cap template axis lands (#491) #491 added `AT_DISPATCH_BOOL_NO_RETURN(params.softcap != 0.f, Softcap, ...)` to the decode/prefill generators, doubling the instantiations inside the existing translation units. The edge and object counts are unchanged (871 edges / 704 objects / 83 .so before and after), so nothing in the build graph flagged it -- only the wall time moved, by +7.6m in a single commit. That left no headroom under the 40m budget. Builds now finish compiling at ~39m43s and get killed during wheel packaging, i.e. after a fully successful compile: PR #398 passed at exactly 40m00s, and the gemma3-bidirectional-image-mask branch has hard-failed four times at 40.1-40.2m without a single test failure. Raise the budget to 75m so a successful compile is not thrown away, and so there is room for the next few quarters of matrix growth. The two downstream steps (Install Dependency 20m, Run Sglang Kernel Cases 70m) are unchanged.

airMeng · 2026-09-22T02:49:42Z

    if (b >= B || kj >= max_seq_len) return;

    int out_idx = b * max_seq_len + kj;
+    int n_vec4 = D / 16;


16 is sg_size?

Thanks. It's the vec len. Replaced with a constexpr var.

airMeng · 2026-09-22T02:52:57Z

+    per_batch_bytes=144128 so chunk_b=7, i.e. two unequal chunks (7 + 1).
+    """
+    tests_dir = os.path.dirname(os.path.abspath(__file__))
+    script = f"""


suggest to use some more smart LLM :)

Thanks. This part is polished.

* [fmha] Select the logit soft-cap at runtime instead of instantiating for it #491 threaded the logit soft-cap through the FMHA kernels as a template bool and dispatched on it with AT_DISPATCH_BOOL_NO_RETURN(params.softcap != 0.f, Softcap, ...) at 17 sites across the eight decode/prefill/split-decode generators. Each of those doubles the instantiations of the branch it wraps, and because the generators are configured over the full AOT matrix (5 QG sizes x 6-8 head dims x 2 page sizes x 2 element types) the cost lands on every translation unit rather than only on the soft-cap path. The build graph does not show it -- edges, objects and .so count are identical before and after -- but the FMHA rebuild went from 31.4m to 40.0m in one commit on bmg-754, with the decode TUs taking the bulk of it (xe_fmha_fwd_decode_page 8.4 -> 11.8m, xe_fmha_fwd_split_decode_page 8.9 -> 11.7m, xe_fmha_fwd_decode_nopage 3.7 -> 5.1m). params.softcap is a kernel-wide scalar, so it does not need to be a template axis: softmax() can duplicate its two inner loops and pick between them with a single uniform branch per softmax block. The uncapped loop body is then byte-for-byte what it was before #491 -- no per-element test -- and one kernel serves both capped and uncapped launches. - collective/xe_fmha_fwd_mainloop.hpp: replace `if constexpr (Softcap)` inside the two softmax() implementations with `if (capped)` around duplicated loops, where `capped = Softcap && softcap != 0`. Softcap is still a template parameter, so `capped` folds to a constant on the Softcap=false instantiations and only the uncapped loops are emitted there. - the eight *_kernel.cpp.in generators: drop the soft-cap dispatch and pass /*Softcap=*/true, restoring the pre-#491 instantiation count exactly (verified per generator against 7d443d1^). - the three relative-attention branches pass /*Softcap=*/false and reject a non-zero cap up front. Both runners already rejected softcap + rel_bias (apply_relative_bias would fold the cap as softcap(QK+bias) while the reference ordering is softcap(QK)+bias), so those Softcap=true instantiations could only ever throw. Soft-cap numerics are unchanged: same apply_logit_softcap on the same natural-scale logit, with masked -INFINITY lanes still passing through untouched. Needs a BMG perf run before merge to confirm the extra branch and the larger Softcap=true kernels do not cost occupancy on the uncapped paths that every non-Gemma model takes. * [ci] Raise the FMHA rebuild step timeout to 75 minutes The "Rebuild sgl-kernel from latest main + PR" step recompiles all ~700 FMHA translation units from scratch on every run (the step deliberately `rm -rf build` so Ninja cannot link a PR against the image-baked .o files), so its wall time tracks the size of the AOT instantiation matrix rather than the size of the PR. That matrix has been growing steadily. Measured on bmg-754, which ran every one of these builds, for the same branches over time: 2026-08-11 27.3m baseline 2026-09-16 31.4m organic growth, ~+1m/week 2026-09-18 40.0m logit soft-cap template axis lands (#491) #491 added `AT_DISPATCH_BOOL_NO_RETURN(params.softcap != 0.f, Softcap, ...)` to the decode/prefill generators, doubling the instantiations inside the existing translation units. The edge and object counts are unchanged (871 edges / 704 objects / 83 .so before and after), so nothing in the build graph flagged it -- only the wall time moved, by +7.6m in a single commit. That left no headroom under the 40m budget. Builds now finish compiling at ~39m43s and get killed during wheel packaging, i.e. after a fully successful compile: PR #398 passed at exactly 40m00s, and the gemma3-bidirectional-image-mask branch has hard-failed four times at 40.1-40.2m without a single test failure. Raise the budget to 75m so a successful compile is not thrown away, and so there is room for the next few quarters of matrix growth. The two downstream steps (Install Dependency 20m, Run Sglang Kernel Cases 70m) are unchanged. * Trigger CI * Update pr-test-xpu.yml * [fmha] Trim the soft-cap comments to 1-3 lines Condense the explanatory blocks added with the runtime soft-cap select down to 1-3 lines per site, and fix the decode softmax() cross-reference that pointed at the decode side instead of the prefill one. Comments only. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * [fmha] Drop the runtime soft-cap select; keep only the unreachable-kernel trim Selecting the soft-cap at runtime meant every FMHA path compiled as Softcap=true and picked between duplicated softmax loops on a uniform branch. The uncapped loop body was arithmetically unchanged, but moving it into a branch shifted codegen enough to break a tight-tolerance case: test_flash_attn_varlen_output[1024-1024-512-...-0.0-...-16-4-dtype1] max |out - out_ref| = 0.0301 vs a 0.0059 budget (softcap=0.0, d=512, causal, fp16); mean error was 2.19e-05, below PyTorch's own 3.51e-05. Only the most accumulation-heavy config crossed its budget -- the other 12935 cases, nearly all uncapped, passed -- so this is rounding, not a wrong result. Rather than chase codegen, restore the Softcap template axis so uncapped launches use bit-identical instantiations to main. What survives is the part that needs no numerical argument: the paged decode and prefill relative-attention branches can hardcode Softcap=false, because DecodeRunner/PrefillRunner already reject softcap with a rel bias, making those Softcap=true kernels unreachable. Those guards are tightened from `softcap > 0` to `softcap != 0` so the hardcode holds for any nonzero cap. Split decode keeps the dispatch: SplitDecodeConfig goes through SplitDecodeKernelRunner, which carries no such guard. The 50-minute rebuild timeout still matters -- the full matrix takes 39m58s, which cleared the old 40m budget by two seconds. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> * [ci] Raise the FMHA rebuild timeout to 60 minutes, and revert the kernel changes The soft-cap instantiation trim was measured on CI and saved 19 seconds of a 40-minute build (39m58s on main -> 39m39s), which is inside run-to-run variance. The relative-attention branch it touched only exists for 2 of the 6 paged head dims and is a leaf of the dispatch tree, so it removed 88 kernel instantiations out of ~700 translation units -- never enough to matter. A change that buys no build time is not worth the review cost of verifying that Softcap=false is unreachable in both runners, so all four kernel files are back to byte-identical with main. What is left is the only thing that actually unblocks CI: the rebuild step now gets 60 minutes instead of 40. The step recompiles the whole AOT matrix from scratch every run (it deliberately rm -rf's build so a SYCL PR cannot be silently linked against image-baked objects), so its wall time tracks the instantiation matrix, not the PR: 27.3m in August, 31.4m in mid-September, 40.0m once the soft-cap template axis landed. PR #501 built in 39m58s on plain main -- two seconds under the old budget. 60 minutes leaves room for the ~1m/week growth instead of needing another bump next month. --------- Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

Add chunking strategy for fp8_paged_mqa_logits

ba52423

polisettyvarma reviewed Aug 20, 2026

View reviewed changes

Merge branch 'main' into fp8_mqa_logits_oom

0cc15ee

Xia-Weiwen added the run-ci label Aug 25, 2026

trigger CI

539c1fc

Xia-Weiwen marked this pull request as ready for review August 25, 2026 05:08

Xia-Weiwen requested review from airMeng and mingfeima and a lite review from Copilot August 25, 2026 05:08

Copilot started reviewing on behalf of Xia-Weiwen August 25, 2026 05:09 View session

Copilot AI reviewed Aug 25, 2026

View reviewed changes

Comment thread src/sycl/Fp8MqaLogitsXe20.cpp Outdated

Comment thread src/sycl/Fp8MqaLogitsXe20.cpp Outdated

airMeng reviewed Aug 25, 2026

View reviewed changes

Refine code

1b59925

arathi-hlab mentioned this pull request Aug 25, 2026

[CI] bench_flash_attn: skip LSE for local/sink combos #405

Merged

2 tasks

Merge branch 'main' into fp8_mqa_logits_oom

e5da11c

Xia-Weiwen marked this pull request as draft August 26, 2026 01:27

Xia-Weiwen added 5 commits September 9, 2026 09:30

Merge branch 'main' into fp8_mqa_logits_oom

44695e5

Merge branch 'main' into fp8_mqa_logits_oom

6fa0ab3

Merge branch 'main' into fp8_mqa_logits_oom

0cfc9fe

Optimize performance of fp8_paged_mqa_logits

5914384

Merge branch 'main' into fp8_mqa_logits_oom

88f3d05

Xia-Weiwen changed the title ~~Add chunking strategy for fp8_paged_mqa_logits~~ Improve fp8_paged_mqa_logits Sep 21, 2026

Xia-Weiwen requested a lite review from Copilot September 21, 2026 07:33

Copilot started reviewing on behalf of Xia-Weiwen September 21, 2026 07:34 View session

Copilot AI reviewed Sep 21, 2026

View reviewed changes

Comment thread src/sycl/kernels/nsa/fp8_mqa_logits_kernel.hpp

Comment thread src/sycl/kernels/nsa/fp8_mqa_logits_kernel.hpp

Comment thread src/sycl/Fp8MqaLogitsXe20.cpp

Add a test case for multi-chunk scenario

a121b98

arathi-hlab mentioned this pull request Sep 21, 2026

[ci] Raise the FMHA rebuild step timeout from 40 to 60 minutes #502

Merged

Xia-Weiwen marked this pull request as ready for review September 21, 2026 14:00

airMeng requested changes Sep 22, 2026

View reviewed changes

Xia-Weiwen added 3 commits September 22, 2026 13:33

Refine UT

6df3e83

Merge branch 'main' into fp8_mqa_logits_oom

704ae22

Avoid magic number

0c0a697

Xia-Weiwen requested a review from airMeng September 22, 2026 05:52

airMeng approved these changes Sep 22, 2026

View reviewed changes

Xia-Weiwen added 2 commits September 23, 2026 14:33

Merge branch 'main' into fp8_mqa_logits_oom

ba70338

Merge branch 'main' into fp8_mqa_logits_oom

cfa22a3

Xia-Weiwen merged commit e872300 into main Sep 24, 2026
6 checks passed

Conversation

Xia-Weiwen commented Aug 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

polisettyvarma left a comment • 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!

Xia-Weiwen commented Aug 21, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

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!

Copilot AI left a comment

Choose a reason for hiding this comment

Copilot review overview

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

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!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Xia-Weiwen commented Aug 20, 2026 •

edited

Loading

polisettyvarma left a comment •

edited

Loading