[roberteg16]

Summary

A Qwen MoE block gates its shared expert with sigmoid(shared_expert_gate(x)), where shared_expert_gate is a ReplicatedLinear(hidden_size, 1). During decode with a single token that is a 1x1xK GEMM, and hipBLASLt answers it with a 64x96x32 macro tile plus a SplitK post-pass in order to produce one scalar. Qwen3.5-35B-A3B has 40 such layers, so this fires 40 times per decode step.

This series replaces that shape with a dot product, then folds the surrounding work into the same kernel. On gfx1151 it is worth 3.2% of decode throughput end to end, with prefill and TTFT unchanged.

Everything is gated on on_gfx1151(), the only architecture measured here. gfx9, gfx1100, gfx12xx and every non-ROCm platform keep the code they have today: the series removes exactly one line of existing upstream code, the shared-expert gate multiply in Qwen2MoeMLP. Everything else is additive.

What this touches

The dispatcher change fires for any unquantized bf16/fp16 Linear(K, 1) on a single token, which in this tree means the shared_expert_gate in qwen2_moe, qwen3_moe, qwen3_next, qwen3_5 and interns2_mobius. The only other 1-wide projection, qwen3_dspark's proj, is fp32 and therefore excluded.

The call-site change is narrower: it lives in Qwen2MoeMLP, so it reaches Qwen2-MoE, Qwen3-Next, Qwen3.5-MoE and InternS2 Mobius. qwen3_moe has its own Qwen3MoeMLP whose forward is line-for-line identical; I left it alone because I have not measured it, but folding the two would be the obvious follow-up.

Checkpoints without a shared expert are untouched, as is any deployment running with VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS, where expert_gate is None.

Commits

commit	what it does
Replace the 1x1xK BLAS call with a fused dot on gfx1151	Routes the m == n == 1 shape to (x*w).sum(dtype=x.dtype) instead of BLAS.
Serve the tiny scalar projection from a Triton kernel	Replaces the two-launch eager chain with a single-workgroup fp32-accumulated kernel, for bias=None and K within the register budget. _tiny_dot holds that choice so callers do not repeat it. Adds the correctness test and the A/B benchmark.
Fold the shared_expert_gate sigmoid into the tiny-dot kernel	APPLY_SIGMOID constexpr collapses mul + sum + sigmoid into one launch at the Qwen2MoeMLP call site, and the scalar output is allocated in the input dtype so the store does the rounding rather than a separate copy. Extends the benchmark with --sigmoid.

End-to-end results

Qwen3.5-35B-A3B-W4A16 on gfx1151 (Strix Halo), concurrency 1, 100-token prompts, 128 output tokens, 10 prompts, max_num_seqs=1. Three runs per arm, medians below.

metric	upstream/main	this branch	change
prefill	612 tok/s	627 tok/s	unchanged
TTFT	163 ms	159 ms	unchanged
decode	43.3 tok/s	44.7 tok/s	+3.2%
TPOT	23.11 ms	22.39 ms	−3.1%
end-to-end latency	3098 ms	3003 ms	−3.1%

Per-run decode, showing the spread: upstream/main 43.3 / 43.3 / 43.2, this branch 44.4 / 44.7 / 44.7. Intra-arm spread is about 0.1%, so the 3.2% separation is roughly 30x the noise.

On the measurement methodology, because the first attempt was wrong and it matters. Running the two arms back to back in a fixed order produced a 17% prefill improvement that cannot exist: prefill submits 100 tokens at once, the fast path requires n == 1, so it never executes during prefill. The cause was thermal. Each arm ran immediately after the other on an already-warm GPU, and whichever arm ran second was penalised. Since the branch always ran first, the bias inflated its result.

The numbers above come from a re-run with the order reversed and a cooldown before every run, with all twelve runs starting at 34-35 °C. Under those conditions prefill and TTFT come out equal between arms, which is what the code requires, and the decode difference survives unchanged. Anyone reproducing this should not interleave the arms without a cooldown.

Kernel microbenchmark

benchmarks/kernels/bench_tiny_dot.py, times in µs. Each cell is the median of five separate process invocations, because the BLAS arm is not repeatable enough for a single reading (see below).

K	dtype	BLAS	eager fused	Triton	Triton vs BLAS
1024	bf16	12.60	9.16	5.24	2.40x
1024	fp16	9.98	9.08	4.88	2.05x
2048	bf16	16.12	8.36	5.22	3.09x
2048	fp16	11.28	8.26	4.88	2.31x
4096	bf16	22.94	8.64	5.20	4.41x
4096	fp16	13.88	8.56	4.92	2.82x

The Triton kernel is flat in K at roughly 5 µs, being one launch over a single workgroup, while BLAS grows with K. Qwen3.5-35B-A3B uses K=2048.

With --sigmoid, which is how the gate is actually used, at K=2048 bf16:

	dot only	with sigmoid
BLAS	16.12 µs	18.92 µs
eager fused	8.36 µs	10.88 µs
Triton	5.22 µs	5.26 µs

Folding it into the kernel is what makes the sigmoid free; the other two arms pay another launch for it, so the Triton margin over the eager chain widens from 1.60x to 2.07x.

