perf: tune Nemotron non-gated MoE on S5000 - #114
Draft
yeahdongcn wants to merge 1 commit into
Draft
yeahdongcn wants to merge 1 commit into
yeahdongcn wants to merge 1 commit into
Conversation
froststeam
reviewed
Sep 16, 2026
| num_tokens = hidden_states.shape[0] | ||
| E, N, _ = w1.shape | ||
| topk = topk_ids.shape[1] | ||
| if (is_gated and activation != "silu") or (not is_gated and activation != "relu2_no_mul"): |
Contributor
There was a problem hiding this comment.
This hard-codes gated=silu / non-gated=relu2_no_mul, so SGLang Gemma (and other gelu/swiglu MoEs) can no longer tune or run through this path. Please keep the Nemotron relu2 path, but don't reject other activations here.
froststeam
reviewed
Sep 16, 2026
| "num_warps": 4, | ||
| "num_stages": 2 | ||
| }, | ||
| "8": { |
Contributor
There was a problem hiding this comment.
This file only covers M=1/4/6/8, so any larger batch falls back to M=8. That is too small for prefill / high-
concurrency decode. Please also tune larger buckets (16, 32, 64, ..., 4096).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add the S5000 Triton 3.2 tuning path for NVIDIA Nemotron 3.5 Lightning's non-gated
relu2MoE and include the first fresh M=8 configuration from the merged vLLM-MUSA performance campaign.Nemotron shape:
E=128K=2688N=1856topk=6relu2_no_mulThe config uses default tiles for M=1/4/6 and a fresh tuned tile for M=8. This prevents the M=8 tile from being selected for MTP2/MTP4 while targeting the M≈7/8 routed-token bucket used by MTP6.
Validation
17.824 ms -> 15.119 ms(15.2%).8.106 ms/token,9.528 sdone latency6.949 ms/token,8.189 sdone latency540ddb72...17 passed.Scope and follow-up
This PR includes the non-gated tuner/runtime semantics required to generate and load the config on the post1 Triton 3.2 stack. MTP2/MTP4 long-run validation with the explicit default buckets remains in progress; no claim is made for those modes yet. The config was generated on the current S5000 host and is not copied from the withdrawn historical Triton 3.2 file.