[ROCm][Bugfix][Perf] Tune multi-stream shared experts use; wvSplitKrc fixes by mawong-amd · Pull Request #56098 · vllm-project/vllm

mawong-amd · 2026-09-09T14:43:48Z

After #55099, multi-stream performance in ROCm has markedly improved and allows new use-cases. In particular, the regime in which multi-stream shared-experts decode (implemented in #52033) brings performance advantages has shifted.

Purpose

This PR expands the use of multi-stream shared-experts decode on ROCm by loosening the gate that enables it and tuning it to where it's measured to improve performance. Furthermore, the correctness issue that led to the original implementation's (#48223) reversion has been identified as a bug in a gfx950-specific skinny GEMM (wvSplitKrc).

wvSplitKrc was not previously multi-stream safe because it uses a statically allocated buffer which 2 or more streams can access at the same time. This has been fixed by enlarging the allocation 8x (7.5 MiB per buffer, so 60 MiB total which is still small) and giving distinct streams distinct slots into the workspace. If more distinct streams exist than can be accommodated, we now fall back to an on-demand allocation, which is safe albeit adding overhead. This also accounts for cases in which wvSplitKrc is graph-captured on stream A but replayed on a different stream B.

In addition, a latent bug in wvSplitKrc that causes it to truncate its readback on certain GEMM shapes (leading to silent incorrectness) is fixed. wvSplitKrc's CHUNKK parameter is also retuned for a bit of performance.

The different memory allocation patterns caused by multi-stream also revealed a latent bug in vllm/v1/sample/ops/topk_topp_triton.py where the buffer cache teardown was not reliably done and so different allocation patterns (e.g. from multi-stream use) could lead to the buffer cache being allocated in the same segment as the KV cache, thus pinning them from being torn down. This has been fixed by doing the buffer cache reset during clean up.

Performance numbers

Output token throughput, mean of 2 runs. gfx950, TP8/DP1, real weights, ISL 1024 / OSL 256.

Model	Concurrency	Overlap off	Overlap on	Change
Qwen3.5-35B-A3B	1	158.0 tok/s	177.4 tok/s	+12.3%
	64	4,566 tok/s	5,967 tok/s	+30.7%
	256	14,399 tok/s	16,006 tok/s	+11.2%
DeepSeek-V4-Pro	1	62.5 tok/s	66.8 tok/s	+6.9%
	64	2,010 tok/s	2,374 tok/s	+18.1%
	256	4,580 tok/s	4,786 tok/s	+4.5%

Output token throughput. DeepSeek-V2-Lite (hidden 2048, 256 routed experts, top-k 8, 40 layers) with shared-expert size overridden for testing and random weights, single-GPU. Concurrency 1, mean of 2 runs.

Parallel config	Shared-expert size	Overlap off	Overlap on	Change
TP1 / DP1	64	140.5 tok/s	146.1 tok/s	+4.0%
	1024	141.0 tok/s	145.8 tok/s	+3.4%
	4096	132.6 tok/s	140.2 tok/s	+5.8%
TP8 / DP1	64	134.9 tok/s	153.8 tok/s	+14.0%
	1024	144.1 tok/s	149.2 tok/s	+3.6%
	4096	144.2 tok/s	150.2 tok/s	+4.2%
TP1 / DP8 + EP	64	95.2 tok/s	96.8 tok/s	+1.7%
	1024	93.9 tok/s	96.8 tok/s	+3.1%
	4096	90.6 tok/s	94.3 tok/s	+4.1%

In particular the numbers at shared-expert size = 64 are quite representative of Qwen3.5-35B at TP8.

Test Plan

2 new tests have been added for wvSplitKrc that exercise shapes it was previously failing on before and that exercise it in a multi-stream concurrent setting. The shared-experts multi-stream correctness (and a further test on wvSplitKrc) are exercised by the LM Eval Qwen3-5 Models test group which was failing earlier on the multi-stream shared experts path.

Test Result

The tests pass, and there are no new test regressions.

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.

The deterministic split-K readback stages one fp32 partial per (k-shard, N-tile) in LDS before summing them. That footprint is 256 * nt * k_rnd float4 and grows with the shard count, but the LDS allocation is fixed, so once it exceeds LDS the global_load_lds writes are dropped and the reduction returns a truncated K-sum -- silently, and only for the N-tiles whose staging fell off the end. On gfx950 this is reachable straight from the dispatch path (10 <= n <= 128, k > 512, m % 16 == 0): n >= 65 breaks from k = 5248 and n = 32 from k = 10368. Every guard passes -- fits_wvsplitkrc bounds the global workspace and CuNeeded bounds the grid, neither bounds LDS -- and the test grid stops at k = 3080, so nothing caught it. The measured error is 2-3x the output standard deviation, confined to exactly the output rows owned by the overflowing N-tile. Stage the readback in batches that fit. The main-loop staging areas are separated from the readback by a __syncthreads(), so a union lets the readback address all of LDS; the first batch is peeled and is the only batch for every shape that fits in one pass, which keeps the batch loop off the common path. Also correct the workspace predicate: the shard count is ceil(k * CHUNKK / 512), not ceil(k / 512), so it was 2x optimistic whenever the kernel halved the K-shard. The host wrapper now checks it rather than overrunning the workspace. Verified on gfx950 (MI355X). The new test fails 8/8 before the kernel change and passes 8/8 after; a 580-configuration sweep over n/k/m/dtype/bias reports no mismatches. Shapes that already worked are within ~1.3% of the previous kernel and remain 1.3-2x faster than hipBLASLt. Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

CHUNKK=2 halves the K-shard and doubles the CUs used, but also doubles the split-K shard count. The readback of those partials is serial in the shard count, so the added parallelism stops paying once the per-block main loop no longer dominates. Cap CHUNKK=2 at 11 shards. Measured on MI355X (bf16, CUDA-graph replay to keep launch overhead from swamping these kernels, min of 11 timings of 100 replays), sweeping 78 shapes over N in {32,64,128}, M in {128..1024} and 4-24 shards, each built twice and forced to each CHUNKK: shards mean CHUNKK=2 vs CHUNKK=1 worst 4 -9.6% -14.0% 6 -7.9% -11.2% 8 -2.8% +3.4% 10 -2.3% +2.5% 12 0.0% +3.4% 16 +6.0% +10.1% 20 +8.9% +12.1% 24 +13.3% +18.2% The crossover is not a constant: it falls as M grows, from ~16 shards at M=128 to ~4 at M=1024, since more output tiles read back at once. A flat cap is therefore a compromise. Scoring each cap against an oracle that picks per shape, caps 10 through 13 are indistinguishable (0.21-0.22% off oracle) and form the optimum; cap 4 costs 1.80% and cap 24 costs 3.71%. This picks the low end of that plateau, since overshooting the crossover costs about twice what undershooting it does. N_p2 == 16 is folded into the CHUNKK expression rather than corrected afterwards; the dispatch below already hardcodes CHUNKK=1 for that case, so this is equivalent and keeps the host rule in one place. The Python mirror moves into wvsplitkrc_dispatch() so the dispatch rule has a single Python definition that tests can consume. No new shape is admitted to the kernel: CHUNKK only ever goes 2 -> 1, which halves k_rnd and can only loosen the workspace guard. Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

wvSplitKrc keeps its split-K partials and reduction counters in function-local static tensors, so one pair serves the whole process. They are recycled by protocol rather than re-zeroed -- the last workgroup to arrive resets them -- which is safe only while every caller is serialised on one stream. Two launches on different streams each satisfy the other's completion test and read the other's partials: no error, no crash, silently wrong results. vLLM's shared-experts overlap does exactly that, running the router gate and the shared expert's gate_up_proj concurrently with both admitted to the kernel. Allocate one pre-partitioned pool per device and hand each stream a slot on first use. Allocating per stream on demand would also fix it, but moves the allocation to an uncontrolled moment: the first qualifying GEMM on a new stream can land inside a cudagraph capture, where the zero-fill becomes a replayed graph node. Every kernel access is relative to the base pointers, so a slot is a pointer offset and the kernel is unchanged. This removes the reason for the ROCm multi-stream safety gate, which attributed the corruption to shared and routed experts aliasing their inputs. They do not alias -- the storage pointers differ at that call site on every launch, and a copy inserted there does not help. Quantized routed experts appeared safe because they take a different kernel, not because a copy broke an alias. gsm8k on DeepSeek-V4-Pro with the overlap enabled: 0.92 with this change against 0.001 without it, in the same batch. Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

… config The overlap was restricted to DP-only deployments on ROCm because it did not pay elsewhere. That finding predates two changes. The runtime now gives a graph's cross-stream dependency a device-resident value word, which removes most of the multi-stream penalty that made the overlap unattractive; and until the preceding commit the overlap produced wrong answers on ROCm, so any timing taken with it enabled was measuring corruption rather than the overlap. Neither reason is specific to a parallel configuration, so the gate goes and ROCm takes the same path as CUDA. The token threshold and the aux-stream availability check still apply to both. Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

claude

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.

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

AndreasKaratzas · 2026-09-09T17:08:46Z

/ci run

github-actions · 2026-09-09T17:09:06Z

✅ Triggered Buildkite CI #87959 for commit 390812e79641.

simondanielsson

Nice work!

simondanielsson · 2026-09-09T16:20:39Z

@@ -112,19 +107,11 @@ def _determine_shared_experts_order(
        if self._mk_can_overlap_shared_experts():
            return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED

-        # On ROCm, empirically only DP-only deployments benefit from the overlap.
-        overlap_is_beneficial = not current_platform.is_rocm() or (


Suggestion: Should we re-run the bench serve experiments from #52033 (or similar) and validate that it's in fact beneficial now also for non-DPA cases?

I did a few sweeps over DSV4 and Qwen-3.5-35B at various concurrencies, TP8. And also did a sweep over DSV2-Lite with synthetic shared-expert sizes and various parallelism configs. It's beneficial now in essentially all cases I tested. Updated the PR body with some numbers.

Very good thanks

simondanielsson · 2026-09-09T16:25:45Z

@@ -283,24 +283,11 @@ def __init__(
        self._shared_experts: SharedExperts | None = None
        if shared_experts is not None:
            can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts
-            # When unquantized, shared expert inputs alias the hidden states,


Suggestion: Can we run gsm8k on gfx950 with the following cfg and check everything is still ok?

vllm serve Qwen/Qwen3.5-35B-A3B \ --max-model-len 4096 \ --data-parallel-size 2 \ --enable-expert-parallel \ --trust-remote-code \ --disable-uvicorn-access-log

(this will anyhow be run in nightly ci)

Yep gsm8k checks out at 0.958-0.970 on Qwen3.5-35B, on all the parallelism configs (TP8, TP1+DP8 [with and without EP], TP2+DP4 [with and without EP] I tested. I also tested DeepSeek V4 Pro on TP8 and it comes out at 0.952-0.962.

mawong-amd · 2026-09-09T22:35:58Z

/ci run

github-actions · 2026-09-09T22:36:14Z

❌ @mawong-amd, A reviewer with write access must run /ci run, approve the PR, or add the ready label first.

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

shen-shanshan · 2026-09-10T02:22:26Z

/ci retry

github-actions · 2026-09-10T02:22:44Z

✅ Queued 5 failed job(s) for retry in Buildkite CI #88008.

…time Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

…/wvsplitkrc-improvements

Fangzhou-Ai · 2026-09-10T03:44:54Z

Great work @mawong-amd ! A quick question do we have any kernels available from AITER directly?

…V cache sometimes Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

AndreasKaratzas · 2026-09-10T18:14:37Z

/ci run

github-actions · 2026-09-10T18:14:55Z

✅ Triggered Buildkite CI #88183 for commit 3d530ec760cb.

mawong-amd · 2026-09-10T21:35:41Z

The CI failures are pre-existing in main:
CPU: Multimodal Processor (2)
MI355: Multimodal Processor
H200 MIG 18GB: Multimodal Processor

… fixes (vllm-project#56098) Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

mawong-amd and others added 5 commits September 8, 2026 04:58

Clean up comments

bef23dd

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

mawong-amd requested review from AndreasKaratzas, dllehr-amd, hongxiayang, mgoin, pavanimajety, tjtanaa and zyongye as code owners September 9, 2026 14:43

claude Bot reviewed Sep 9, 2026

View reviewed changes

mergify Bot added rocm Related to AMD ROCm bug Something isn't working labels Sep 9, 2026

github-project-automation Bot added this to AMD Sep 9, 2026

github-project-automation Bot moved this to Todo in AMD Sep 9, 2026

Add multi-stream wvSplitKrc tests

390812e

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

mawong-amd force-pushed the perf/wvsplitkrc-improvements branch from ffe1c08 to 390812e Compare September 9, 2026 14:48

simondanielsson reviewed Sep 9, 2026

View reviewed changes

Harden wvSplitKrc workspace handling during graph capture

1b2fad6

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

mawong-amd force-pushed the perf/wvsplitkrc-improvements branch from bd9ecf9 to 1b2fad6 Compare September 9, 2026 22:37

shen-shanshan self-assigned this Sep 10, 2026

Relax workspace safety based on vLLM only dispatching one graph at a …

3f58d00

…time Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

Merge branch 'main' of https://github.com/vllm-project/vllm into perf…

4b522c2

…/wvsplitkrc-improvements

Empty triton sampler buffer caches during teardown to avoid pinning K…

3d530ec

…V cache sometimes Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>

AndreasKaratzas approved these changes Sep 10, 2026

View reviewed changes

ywang96 merged commit 9163190 into vllm-project:main Sep 10, 2026
171 of 175 checks passed

github-project-automation Bot moved this from Todo to Done in AMD Sep 10, 2026

shen-shanshan mentioned this pull request Sep 11, 2026

[RFC]: DeepSeek-V4 Performance Optimization on ROCm (Phase Two) #52911

Open

17 tasks

djramic mentioned this pull request Sep 12, 2026

[ROCm][Bugfix] Restrict multi-stream shared experts to gfx950 #56630

Closed

Uh oh!

Uh oh!

Conversation

mawong-amd commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Performance numbers

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

AndreasKaratzas commented Sep 9, 2026

Uh oh!

github-actions Bot commented Sep 9, 2026

Uh oh!

simondanielsson left a comment

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

mawong-amd Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

simondanielsson Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

mawong-amd Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mawong-amd commented Sep 9, 2026

Uh oh!

github-actions Bot commented Sep 9, 2026

Uh oh!

shen-shanshan commented Sep 10, 2026

Uh oh!

github-actions Bot commented Sep 10, 2026

Uh oh!

Fangzhou-Ai commented Sep 10, 2026

Uh oh!

AndreasKaratzas commented Sep 10, 2026

Uh oh!

github-actions Bot commented Sep 10, 2026

Uh oh!

mawong-amd commented Sep 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

mawong-amd commented Sep 9, 2026 •

edited

Loading

mawong-amd Sep 9, 2026 •

edited

Loading

mawong-amd Sep 9, 2026 •

edited

Loading