Replace chunk prefill with AppendKV FMHA prefill by yuankuns · Pull Request #307 · sgl-project/sgl-kernel-xpu

yuankuns · 2026-07-17T04:35:01Z

Summary

Replace the old two-launch chunkprefill dispatcher with the FMHA prefill-with-KVcache AppendKV path.
Route mixed/prefill batches through prefill::mha_fwd; keep pure paged decode on decode::mha_fwd.
Wire existing Python k / v / cu_seqlens_k_new inputs through to native AppendKV params.
Enable AppendKV coverage in test_flash_attn_kvcache.
Keep FP8 e5m2 support on decode/split-decode; prefill/AppendKV remains unsupported for e5m2.
Pull in the latest hd128 mixed-batch optimization from fmha-cri: remove the post-AppendKV device fence and add small/underutilized hd128 paged prefill tile dispatch.

Validation

MAX_JOBS=4 CMAKE_BUILD_PARALLEL_LEVEL=4 python -m pip install -v . passed.
pytest -q tests/test_flash_attention.py::test_flash_attn_kvcache: 1200 passed, 1872 skipped.
pytest -q tests/test_flash_attention.py::test_flash_attn_decode_kvcache: 640 passed, 384 skipped.
python -m pre_commit run --all-files --show-diff-on-failure passed.

benchmark/bench_flash_attn.py Paged Prefill Sweep

Measured with benchmark/bench_flash_attn.py filtered to the paged prefill path.

Command ran in the sglang-syk container with ZE_AFFINITY_MASK=0 and ONEAPI_DEVICE_SELECTOR=level_zero:gpu.

Filtered configs:

head_dim={64,128}
q_seq_length=128, kv_seq_length=4096
paged KV: page_size=128
dtype=bf16, local=False, use_sinks=False
batch={1,8,16}, num_heads_q=16, num_heads_kv={4,8}, causal={True,False}

Baseline: 590d8d1 prefill kernel .so.
Current: this PR's rebuilt prefill kernel .so.

head_dim	configs	base avg ms	current avg ms	avg latency change	mean per-config speedup	best	worst
64	12	0.4024	0.3831	-4.8%	+4.3%	+7.4%	+1.2%
128	12	0.8880	0.8992	+1.3%	-0.5%	+1.5%	-2.5%

Notes:

hd64 improved across all measured configs.
hd128 is effectively flat: small configs slightly improve, larger configs are slightly slower.
benchmark/bench_flash_attn.py does not exercise AppendKV or sparse varlen mixed-query cases, so this sweep mainly measures the paged prefill/page-table path.

Benchmark Notes

The main baseline is the complete dfa5630 (origin/main) build. Its timed callable consumes new_k/new_v: main's fused store_cache_xpu writes both tensors to their physical paged-cache slots in one launch, then the old chunkprefill path runs decode::mha_fwd plus prefill::mha_fwd. The baseline therefore measures one cache-store launch plus the old two attention launches.

Main's flash_attn_with_kvcache Python signature accepts k/v, but the main native call does not consume them. Validation confirmed that passing k/v alone leaves the cache unchanged; the composite baseline cache and attention output match explicit reference cache writes.

AppendKV timing is measured on the current PR build through torch.ops.sgl_kernel.fwd_appendkv with old KV cache plus k_new/v_new/cu_seqlens_k_new. AppendKV ms and speedup ratio are filled for head_dim={64,96,128,192,256}.

Benchmarks use bf16 KV, heads q/kv = 16/4, ZE_AFFINITY_MASK=0, ONEAPI_DEVICE_SELECTOR=level_zero:gpu, and triton.testing.do_bench(warmup=100, rep=500, return_mode="mean"). Each table cell is the median of three full sweeps. Speedup ratio is main store + double-launch ms / AppendKV ms, shown as a percentage.

AppendKV Size Sweep

Cache old length is 4096, page size is 64, causal/local are both false.

