Repository navigation
[Nemotron][CUDA] Add deterministic moe_router_dispatch - #504
Conversation
Signed-off-by: RichApple123 <134492660+RichApple123@users.noreply.github.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configuration
You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
Signed-off-by: RichApple123 <134492660+RichApple123@users.noreply.github.com>
|
This PR is now ready for review. The CI run is still marked |
|
Thank you for your contribution. CC @KJLdefeated PTAL. |
|
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
left a comment
There was a problem hiding this comment.
Looks good to me. @KJLdefeated, could you confirm the proposed routing order and arithmetic contract before we merge into test-nemotron?
KJLdefeated
left a comment
There was a problem hiding this comment.
LGTM! I will confirm the routing order and arithmetic contract before merging into test-nemotron in next few days.
|
Hi @KJLdefeated, just following up on the routing order and arithmetic contract review. Is there anything I can clarify or provide to help? Thanks! |
Implements the SM90 CUDA
moe_router_dispatchwork 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-nemotronlayout 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 inrl_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:
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.