On the variance, since it decides how these numbers should be read. Across the five invocations the Triton arm holds within 2-7% and the eager arm within 1-2% on most cells, but the BLAS arm swings by up to 49% and in one invocation returned 4471 µs against a median of 18.92. Single readings of the BLAS arm are therefore not reportable, and an earlier revision of this description quoted one such outlier. Anyone reproducing this should take medians across processes rather than trusting one run, do_bench's internal median notwithstanding.

Why K is bounded, and where

Doing the reduction in one workgroup is what makes it one launch: there is no cross-workgroup reduction inside a kernel, so splitting K would cost either a second launch or atomics, and the launch is the thing being eliminated. The price is that register pressure grows with K, and K <= 4096 was inherited from the downstream branch without a measurement behind it.

Reading VGPR counts out of the compiled kernel shows the real boundary is a function of elements per lane alone, identical across different combinations of K and num_warps that share it:

elements/lane	VGPRs	spills	time
32	65	0	5.4-5.7 µs
64	99	0	5.8-7.7 µs
128	175	0	6.3-8.1 µs
256	256	257	30.3-32.1 µs
512	256	829	69.2 µs

The allocator hits the 256-VGPR ceiling at 256 elements per lane and spills, costing a factor of 5. So the kernel now sizes the workgroup to hold each lane at 128 elements, and the bound becomes a property of the hardware: 32 warps per workgroup times 32 lanes times 128 elements, i.e. 131072.

Measured with that heuristic, against the eager fallback:

K	warps	elements/lane	spills	Triton	eager
16384	4	128	0	6.30 µs	9.44 µs
65536	16	128	0	8.04 µs	12.76 µs
131072	32	128	0	11.26 µs	19.62 µs
262144	32	256	343	47.00 µs	25.96 µs

131072 is the last K that fits; at 262144 the warp count is already at the hardware maximum, elements per lane double, and eager wins. Nothing changes for the models in this series, which use K=2048 and sit a factor of 16 below the ceiling; the point is that the limit is now derived rather than guessed, and large-hidden MoEs stop falling back for no reason.

Accuracy

GSM8K, 5-shot, temperature 0, 400 questions per arm, no invalid responses on either side.

arm	accuracy
upstream/main	0.9225 (369/400)
this branch	0.9125 (365/400)

The four-question gap is 0.51 standard errors of the difference, so the two arms are indistinguishable. That is the expected outcome rather than a lucky one: the Triton path accumulates in fp32 and rounds once on the store, which is what (x*w).sum(dtype=x.dtype) does, so it is numerically equivalent to the code it replaces rather than an approximation of it. tests/kernels/test_tiny_dot_triton.py checks that equivalence directly across 24 combinations of K, dtype and sigmoid.

400 rather than the full 1319 because this box sustains about 40 tok/s aggregate, which puts the full set at roughly 100 minutes per arm. At n=400 the standard error on each arm is 1.4%, enough to catch a gross regression and not enough to resolve a sub-1% one.

Testing

pytest tests/kernels/test_tiny_dot_triton.py          24 passed
pre-commit run --files <the four changed files>       clean (ruff, ruff-format, mypy, SPDX)

The new test compares the kernel against the eager reference across K = 32 / 1024 / 2048 / 4096 / 65536 / 131072, both dtypes, with and without the folded sigmoid. The two largest K exercise the multi-warp launches.

Software stack

Measured on ROCm 10.1.0a20260822 (torch.version.hip 7.16), PyTorch 2.13.0, Triton 3.8.0. This is newer than what docker/Dockerfile.rocm_base pins for the shipped image (ROCm 7.2.3, PyTorch 2.12, Triton 3.7), so the ratios against BLAS are specific to this stack, hipBLASLt being exactly the component that moves between them. The kernel itself uses only long-stable Triton primitives (tl.arange, tl.load, tl.sum, tl.exp, tl.store) and carries no version requirement. gfx1151 is a shipped target in that image's PYTORCH_ROCM_ARCH, but no gfx11 runner exists in CI, so these numbers cannot be reproduced by CI.

Not a duplicate

• [ROCm][Perf] Route the fused shared-expert gate GEMM through the platform dispatcher vllm-project/vllm#54185 touches the same conceptual area and is complementary, not overlapping. It routes the fused shared-expert gate through dispatch_unquantized_gemm() in moe_runner.py; that path concatenates the router gate with the shared gate into a [num_experts + 1, hidden] weight, so it is an m=513 GEMM. This series handles the unfused path, where the gate is a standalone Linear(hidden, 1) and m == 1. Different file, disjoint conditions. When VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS is on, Qwen2MoeMLP.expert_gate is None and none of this runs at all.
• [ROCm][Perf] Encoder prefill attention on gfx1151: split-D head dims, alignment hint, launch config vllm-project/vllm#56450 is mine and also targets gfx115x. An earlier revision of this branch shared its on_gfx115x() predicate commit; that commit was dropped here in favour of the existing on_gfx1151(), so the two PRs no longer touch a common line. gfx1151 is also the only part of gfx115x measured in this series.

AI assistance

This change was prepared with AI assistance (Claude). Every line was reviewed by the submitter, and the measurements above were run and checked by hand.

[roberteg16]

Upstream PR: vllm-project#56276