Skip to content

[Nemotron][CUDA] Add deterministic moe_router_dispatch - #504

Merged
KJLdefeated merged 2 commits into
RL-Align:test-nemotronfrom
RichApple123:feat/nemotron-router-sm90
Oct 26, 2026
Merged

KJLdefeated merged 2 commits into
RL-Align:test-nemotronfrom
RichApple123:feat/nemotron-router-sm90

Conversation

@RichApple123

@RichApple123 RichApple123 commented Oct 10, 2026 •

Copy link
Copy Markdown

Implements the SM90 CUDA moe_router_dispatch work item of #434: FP32 projection, sigmoid, bias-corrected top-six selection, normalization and deterministic token packing, with backward for routing weights and packed payloads. Supports BF16/FP32 inputs with FP32 router weights at Nano's 2688-hidden, 128-expert geometry.

Merges the latest test-nemotron layout refactor and places the implementation, reference, tests and benchmark under the canonical directories. The operator and reference are byte-identical to the previous PR version; test and benchmark computations are unchanged. Registration lives in rl_engine.runtime.registry, preserving the upstream compatibility entry point.

Adds a numerical reference, CPU/GPU correctness and batch-invariance tests, a cuBLAS + Transformer Engine benchmark, documentation and CI wiring for test-nemotron. The ordering and fixed arithmetic policy remain proposed for section 4.4 review. Expert MLP/combine and distributed communication are outside this operator.

Local migration validation:

  • Focused router/reference/registry/layout tests: 41 passed, 81 CUDA skipped; upstream layout suite: 32 passed, 1 GPU extension skipped. These sets overlap.
  • Additional dispatch, stateless-executor, cross-configuration and tolerance tests: 87 passed.
  • Strict documentation build, scoped formatting/lint/types, wheel build and imports outside the source checkout pass.
  • Full-package MyPy reports the same 14 errors as the target branch in the same local environment, with no new diagnostics.

Under the canonical layout, all 106 single-H20 router/reference/registry tests passed without skips. A fresh cuBLAS + TE 2.19 index/eager replay completed 256 route/full measurements with accuracy qualification; its 128 full-operator rows retain paired samples and regressions. BF16 random forward+backward at 8192/32768 tokens measured 1.087–1.092x; FP32 remained near parity and small-token regressions persist. Raw measurements are retained separately for review. The unchanged operator also has historical 36-case bitwise comparisons and four-tool Compute Sanitizer qualification; those checks were not rerun during this layout replay.

Target-branch CI and maintainer approval of the proposed contract remain pending.

Signed-off-by: RichApple123 <134492660+RichApple123@users.noreply.github.com>
@coderabbitai

coderabbitai Bot commented Oct 10, 2026 •

Copy link
Copy Markdown

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration
  • Configuration used: defaults
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 998bd0a1-83de-44d3-9c79-8678ac87f0e2

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

  • 🔍 Trigger review
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

Signed-off-by: RichApple123 <134492660+RichApple123@users.noreply.github.com>
@RichApple123
RichApple123 marked this pull request as ready for review October 14, 2026 02:57
@RichApple123

Copy link
Copy Markdown
Author

This PR is now ready for review. The CI run is still marked action_required. Could you approve the run and help review the PR?

@Flink-ddd

Copy link
Copy Markdown
Collaborator

Thank you for your contribution. CC @KJLdefeated PTAL.

@RichApple123

Copy link
Copy Markdown
Author

Hi @KJLdefeated, a gentle follow-up on this PR. It has been updated with test-nemotron, has no merge conflicts, and all CI checks are passing. Could you take a look when you have time? Thanks!

@Flink-ddd Flink-ddd left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me. @KJLdefeated, could you confirm the proposed routing order and arithmetic contract before we merge into test-nemotron?

@KJLdefeated KJLdefeated left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! I will confirm the routing order and arithmetic contract before merging into test-nemotron in next few days.

@Flink-ddd Flink-ddd added platform: cuda Specific optimizations or bugs in NVIDIA graphics cards (such as FlashInfer, TMA optimizations) new-model Nemotron labels Oct 22, 2026
@RichApple123

Copy link
Copy Markdown
Author

Hi @KJLdefeated, just following up on the routing order and arithmetic contract review. Is there anything I can clarify or provide to help? Thanks!

@KJLdefeated
KJLdefeated merged commit 3c34107 into RL-Align:test-nemotron Oct 26, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Nemotron new-model platform: cuda Specific optimizations or bugs in NVIDIA graphics cards (such as FlashInfer, TMA optimizations)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants