Skip to content

perf(musa): optimize Qwen2 decode for CosyVoice3#117

Merged
yeahdongcn merged 5 commits into
v0.24.0-devfrom
perf/musa-60029-qwen2-llm
Jul 24, 2026
Merged

perf(musa): optimize Qwen2 decode for CosyVoice3#117
yeahdongcn merged 5 commits into
v0.24.0-devfrom
perf/musa-60029-qwen2-llm

Conversation

@yeahdongcn

@yeahdongcn yeahdongcn commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • tune the MP31 small-hidden add-RMSNorm launch used by the Qwen2 talker;
  • fuse the exact Qwen2 BF16 RoPE + NHD KV-cache update chain into one MUSA
    Triton launch;
  • apply the graph rewrite atomically before vLLM splitting so layer-name
    strings remain outside Inductor graphs and compiled pieces stay reusable;
  • keep the upstream KV split and all unsupported model, graph, dtype, prefill,
    and tensor-layout paths on the existing implementation.

This PR changes vLLM-MUSA only. vLLM-Omni is an integration test path, not an
implementation dependency or patch target.

Why the rewrite is pre-split

The first post-grad prototype fused correctly but caused about 95.6 seconds of
cold compilation because 24 layer-name strings produced non-reusable pieces.
The accepted raw-FX rewrite requires all 24 exact sites to match before any
mutation, adds the fused outer op as a split, and retains the original KV split
as a fail-closed fallback.

The accepted path restores compile plus initial warmup to 5.90 seconds on the
development container and produces only four reusable compile artifacts. It is
exact-gated and default-on; no serving flag is required. The environment switch
is retained only as an A/B and emergency-disable escape hatch.

Performance

Single S5000, BF16, FlashAttention-3, normal vLLM compile/FULL graph capture,
357 input tokens, fixed 1024 output tokens (40.96-second audio proxy),
temperature 1.0, repetition penalty 1.0001. Two alternating same-host rounds,
three measured requests per cell after one warmup:

Variant Runs Mean RTF Median RTF
RoPE+KV fusion 6 0.141548 0.141559
baseline on the same branch 6 0.145836 0.145760

Combined mean gain: 2.94%; combined median gain: 2.88%. Both paired rounds are
positive. A deterministic token-identical A/B improves mean RTF from 0.077270
to 0.067954 (12.06%), isolating launch-heavy LLM decode from sampling cost.

A candidate MUPTI activity trace confirms the production path retains all
three native BF16 TCE GEMM families with the expected 2:1:1 call ratio. The
accepted fused kernel averages 4.445 us over the captured startup, graph,
warmup, and measured-request scope. Generated Inductor artifacts use external
addmm/mm kernels and do not Triton-lower the GEMMs.

Validation

  • focused FX/backend/config/real-S5000 kernel suite: 19 passed;
  • default-on integration: matched 24/24 layers, exercised a 357-token prefill
    fallback, and completed 256 deterministic decode tokens with checksum
    953984;
  • deterministic 1024-token candidate and baseline: identical checksum 15219200
    in all six measured requests;
  • broader Python slice: 117 passed, 2 skipped, 3 failed; the same three
    unrelated FP8-MoE/format-patch failures reproduce from the pre-overlay source
    archive;
  • formatting and lint pass for all new/changed files except three unchanged
    baseline findings in flash_attn.py.

Not run / Draft status

  • full synthesized-waveform listening and quality validation belongs to the
    downstream CosyVoice integration provider;
  • project CI and maintainer review are still pending.

Keep this PR in Draft until those external gates complete.

@yeahdongcn yeahdongcn changed the title perf(musa): tune Qwen2 small-hidden RMSNorm launch perf(musa): optimize Qwen2 decode for CosyVoice3 Jul 23, 2026
@yeahdongcn
yeahdongcn marked this pull request as ready for review July 24, 2026 06:01
@yeahdongcn
yeahdongcn merged commit 3e7c503 into v0.24.0-dev Jul 24, 2026
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