head_dim	batch	new_k/v per row	total new_k/v	new_k/v shape	main store + double-launch ms	AppendKV ms	speedup ratio
64	8	16	128	`[8,16,4,64]`	0.6378	0.2684	237.6%
64	8	64	512	`[8,64,4,64]`	0.3395	0.3213	105.7%
64	8	128	1024	`[8,128,4,64]`	0.3631	0.3606	100.7%
64	8	256	2048	`[8,256,4,64]`	0.6651	0.6713	99.1%
64	16	16	256	`[16,16,4,64]`	0.6369	0.3923	162.3%
64	16	64	1024	`[16,64,4,64]`	0.6343	0.5418	117.1%
64	16	128	2048	`[16,128,4,64]`	0.6603	0.6600	100.0%
64	16	256	4096	`[16,256,4,64]`	1.1625	1.1795	98.6%
96	8	16	128	`[8,16,4,96]`	0.6424	0.6928	92.7%
96	8	64	512	`[8,64,4,96]`	0.6895	0.6640	103.8%
96	8	128	1024	`[8,128,4,96]`	0.7519	0.6938	108.4%
96	8	256	2048	`[8,256,4,96]`	1.3274	1.2324	107.7%
96	16	16	256	`[16,16,4,96]`	1.1030	1.2131	90.9%
96	16	64	1024	`[16,64,4,96]`	1.1943	1.1705	102.0%
96	16	128	2048	`[16,128,4,96]`	1.3155	1.2227	107.6%
96	16	256	4096	`[16,256,4,96]`	2.4015	2.3559	101.9%
128	8	16	128	`[8,16,4,128]`	1.4784	0.5075	291.3%
128	8	64	512	`[8,64,4,128]`	1.4368	0.6469	222.1%
128	8	128	1024	`[8,128,4,128]`	1.4033	0.9119	153.9%
128	8	256	2048	`[8,256,4,128]`	1.4271	1.3852	103.0%
128	16	16	256	`[16,16,4,128]`	2.6168	0.8151	321.0%
128	16	64	1024	`[16,64,4,128]`	2.5354	1.0319	245.7%
128	16	128	2048	`[16,128,4,128]`	2.4691	1.5883	155.5%
128	16	256	4096	`[16,256,4,128]`	2.5156	2.4265	103.7%
192	8	16	128	`[8,16,4,192]`	2.5809	1.4430	178.9%
192	8	64	512	`[8,64,4,192]`	2.6133	1.5128	172.7%
192	8	128	1024	`[8,128,4,192]`	2.7192	1.6169	168.2%
192	8	256	2048	`[8,256,4,192]`	2.9481	2.8980	101.7%
192	16	16	256	`[16,16,4,192]`	4.7736	2.5229	189.2%
192	16	64	1024	`[16,64,4,192]`	4.8360	2.6651	181.5%
192	16	128	2048	`[16,128,4,192]`	5.0535	2.8582	176.8%
192	16	256	4096	`[16,256,4,192]`	5.5113	5.4269	101.6%
256	8	16	128	`[8,16,4,256]`	3.0493	1.8827	162.0%
256	8	64	512	`[8,64,4,256]`	3.0227	1.8817	160.6%
256	8	128	1024	`[8,128,4,256]`	3.0448	1.9531	155.9%
256	8	256	2048	`[8,256,4,256]`	3.1181	3.1183	100.0%
256	16	16	256	`[16,16,4,256]`	5.6641	3.2809	172.6%
256	16	64	1024	`[16,64,4,256]`	5.6140	3.3701	166.6%
256	16	128	2048	`[16,128,4,256]`	5.6500	3.5085	161.0%
256	16	256	4096	`[16,256,4,256]`	5.8156	5.8147	100.0%

Mixed Varlen Chunk-Prefill

The first row mirrors the standalone window-3 style shape. The remaining mixed rows cover max new_k/v <= 256 for decode-heavy, balanced, and prefill-heavy mixes.

Current mixed-table AppendKV values were refreshed at a04b0b6.

