Skip to content

Qwen3.5 moe optimizations#113

Merged
zichen-kuuga merged 16 commits into
MooreThreads:v0.24.0-devfrom
zichen-kuuga:qwen35_moe-merged
Jul 21, 2026
Merged

Qwen3.5 moe optimizations#113
zichen-kuuga merged 16 commits into
MooreThreads:v0.24.0-devfrom
zichen-kuuga:qwen35_moe-merged

Conversation

@zichen-kuuga

@zichen-kuuga zichen-kuuga commented Jul 21, 2026

Copy link
Copy Markdown
Collaborator

These 15 commits optimize two hot subsystems for Qwen3.5-35B-A3B — GDN (the gated-delta-net linear attention layer) and MoE — plus KV-cache handling and build fixes.

GDN linear attention

373565d tilelang gated RMSNorm
The GDN output gating (RMSNorm + SiLU gate) was a chain of elementwise kernels; replaced with a single fused tilelang kernel.

7fa8fe4 fused zba split
The GDN forward splits z/b/a out of the packed projection. Replaces split + reshape + two contiguous copies with one kernel that slices contiguous z/b/a directly.

9e61f28 mate decode
Single-token GDN decode now runs mate's fp32 VK recurrent kernel (was the slower Triton fused-recurrent update), with automatic fallback.

f046133 tilelang prefill causal conv1d
The causal 1D convolution in GDN (prefill) moved to a tilelang kernel.

98a253d strided q/k/v + fused gating
Feeds mate strided q/k/v views instead of materializing contiguous copies (zero-copy), and fuses the log-space GDN gating (g, beta) into one kernel — replacing ~9 elementwise kernels (ported from SGLang).

b7d72d4 drop D2H + 3D→2D
Removes a per-call device→host sync in conv1d (reads a CPU-side value already carried on the metadata); routes the q/k GemmaRMSNorm through the kernel for 3D inputs (which previously fell back to the slow path).

fd9ad37 mate output buffer + int32
Lets mate write straight into the caller's buffer, saving one D2D copy; switches cu_seqlens from int64 to int32 (mate specializes on int32 internally).

MoE

b03acb0 generalize DeepGEMM MoE prefill
Extends the DeepGEMM fused-prefill path from DeepSeek-V4-only to any qualifying block-FP8 MoE ( scales, etc.). Default-on, gated by VLLM_MUSA_MOE_DEEPGEMM_PREFILL. This is the main TTFTlever.

8af8300 shared-expert fold
Qwen3.5's shared expert ran as a separate serial MLP, adding a chain of kernels per layer. Folds it into the routed grouped GEMM as an always-selected extra expert (slot 256), producing the result in one kernel. Handles both unquantized and FP8 weights; disabled under expert parallelism.

1f369d9 row-tiled fused silu+mul+FP8 quant
Fuses the MoE silu_and_mul and per-token-group FP8 quantization into one row-tiled kernel.

5f8dcf6 per-token-group vec quant
group-128 bf16/fp16 activation quant now uses a register-resident vec kernel (was routing through the slower upstream path); caches a dummy tensor to avoid a per-call torch.empty.

KV cache

6558174 + a4b544c separate mamba state pool
Gives the mamba (GDN) state its own contiguous BlockPool with a private block-id space, so mate's GDN decode updates the recurrent state in place (no gather/scatter). This also fixes an overlap bug where attention KV and mamba state shared one memory tensor (garbage decode output under high concurrency), so it is defaulted on.

QK-norm + positional encoding

7732461 fused QK-RMSNorm + MRoPE
For interleaved-MRoPE gated attention, fuses QK-RMSNorm and MRoPE into one kernel. These modelsry path and emit a long tail of elementwise kernels. Dispatched ahead of the IR op, which staysas the fallback. Gated by VLLM_MUSA_FUSED_QK_MROPE.

Build fixes

a90b1fc rebase two build patches
0060 and 0085 no longer matched their surrounding series after upstream changes, breaking a clean install; their context was re-aligned.

Perf test

Qwen3.5-35B-A3B-FP8 · TP1 · FULL_DECODE_ONLY · 8 prompts (~2.5k in / 1.5k out)

Metric Before After Δ
Output token throughput (tok/s) 208.53 396.12 +90.0%
Total token throughput (tok/s) 557.51 1059.05 +90.0%
Benchmark duration (s) 57.55 30.29 −47.4%
Mean TTFT (ms) 913.54 536.78 −41.2%
Mean TPOT (ms) 37.76 19.84 −47.5%
Mean ITL (ms) 37.73 19.82 −47.5%
Raw benchmark output

Before

============ Serving Benchmark Result ============
Successful requests:                     8
Failed requests:                         0
Maximum request concurrency:             64
Benchmark duration (s):                  57.55
Total input tokens:                      20083
Total generated tokens:                  12000
Request throughput (req/s):              0.14
Output token throughput (tok/s):         208.53
Peak output token throughput (tok/s):    216.00
Peak concurrent requests:                8.00
Total token throughput (tok/s):          557.51
Mean TTFT (ms):                          913.54
Median TTFT (ms):                        957.80
P99 TTFT (ms):                           1133.53
Mean TPOT (ms):                          37.76
Median TPOT (ms):                        37.73
P99 TPOT (ms):                           38.03
Mean ITL (ms):                           37.73
Median ITL (ms):                         37.45
P99 ITL (ms):                            41.27
==================================================

After

============ Serving Benchmark Result ============
Successful requests:                     8
Failed requests:                         0
Maximum request concurrency:             64
Benchmark duration (s):                  30.29
Total input tokens:                      20083
Total generated tokens:                  12000
Request throughput (req/s):              0.26
Output token throughput (tok/s):         396.12
Peak output token throughput (tok/s):    448.00
Peak concurrent requests:                8.00
Total token throughput (tok/s):          1059.05
Mean TTFT (ms):                          536.78
Median TTFT (ms):                        557.22
P99 TTFT (ms):                           680.83
Mean TPOT (ms):                          19.84
Median TPOT (ms):                        19.82
P99 TPOT (ms):                           19.99
Mean ITL (ms):                           19.82
Median ITL (ms):                         19.92
P99 ITL (ms):                            20.72
==================================================

(cherry picked from commit 58d9e4a)
(cherry picked from commit 8b463e9)
(cherry picked from commit e947cdb)
Series patch renumbered 0086 -> 0081; v0.24.0-dev already uses 0086.
(cherry picked from commit 088e0fb)
Keep a default-on VLLM_MUSA_MOE_DEEPGEMM_PREFILL kill switch so the
generalized prefill path can be turned off without a code change.
(cherry picked from commit 7312722)
Widen the existing fold gate to accept FP8 [128,128] block weights in
addition to unquantized, keep it off under expert parallelism / EPLB,
and add the VLLM_MUSA_MOE_SHARED_EXPERT_FUSION switch.
…ken_group_fp8_quant

(cherry picked from commit 44268a4)
(cherry picked from commit 7e8c9b3)
Only the dedicated BlockPool half: the contiguous per-state striding it
also carried is already unconditional in the gpu_model_runner patch.
Regenerated the kv_cache_utils patch, whose upstream copy was truncated.
Sharing one KV tensor between a mamba layer and an attention layer is only
safe while both index it as block_id * page_size. The segregated per-state
mamba layout breaks that congruence, so without the dedicated pool the two
overlap and decode output degrades past ~9 concurrent requests.

Verified on Qwen3.5-35B-A3B-FP8, tp1, 16 capture sizes: shared pool is clean
at 4 and 8 concurrent and garbage at 16; the dedicated pool is 16/16 clean.
(cherry picked from commit 8eae317)
Series patch renumbered 0090 -> 0085. Also aligns the GDN decode env
default with the kv-cache one so the dedicated pool takes the in-place
path instead of falling back to gather/scatter.
0060 carried pre-MooreThreads#108 context for the PrivateUse1 registration block that
0044 now emits two lines longer, so the final hunk no longer matched.

0085 re-applied the is_musa() branch that 0012 already adds to the
per-token-group quant dispatch; it now only introduces the vec fast path.

Both were caught by running vllm_musa/patches/build_apply.py, which is what
setup.py invokes; a plain three-way apply absorbs the duplicate and hides it.
Ports the MUSA kernel from the qwen35_moe branch. The IR provider registered
for gated_qkv_rms_norm_rope only enters the priority list under Inductor and
excludes routed MoE, so Qwen3.5 ran the eager rotary path; this dispatches the
fused kernel ahead of the IR call and leaves it as the fallback.

Measured on Qwen3.5-35B-A3B-FP8 tp1, 2500/1500 x16 @ concurrency 64:
output throughput 509.45 -> 557.54 tok/s, mean TPOT 30.65 -> 27.99 ms.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant