perf(moe): fuse router weight multiply with Triton kernel

[piyushumate]

What does this PR do ?

Implements a fused Triton forward and backward kernel for MoE router-weight multiply (_apply_router_weight_fp32), replacing the CPU Python chunking loop with in-register FP32 computation and reduction to eliminate kernel launch overhead and intermediate HBM memory traffic.

Changelog

• nemo_automodel/components/moe/optimized_ops.py:
  • Added fused Triton forward kernel _router_weight_fwd_kernel computing out = (x * prob).to(out_dtype) in-register in FP32.
  • Added fused Triton backward kernel _router_weight_bwd_kernel computing grad_x and grad_p (in-register reduction along hidden dimension) in a single pass.
  • Added custom autograd Function _TritonRouterWeightMulFunction with support for save_x=False optimization when probs requires no gradient.
  • Updated _apply_router_weight_fp32 to dispatch to Triton on CUDA and safely fall back to chunked custom autograd on CPU/non-CUDA.
  • Guarded imports via safe_import("triton") and safe_import("triton.language") per repository review guidelines.
• tests/unit_tests/moe/test_router_weight_mul.py:
  • Added Triton forward and backward parity tests against eager PyTorch ground truth.
  • Verified save_x=False behavior when probs.requires_grad=False.
  • Updated autograd function class name check to accept both _TritonRouterWeightMulFunction and _RouterWeightMulFunction.

1. Forward Pass

Configuration	Shape [tokens, hidden]	Chunked (main)	Fused Triton (PR #3800)	Speedup	Latency Reduction
Nemotron / Qwen-MoE	[16384, 4096]	0.671 ms	0.092 ms	7.31x	-86.3%
Kimi K3 / DeepSeek V3	[16384, 7168]	1.147 ms	0.158 ms	7.24x	-86.2%
Large Batch / SP	[65536, 7168]	4.589 ms	0.622 ms	7.38x	-86.4%

2. End-to-End Training (Forward + Backward)

Configuration	Shape [tokens, hidden]	Chunked (main)	Fused Triton (PR #3800)	Speedup vs main	Latency Reduction
Nemotron / Qwen-MoE	[16384, 4096]	2.232 ms	0.319 ms	6.99x	-85.7%
Kimi K3 / DeepSeek V3	[16384, 7168]	3.815 ms	0.548 ms	6.96x	-85.6%
Large Batch / SP	[65536, 7168]	15.241 ms	2.150 ms	7.09x	-85.9%

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Additional Information

• cc @akoumpa

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[HuiyingLi]

Thank you @piyushumate ! I am wondering if you have any speed comparisons?

[piyushumate]

@HuiyingLi yes, i ran a benchmark and numerical parity audit on H100 , comparing chunked autograd implementation currently on mainline against this PR

1. Forward Pass

Configuration	Shape [tokens, hidden]	Chunked (main)	Fused Triton (PR #3800)	Speedup	Latency Reduction
Nemotron / Qwen-MoE	[16384, 4096]	0.671 ms	0.092 ms	7.31x	-86.3%
Kimi K3 / DeepSeek V3	[16384, 7168]	1.147 ms	0.158 ms	7.24x	-86.2%
Large Batch / SP	[65536, 7168]	4.589 ms	0.622 ms	7.38x	-86.4%

2. End-to-End Training (Forward + Backward)

Configuration	Shape [tokens, hidden]	Chunked (main)	Fused Triton (PR #3800)	Speedup vs main	Latency Reduction
Nemotron / Qwen-MoE	[16384, 4096]	2.232 ms	0.319 ms	6.99x	-85.7%
Kimi K3 / DeepSeek V3	[16384, 7168]	3.815 ms	0.548 ms	6.96x	-85.6%
Large Batch / SP	[65536, 7168]	15.241 ms	2.150 ms	7.09x	-85.9%

[piyushumate]

@HuiyingLi any suggestions (/next steps) based on the benchmark numbers?