head_dim	pattern	rows decode/prefill	q_lens	old_lens	total new_k/v	max new_k/v	new_k/v shape	page	mask	main store + double-launch ms	AppendKV ms	speedup ratio
64	window3_q1_q512	1/1	`1/512`	`8191/7680`	513	512	`[513,4,64]`	128	causal	0.3973	0.3702	107.3%
64	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,64]`	64	none	0.4491	0.3252	138.1%
64	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,64]`	64	causal	0.4496	0.3689	121.9%
64	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,64]`	64	none	0.4408	0.4254	103.6%
64	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,64]`	64	causal	0.4395	0.4152	105.8%
64	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,64]`	64	local	0.3732	0.3053	122.2%
64	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,64]`	64	none	0.5864	0.5980	98.1%
64	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,64]`	64	causal	0.5843	0.6071	96.2%
96	window3_q1_q512	1/1	`1/512`	`8191/7680`	513	512	`[513,4,96]`	128	causal	0.7777	0.7952	97.8%
96	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,96]`	64	none	0.8322	0.8132	102.3%
96	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,96]`	64	causal	0.7925	0.8290	95.6%
96	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,96]`	64	none	0.8231	0.8282	99.4%
96	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,96]`	64	causal	0.7867	0.8375	93.9%
96	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,96]`	64	local	0.2944	0.2625	112.1%
96	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,96]`	64	none	1.1185	1.1098	100.8%
96	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,96]`	64	causal	1.0580	1.1190	94.5%
128	window3_q1_q512	1/1	`1/512`	`8191/7680`	513	512	`[513,4,128]`	128	causal	0.7349	0.7532	97.6%
128	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,128]`	64	none	1.3376	1.0052	133.1%
128	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,128]`	64	causal	1.3077	0.9893	132.2%
128	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,128]`	64	none	1.6380	1.1030	148.5%
128	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,128]`	64	causal	1.5813	1.0941	144.5%
128	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,128]`	64	local	0.2656	0.1891	140.5%
128	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,128]`	64	none	2.2697	1.4583	155.6%
128	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,128]`	64	causal	2.2178	1.4433	153.7%
192	window3_q1_q512	1/1	`1/512`	`8191/7680`	513	512	`[513,4,192]`	128	causal	1.7540	1.7988	97.5%
192	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,192]`	64	none	2.3294	1.7715	131.5%
192	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,192]`	64	causal	2.3111	1.7127	134.9%
192	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,192]`	64	none	2.9521	2.0641	143.0%
192	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,192]`	64	causal	2.9372	1.9969	147.1%
192	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,192]`	64	local	0.2110	0.2196	96.1%
192	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,192]`	64	none	4.3535	2.9114	149.5%
192	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,192]`	64	causal	4.3125	2.7751	155.4%
256	window3_q1_q512	1/1	`1/512`	`8191/7680`	513	512	`[513,4,256]`	128	causal	1.9695	1.9794	99.5%
256	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,256]`	64	none	2.7114	2.1799	124.4%
256	decode_heavy	10/6	`1/1/1/1/1/1/1/1/16/32/64/128/1/1/8/256`	`4096 x16`	514	256	`[514,4,256]`	64	causal	2.7019	2.1654	124.8%
256	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,256]`	64	none	3.4770	2.5767	134.9%
256	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,256]`	64	causal	3.4419	2.5562	134.7%
256	balanced	8/8	`1/16/1/32/1/64/1/128/1/256/8/1/64/1/32/1`	`4096 x16`	608	256	`[608,4,256]`	64	local	0.2153	0.2047	105.2%
256	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,256]`	64	none	5.0306	3.5267	142.6%
256	prefill_heavy	3/13	`64/128/256/128/64/32/16/8/1/1/64/128/256/32/1/16`	`4096 x16`	1195	256	`[1195,4,256]`	64	causal	4.9867	3.4919	142.8%

Copilot

Pull request overview

This PR replaces the prior two-launch “chunkprefill” mixed-batch path with a single prefill kernel path that supports in-kernel KV append (AppendKV), routing mixed/prefill batches through prefill::mha_fwd while keeping pure paged decode on decode::mha_fwd. It also wires optional k_new/v_new/cu_seqlens_k_new through the Python → Torch op → native SYCL layers and expands tests/benchmarks accordingly.

Changes:

Extend the public op/API surface to accept optional AppendKV inputs (k_new, v_new, cu_seqlens_k_new) and plumb them into the prefill mainloop.
Update kernel mainloop/epilogue to support AppendKV prefill (and remove/limit unsupported fp8 prefill variants as described).
Enable AppendKV coverage in test_flash_attn_kvcache and align benchmark config with supported fp8 paths.

Reviewed changes

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

Show a summary per file

