[Perf] dsv41: Serve token-interleaved V4.1 KV caches with sgl-flash-mla by DarkSharpness · Pull Request #45893 · sgl-project/sglang

DarkSharpness · 2026-10-22T20:24:40Z

(auto generated by Claude)

Motivation

FlashMLA's 2026.09.30 V4.1 kernels read a token-interleaved KV cache and are faster at decode. They ship in the sgl-flash-mla wheel (release workflow #45651), which also carries every pre-release kernel sglang uses behind use_legacy=True. This moves all CUDA FlashMLA calls to that wheel and adds the interleaved layout as an opt-in, so existing deployments are unchanged.

Supersedes the serving side of #45635 (co-authored), keeping the block-split v41 layout so gfx950 / aiter_sparse keeps working.

Modifications

New SGLANG_DSV4_KV_LAYOUT=v41_new (KVLayout.V41_NEW / V41_FP4_NEW): the V4.1 bytes interleaved per token, unpadded. v41 keeps the block-split layout. Writers, dequant and the request-window copy address both through row_addressing; the PD layout string names the arrangement, so mismatched prefill / decode pairs are refused.
sglang.kernels.ops.attention.flash_mla_interface wraps the wheel with the call signatures of the former sgl_kernel.flash_mla / flash_mla_mega ops. Old kernels default to use_legacy=True; only V41_NEW uses the current kernels. The V3/R1 dense, V3.2 DSA and DSv4 backends import from it; HIP / XPU keep their sgl-kernel ops.
Decode skips the schedule precompute when the current kernel does not split the KV (topk + extra_topk <= 640).
KVLayout properties are functools.cached_property.
Adds sgl-flash-mla to the dependencies.

Accuracy and speed

DeepSeek-V4.1-Flash, 4x B200, TP4 / EP4 / DSpark, sgl-eval with temperature 1.0 / top_p 0.95:

	v41	v41_new
gsm8k (1319)	96.66	96.51
ruler2 64K (240 x 3)	86.31 / 87.03 (two servers)	87.57
decode TPOT, 8K in / 1K out, concurrency 16	2.50 / 2.46 ms	2.41 / 2.35 ms

Greedy ruler2 outputs agree as often across layouts (65.0%) as across two v41 servers (66.7%). On the same tokens the current and legacy V4.1 sparse decode kernels give bitwise-identical outputs.

Checklist

Pin sgl-flash-mla to the first published version (1.0.0 is not on PyPI yet)

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): 🚫 Run #39877090044
Latest PR Test (Extra): 🚫 Run #39877090183
Latest PR Test (AMD ROCm 10): 🚫 Run #39877090151

…la and drop sgl-kernel's FlashMLA Add the opt-in SGLANG_DSV4_KV_LAYOUT=v41_new layout (KVLayout.V41_NEW / V41_FP4_NEW): the V4.1 bytes interleaved per token, read by FlashMLA's 2026.09.30 kernels. v41 keeps the block-split layout. Route every CUDA FlashMLA call through sglang.kernels.ops.attention.flash_mla_interface over the sgl-flash-mla wheel, with use_legacy=True for the pre-release kernels, and add the wheel to the dependencies. Co-Authored-By: kpham-sgl <khoa.pham@radixark.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…missing legacy library Pass `is_fp8_kvcache` through on the scheduler-object path instead of forcing it to False when `indices` is None, so a dense fp8 request without tensor metadata hits the fork's assertion rather than running the bf16 kernel on fp8 bytes. The tensor-metadata path takes its ops from `flash_mla._loader.get_ops(use_legacy=True)`, which raises the wheel's own message when the legacy extension is absent. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

The layout flag, the interface import and the use_legacy partial are constant for the backend's life; cache them on the instance, and the split threshold import on the module. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Prefill reads unpaged bf16 rows, so the layout did not force the choice; one kernel set per mode is the stated one. The generic sparse prefill stays on the legacy kernel, whose head-count contract its callers rely on. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

The two callers that hold raw (metadata, num_splits) tensors now get a fresh scheduler object carrying them, and the fork's own flash_mla_with_kvcache dispatches; the private reimplementation with its raw positional op calls goes. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

