Repository navigation
[Perf] [AMD] dsv4: fp8 unified_kv decode/verify on FlyDSL v4 nm decode with a latency-model split plan (1/8) - #44016
Merged
HaiShaw merged 10 commits intoOct 16, 2026
Conversation
…ersistent kernel (aiter sgl-project#6126); FlyDSL H128 opt-in
… until aiter ships it The official aiter pin (ROCm/aiter main 02575f1c3a) has no FlyDSL v4 nm decode, so sgl-project#41383's fp8 unified_kv decode/verify fell back to the asm kernel. Carry ROCm/aiter#5882 at its sgl-project#6377-adapted head 03359c43df in-tree: - python/sglang/kernels/ops/attention/dsv4/flydsl_v4nm/: mla_v4_decode.py, mla_v4_grouped.py, kernels/mla_v4_decode_{common,gfx950}.py copied from the PR (license header kept, source PR / commit / file noted); only the import paths changed (helpers come from aiter's public FlyDSL utils: communication_ops_utils, buffer_ops, kernels_common, tensor_shim) plus isort / ruff format. Importing it on an aiter before sgl-project#6377 raises ImportError. - runtime._flydsl_v4_decode(): aiter.ops.flydsl when the installed aiter has the API, else the in-tree copy; logs which one once per process. SGLANG_DSV4_FLYDSL_DECODE_SOURCE=auto|aiter|intree (default auto) forces one copy; SGLANG_DSV4_FLYDSL_DECODE=0 still turns FlyDSL off (asm). SGLANG_DSV4_FLYDSL_DECODE_H128 gating is unchanged. - test/registered/kernels/ops/attention/dsv4/test_flydsl_v4nm_decode.py: the PR's op_test ported to pytest, trimmed to the DSv4 shapes (H16/32/64/128 msq1, H16 msq4, fused wo_a epilogue, grouped verify at 4/7 drafts), plus the source-selection contract. - Kernel inventory entries for the two entry points; codespell skips the vendored directory (helper name `padd`). Every piece carries TODO(aiter#5882): remove once ROCm/aiter#5882 is in the aiter pin.
kkHuang-amd
requested review from
BBuf,
DarkSharpness,
Fridge003,
HaiShaw,
HydraQYH,
Qiaolin-Yu,
celve,
ch-wan,
fzyzcjy,
hebiao064,
ispobock,
merrymercy and
yuan-luo
as code owners
October 13, 2026 08:21
5 tasks
HaiShaw
approved these changes
Oct 16, 2026
HaiShaw
left a comment
Collaborator
There was a problem hiding this comment.
HIP, unified kv, gfx95, and flyDSL specific
4 of 5 tasks
This was referenced Oct 16, 2026
This was referenced Nov 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
DeepSeek-V4 fp8 unified_kv decode/verify on gfx950 runs the asm v4 nm kernel per token plus a separate merge and inverse-RoPE quant. This moves it to aiter's FlyDSL v4 nm fused decode and sizes the asm fallback's KV split from aiter's latency model. Includes #41383 (not merged yet).
Depends on: #41383 (open, included here); ROCm/aiter#6126 (in aiter main, needs the aiter pin bump in #44035); ROCm/aiter#5882 (open, carried in-tree until it lands).
Modifications
SGLANG_DSV4_FLYDSL_DECODE, default on;SGLANG_DSV4_FLYDSL_DECODE_SOURCE=auto|aiter|intree, default auto; 128 local heads opt-in viaSGLANG_DSV4_FLYDSL_DECODE_H128, default off).SGLANG_DSV4_FP8_DECODE_SPLIT_MODEL, default on;SGLANG_DSV4_FP8_DECODE_SPLIT_CTX, default 200000).mla_decode_fwd_v4_nm_ps([ASM] [HIP] [MLA v4 nm] Split plan helper and persistent decode kernel with in-kernel split and merge ROCm/aiter#6126) from 15 rows (SGLANG_DSV4_FP8_HCA_PERSIST, default on).test_flydsl_v4nm_decode.py, plus unified fp8 decode tests.Accuracy Tests
GSM8K: main 0.945 -> 0.951 (current aiter pin); 0.953 -> 0.948 (aiter 02575f1c3a).
Speed Tests and Profiling
From #41383 (8K in / 1K out, step time stock asm -> FlyDSL): TP8 c8 13.28 -> 12.75 ms (-3.9%), TP8 c32 17.13 -> 16.61 ms (-3.1%), DP8 attention c64 21.24 -> 21.01 ms (-1.1%).
Split plan + HCA persistent (asm path, used with DP attention), DP8 attention + EAGLE, median ITL on vs off: 8K in / 1K out c16 6.71 vs 6.76 ms, c64 9.32 vs 9.32 ms; 64K in / 1K out (context length 200000) c8 6.26 vs 6.26 ms, c32 7.60 vs 7.53 ms. No measurable difference in these cells.
Co-authored-by: @AMD-yanfeiwang
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): 🚫 Run #38834001002
Latest PR Test (Extra): ✅ Run #38834000818
Latest PR Test (AMD ROCm 10): ⏳ Run #38834000890