Skip to content

feat: support Nemotron non-gated MoE tuning on Triton 3.6 - #112

Draft
yeahdongcn wants to merge 22 commits into
mainfrom
xd/nemotron35-triton36
Draft

feat: support Nemotron non-gated MoE tuning on Triton 3.6#112
yeahdongcn wants to merge 22 commits into
mainfrom
xd/nemotron35-triton36

Conversation

@yeahdongcn

@yeahdongcn yeahdongcn commented Sep 13, 2026

Copy link
Copy Markdown
Collaborator

This draft adds the torchada-side support needed to tune Nemotron 3.5 Lightning's non-gated MoE under Triton 3.6.

Changes:

  • recognize Triton 3.6.0 when selecting fused-MoE configuration directories;
  • model Nemotron's relu2_no_mul route as non-gated, with E=128/N=1856 projection and quantized-K shapes;
  • keep routed weighting, activation, reduction, and graph-capture metadata consistent with the runtime path;
  • add focused configuration and pipeline tests;
  • add S5000 Triton 3.6 configs for Nemotron E=128/N=1856 MoE and Mamba SSU shapes.

Validation:

  • py_compile, ruff, JSON validation, and git diff --check pass.
  • On leased S5000 GPU0 (driver 3.3.5-server), vLLM-MUSA post2/Triton 3.6 explicitly loaded both config files and completed PCG + decode FULL graph capture.
  • Same-host target-only E2E improved from B8 53.693 tok/s without configs to 83.012 tok/s with configs (three measured runs, two warmups, fixed TP1/BF16 workload); B1 was 27.542 tok/s.
  • MTP2 graph/request smoke completed at B1 24.079 / B8 41.468 tok/s, but repeated B8 output was not deterministic, so this does not qualify MTP2 for production.
  • The candidate config files were initially generated on a rejected host and then revalidated by exact config-load and E2E replay on the correct leased S5000. A clean current-host autotuner rerun was started but interrupted when SOL reclaimed the lease; its incomplete output is not used as evidence.

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