File	Description
`tests/test_flash_attention.py`	Enables AppendKV test coverage and adds guardrails for cache capacity when appending.
`src/torch_extension_sycl.cc`	Extends Torch library schema to accept optional `k_new/v_new/cu_seqlens_k_new`.
`src/sycl/xe_fmha_fwd_prefill_kernel.cpp.in`	Adds hd128 mixed/small paged prefill tile dispatch and removes legacy fp8 prefill handling.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_runner.hpp`	Adds AppendKV argument wiring into the prefill kernel mainloop args.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_prefill_dispatch.hpp`	Fully-qualifies the dispatch macro to avoid namespace ambiguity.
`src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_decode_runner.hpp`	Updates fp8 scale pointer types and split-decode mainloop arg wiring (non-split needs a fix).
`src/sycl/kernels/flash_attention_v2/kernel/xe_fhma_fwd_kernel.hpp`	Implements effective-K length handling and append-store length logic for AppendKV.
`src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp`	Introduces AppendKV params/storage and refactors decode fp8 scaling + masking/prefetch logic.
`src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_epilogue.hpp`	Refactors sink handling and softmax normalization paths (and removes fp8 scale folding in epilogue).
`src/sycl/flash_attention.cpp`	Removes chunkprefill dispatcher, routes mixed batches to prefill, and wires AppendKV inputs/validation.
`python/sgl_kernel/flash_attn.py`	Forwards `k/v/cu_seqlens_k_new` to the native op and computes an appropriate `max_seqlen_k` for native.
`include/sgl_flash_kernel_ops.h`	Updates C++ interface for `mha_fwd` to accept AppendKV inputs.
`benchmark/bench_flash_attn.py`	Narrows fp8 benchmarking to supported decode path and adjusts local window sampling.

Comments suppressed due to low confidence (1)

src/sycl/kernels/flash_attention_v2/xe_fmha_fwd_decode_runner.hpp:331

Non-split decode runner is not forwarding fp8 KV descale pointers (k_scale_ptr/v_scale_ptr) and total KV length into the decode mainloop arguments. After the DecodeFwdMainloop argument layout change, this will either mis-initialize the mainloop args (wrong field order) or drop fp8 scaling entirely.

        {params.softmax_scale,
         params.page_table,
         params.page_size,
         params.max_num_pages_per_seq,
         params.window_size_left,
         params.window_size_right},

+    int total_knew = 0;
+    int seqlen_knew = max_seqlen_k > 0 ? max_seqlen_k : max_seqlen_q;
+    if (k_new.dim() == 3) {
+      total_knew = k_new.size(0);
+      CHECK_SHAPE(k_new, total_knew, num_heads_k, head_size);
+      CHECK_SHAPE(v_new, total_knew, num_heads_k, head_size_v);
+      TORCH_CHECK(
+          cu_seqlens_k_new_.has_value() || seqlen_knew > 0,
+          "ragged k_new requires cu_seqlens_k_new or positive max_seqlen_k");
+    } else {
+      TORCH_CHECK(k_new.size(0) == batch_size, "batched k_new first dimension must match batch size");
+      int const k_new_seqlen = k_new.size(1);
+      total_knew = batch_size * k_new_seqlen;
+      seqlen_knew = max_seqlen_k > 0 ? max_seqlen_k : k_new_seqlen;
+      CHECK_SHAPE(k_new, batch_size, k_new_seqlen, num_heads_k, head_size);
+      CHECK_SHAPE(v_new, batch_size, k_new_seqlen, num_heads_k, head_size_v);
+    }


  params.o_row_stride = out.stride(-3);
  params.o_head_stride = out.stride(-2);



Valentine233 · 2026-07-17T08:40:17Z

Hi @yuankuns, I suppose this PR mainly targets to support the new kv cache in addition to the current one. Firstly, could you please answer several questions:

What is the scenario or model to use new kv cache?
Why is the chunk prefill launcher replaced?
Why is FP8 e5m2 removed?

yuankuns · 2026-07-17T19:17:32Z

@Valentine233 This kernel aims to replace chunk prefill, to merge 3 kernel launch (appendkv, prefill, decode) into 1 kernel launch (chunk-prefill) and according to the measurement, it's faster. I only tuned hddim=128 cases, will add more headdim tuning. If you believe there are some regression in performance, please also provide the shape so that I can tune these shapes.
The fp8 removal is by accidental, already recovered.

