NVIDIA / NVIDIA/Megatron-LM

Optimize MoE routing-map construction to avoid int32 intermediates

Open
#4,965 0 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 3h
Merged PRs (30d)
272

Description

**Is your feature request related to a problem? Please describe.**

The unfused MoE router paths build routing maps as boolean masks, but some paths first allocate logits-shaped integer tensors, scatter integer values into them, and then convert the result to `bool`.

@NVIDIA/mcore-reviewers @NVIDIA/mcore-oncall

This is correct behaviorally, but it introduces avoidable temporary memory traffic in a repeated MoE router path. The cost can compound across MoE-heavy models where routing-map construction runs in many layers.

**Describe the solution you'd like**

Build the routing map directly as a boolean tensor and scatter `True` values into it. The routing decisions and routing probability tensors should remain unchanged; only the intermediate dtype/allocation pattern changes.

**Describe alternatives you've considered**

The existing implementation can be left as-is if the integer intermediate is needed for compatibility with older PyTorch/CUDA behavior. In local equivalence checks on the current environment, direct boolean `scatter`/`index_put_` produced the same routing maps across CPU/CUDA and fp32/fp16/bf16 logits.

**Additional context**

Related PR: #4964

This is intended as a router preprocessing micro-optimization, not an end-to-end training throughput claim. Local standalone CUDA timings showed lower routing-map construction time for larger token/expert shapes, but full training impact depends on workload, MoE configuration, expert compute, and communication overhead.

Contributor guide

Open the contributing guide

Research direction

Locate the unfused MoE router paths that construct routing maps and compare the intended change with related PR #4964. Verify that direct boolean scatter or index assignment preserves routing maps across CPU/CUDA and fp32, fp16, and bf16 logits while leaving routing decisions and probabilities unchanged.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning, performance
Issue type
Refactor
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.