[DSA] Support flashmla_sparse_q8 prefill backend on CUDA by yzk1989 · Pull Request #44217 · sgl-project/sglang

yzk1989 · 2026-10-14T06:27:34Z

Motivation

Enable flashmla_sparse_q8 DSA backend for prefill on CUDA, import performents for models such as GLM-5.3-Flash.

Modifications

Enable flashmla_sparse_q8 prefill with TileLang decode for GLM-5.3-Flash’s raw FP8 KV layout on SM90, restricted to NoPE models with kv_lora_rank=512.
Allow Q8 prefill to consume KPool top-k indices, including appended live-tail tokens.
Route raw FP8 paged KV directly to the Q8 kernel using physical KV-pool indices.
Pad top-k indices to a multiple of 128 with -1 sentinels, preserving all valid tail tokens.
Handle zero-dimensional RoPE queries without invoking the concat kernel.
Preserve the raw 512-byte KV row layout shared by Q8 prefill and TileLang decode.

Depends on CUDA TileLang FP8 KV support from #42957.

Accuracy Tests

Validated GLM-5.3-Flash on 4 × H200 with flashmla_sparse_q8 and TileLang prefill .

sglang serve \
  --model-path /path/to/GLM-5.3-Flash \
  --served-model-name llm \
  --tp-size 4 \
  --mem-fraction-static 0.75 \
  --dsa-prefill-backend flashmla_sparse_q8 \
  --dsa-decode-backend tilelang \
  --kv-cache-dtype fp8_e4m3 \
  --moe-runner-backend triton \
  --ep-size 4 \
  --speculative-algorithm EAGLE \
  --speculative-num-steps 5 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 6 \
  --host 0.0.0.0 \
  --port 8009 \
  --cuda-graph-backend-prefill breakable \
  --reasoning-parser auto \
  --tool-call-parser auto

GSM8K: 1,314 scored questions per run, 5-shot, temperature 0

dsa-prefill-backend	Run 1 (rounded)	Run 2 (rounded)	Mean
tilelang	95.3%	95.7%	95.5%
flashmla_sparse_q8	95.7%	95.5%	95.6%

Speed Tests and Profiling

4 × H200, TP4/EP4, TileLang DSA, triton MoE, EAGLE 5/1/6 and Prefill BCG,kv-cache-dtype fp8_e4m3. Only --dsa-prefill-backend changes between cases.

dsa-prefill-backend	TTFT (ms)	TPOT (ms)	Output throughput (tok/s)
tilelang	2437.68	24.68	859.40
flashmla_sparse_q8	1922.94	23.40	957.61

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #40813492287
Latest PR Test (Extra): ❌ Run #40813492239
Latest PR Test (AMD ROCm 10): ❌ Run #40813492251

yzk1989 · 2026-10-16T11:54:15Z

@Ying1123 @merrymercy @hnyls2002 Could you please take a look at this PR when you have a chance? Thanks!

mmangkad

I think this needs a rebase onto #41790. After that the fix might be pretty small: the gather kernel already copies any dtype, so maybe just relax the bf16 assert in gather_cast_kv_fp8_pad_paged for NoPE, and let _validate_flashmla_sparse_q8_backend accept the unpacked fp8 pool when rope=0 and decode is tilelang.
Maybe update the GLM-5.3-Flash cookbook too, since it still says FP8 KV doesn't work with q8 prefill.
Could you rerun GSM8K with sgl-eval, with prefix caching on?

mmangkad · 2026-10-16T12:38:47Z

+            if (
+                dsa_impl == "flashmla_sparse_q8"
+                and topk_transform_method == TopkTransformMethod.PAGED
+                and not self.dsa_kv_cache_store_fp8
+                and kv_cache.shape[-1] == 512
+            ):


I think after #41790 this never runs, since q8 prefill is always RAGGED now, and the launch command actually fails at startup in _validate_flashmla_sparse_q8_backend. Maybe move this into the RAGGED path?

I think after #41790 this never runs, since q8 prefill is always RAGGED now, and the launch command actually fails at startup in _validate_flashmla_sparse_q8_backend. Maybe move this into the RAGGED path?

Done. Removed the PAGED fallback. Raw FP8 Q8 prefill now reuses the RAGGED path introduced in #41790.

Updated _validate_flashmla_sparse_q8_backend to accept only the supported raw FP8 configuration: NoPE, kv_lora_rank=512, TileLang decode, and SM90.

mmangkad · 2026-10-16T12:39:21Z

+                return self._forward_flashmla_sparse_q8kv8(
+                    q_nope=q_nope,
+                    q_rope=q_rope,
+                    kv_bf16=kv_cache.view(-1, 1, 512),


I guess this copies the whole layer pool into _q8kv8_kv_buf on every prefill layer? Maybe gather just the needed slots via page_table_1_flattened?

I guess this copies the whole layer pool into _q8kv8_kv_buf on every prefill layer? Maybe gather just the needed slots via page_table_1_flattened?

Done. Removed the whole-pool copy.

Prefix requests now use page_table_1_flattened to gather only the referenced context slots through gather_cast_kv_fp8_pad_paged. For raw FP8 NoPE rows, this copies the FP8 values without scale conversion.

mmangkad · 2026-10-16T12:40:21Z

        # concat + .to(fp8)); it requires power-of-two head/dim counts (a
        # tl.arange constraint), so non-power-of-two head counts fall back to
        # the generic concat + cast.
+        elif q_rope.shape[-1] == 0:


I think concat_and_cast_q_fp8_pad already handles ROPE=0 on main, so maybe this isn't needed?

I think concat_and_cast_q_fp8_pad already handles ROPE=0 on main, so maybe this isn't needed?

Done. Removed the separate NoPE query-preparation branch and reused concat_and_cast_q_fp8_pad.

The ROPE=0 and ROPE>0 GPU regression tests both pass.

mmangkad · 2026-10-16T12:40:41Z

            topk_indices is None
            or self.dsa_index_kpool <= 1
            or dsa_impl in ("fa3", "tilelang", "trtllm")
+            or (dsa_impl == "flashmla_sparse_q8" and phase == "prefill")


I think this opens up packed fp8 too, which main keeps closed on purpose, and it breaks test_kpool_tail_guard_still_rejects_q8_over_a_packed_fp8_pool. Maybe drop it?

I think this opens up packed fp8 too, which main keeps closed on purpose, and it breaks test_kpool_tail_guard_still_rejects_q8_over_a_packed_fp8_pool. Maybe drop it?

Done. Removed the unconditional Q8 prefill allowance and restored the packed FP8 KPool guard.

test_kpool_tail_guard_still_rejects_q8_over_a_packed_fp8_pool now passes.

mmangkad · 2026-10-16T12:41:02Z

+        raw_fp8_nope=(
+            model_arch == "Glm5NextForConditionalGeneration"
+            and model_config_of(view).qk_rope_head_dim == 0
+            and model_config_of(view).kv_lora_rank == 512
+        ),


Maybe check the layout instead of the model name, and share one helper with calculate_mla_kv_cache_dim?

Maybe check the layout instead of the model name, and share one helper with calculate_mla_kv_cache_dim?

Done. Added is_q8_raw_fp8_nope_layout as a shared predicate for argument validation, backend validation, and KV pool allocation.

It checks the backend pairing and layout geometry rather than the model architecture name.

yzk1989 · 2026-10-19T06:03:42Z

I think this needs a rebase onto [DSA] Q8KV8 FP8 Sparse Prefill on GLM-5.3-Flash with BF16 KV Cache #41790. After that the fix might be pretty small: the gather kernel already copies any dtype, so maybe just relax the bf16 assert in gather_cast_kv_fp8_pad_paged for NoPE, and let _validate_flashmla_sparse_q8_backend accept the unpacked fp8 pool when rope=0 and decode is tilelang.

Maybe update the GLM-5.3-Flash cookbook too, since it still says FP8 KV doesn't work with q8 prefill.

Could you rerun GSM8K with sgl-eval, with prefix caching on?

Thanks for the detailed review! I’ve addressed the comments and rebased this branch onto main, including #41790.

Changes:

Added a shared layout predicate for argument validation, backend validation, and KV pool allocation, checking backend pairing and geometry rather than model architecture.
Restricted raw FP8 support to 512-dim NoPE KV with Q8 prefill and TileLang decode on SM90.
Removed the PAGED fallback and whole-pool copy. Prefix requests now reuse the RAGGED path and gather the referenced context slots via page_table_1_flattened.
Removed the redundant top-k padding and NoPE query-preparation branches, reusing the implementations from [DSA] Q8KV8 FP8 Sparse Prefill on GLM-5.3-Flash with BF16 KV Cache #41790.
Restored the packed FP8 KPool guard.
Updated the GLM-5.3-Flash cookbook and added regression tests.

Validation:

49 CPU tests and 29 GPU tests passed, including prefix gathering, live KPool tail tokens, and unsupported-layout rejection.
The 8 Q8 helper cases also passed after rebuilding with an isolated JIT cache.
isort and Ruff checks passed.

I also reran GSM8K using sgl-eval with prefix caching enabled on the server:

1,319 examples × 2 repeats
Concurrency: 8
temperature=0, top_p=1, max_tokens=4096, reasoning_effort=low
pass@1 (averaged over two repeats): 92.53%
pass@2: 92.57%
No-answer, truncation, and request error rates: all 0%

The average input length was 94 tokens, so this GSM8K run is an end-to-end accuracy sanity check rather than evidence of long-prefix Q8 sparse-prefill coverage. The GPU regression tests separately cover prefix gathering and numerical agreement with the reference.

Could you please take another look when you have time? Thanks again for your feedback!

yzk1989 · 2026-10-19T06:06:11Z

@mmangkad could you help me take a look at this PR again?Thanks!

yzk1989 · 2026-10-19T06:06:22Z

@mmangkad could you help me take a look at this PR again?Thanks!

mmangkad · 2026-10-19T13:23:38Z

+# Generated by ACS2 SKILL - Ant International Security
+"""Shared layout checks for DSA backend validation and KV pool allocation."""
+
+from typing import Optional
+
+
+def is_q8_raw_fp8_nope_layout(


Move this helper into configs/model_config.py next to is_deepseek_dsa. A whole new file for one small function is a bit much

Move this helper into configs/model_config.py next to is_deepseek_dsa. A whole new file for one small function is a bit much

Done. Moved is_q8_raw_fp8_nope_layout next to is_deepseek_dsa in configs/model_config.py and removed kv_layout.py. Updated all imports; overrides.py uses a function-local import to avoid a circular dependency. Both import orders were checked successfully.

mmangkad · 2026-10-19T13:24:25Z

+        kv_lora_rank=getattr(model_config_of(view), "kv_lora_rank", None),
+        qk_rope_head_dim=getattr(model_config_of(view), "qk_rope_head_dim", None),


Do not use getattr. I think DSA models always have these, so maybe just access them directly?

Do not use getattr. I think DSA models always have these, so maybe just access them directly?

Done. Replaced getattr with direct access to kv_lora_rank and qk_rope_head_dim in both the normal backend-resolution path and the HiSparse path.

mmangkad · 2026-10-19T13:24:59Z

+    kv_lora_rank: Optional[int] = None,
+    qk_rope_head_dim: Optional[int] = None,


Same here, make these plain required ints (and in the validator at dsa_backend.py L423)

Same here, make these plain required ints (and in the validator at dsa_backend.py L423)

Done. Made kv_lora_rank and qk_rope_head_dim required int parameters in the shared layout predicate and backend validators. Updated the call sites and test mocks to provide the geometry explicitly.

mmangkad · 2026-10-19T13:25:42Z

+        (512, 64, "tilelang", False),
+        (512, 0, "trtllm", False),
+        (512, 0, "flashmla_kv", False),
+        (None, 0, "tilelang", False),


This (and the (None, None) case in test_dsa_tilelang_fp8_validation.py) can go once the args are required

This (and the (None, None) case in test_dsa_tilelang_fp8_validation.py) can go once the args are required

Done. Removed the None geometry cases from both test files. Kept the rejection coverage for unsupported dimensions, nonzero RoPE, incompatible decode backends, and unsupported serving combinations.

mmangkad · 2026-10-19T13:26:01Z

    ) -> None:
        # q8 prefill consumes the tail columns as ragged ids, like the history
-        # picks. Only over a bfloat16 pool: gathering packed fp8 rows through
+        # picks over a bfloat16 or raw NoPE fp8 pool: gathering packed fp8 rows through


Nit: I think this lost the "Only", so it reads a bit off now. Maybe # picks. Only over a bf16 or raw NoPE fp8 pool: gathering packed?

Nit: I think this lost the "Only", so it reads a bit off now. Maybe # picks. Only over a bf16 or raw NoPE fp8 pool: gathering packed?

Fixed the wording to explicitly say “Only over a bf16 or raw NoPE fp8 pool.” The packed FP8 guard is unchanged.

mmangkad · 2026-10-19T13:26:39Z

-On Hopper, BF16 KV can also run the sparse prefill in FP8 with `--dsa-prefill-backend flashmla_sparse_q8` (decode stays on `tilelang` and the cache stays BF16): on H200 this pairing measured 5.8–12.1% higher throughput on prefill-only requests, with GSM8K accuracy within noise of BF16 + TileLang. It is SM90-only and requires a build that includes [PR #41790](https://github.com/sgl-project/sglang/pull/41790); FP8 KV is not available with it for this model.
+On Hopper, BF16 KV can also run the sparse prefill in FP8 with `--dsa-prefill-backend flashmla_sparse_q8` (decode stays on `tilelang` and the cache stays BF16): on H200 this pairing measured 5.8–12.1% higher throughput on prefill-only requests, with GSM8K accuracy within noise of BF16 + TileLang. It is SM90-only and requires a build that includes [PR #41790](https://github.com/sgl-project/sglang/pull/41790).
+
+With the raw FP8 KV support in [PR #42957](https://github.com/sgl-project/sglang/pull/42957) and [PR #44217](https://github.com/sgl-project/sglang/pull/44217), you can also pair Q8 prefill with `--kv-cache-dtype fp8_e4m3` and TileLang decode on CUDA H100/H200 (SM90). This path requires NoPE MLA (`qk_rope_head_dim=0`, `kv_lora_rank=512`) and gathers only the requested prefix rows from the raw FP8 pool. It does not use the scaled/packed FP8 layout. Keep prefix caching enabled; HiSparse, mixed chunks, and `--dcp-size > 1` are not supported with this pairing. The BF16 benchmark results above do not measure this FP8 KV configuration.


Why keep prefix caching on? I think it works fine without cache hits too, so maybe drop that part or say why

Why keep prefix caching on? I think it works fine without cache hits too, so maybe drop that part or say why

Agreed. Removed the mandatory prefix-caching instruction from the cookbook. Prefix caching was enabled for validation, but that should not have been presented as a requirement for using this backend.

mintlify · 2026-10-21T02:57:49Z

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit 6b1ac36.

Approve preview

yzk1989 · 2026-10-21T03:06:09Z

Thanks for the follow-up review! All six comments have been addressed.

Could you please take another look when you have a chance? Thanks!@mmangkad

github-actions · 2026-10-29T03:44:32Z

Required checks failed on 6b1ac36; this PR cannot be merged until they pass.

PR title (run)

PR title must use "[Type] scope: title" or "[Type] [Hardware] scope: title". Type must be one of: Feat, Fix, Refactor, Perf, Doc, Test, CI, Chore. Scope is required and may be any module or model name, with unrestricted letter case (letters, numbers, dots, underscores, slashes, or hyphens). Hardware is optional; use a single tag without spaces. Use uppercase hardware tags (AMD, NVIDIA, NPU, MI300X), except for exactly these four spellings: Intel, Apple, Huawei, Ascend. Reverts must use GitHub's default format: Revert "original title". Use one space between fields and after the colon, with a non-empty title and no leading or trailing whitespace. Examples: "[Fix] attention: Fix KV cache indexing", "[Feat] qwen3: Add model support", "[Perf] [AMD] flash-attn: Optimize decode performance".
Edit the title; the check reruns on its own.

yzk1989 requested review from 1am9trash, Fridge003, HaiShaw, Qiaolin-Yu, YAMY1234, Ying1123, alphabetc1, hanming-lu, hebiao064, hnyls2002, huangtingwei9988, hubertlu-tw, hzh0425, ispobock, kkHuang-amd, merrymercy, rainj-me, xiezhq-hermann and yizhang2077 as code owners October 14, 2026 06:27

yzk1989 changed the title ~~support flashmla_sparse_q8 backend for dsa-prefill-backend in GLM5.3-Flash~~ [DSA] Support TileLang flashmla_sparse_q8 backend on CUDA Oct 16, 2026

This comment has been minimized.

Sign in to view

yzk1989 changed the title ~~[DSA] Support TileLang flashmla_sparse_q8 backend on CUDA~~ [DSA] Support flashmla_sparse_q8 prefill backend on CUDA Oct 16, 2026

This comment has been minimized.

Sign in to view

mmangkad suggested changes Oct 16, 2026

View reviewed changes

yangzhenkun.yzk added 2 commits October 19, 2026 11:22

support flashmla_sparse_q8 for dsa-prefill-backend in GLM5.3-Flash

0944fe9

[Fix] dsa: Address raw FP8 Q8 prefill review feedback

06404b9

yzk1989 force-pushed the glm-flashmla-sparse-q8 branch from 667c4af to 06404b9 Compare October 19, 2026 03:46

yzk1989 requested a review from DarkSharpness as a code owner October 19, 2026 03:46

yzk1989 requested review from BBuf, HydraQYH, JustinTong0323, celve, sogalin, wisclmy0611, yuan-luo and zijiexia as code owners October 19, 2026 03:46

github-actions Bot added documentation Improvements or additions to documentation quant LLM Quantization jit-kernel labels Oct 19, 2026

This comment has been minimized.

Sign in to view

mmangkad reviewed Oct 19, 2026

View reviewed changes

JackChuang mentioned this pull request Oct 20, 2026

[Tracking] GLM-5.3-Flash performance optimization on Hopper (SM90) #45373

Open

14 of 49 tasks

[Refactor] dsa: Address Q8 prefill review feedback

efb7b86

This comment has been minimized.

Sign in to view

Merge upstream main and resolve Q8 prefill conflicts

6b1ac36

This was referenced Nov 2, 2026

[Fix] dsa: Reject unsupported k-pool backends early #41898

Open

[Doc] cookbook: Correct the GLM-5.3-Flash note on FP8 KV with TileLang DSA #48179

Closed

airMeng mentioned this pull request Nov 3, 2026

[Feat] xpu: Consolidate simple XPU enablements #47956

Open

		kv_lora_rank=getattr(model_config_of(view), "kv_lora_rank", None),
		qk_rope_head_dim=getattr(model_config_of(view), "qk_rope_head_dim", None),

		kv_lora_rank: Optional[int] = None,
		qk_rope_head_dim: Optional[int] = None,

Conversation

yzk1989 commented Oct 14, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

yzk1989 commented Oct 16, 2026

Uh oh!

mmangkad 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment has been minimized.

yzk1989 commented Oct 19, 2026

Uh oh!

yzk1989 commented Oct 19, 2026

Uh oh!

yzk1989 commented Oct 19, 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!

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!

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

Uh oh!

Preview approval required

Uh oh!

This comment has been minimized.

yzk1989 commented Oct 21, 2026

Uh oh!

github-actions Bot commented Oct 29, 2026

yzk1989 commented Oct 14, 2026 •

edited by github-actions Bot

Loading

mintlify Bot commented Oct 21, 2026 •

edited

Loading