Valentine233 · 2026-07-20T08:23:35Z

Thanks, @yuankuns .

Could you elaborate on the scenario or model which requires new kv cache?
For benchmark, please compare the perf without new kv cache. Run https://github.com/sgl-project/sgl-kernel-xpu/blob/main/benchmark/bench_flash_attn.py, and provide the result before/after this PR.

airMeng · 2026-07-22T16:33:54Z

Thanks, @yuankuns .

Could you elaborate on the scenario or model which requires new kv cache?

@Valentine233 I think Yuankun means fusing storing kv cache into attention. In current intel_xpu backend, storing kv cache is a saparete kernel

yuankuns · 2026-07-22T20:03:46Z

@Valentine233 @airMeng This pr also merges 2 fmha kernel launch into 1 kernel launch in chunk prefill cases. The master branch launchs 1 prefill + 1 decode for chunk prefill. The script benchmark/bench_flash_attn.py didn't include chunk prefill cases, either.

Valentine233 · 2026-07-23T02:08:47Z

With benchmark/bench_flash_attn.py, we are able to compare the perf of prefill/decode with/without paged cache. It doesn't seem the result in the description complete for all these configs. For the varlen chunk-prefill scenario, there is still performance drop via description, for example window-3.
We do see the perf gain to fuse appendKV into attention, but we still need to validate the affect on the cases without appendKV.

sunjiweiswift · 2026-07-27T05:27:30Z

sycl-tla example06 supports both kvcache and kvnew, you can refer to that for reference.
The code for kv_new and kv_cache should maintain a consistent style as much as possible.

sgl-kernel was derived by stripping down code that originally supported the coexistence of both kvnew and kvcache; the kvnew code can be restored.

Change the existing kv in the current code to kv_cache, and then add kv_new. The code for kv_new and kv_cache should follow a consistent coding style.

sunjiweiswift · 2026-07-27T05:29:19Z

-          constexpr int elems_per_n = tSrS.size() / n_reps;
-          int k_base = K * kTileK;
+          // Need to get global col and row indices to mask the elements
+          Tensor cPgP = make_identity_tensor(make_shape(seq_len, seq_len));


I’ve tested this implementation before; TFLOPS drop by 3–5 T. Although this modification looks better
@Valentine233 I remember it being like this. right？

Yes, please keep the original one, referring to #260.

I have validate the pr260 change, no significant performance difference in those 2 style. <1%. Which shape you've measured results in 3-5Tops dropping?

Have you tried the shape in pr260's description? What is the tflops now?

sunjiweiswift · 2026-07-27T05:32:48Z

+  params.knew_ptr = nullptr;
+  params.vnew_ptr = nullptr;
+  params.cu_seqlens_knew = nullptr;
+  params.seqlen_knew = 0;


Change the original total in the code to seqlen_kvcache.

seqlen_knew is max new_kv len in single request, and seqlen_kvcache is the maximum size of current page table.

sunjiweiswift · 2026-07-27T05:33:55Z

