Skip to content

perf: tune Nemotron MTP MoE decode on S5000 - #115

Draft
yeahdongcn wants to merge 13 commits into
mainfrom
codex/nemotron35-moe-tpot-20260918
Draft

yeahdongcn wants to merge 13 commits into
mainfrom
codex/nemotron35-moe-tpot-20260918

Conversation

@yeahdongcn

@yeahdongcn yeahdongcn commented Sep 18, 2026

Copy link
Copy Markdown
Collaborator

Status

This draft carries the shape map that produced the historical Nemotron MTP
result, plus a Triton capability guard. A clean vLLM-MUSA integration gate is
still required before merge.

Changes

The PR packages the S5000 Nemotron non-gated MoE autotune map for
E=128,N=1856,device_name=MTT_S5000 with measured runtime and graph buckets:

16, 49, 80, 113, 144, 511, 512, 513

The serving routes are M≈115/120 -> 113 and M=512 -> 512; values above
the measured range use the generic heuristic. The packaged JSON SHA256 is
a0b192a4625e71cf46c52cca3bf5f042d9988da927524d7c9e4d10a0bbabc983.

The descriptor capability check now requires the actual Triton language API
triton.language.make_tensor_descriptor. Triton 3.2 can expose the tools
module without exposing that constructor; it must use the fallback path. Triton
3.6 continues to enable the descriptor path when the API is present.

Validation

  • Candidate commits: 668bd68, f4e2ffb.
  • Clean wheel builds succeed and package the JSON above.
  • Focused config/resolver and descriptor tests: 10 passed on the MUSA image.
  • The full torchada suite reached 508 passed, 19 skipped; the only reported
    failure was the project-version check from a source-only staging harness that
    did not include project metadata.
  • The historical manual-injection A/B on the original post1/Triton 3.2 setup
    produced exact-output results:
path TPOT
target-only 19.6128 ms/token
MTP1 11.4567 ms/token
MTP2 8.8937 ms/token
MTP4 7.1127 ms/token
MTP6 6.6091 ms/token

Those numbers came from the eight-key JSON injected into the installed runtime;
they are historical evidence, not yet a clean torchada==0.1.88 claim.

A diagnostic clean-source run confirmed that the candidate JSON was loaded and
that runtime M=112/126 selected key 113. Its warm-cache result was about
9.46 ms/token on a different S5000 host. A fresh post1/Triton-3.2 compile path
also exposed the descriptor compatibility issue now fixed by f4e2ffb. Neither
diagnostic number is accepted as the final performance result.

The PR should remain Draft until the packaged wheel is run with current
vLLM-MUSA source on the same hardware under a clean, interleaved
baseline-vs-candidate gate, with output hashes, selected MoE buckets,
acceptance, and TPOT recorded.

@yeahdongcn
yeahdongcn force-pushed the codex/nemotron35-moe-tpot-20260918 branch from a288538 to 364b0ec Compare September 18, 2026 13:04
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