kpham-sgl · 2026-10-22T23:49:11Z

Three more commits, per review discussion:

fb8eb07: the V4.1 decode path's layout flag, interface import and use_legacy partial are cached on the backend instance (functools.cached_property), and the split-threshold import on the module, instead of being recomputed per layer per step.
9c3e121: MegaAttention prefill also runs on the current kernels under v41_new, so each mode uses one kernel set. The generic flash_mla_sparse_fwd prefill stays on the legacy kernel, whose head-count contract its callers rely on.
8df6313: the two tensor-metadata callers (dsa_backend, flashmla_backend) get a fresh FlashMLASchedMeta carrying their tensors and the fork's own flash_mla_with_kvcache dispatches; _flash_mla_with_kvcache_tensor_meta and its raw positional op calls are gone. The V3.2-no-RoPE kv_format inference and the dense-fp8-by-q-dtype selection are kept on that path.

Validation in progress: the PR's kernel and pool tests on GB300 against the wheel, and a DEP4 GSM8K A/B of v41_new against v41 on the same head.

Plain attributes set under is_cuda() replace the cached properties and the cached import helper. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

kpham-sgl · 2026-10-23T00:41:45Z

GB300 serving check on head bd3b828, with the published sgl-flash-mla 1.0.0 wheel: DeepSeek-V4.1-Flash, one node DEP4 (TP4/DP4/EP4), MegaMoE, flashmla_mega, DSpark 3x3 with native acceptance, GSM8K 1,314 x 5-shot, thinking off, temperature 0. Same recipe, only the layout differs.

`SGLANG_DSV4_KV_LAYOUT`	kernels (prefill + decode)	GSM8K	latency
`v41_new`	current	0.9718	114.9 s
`v41`	legacy	0.9749	111.8 s

4 answers apart out of 1,314, inside the run-to-run spread of this recipe (0.971 to 0.976 across the earlier runs on this node type). Kernel and pool tests on the same head pass on GB300: test_flashmla_v41 7/7, test_mega_attention 2/2, test_v41_kv_store 8 (1 skip), test_dsv41_request_window_copy 4/4, test_dsv4_compressed_pools 19/19, test_flashmla_sched_meta 41/41, test_dsa_sink_pad_cache 1/1.

Every CUDA FlashMLA call goes through the sgl-flash-mla wheel since #45893, so sgl-kernel stops fetching and compiling the fork: the cmake fragment, the two extension sources, the ops declarations, the Python wrappers and their tests and benchmark. XPU keeps its own ops. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…uses SGLang v0.5.23 (sgl-project/sglang#45893) depends on sgl-flash-mla==1.0.0 and calls decode APIs only that build has (get_mla_decoding_metadata, fused_norm_rope_attn_rope_cast, the legacy SM90 kernels). The image installed the wheels release's FlashMLA nv_dev over it under the same `flash_mla` name, so SGLang's DSA paths failed at warmup: "module 'flash_mla' has no attribute 'get_mla_decoding_metadata'" (run 41189133939, both stage-c-b200 GLM-5.2 NVFP4 jobs) and "flash_mla_sparse_fwd() got an unexpected keyword argument 'use_legacy'" (the H200 DeepSeek-V4.1 job). Dropping nv_dev is not an option: Megatron's CSA forward passes `indexer_topk` on every call, and miles' DSA kernel runs the sparse prefill on SM90 and with 576-dim heads. sgl-flash-mla has neither (its current kernels are V4.1 on SM100 only), so the trainer would fail instead. Keep nv_dev as `flash_mla`, so the trainer stays exactly as on main, and move the base image's sgl-flash-mla package to `sgl_flash_mla`, which sglang-miles imports first. The two builds load their kernels differently (a CPython extension vs STABLE_TORCH_LIBRARY ops), so both can live in one process. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…can ship nv_dev as flash_mla v0.5.23 moved SGLang's FlashMLA ops to the sgl-flash-mla wheel (#45893), which installs as `flash_mla`. miles' image must also ship the FlashMLA nv_dev build under that name, because Megatron's CSA forward and miles' DSA kernel call its `indexer_topk` and SM90 sparse-prefill paths, which sgl-flash-mla lacks. The image therefore moves sgl-flash-mla's package to `sgl_flash_mla`; this module, SGLang's only FlashMLA import site, takes it from there when present and falls back to `flash_mla` everywhere else. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

