Skip to content

[AMD] Dsv4/pr2 compressor opt - #26208

Merged
HaiShaw merged 11 commits into
sgl-project:mainfrom
HaiShaw:dsv4/pr2-compressor-opt
May 26, 2026
Merged

[AMD] Dsv4/pr2 compressor opt#26208
HaiShaw merged 11 commits into
sgl-project:mainfrom
HaiShaw:dsv4/pr2-compressor-opt

Conversation

@kkHuang-amd

@kkHuang-amd kkHuang-amd commented May 24, 2026

Copy link
Copy Markdown
Collaborator

Co-authored-by: @1am9trash @RolaoDenthu @yichiche @hubertlu-tw @sogalin @HaiShaw @yctseng0211 @bingxche @amd-danli103 @Raiden-Makoto

Motivation

This PR improves DeepSeek-V4 inference performance on AMD ROCm by reducing decode/prefill hot-path overhead in compressor, indexer, and fused attention execution.
It also consolidates kernel options so we can enable high-performance fused paths with clearer runtime flags while maintaining numerical correctness checks.

Modifications

  • Add DSV4 fused compress implementations (fused_compress_kernel.py, fused_compress_triton.py) and integrate them into compressor flows.
  • Introduce Triton NSA decode kernel stack (triton_decode/*) for optimized/fused/split-k decode paths.
  • Optimize DSV4 compressor/indexer/metadata and HIP radix backend plumbing to support new fused execution paths.
  • Add fused QK norm support (fused_qk_norm.py) and wire related model-layer updates (deepseek_v4.py, RoPE, top-k, fp8, mem cache, MoE runner).
  • Expose/configure new runtime environment toggles in environ.py for selecting fused kernels and indexer behaviors.
  • Add validation and benchmark artifacts, including manual DSV4 fused-compress attention test and sgl-kernel norm+rope benchmark/test.
  • Add engineering notes documenting DSV4 HIP kernel fusion decisions.

AMD Accuracy Tests

server command

export SGLANG_OPT_USE_JIT_KERNEL_FUSED_TOPK=true

export SGLANG_OPT_DEEPGEMM_HC_PRENORM=false

export SGLANG_ENABLE_THINKING=1
export SGLANG_USE_AITER=1
export SGLANG_USE_ROCM700A=1
export SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1

export SGLANG_OPT_USE_OVERLAP_STORE_CACHE=false
export SGLANG_OPT_USE_FUSED_STORE_CACHE=true

# changed
export SGLANG_OPT_USE_FUSED_COMPRESS=true
export SGLANG_TOPK_TRANSFORM_512_TORCH=0
export SGLANG_OPT_USE_TILELANG_INDEXER=true
export SGLANG_OPT_USE_AITER_INDEXER=false
export SGLANG_HACK_FLASHMLA_BACKEND=triton
export SGLANG_REASONING_EFFORT=max
export SGLANG_OPT_USE_AITER_MHC_PRE=true
export SGLANG_OPT_USE_AITER_MHC_POST=true
export SGLANG_OPT_USE_FUSED_HASH_TOPK=true
export SGLANG_OPT_FUSE_WQA_WKV=true

export SGLANG_OPT_USE_FUSED_QK_NORM_ROPE=true
export SGLANG_OPT_USE_FUSED_CLAMP_ACT_MUL=true
export SGLANG_OPT_USE_FUSED_COMPRESS_TRITON=true

export SGLANG_OPT_FP8_WO_A_GEMM=false
export SGLANG_OPT_USE_TOPK_V2=false

export SGLANG_OPT_USE_COMPRESSOR_V2=false
export SGLANG_OPT_USE_JIT_NORM=false

export SGLANG_OPT_DPSK_V4_RADIX=1
export SGLANG_OPT_C4_SPARSE_TOPK=512

export SGLANG_OPT_USE_JIT_INDEXER_METADATA=false
export SGLANG_OPT_USE_MULTI_STREAM_OVERLAP=false
#export SGLANG_ROCM_USE_MULTI_STREAM=1

export AITER_BF16_FP8_MOE_BOUND=1
export GPU_MAX_HW_QUEUES=5

MODEL=/dockerx/data/models/DeepSeek-V4-Pro/

python3 -m sglang.launch_server \
    --model-path ${MODEL} \
    --trust-remote-code \
    --tp 8 \
    --dp 8 --enable-dp-attention --enable-prefill-delayer \
    --disable-radix-cache \
    --attention-backend compressed \
    --max-running-request 256 \
    --page-size 256 \
    --chunked-prefill-size 8192 \
    --port 8000 \
    --disable-shared-experts-fusion \
    --tool-call-parser deepseekv4 \
    --reasoning-parser deepseek-v4

client command

python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319 --num-shots 5 --port 8000

Result

/sgl-workspace/sglang# python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319 --num-shots 5 --port 8000
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [01:19<00:00, 16.68it/s]
Accuracy: 0.951
Invalid: 0.000
Latency: 79.067 s
Output throughput: 1503.000 token/s

NVIDIA Accuracy Tests

server command

export SGLANG_JIT_DEEPGEMM_PRECOMPILE=0
export SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT=1
export SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1
export SGLANG_OPT_USE_JIT_INDEXER_METADATA=1
export SGLANG_OPT_USE_JIT_NORM=1
export SGLANG_OPT_USE_TOPK_V2=1

export SGLANG_OPT_USE_COMPRESSOR_V2=false

python3 -m sglang.launch_server --model-path /dockerx/mnt/models/deepseek-ai/DeepSeek-V4-Pro --host 0.0.0.0 --port 8000 --trust-remote-code --tp 8 --disable-radix-cache --max-running-requests 256 --mem-fraction-static 0.90 --swa-full-tokens-ratio 0.1 --moe-runner-backend flashinfer_mxfp4 --chunked-prefill-size 8192 --disable-flashinfer-autotune

Result

/sgl-workspace/sglang# python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --parallel 1319 --num-shots 5 --port 8000
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [01:37<00:00, 13.56it/s]
Accuracy: 0.945
Invalid: 0.000
Latency: 97.288 s
Output throughput: 1207.229 token/s

Speed Tests and Profiling

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. 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
  4. 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 #26403671212
Latest PR Test (Extra): ❌ Run #26403670966

@github-actions github-actions Bot added documentation Improvements or additions to documentation amd deepseek sgl-kernel labels May 24, 2026

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces fused Triton kernels and performance optimizations for DeepSeek V4 on HIP, including fused compressor attention, Q/K RMSNorm, and sparse attention decode, alongside multi-stream overlap and aiter integration. The review highlights critical bugs regarding numerical stability and memory safety: online softmax implementations require guards against -inf to prevent NaN propagation, and several kernels need improved masking to avoid unsafe memory overwrites or incorrect reduction sums when handling invalid indices.

Comment thread python/sglang/srt/layers/attention/dsv4/compressor_v2.py Outdated
Comment thread python/sglang/srt/layers/attention/dsv4/compressor_v2.py Outdated
Comment thread python/sglang/srt/layers/attention/dsv4/compressor_v2.py Outdated
Comment thread python/sglang/srt/layers/attention/dsv4/compressor_v2.py Outdated
Comment thread python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py Outdated
Comment thread python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py
Comment thread python/sglang/srt/layers/attention/dsv4/fused_compress_triton.py
wunhuang and others added 2 commits May 24, 2026 09:18
Remove the local fused-compress notes file from this branch to keep the PR focused on executable DSV4 optimization changes.

Co-authored-by: Cursor <cursoragent@cursor.com>
Fix online softmax NaN guards and invalid write/index masking in HIP compressor paths, and keep fused kernel import paths consolidated after merging fused_compress_kernel into fused_compress_triton.
@kkHuang-amd

Copy link
Copy Markdown
Collaborator Author

Addressed review feedback in commit 41cd3691f:

  • Added -inf guards for online softmax exp_cur in HIP c128 decode/prefill kernels (compressor_v2.py).
  • Fixed fallback write path to only scatter valid indices (avoid writing invalid entries to slot 0).
  • Added invalid-index store mask for c4 decode state write (fused_compress_triton.py).
  • Masked exp_s reductions by validity in c128 decode/prefill chunked kernels to avoid incorrect running_sum updates on invalid tiles.
  • Kept fused-kernel consolidation by moving fused_ape_pool_norm_rope into fused_compress_triton.py and leaving a compatibility shim in fused_compress_kernel.py.

wunhuang added 2 commits May 24, 2026 11:57
Delete fused_compress_kernel.py now that fused kernel implementation has been merged into fused_compress_triton.py.
Inline minimal test-only plan/write/compress helpers in test_fused_compress_attn_hip.py so the manual test no longer depends on the removed fused_compress_kernel module.
wunhuang and others added 2 commits May 24, 2026 13:11
Preserve input-path contributions in c128 prefill softmax masking and only gate state-path entries by src_pos validity to avoid regression in GSM8K accuracy.
wunhuang and others added 2 commits May 25, 2026 01:11
Drop duplicated DeepSeek-V4 norm/rope operator defs in common_extension_rocm to prevent c10 duplicate-registration aborts in AMD CI.

@HaiShaw HaiShaw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@yctseng0211 yctseng0211 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@HaiShaw

HaiShaw commented May 26, 2026

Copy link
Copy Markdown
Collaborator

@amd-bot ci-status

@amd-bot

amd-bot commented May 26, 2026

Copy link
Copy Markdown

@HaiShaw

CI Status for PR #26208

PR: [AMD] Dsv4/pr2 compressor opt
Changed files: python/sglang/srt/layers/attention/dsv4/* (massive additions), python/sglang/srt/layers/attention/nsa/triton_decode/* (~6000 lines added), python/sglang/srt/models/deepseek_v4.py (+152/-10), python/sglang/srt/layers/{activation,fused_qk_norm,deepseek_v4_rope,moe/topk,quantization/fp8}.py, sgl-kernel/* (DSv4 norm+rope kernel), and 6 more files.

AMD: 3 failures (0 likely related) | Others: 8 failures (0 related)

AMD CI Failures

Job Test File Test Function Error Related? Explanation Log
stage-b-test-1-gpu-small-amd (7) test/registered/perf/test_vlm_perf_5090.py test_vlm_offline_throughput aiohttp.ClientPayloadError: Response payload is not completed (server-side disconnects during bench) 🟢 Unlikely VLM bench-serving on a non-DSv4 model (minicpmv). PR's DSv4/NSA/HIP kernels and deepseek_v4* paths are not exercised by this test. The shared edits (activation.py, moe/topk.py, fp8.py) could in theory touch generic paths, but the failure mode is HTTP-level (server dropping the stream), which is a known infra/runner flake on the 5090-perf bench. Log
stage-b-test-1-gpu-large-amd (1) test/registered/perf/test_bench_serving_1gpu_part2.py test_embeddings_api_batch_scaling AssertionError: 226.13… not less than 80 (latency threshold) 🟢 Unlikely Embeddings API perf threshold. The failing path is the embeddings/HTTP server, not DSv4/MoE/fp8. Numbers consistent with a perf-regression flake (212 → 226 across attempts). PR changes do not touch embedding code. Log
stage-b-test-1-gpu-small-amd-nondeterministic test/registered/models/test_vlm_models.py TestVLMModels.test_vlm_mmmu_benchmark RuntimeError: split_with_sizes expects split_sizes to sum exactly to 2112 … got 12×64=768 in mm_utils._get_chunked_prefill_embedding → server crash → 45-min timeout 🟡 Possibly Crash is in generic managers/mm_utils.py:708 chunked-prefill embedding path on minicpmv (not DSv4). The PR does not modify mm_utils.py or models/minicpmv.py, and the mismatch (2112 vs 12·64) looks like a multimodal-token-bookkeeping bug independent of this PR. Marked Possibly only because PR's moe/topk.py / activation.py edits run for every model and a downstream side-effect cannot be fully ruled out from the log alone. Re-run after a rebase on main to disambiguate. Log

The three finish/gate jobs (pr-test-amd-finish, wait-for-stage-b-amd, call-gate / pr-gate, pr-test-extra-finish) are aggregator jobs that failed only because of the upstream stage-b failures above — no independent failure.

Other CI Failures

All 7 NPU jobs failed at workflow startup with the same infrastructure error — none ran any PR code.

Job Test File Test Function Error Related? Explanation Log
stage-b-test-1-npu-a2 (0) N/A N/A ImportError: …libnpu_profiler.so: undefined symbol: _ZNK3c1010TensorImpl15decref_pyobjectEvFailed to load backend extension: torch_npu 🟢 Unlikely torch_npu ABI mismatch against torch — environment/runner issue, fires before any sglang test executes. Log
stage-b-test-1-npu-a2 (1) N/A N/A same torch_npu ImportError 🟢 Unlikely Same infra failure. Log
stage-b-test-2-npu-a2 (1) N/A N/A same torch_npu ImportError 🟢 Unlikely Same infra failure. Log
stage-b-test-4-npu-a3 N/A N/A same torch_npu ImportError 🟢 Unlikely Same infra failure. Log
multimodal-gen-test-1-npu-a3 N/A N/A same torch_npu ImportError 🟢 Unlikely Same infra failure. Log
multimodal-gen-test-2-npu-a3 N/A N/A same torch_npu ImportError 🟢 Unlikely Same infra failure. Log
multimodal-gen-test-8-npu-a3 N/A N/A same torch_npu ImportError 🟢 Unlikely Same infra failure. Log

Details

Bottom line: none of the failures clearly trace to this PR's code.

  • All NPU failures (7 jobs) and the pr-test-npu-finish aggregator are a single infrastructure problem (torch_npu ABI symbol _ZNK3c1010TensorImpl15decref_pyobjectEv missing) hit before any sglang import — affects every PR running on that runner image today. Safe to ignore from this PR's perspective; needs an NPU image/torch_npu version pin.
  • The two perf tests (test_vlm_perf_5090, test_bench_serving_1gpu_part2::test_embeddings_api_batch_scaling) are non-DSv4 perf/throughput checks. PR's DSv4 backend, NSA decode kernels, DSv4 RoPE/qk_norm, and DSv4 fused-compress Triton work are not on the code path of either test. The cross-cutting edits (moe/topk.py, activation.py, quantization/fp8.py) are small and the failures are HTTP/perf-threshold flavored, not numerical or import errors.
  • The one failure worth a second look is test_vlm_mmmu_benchmark on minicpmv: the crash is split_with_sizes in python/sglang/srt/managers/mm_utils.py:708 (_get_chunked_prefill_embedding) — completely outside the files this PR touches, and the symptom (12 chunks of 64 vs total 2112) looks like an existing multimodal-token-bookkeeping issue. Re-run after rebasing on a recent main to confirm it's pre-existing.

Suggested action: rebase on latest main and re-trigger AMD CI; if the mm_utils crash persists on a fresh run, file as a separate VLM bug rather than blocking this PR.

Generated by amd-bot using Claude Code CLI

@yctseng0211

Copy link
Copy Markdown
Collaborator

@HaiShaw

CI Status for PR #26208

PR: [AMD] Dsv4/pr2 compressor opt Changed files: python/sglang/srt/layers/attention/dsv4/* (massive additions), python/sglang/srt/layers/attention/nsa/triton_decode/* (~6000 lines added), python/sglang/srt/models/deepseek_v4.py (+152/-10), python/sglang/srt/layers/{activation,fused_qk_norm,deepseek_v4_rope,moe/topk,quantization/fp8}.py, sgl-kernel/* (DSv4 norm+rope kernel), and 6 more files.

AMD: 3 failures (0 likely related) | Others: 8 failures (0 related)

Suggested action: rebase on latest main and re-trigger AMD CI; if the mm_utils crash persists on a fresh run, file as a separate VLM bug rather than blocking this PR.

Generated by amd-bot using Claude Code CLI

No need to rebase and re-trigger AMD CI, the mm_utils crash in est/registered/models/test_vlm_models.py is a known issue, we also see it in scheduled run : https://github.com/sgl-project/sglang/actions/runs/26389348481

@HaiShaw
HaiShaw merged commit 3f5e2c7 into sgl-project:main May 26, 2026
230 of 267 checks passed
bingxche added a commit that referenced this pull request Jun 1, 2026
Resolve conflict in AiterMoeQuantInfo (keep both swiglu_limit and
gate_mode fields) and reconcile AiterRunnerCore.run() into a single
gate_mode source.

main #26208 set extra["gate_mode"]=INTERLEAVE when swiglu_limit>0, while
this branch passes gate_mode=quant_info.gate_mode directly to fused_moe.
A naive merge would pass gate_mode twice (TypeError) on the FP8/DSv4 path
where swiglu_limit>0. Now gate_mode is resolved once: it comes from
quant_info.gate_mode (MXFP4/GPT-OSS sets it from activation), and
swiglu_limit>0 forces INTERLEAVE to preserve the original DSv4 behavior.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
mqhc2020 pushed a commit to mqhc2020/sglang that referenced this pull request Jun 2, 2026
Co-authored-by: wunhuang <wunhuang@amd.com>
Co-authored-by: Thomas Wang <1am9trash@gmail.com>
Co-authored-by: Xinyi Song <86638975+RolaoDenthu@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: amd-danli103 <danli103@amd.com>
Co-authored-by: Lin, Soga <soga.lin@amd.com>
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Co-authored-by: yichiche@amd.com <jacky.cheng>
Co-authored-by: yctseng0211 <yctseng@amd.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>
zijiexia added a commit to zijiexia/sglang that referenced this pull request Jun 4, 2026
…iffusion matrix

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

amd deepseek documentation Improvements or additions to documentation run-ci sgl-kernel

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants