Skip to content

[ROCm][Perf] Enable CSA multi-stream overlap for DeepSeek-V4 - #51794

Draft
shen-shanshan wants to merge 13 commits into
vllm-project:mainfrom
shen-shanshan:rocm-dsv4-csa-multi-stream
Draft

[ROCm][Perf] Enable CSA multi-stream overlap for DeepSeek-V4#51794
shen-shanshan wants to merge 13 commits into
vllm-project:mainfrom
shen-shanshan:rocm-dsv4-csa-multi-stream

Conversation

@shen-shanshan

@shen-shanshan shen-shanshan commented Aug 11, 2026

Copy link
Copy Markdown
Collaborator

Purpose

Part of #41820.
Continue working on #43718.
Based on #50866 (This one should be merged before this PR).
Added #52093 and #52094.

Enable kernel-level multi-stream overlap for DeepSeek-V4 CSA (Compressed Sparse Attention, compress_ratio=4) attention layers on ROCm, gated by a new environment variable VLLM_ROCM_DSV4_CSA_MULTI_STREAM (default off, serial fallback when unset). When enabled, the attention forward forks three HIP streams before the input GEMMs:

  • default stream: fused wqa+wkv GEMM → q/kv RMSNorm → wq_b → qnorm/RoPE → SWA KV insert
  • aux stream 0: main-compressor wkv_gate GEMM → compress kernels (compressed KV cache write)
  • aux stream 1: indexer-compressor wkv_gate GEMM → K-cache write

After the join, the indexer weights GEMM, the indexer q-side (wq_b + fused q rope+quant), the sparse indexer op (top-k selection), and the MLA attention run serially on the main stream. Moving the two compressor wkv_gate GEMMs onto side streams takes them off the serial critical path; each stream runs a self-contained chain straight from
hidden_states, so no fine-grained cross-stream dependencies are required.

Before:

before

After:

Test Plan

  • Benchmark with SA InferenceX 8k1k and agentic workload.
  • Acc test with gsm8k dataset.

Test Result

Benchmark (workload=8k1k tp8)

Benchmark (workload=agentic tp8)

To be completed...

Acc Test

To be completed...


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@shen-shanshan
shen-shanshan marked this pull request as draft August 11, 2026 09:49
@mergify mergify Bot added deepseek Related to DeepSeek models nvidia rocm Related to AMD ROCm labels Aug 11, 2026
@github-project-automation github-project-automation Bot moved this to Todo in AMD Aug 11, 2026
@mergify

mergify Bot commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @shen-shanshan.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Aug 13, 2026
@shen-shanshan
shen-shanshan force-pushed the rocm-dsv4-csa-multi-stream branch from 796108d to 5f88b24 Compare August 13, 2026 03:24
@mergify mergify Bot removed the needs-rebase label Aug 13, 2026
@shen-shanshan
shen-shanshan force-pushed the rocm-dsv4-csa-multi-stream branch from 48e4031 to 06d206b Compare August 13, 2026 06:35
@mergify

mergify Bot commented Aug 16, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @shen-shanshan.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

@mergify mergify Bot added the needs-rebase label Aug 16, 2026
shen-shanshan and others added 10 commits August 17, 2026 14:53
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
…ct#51430)"

Restore the wider eager region because the narrowed capture corrupts ROCm output and regresses long-context serving performance.

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>

Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
…roject#51768)"

Restore Model Runner V1 after the prerequisite eager-region revert makes piecewise graph execution correct again and recovers ROCm decode throughput.

Co-authored-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Fangzhou Ai <fangzhou.ai@amd.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
@shen-shanshan
shen-shanshan force-pushed the rocm-dsv4-csa-multi-stream branch from 06d206b to d314c00 Compare August 17, 2026 08:06
@mergify mergify Bot removed the needs-rebase label Aug 17, 2026
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
Signed-off-by: Shanshan Shen <87969357+shen-shanshan@users.noreply.github.com>
@mergify

mergify Bot commented Aug 19, 2026

Copy link
Copy Markdown
Contributor

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @shen-shanshan.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

deepseek Related to DeepSeek models DSv4 needs-rebase nvidia rocm Related to AMD ROCm

Projects

Status: Todo
Status: No status

Development

Successfully merging this pull request may close these issues.

2 participants