Skip to content

perf: add chunked min-p sampling on MUSA#111

Closed
yeahdongcn wants to merge 10 commits into
xd/musa-0760-rms-deepgemmfrom
xd/musa-0763-chunked-min-p
Closed

perf: add chunked min-p sampling on MUSA#111
yeahdongcn wants to merge 10 commits into
xd/musa-0760-rms-deepgemmfrom
xd/musa-0763-chunked-min-p

Conversation

@yeahdongcn

@yeahdongcn yeahdongcn commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

This is PR 3/4 in the MUSA kernel stack and is based on the MUSA-0760
RMS/DeepGEMM PR.

Review/merge order: #109 -> #110 -> #111 -> #112.

It adds a chunked MUSA min-p sampler and automatic guarded dispatch for the
validated Qwen dense sampling contract. The fast path requires S5000/mp31,
contiguous FP32 probabilities, a supported Qwen vocabulary (151936 or
248320), batch 1-64, deterministic sampling, and a matching MUSA
generator/device contract. Device capability is checked through the logical
MUSA ordinal, and query failure falls back safely. Philox state is unpacked in
the kernel so graph replay does not capture stale host-side state. Every
unsupported case uses the existing native sampler. No campaign- or
kernel-specific environment control is introduced or read to select this
path.

Exact source:

  • base: xd/musa-0760-rms-deepgemm at
    ee367bc5a3b6e2704e1a177a0bc214f8d511f003
  • head: c8ade1f9f92189ea18368912165f4d9fbb79a219
  • head ref: xd/musa-0763-chunked-min-p
  • pinned vLLM: ee0da84ab9e04ac7610e28580af62c365e898389

Validation

Historical per-ticket performance evidence

The figures below are retained from the earlier MUSA-0763 incremental
campaign. They predate the current graph-safe Philox and
environment-independent follow-ups, so they are directional historical
evidence, not a result from the current exact head and not the final combined
four-PR A/B. The A/B/A used Qwen3-8B-FP8, TP8 on an isolated S5000 node,
top_k=50, min_p=0.05, normal Inductor, FULL_AND_PIECEWISE graph capture,
no MTP/speculative decoding, and separately built candidate/native-control
trees. Every request completed at 4104 input and 1024 output tokens.

concurrency mean TTFT change mean TPOT reduction
1 +0.64% +1.06%
4 +8.24% +1.13%
16 -3.20% +1.46%
64 +3.18% +1.19%

Median TTFT improved in all four historical cells; mean TTFT at concurrency
16 was noisy. TPOT improved in all four cells.

Historical checks also covered min-p correctness at vocabularies 151936 and
248320 for batch 1/4/16/64 and guarded fallback behavior. These checks remain
useful background but do not replace native/graph tests at the current exact
head.

Superseded descendant-stack evidence

The former descendant head f53cf6cbbd7f4d85e82270a86db1938c713d948a
passed a non-identity visible-device smoke and dense, MoE, and MLA functional
serving checks. That SHA predates the current environment-independent and
graph-safe sampling head. The associated host-81 combined latency run also had
a foreign ExecBench workload, so it is rejected for performance and is not an
exact-head readiness claim.

Exact-final MUSA-0767 gate

Status: EXACT-HEAD FUNCTIONAL REGRESSION COMPLETE; PERFORMANCE GATES
PENDING -- keep this PR in draft.

  • PR-specific head: c8ade1f9f92189ea18368912165f4d9fbb79a219
  • final stacked head: 95c3661c8658a861a1891f981c81dcadabea0654
  • final source tree: a766855ed712d61f81acb6b11d456ea5a128cac5
  • source archive SHA-256:
    03dbe0526f72d4dcba6b743c500d4d5eaaf9676e263f9bd91cbfd265db140feb
  • exact environment and build: 8x MTT S5000, driver 3.3.5-server,
    MUSA runtime 50200, torch==2.9.1.post1,
    torch_musa==2.9.1.post1+4c646b6, torchada==0.1.75,
    MATE/FlashAttention3 0.2.4, candidate extension SHA-256
    d0baaa8ba11981438a549115d50e89c10ec566dfac1ba27c0ccdc8364c5c2bd7
  • focused min-p correctness, fallback, RNG, and graph-replay tests:
    PASS as part of the 60 passed in 44.72s final-stack suite
  • exact-head non-identity logical-device sampling smoke: PASS with
    MUSA_VISIBLE_DEVICES=4,5; physical devices 4 and 5 were exposed as
    logical devices 0 and 1, and optimized min-p selected on logical device 0
  • compiled/captured Qwen dense, Qwen MoE, and DeepSeek MLA regression:
    PASS, 3/3 rows and 10/10 semantic requests (Qwen3-8B-FP8 TP1
    FLASH_ATTN, Qwen3.5-35B-A3B-FP8 TP4 FLASH_ATTN, and
    DeepSeek-V2-Lite-Chat-FP8 TP1 FLASHMLA; Inductor,
    FULL_AND_PIECEWISE, MTP off)
  • clean-node combined Qwen3-8B-FP8 candidate/control result at exact
    4104/1024 and concurrency 1/4/16/64, without MTP:
    PENDING -- not present in this regression-only bundle
    • the first control-only attempt is rejected because default prefix caching
      plus repeated deterministic prompts invalidated cold-prefill TTFT; every
      A/B leg will restart with prefix caching disabled
  • Qwen3.5-2B and Qwen3.6-27B exact-token widening A/B evidence:
    PENDING -- not present in this regression-only bundle
  • regression evidence bundle SHA-256:
    ea1b38559be0458343b85f3cd6bba275d320d2b1c729071f455e2dccb04c5272;
    its focused monitor verdict was FAIL because foreign ExecBench GPU work
    was present, so the clean performance-monitor gate remains PENDING

Remaining gate

The exact-head functional regression is complete. This PR remains draft until
the combined and widening A/B plus clean performance-monitor gates are
complete and accepted. Qwen3.5/Qwen3.6 historical runs remain compatibility
evidence, not current candidate/control performance claims.

@yeahdongcn

Copy link
Copy Markdown
Collaborator Author

Superseded by #115, which consolidates the four logical kernel changes on the current v0.24.0-dev head (4416dc24) while preserving per-change commits. This draft and its historical evidence remain available for traceability; it is not a current-head readiness claim.

@yeahdongcn yeahdongcn closed this Jul 21, 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