[Perf] Re-enable dual-stream input projection for Qwen3/Qwen3.5 GDN#39748
[Perf] Re-enable dual-stream input projection for Qwen3/Qwen3.5 GDN#39748jhsmith409 wants to merge 1 commit intovllm-project:mainfrom
Conversation
jhsmith409
requested review from
ZJY0516,
tdoublep and
vadiklyutiy
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces dual-stream parallel execution for input projections in the GDN linear attention layer to improve performance and torch.compile compatibility. A bug was identified in the initialization of synchronization events where the platform check for CUDA might cause failures on ROCm systems; it is recommended to check for the existence of the auxiliary stream instead and disable event timing to reduce overhead.
| self.events = ( | ||
| [torch.cuda.Event(), torch.cuda.Event()] | ||
| if current_platform.is_cuda() | ||
| else [None, None] | ||
| ) |
There was a problem hiding this comment.
The initialization of self.events uses current_platform.is_cuda(), which returns False on ROCm (HIP) platforms. However, self.aux_stream is initialized using aux_stream(), which returns a valid stream on ROCm (as it checks is_cuda_alike()). This mismatch will cause a crash in maybe_execute_in_parallel when it attempts to call .record() on a None event.
Additionally, for synchronization events where timing is not required, it is recommended to use enable_timing=False to reduce overhead.
| self.events = ( | |
| [torch.cuda.Event(), torch.cuda.Event()] | |
| if current_platform.is_cuda() | |
| else [None, None] | |
| ) | |
| self.events = ( | |
| [torch.cuda.Event(enable_timing=False), | |
| torch.cuda.Event(enable_timing=False)] | |
| if self.aux_stream is not None else [None, None] | |
| ) |
|
cc @xyang16 @zou3519 — this uses the |
Cold-compile timing (response to #38152's concern)Measured on the same RTX PRO 6000 Blackwell setup,
Difference is ~2 s (≈1.8 %), smaller than the run-to-run noise (~25 s range within each config). The original concern in #38152 was a ~4× regression — that would have been ~400-500 s here, which is clearly not what we're seeing. The Torch version in image: Happy to rerun with more iterations or a different test harness if reviewers want tighter bounds. |
jhsmith409
force-pushed
the
reenable-qwen3-dual-stream
branch
from
73bed11 to
5db5195
Compare
TTFT (streaming, short prompts)Same hardware (RTX PRO 6000 Blackwell,
TTFT is ~1.5 ms higher with the patch — worth calling out honestly. At 8 prompt tokens the two input projections are tiny, and the Net picture:
If reviewers want TTFT to be neutral or positive, one option is to gate the dual-stream dispatch on a token-count threshold (only use aux_stream when |
jhsmith409
force-pushed
the
reenable-qwen3-dual-stream
branch
from
5db5195 to
ed41543
Compare
|
Applied both suggestions in Thanks for the review. |
|
cc @tdoublep @ZJY0516 @vadiklyutiy as CODEOWNERS for |
Output equivalence check (re: correctness of the stream dispatch)5 fixed prompts,
The three long-output prompts also differ across two repeat calls to the same unchanged patched server — prompts 1/2/4 produce different hashes across runs with identical inputs. Short deterministic prompts (3, 5) match byte-for-byte both within- and across-config. Root cause of the drift is not this PR but the serving config ( The two short-prompt matches confirm that when the server is deterministic, the patch preserves outputs exactly — which is what you'd expect: the dual-stream dispatch changes when the two projections run (parallel vs sequential) but not what they compute. Happy to rerun with a deterministic server config if reviewers want a tighter bound. |
|
I remember we want to wait for torch 2.12 which can hadle multi stream so we dodn't need to wrap it as a custom op |
Update on the equivalence checkTried to tighten the bound by serving with:
Two consecutive calls to the same unchanged server, identical inputs:
So even with the most invariant config vLLM offers on this model, longer completions are still non-deterministic on Qwen3.5-35B-A3B (fp8 KV, compressed-tensors AWQ, hybrid GDN+MoE). Likely sources are the MoE router and chunked-prefill chunking boundaries — neither of which this PR touches. The short-output prompts that are byte-reproducible on the base server stay byte-reproducible with the patch (see my earlier equivalence comment). Happy to pursue this further if a reviewer points at a config that does produce byte-exact output for this model. Otherwise this is where I'll leave the equivalence evidence — the two-short-prompts-match + the mechanical argument (dual-stream changes when the two projections run, not what they compute) is the strongest claim I can defend without a deterministic base to compare against. |
|
@ZJY0516 thanks for the context — I hadn't realized torch 2.12's native multi-stream handling was the intended path. A clarifying question: is torch 2.12 expected to land soon enough that a few weeks of the current ~8 % Qwen3.5/Qwen3-Next decode regression is preferable to an interim custom-op? For what it's worth, I'm personally fine running a local overlay until 2.12 is available — so if you'd rather not carry this code, I'm happy to withdraw. But I'm also happy to keep iterating on this PR so other users of the nightly don't have to hold the regression. If we go that route I can add an explicit `# TODO(torch>=2.12): remove this custom op and call self.in_proj_qkvz / self.in_proj_ba directly once torch.compile handles parallel CUDA streams natively` marker on both the registration and the call site so cleanup after 2.12 is a quick grep. Which would you prefer? |
Independent bench on a different config — directionally consistentRe-ran the A/B on RTX 5090 (sm_120, 32 GB) + Baseline image:
Smaller magnitude than the PR's +8 % PRO 6000 / AWQ / fp8 result but sign-consistent, and the curve makes sense — at par=1 a single batch of
(AI-assisted verification run. Human submitter reviewed the overlay and ran both A/B configurations.) |
|
Hi @jhsmith409, running a variant of this PR as a monkey-patch locally on A5000 + Qwen3.6-35B-A3B-FP8 — wanted to flag one thing I ran into. (Small disclaimer: I'm from Ukraine and my English is still a work in progress, so I'm using AI to help with translation. Hope it reads okay!) For the fake-tensor sizes in In my patch I pre-compute the per-TP output sizes at # at init, once
_qkvz_per_tp = (
self.head_k_dim * self.num_k_heads +
self.head_v_dim * self.num_v_heads * 2
) // self.tp_size
_ba_per_tp = 2 * self.num_v_heads // self.tp_sizeand use those in the fake impls. This is robust against weight repacking regardless of backend (Marlin FP8, AWQ, GPTQ) because it never touches Happy to put together a small delta on top of this PR if you'd like — let me know. Otherwise just flagging in case someone hits it later. Performance-wise your PR's claim holds on A5000 too — I see a similar +4.x% decode improvement on our 2× A5000 TP=2 setup vs. the no-dual-stream baseline (though my measurement is the monkey-patch, not a clean cherry-pick of this exact PR). |
|
@jhsmith409 sorry, but here posted a lot of AI generated text. Pls make concise and informative description |
The short description is that if you merge this PR to re-enable dual stream input processing, you get a ~10% gain in tokens/sec. I understand the next release of pytorch will have support for implementing this a different way. In the mean time I wanted the 15% back. I'm using it as a patch, easier for me if you merge it back into main. You can take it back out whenever someone writes a new method using the next version of pytorch. Or @vadiklyutiy are you talking about a different PR like the OOM PR on JartX's repo? I referenced this patch in that PR so it was cross-posted here. |
I can't find this info in PR descr |
@vadiklyutiy , thanks for looking at this PR. You are right I did not specify the gain, instead I gave a table. More difficult to interpret that way. On the first post in this PR, it shows token generation dropping from 198 to 183 (which is an 8% drop not 10%) after March 23 (and PR #38152 merge to the nightly). This dropped dual stream processing and slowed token generation. One of the maintainers commented that they wanted to wait til Pytorch 2.12 dropped then fix this a different way. I wanted the token speed back for one of my applications that is running about an hour on this model so I created a PR (or a patch depending on how you look at it) that puts that speed back (which is an 10% gain). Funny how an 8% loss becomes a 10% gain but that's because the base reference changes from close to 200 to 180 (approximate numbers). See the chart in the first entry of this PR for the drop in throughput. |
|
BTW, vllm PR#38123 — introduced the LayerName opaque type that lets the dual-stream dispatch co-exist with |
|
Heads-up for anyone applying this as an out-of-tree overlay (i.e patch) rather than as a merged upstream as part of a PR. Notes below not relevant to PR itself: |
|
pls provide command lines for perf testing and result with and without this PR |
Here is the docker-compose.yml used to run the model on an RTX 6000 Pro Blackwell QMax. It uses last night's cu130 build of vllm. Simply comment out the patch to test the version without the PR implemented. (Forced load of Transformers to prevent race condition on patching.): ***docker-compose.yml This is the benchmark used: #!/usr/bin/env python3 import aiohttp HOST = "http://localhost:8006" def _build_payload(animal: str) -> dict: async def _run_one(session: aiohttp.ClientSession, animal: str) -> dict: async def _sweep(n: int) -> dict: async def _warmup() -> None: async def main() -> None: if name == "main": These are the results. Gains vary. Longer generation sequences typically see gains as high as 8%.: Qwen3.5-122B-A10B-AWQ-4bit · cu130-nightly commit gb47840019 · max-num-seqs=16 · 512 decode tokens · Dual-stream ON (patched): ┌─────┬──────────┬──────────┬──────────────────────┬────────────────┐ Dual-stream OFF (upstream): ┌─────┬──────────┬──────────┬──────────────────────┬────────────────┐ Delta (ON vs OFF): ┌─────┬────────────────┬──────────────────┬────────┐ Overlay wins on per-request decode at every concurrency (1.8%–6.2%, strongest at N=2–4), matches or |
vadiklyutiy
left a comment
vadiklyutiy
left a comment
There was a problem hiding this comment.
@jhsmith409 did you try to read your post by yourself? Do you think is it proper way to write?
|
Results for Qwen36-35B in awq quant from cyankiwi on RTX 5090. Same benchmark, same docker compose setup. Similar results. WITH dual-stream overlay WITHOUT (stock upstream) Per-request decode tok/s delta (WITH − WITHOUT) / WITHOUT:
Decode-throughput signal is consistent with the handoff's expected range (2–6%, strongest at N=2–4, Caveat on the aggregate / TTFT numbers: The WITH run shows anomalously high TTFT at N=4 and N=8 |
I'm sorry you have a problem with the way I write. I relied more on AI in prior posts and I think that writing is more udnerstandable, but you asked me to write it myself. In any case, you are welcome to implement this or to close it out. The patch is working for me and I look forward to someone else picking up the task of re-writing the multi stream process when the next version of pytorch is released. Thank you for your time in looking into this. |
adequate formatting is mandatory. Pls look on tables in last 2 messages. It is very hard to understand the results. |
|
Apologies for the mess on the recent perf-data post. I owe you a clearer explanation. When you said "posted a lot of AI generated text. Pls make concise and informative description" earlier in this thread, I read that as "don't use AI to write the post" and tried to follow it literally. I don't actually know how to convert benchmark-script output into proper Markdown tables by hand, so I dumped the raw column-aligned text — which I now agree looks awful as a comment. I was overcorrecting away from AI when you'd really only flagged AI-generated prose, not formatting. So this time I asked an AI just to convert the same numbers I already collected into proper Markdown tables — no analysis, no rewriting, only RTX 5090 — Qwen3.5-35B-A3B (cyankiwi AWQ)Same benchmark and docker-compose I posted above; this is the 35B variant on a single RTX 5090. WITH dual-stream overlay (this PR)
WITHOUT (stock upstream)
Per-request decode tok/s delta (WITH − WITHOUT)
Decode-throughput signal is in the handoff's expected 2–6 % range, strongest at N=2–4. Sign-consistent at every sweep point, no regression on the per-request decode metric. Caveat on the aggregate / TTFT numbersThe WITH row at N=4 and N=8 shows anomalously high TTFT (429 ms / 395 ms vs 53 ms / 95 ms on WITHOUT). That's a scheduler-warmup artifact from running only one warmup-per-sweep — the first batches after cold start catch a cudagraph capture / prefill queue backlog rather than anything attributable to the overlay. The aggregate Happy to gather more data if anything else would help. |
|
What is the shape of weight of |
|
They're That said, my understanding from earlier in this thread (@ZJY0516's comment) was that the team plans to rewrite this path once PyTorch 2.12 lands — which on the current schedule is roughly two weeks away — because 2.12's native multi-stream support removes the need for any custom-op wrapper here at all. The intent of this PR is narrower than that: it's a straight revert to dual-path execution to recover the ~5–8% throughput that was lost when #38152 disabled it to work around a torch.compile deduplication bug. That bug is now fixed upstream by #38123 ( Happy to flag both the registration and the call site with |
Benchmark results — Qwen3.5-122B (PRO 6000 Blackwell) and Qwen3.6-35B (RTX 5090)A/B measured on 2026-04-30 against vLLM v0.20.0 stable (cu130). Patch applied as a file overlay built by running the PR diff (head SHA Hardware / software
Common: vLLM v0.20.0, torch 2.11.0+cu130, transformers 5.6.2, Python 3.12.13. Results — Qwen3.5-122B-A10B-AWQ on RTX PRO 6000 Blackwell
Results — Qwen3.6-35B-A3B-AWQ on RTX 5090
Side-by-side
InterpretationBoth model+GPU combinations show the same shape (peak gain at mid-concurrency, falling off at very low N and at saturation), but the magnitude and width of the sweet spot differ. The 35B/5090 combination wins at every concurrency, peaks at +6.7%, and holds +5.6% gain out to N=8; the 122B/PRO 6000 combination is narrower (only N=2–4 see meaningful gains) and lower amplitude. The pattern is consistent with the underlying mechanism: dual-stream overlap of No regressions observed at any concurrency on either model. Output token counts and content match between unpatched and patched runs (the change is a parallelism rearrangement, not a math change). Notes for reviewers
Benchmark sourceIdentical bash script used for both models, parameterized by
|
Restores the parallel execution of in_proj_qkvz and in_proj_ba that was introduced in vllm-project#36795 and reverted in vllm-project#38152. The revert was necessary because the original custom op took the layer name as a string, which caused torch.compile to produce one compiled artifact per GDN layer (~4x cold-compile regression). vllm-project#38123 added the LayerName/OpaqueObject mechanism so custom ops can accept layer-name arguments without baking them into the graph. This PR rewires the gdn_in_proj custom op to use LayerName via the existing _encode_layer_name / _resolve_layer_name helpers, matching the pattern already used by mamba_mixer2, gdn_attention_core, and other mamba ops. On the Qwen3.5-35B-A3B decode path the two projections can now overlap on aux_stream again; the serial fallback is preserved for non-CUDA platforms (XPU forward path unchanged) and when aux_stream() returns None. Signed-off-by: Jim Smith <jim@joshua8.ai> Co-authored-by: Claude
jhsmith409
force-pushed
the
reenable-qwen3-dual-stream
branch
from
70b4c7c to
ea81e43
Compare
No, I think he didn't say it |
I think this is the right approach that should be implemented. |
|
Superseded by #41457 — "[Perf] Fuse Qwen3.5 GDN in_proj_ba into 6-way in_proj MergedColumnParallelLinear." Per @vadiklyutiy's suggestion, I prototyped the fused-
It's just one bigger GEMM and a Apples-to-apples on v0.20.0 + RTX 5090 + Qwen3.6-35B-A3B-AWQ:
Closing this PR. Discussion continues on #41457. |
|
Closing in favor of #41457. |
…allelLinear Replaces the prior `in_proj_qkvz` (4-way merge: q, k, v, z) plus separate `in_proj_ba` (2-way merge: b, a) with a single 6-way `in_proj` MergedColumnParallelLinear of output_sizes=[k, k, v, v, num_v_heads, num_v_heads]. The in_proj_ba projection has output dim ~0.5% of in_proj_qkvz (e.g. 64 vs 12288 on Qwen3.5-35B-A3B), so concatenating along the output axis costs essentially nothing extra and lets a single GEMM cover both projections. This recovers the projection-overlap throughput win that vllm-project#36795 chased with dual-stream dispatch (and that vllm-project#38152 had to disable due to a torch.compile per-layer subgraph regression) without adding any custom-op wrapper, CUDA event synchronization, or `LayerName` / `OpaqueObject` plumbing. The mechanism stays vanilla `MergedColumnParallelLinear` -> single matmul -> `torch.split`. Approach suggested by @vadiklyutiy on vllm-project#39748. Supersedes vllm-project#39748. Layout matrix (handled by the existing `gqa_interleaved_layout` / `create_in_proj_qkvz` flags plus the new `fuse_in_proj_ba` flag): - Qwen3.5 non-LoRA (this PR's hot path): single 6-way `in_proj`. - Qwen3.5 LoRA: in_proj_qkv + in_proj_z + in_proj_ba kept separate so adapters attach independently. Unchanged. - Qwen3-Next (gqa_interleaved_layout=True): unchanged dual-module path. The 6-way fuse is gated on non-interleaved layout. Benchmarks (RTX PRO 6000 Blackwell + Qwen3.5-35B-A3B-AWQ-4bit on cu130-nightly, identical config: --gpu-memory-utilization 0.5 --max-num-seqs 8). Decode tok/s, 1024-token completions: | Config | Seq | Par 2 | Par 4 | Par 8 | |---------------------|----:|------:|------:|-------:| | stock latest | (regression vs mar23 -- see vllm-project#39748) | | vllm-project#39748 dual-stream | 197.9 | 342.2 | 654.9 | 1123.1 | | this PR (fused 6-way) | 197.7 | 343.0 | 659.0 | 1100.8 | | mar23 + vllm-project#37700 ref | 198.3 | 342.3 | 653.9 | 1119.1 | Apples-to-apples on RTX 5090 + Qwen3.6-35B-A3B-AWQ + v0.20.0 image (stock vs dual-stream vs fused all freshly measured today): | Config | Seq | Par 2 | Par 4 | Par 8 | |----------------|------:|------:|------:|-------:| | stock | 209.7 | 317.7 | 617.5 | 1083.7 | | dual-stream | 212.4 | 337.2 | 652.4 | 1133.8 | | fused 6-way | 208.4 | 324.5 | 661.9 | 1129.2 | Fused is within run-to-run noise of dual-stream on every working point on both GPUs and both models, with the cleanup of removing all custom-op / event / LayerName machinery. AWQ packing is not affected: the GDN input projections are in the cyankiwi AWQ checkpoints' `ignore` list, so they are stored as bf16 and concatenation along the output axis is mechanically a no-op for the quant config. Test plan: - Loaded `cyankiwi/Qwen3.5-35B-A3B-AWQ-4bit` end-to-end on cu130-nightly, weights loaded in 11.7s / 22.39 GiB; all 79 GDN layers' `in_proj` parameters populated correctly via the new `("in_proj", "in_proj_qkv|z|b|a", shard_id)` stacked_params_mapping. - Loaded `cyankiwi/Qwen3.6-35B-A3B-AWQ-4bit` end-to-end on v0.20.0 + RTX 5090; same path. - Generation deterministic on temperature=0 seed=0 short prompts: `"List the first ten primes, comma separated."` -> `"2, 3, 5, 7, 11, 13, 17, 19, 23, 29"` byte-identical across runs. - Throughput sweep above (1, 2, 4, 8 concurrent decode streams). - LoRA path is unchanged: `create_in_proj_qkvz=False` still produces separate `in_proj_qkv` / `in_proj_z` / `in_proj_ba` modules with the prior `packed_modules_mapping` and `stacked_params_mapping` entries, exercised via the `if self.enable_lora` branch. AI assistance was used in drafting the patch and benchmark harness; the human submitter (@jhsmith409) reviewed every changed line, ran the tests above, and is accountable for the change. Signed-off-by: Jim Smith <jim@joshua8.ai>
5 participants
Re-enable dual-stream execution of input projection for Qwen3/Qwen3.5 GDN
Purpose
Restores the parallel execution of
in_proj_qkvzandin_proj_bathat wasintroduced in #36795 and reverted in #38152. The revert was necessary because
the original
gdn_in_projcustom op took the layer name as a string, whichcaused torch.compile to produce one compiled artifact per GDN layer (~4×
cold-compile regression, see #38152's description).
#38123 added the
LayerName/OpaqueObjectmechanism precisely so custom opscan accept layer-name arguments without baking them into the graph. This PR
rewires the
gdn_in_projcustom op to useLayerNamevia the existing_encode_layer_name/_resolve_layer_namehelpers — matching the patternalready in use by
mamba_mixer2,gdn_attention_core, and other mamba opsthat were updated in #38123.
The serial fallback path is preserved when
aux_stream()returnsNone(non-CUDA platforms);
forward_xpuis unchanged.Why this is not duplicating an existing PR
Checked on 2026-04-13:
No open PR re-enables dual-stream for Qwen3/Qwen3.5. #38123's TODO in #38152
is still outstanding.
Test plan
Tested on a workstation with an RTX PRO 6000 Blackwell Max-Q (SM 12.0, 97 GB),
vllm/vllm-openai:cu130-nightlyimage (vllm0.19.1rc1.dev231+g9dd5ee011), againstcyankiwi/Qwen3.5-35B-A3B-AWQ-4bitwith
--quantization compressed-tensors --kv-cache-dtype fp8 --enable-prefix-caching --max-num-seqs 8 --gpu-memory-utilization 0.5.Three server configurations, GPU held constant, only image/overlay
changing:
cu130-nightly-mar23(pre-#38152 regression)cu130-nightly(current, #38152 in effect)cu130-nightly+ this PRAll figures in tokens/second. Measured with an 8-request rolling decode
workload (1024 tokens/request, unique prompts to bypass prefix cache).
Regression suite (
vllm_tests.py, 36 tests covering chat completions, toolcalls, vision, logprobs, streaming, determinism, max_tokens, error handling):
32 passed / 4 failed. The 4 failures are all in the "reasoning separated"
family and reproduce on the unpatched baseline — they are driven by
enable_thinking=Falsein the chat-template kwargs of this server config,not by this PR.
Summary of the diff
vllm/model_executor/layers/mamba/gdn_linear_attn.py(+68 / -2)maybe_execute_in_parallelandaux_stream.self.aux_streamand twotorch.cuda.Events inGatedDeltaNetAttention.__init__(no-op on non-CUDA).in_proj_qkvz+in_proj_bacalls with atorch.ops.vllm.gdn_in_proj(...)call using_encode_layer_name._forward_in_projmethod that dispatches viamaybe_execute_in_parallelonaux_stream.gdn_in_projandgdn_in_proj_fakecustom ops withlayer_name: LayerNameType(mirrors the existinggdn_attention_coreregistration pattern).
XPU path (
forward_xpu) and the LoRA fallback path (in_proj_qkv+in_proj_z+in_proj_ba) are unchanged.AI-assisted contribution disclosure
Per AGENTS.md: Claude (Anthropic) assisted with writing this PR. The human
submitter has reviewed every changed line, understands the LayerName
plumbing and multi-stream-utils mechanism, and can defend the change
end-to-end. Co-authored-by trailer is included in the commit.