DarkSharpness requested review from BBuf, Fridge003, HaiShaw, HydraQYH, Qiaolin-Yu, Ying1123, alphabetc1, celve, hanming-lu, hebiao064, hnyls2002, huangtingwei9988, hzh0425, ispobock, merrymercy, xiezhq-hermann, yizhang2077 and yuan-luo as code owners October 22, 2026 20:24

github-actions Bot added documentation Improvements or additions to documentation quant LLM Quantization dependencies Pull requests that update a dependency file deepseek jit-kernel memory-pool labels Oct 22, 2026

This comment has been minimized.

Sign in to view

kpham-sgl mentioned this pull request Oct 22, 2026

[Perf] dsv41: Token-interleaved V4.1 KV caches for the FlashMLA V4.1 kernels #45635

Closed

kpham-sgl added the run-ci CI: run the baseline test suite on this PR label Oct 22, 2026

Format with isort and ruff

ea11ee1

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

This comment has been minimized.

Sign in to view

kpham-sgl added run-ci-extra CI: also run the extra suite (requires run-ci) highest-priority CI: all three control labels, plus never batch-cancelled or stale-closed bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) labels Oct 22, 2026

kpham-sgl and others added 4 commits October 22, 2026 16:36

kpham-sgl and others added 3 commits October 22, 2026 17:01

Bind the FlashMLA interface once in the backend constructor

38c4714

Plain attributes set under is_cuda() replace the cached properties and the cached import helper. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Import the FlashMLA interface as a module in the mega backend

d7b0e5a

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Pass the split threshold to _flashmla_v41_splits_kv in its test

bd3b828

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

kpham-sgl approved these changes Oct 23, 2026

View reviewed changes

kpham-sgl merged commit e83f585 into main Oct 23, 2026
344 of 368 checks passed

kpham-sgl deleted the dsv41-flashmla-wheel branch October 23, 2026 05:02

akhilg-nv mentioned this pull request Oct 23, 2026

[Perf] [NVIDIA] deepseek-v4.1: Keep one zeroed MegaAttention packed-query buffer instead of zeroing it per layer #46062

Open

4 of 5 tasks

zcnrex mentioned this pull request Oct 29, 2026

[Perf] [NVIDIA] sgl-kernel: Build FlashMLA natively for sm_107a #45884

Closed

2 of 5 tasks

kpham-sgl mentioned this pull request Oct 31, 2026

[Refactor] [NVIDIA] attention: Migrate DSA/DSV4 indexers to DeepGEMM 0.2.5.post2 #44350

Open

4 of 5 tasks

Rockdu mentioned this pull request Oct 31, 2026

build(docker): bump SGLang to v0.5.23 radixark/miles#4548

Merged

6 tasks done

Navigation Menu

Repository navigation

[Perf] dsv41: Serve token-interleaved V4.1 KV caches with sgl-flash-mla - #45893

[Perf] dsv41: Serve token-interleaved V4.1 KV caches with sgl-flash-mla#45893
kpham-sgl merged 11 commits into
mainfrom
dsv41-flashmla-wheel

DarkSharpness commented Oct 22, 2026 •

edited by github-actions Bot

Loading

Uh oh!

This comment has been minimized.

This comment has been minimized.

kpham-sgl commented Oct 22, 2026

Uh oh!

kpham-sgl commented Oct 23, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Motivation

Modifications

Accuracy and speed

Checklist

CI States

Uh oh!

This comment has been minimized.

This comment has been minimized.

kpham-sgl commented Oct 22, 2026

Uh oh!

kpham-sgl commented Oct 23, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

DarkSharpness commented Oct 22, 2026 •

edited by github-actions Bot

Loading