-    std::optional<at::Tensor> skip_batch_mask_opt = std::nullopt) {
+    std::optional<at::Tensor> skip_batch_mask_opt = std::nullopt,
+    std::optional<const at::Tensor> k_new_ = std::nullopt,
+    std::optional<const at::Tensor> v_new_ = std::nullopt,


It is adjacent to kvche and is a mandatory input parameter.

sunjiweiswift · 2026-07-27T05:36:44Z

+    int const sm_margin,
+    std::optional<at::Tensor> out_opt = std::nullopt,
+    std::optional<at::Tensor> skip_batch_mask_opt = std::nullopt) {
+  return mha_fwd_appendkv(


Why make this modification? The appendkv logic is only entered when kvnew exists. If you need to modify the prefill or decode phases, please make those changes within the decode or prefill functions themselves.

sunjiweiswift · 2026-07-27T05:42:36Z

@Valentine233 This kernel aims to replace chunk prefill, to merge 3 kernel launch (appendkv, prefill, decode) into 1 kernel launch (chunk-prefill) and according to the measurement, it's faster. I only tuned hddim=128 cases, will add more headdim tuning. If you believe there are some regression in performance, please also provide the shape so that I can tune these shapes. The fp8 removal is by accidental, already recovered.

mha_fwd() {
if (no_page)
   no_page(decode/prefill)
elseif (page)
   if (kv_new) 
        decode/prefill
    if (!kv_new)
       decode/prefill

}

or

mha_fwd() {
  if decode {
     no_page/page/page_kv_new
  } else prefill {
     no_page/page/page_kv_new
  } else chunkprefill{
     no_page/page/page_kv_new
  }
}

The unified interface should be named mha_fwd, with conditional logic based on whether it is a page or kvnew.

sunjiweiswift · 2026-07-27T05:50:46Z

+      "    int      sm_margin,"
+      "    Tensor(a!)?  out=None,"
+      "    Tensor?  k_new=None,"
+      "    Tensor?  v_new=None,"


up to line 180

Port the q128/k64/sg16 paged tile configuration from cutlass-sycl commit 9bf71715.

yuankuns · 2026-07-28T17:36:07Z

xe_fmha_fwd_mainloop.hpp

The appendKV in example06 is slower. So not going to copy the code directly. Already changed the coding style

yuankuns · 2026-07-29T02:14:49Z

@Valentine233 This kernel aims to replace chunk prefill, to merge 3 kernel launch (appendkv, prefill, decode) into 1 kernel launch (chunk-prefill) and according to the measurement, it's faster. I only tuned hddim=128 cases, will add more headdim tuning. If you believe there are some regression in performance, please also provide the shape so that I can tune these shapes. The fp8 removal is by accidental, already recovered.
mha_fwd() {
if (no_page)
   no_page(decode/prefill)
elseif (page)
   if (kv_new) 
        decode/prefill
    if (!kv_new)
       decode/prefill

}
or
mha_fwd() {
  if decode {
     no_page/page/page_kv_new
  } else prefill {
     no_page/page/page_kv_new
  } else chunkprefill{
     no_page/page/page_kv_new
  }
}
The unified interface should be named mha_fwd, with conditional logic based on whether it is a page or kvnew.

change

[2d09c4a](/sgl-project/sgl-kernel-xpu/pull/307/commits/2d09c4af6f47d5adf182dd76ec249844a9787954)

There are many code blocks are shared between prefill/chunkprefill. If go this branch style, there will be more dupilicated code block

jmunetong · 2026-07-31T22:00:07Z

I am providing results for Gemma-4-31b-it. The below table simulates the dump shapes of the model in the SWA layers where head_dim=256.

Sliding family — nkv=4, hd=256, window (1023,0):

batch	new/row	baseline ms	PR307 ms	speedup
1	16	0.1329	0.1130	1.18×
1	64	0.1328	0.1003	1.32×
1	128	0.1402	0.1112	1.26×
1	256	0.1612	0.1208	1.33×
1	512	0.1678	0.2145	0.78×
8	16	0.5131	0.3168	1.62×
8	64	0.5128	0.3125	1.64×
8	128	0.5460	0.3648	1.50×
8	256	0.6388	0.6524	0.98×
8	512	1.0769	1.2123	0.89×
16	16	0.8824	0.5994	1.47×
16	64	0.8814	0.5743	1.54×
16	128	0.9389	0.6590	1.42×
16	256	1.0943	1.1369	0.96×
16	512	1.9793	2.2595	0.88×

with head_dim=512, there are no improvements but #342 will be taking care of this:

Full-attn family — nkv=1, hd=512, full window:

batch	new/row	baseline ms	PR307 ms	speedup
1	16	0.7150	0.7464	0.96×
1	64	0.7099	0.7291	0.97×
1	128	0.7512	0.7807	0.96×
1	256	0.7991	0.8537	0.94×
1	512	1.5230	1.7505	0.87×
8	16	4.7870	5.0250	0.95×
8	64	4.8230	4.8589	0.99×
8	128	4.9941	5.1526	0.97×
8	256	5.0158	5.3635	0.94×
8	512	9.2610	10.5985	0.87×
16	16	8.9409	9.3249	0.96×
16	64	9.0193	9.0999	0.99×
16	128	9.3618	9.6676	0.97×
16	256	9.5051	10.1424	0.94×
16	512	18.3151	20.9391	0.87×

Copilot AI review requested due to automatic review settings July 17, 2026 04:35

yuankuns mentioned this pull request Jul 17, 2026

Replace chunk prefill with AppendKV FMHA prefill #306

Closed

Copilot started reviewing on behalf of yuankuns July 17, 2026 04:35 View session

yuankuns added the run-ci label Jul 17, 2026

yuankuns requested a review from airMeng July 17, 2026 04:38

Copilot AI reviewed Jul 17, 2026

View reviewed changes

airMeng requested review from Valentine233 and sunjiweiswift July 17, 2026 05:00

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch from 45df940 to ca1eb26 Compare July 17, 2026 16:25

yuankuns requested a review from kareemshaik80 as a code owner July 17, 2026 16:25

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch 2 times, most recently from c91f681 to 365294c Compare July 17, 2026 19:17

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch from cf2c45f to 590d8d1 Compare July 22, 2026 22:28

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch from fb64b73 to 4129ac7 Compare July 27, 2026 02:44

sunjiweiswift reviewed Jul 27, 2026

View reviewed changes

yuankuns added 6 commits July 28, 2026 10:29

Replace chunk prefill with AppendKV FMHA prefill

06cbec0

Tune hd64 FMHA prefill mixed batches

8b0f5cc

Tune hd64 sparse mixed KV tile

e311036

Drop unstable hd64 balanced mixed tile path

f107c1e

Tune hd64 sparse mixed subgroup layout

aa07ff7

Tune hd256 FMHA prefill tile

574851a

Port the q128/k64/sg16 paged tile configuration from cutlass-sycl commit 9bf71715.

yuankuns added 4 commits July 28, 2026 10:30

Tune hd192 FMHA prefill tile

d41ee92

Tune hd96 FMHA prefill tiles

4aa21ca

Optimize q256 AppendKV prefill

49183d2

Optimize q256 AppendKV cache access

2d09c4a

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch 2 times, most recently from 2554979 to 2d09c4a Compare July 28, 2026 23:32

sunjiweiswift reviewed Jul 29, 2026

View reviewed changes

Comment thread src/sycl/kernels/flash_attention_v2/collective/xe_fmha_fwd_mainloop.hpp Outdated

yuankuns added 2 commits July 28, 2026 21:51

Add FMHA prefill AppendKV wiring

b6c23bd

Fix FMHA prefill AppendKV wiring

e23eec7

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch from 391f541 to 06961ea Compare July 30, 2026 00:51

Consolidate prefill appendkv modes

0d94217

yuankuns force-pushed the fmha-prefill-kvcache-appendkv branch from 06961ea to 0d94217 Compare July 30, 2026 03:01

yuankuns added 2 commits July 30, 2026 10:47

Add wide AppendKV enum mode

1c4b521

Restore AppendKV mode dispatch

b1fe867

yuankuns requested a review from sunjiweiswift August 3, 2026 16:45

jmunetong mentioned this pull request Aug 3, 2026

[XPU] Enable AppendKV FMHA prefill in intel_xpu attention backend sgl-project/sglang#33430

Draft

Use a standalone prepass for large AppendKV prefill

67e48f6

		params.o_row_stride = out.stride(-3);
		params.o_head_stride = out.stride(-2);

Uh oh!

Conversation

yuankuns commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

benchmark/bench_flash_attn.py Paged Prefill Sweep

Benchmark Notes

AppendKV Size Sweep

Mixed Varlen Chunk-Prefill

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Valentine233 commented Jul 17, 2026

Uh oh!

yuankuns commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Valentine233 commented Jul 20, 2026

Uh oh!

airMeng commented Jul 22, 2026

Uh oh!

yuankuns commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

sunjiweiswift commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

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!

sunjiweiswift commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuankuns commented Jul 28, 2026

Uh oh!

yuankuns commented Jul 29, 2026

Uh oh!

Uh oh!

jmunetong commented Jul 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

yuankuns commented Jul 17, 2026 •

edited

Loading

yuankuns commented Jul 17, 2026 •

edited

Loading

yuankuns commented Jul 22, 2026 •

edited

Loading

Valentine233 commented Jul 23, 2026 •

edited

Loading

sunjiweiswift commented Jul 27, 2026 •

edited

Loading

sunjiweiswift commented Jul 27, 2026 •

edited

Loading