`core.transformer.moe.experts`#

Module Contents#

Classes#

`GroupedLinearFc1Interface`	Interface for linear_fc1 module in TEGroupedMLP.
`GroupedLinearFc1Builder`	Protocol describing how to build a linear_fc1 layer in TEGroupedMLP.
`GroupedLinearFc2Interface`	Protocol for linear_fc2 module in TEGroupedMLP.
`GroupedLinearFc2Builder`	Protocol describing how to build a linear_fc2 layer in TEGroupedMLP.
`GroupedMLPSubmodules`	The dataclass for ModuleSpecs of TEGroupedMLP submodules including linear fc1, activation function, linear fc2.
`_VirtualExpertFC2WgradStore`	Start the FC2 virtual-expert gradient reduction when TE enqueues FC2’s wgrad GEMM.
`TEGroupedMLP`	An efficient implementation of the Experts layer using TE’s GroupedLinear.
`InferenceGroupedMLP`	Inference-optimized GroupedMLP with GPU-resident offsets.
`SequentialMLP`	An implementation of the Experts layer using a sequence of MLP layers.

Functions#

`_te_supports_scaled_tanh_srelu`	Whether the installed TE provides `ScaledTanhSReLU` accepting `tanh_clamp_scale`.
`_require_te_tanh_clamp_support`	Raise an actionable error if the op fuser is asked to clamp and the installed TE cannot.
`_te_grouped_tensor_supports_sharded_weights`	Whether TE threads the sharded parameters through its grouped-tensor path.

Data#

logger

API#

core.transformer.moe.experts.logger#: ‘getLogger(…)’

core.transformer.moe.experts._te_supports_scaled_tanh_srelu() → bool#

Whether the installed TE provides ScaledTanhSReLU accepting tanh_clamp_scale.

Checked by capability rather than version so it works against a development TE, and not cached because tests swap the TE module out.

core.transformer.moe.experts._require_te_tanh_clamp_support(config) → None#: Raise an actionable error if the op fuser is asked to clamp and the installed TE cannot.

class core.transformer.moe.experts.GroupedLinearFc1Interface#

Bases: typing.Protocol

Interface for linear_fc1 module in TEGroupedMLP.

forward( permuted_local_hidden_states: torch.Tensor, tokens_per_expert: list[int], /, ) → tuple[torch.Tensor, torch.Tensor | None]#: Forward method for linear_fc1 module.

backward_dw() → None#: Backward method for linear_fc1 module.

class core.transformer.moe.experts.GroupedLinearFc1Builder#

Bases: typing.Protocol

Protocol describing how to build a linear_fc1 layer in TEGroupedMLP.

__call__( num_local_experts: int, input_size: int, output_size: int, /, *, num_virtual_experts: int = 0, config: megatron.core.transformer.transformer_config.TransformerConfig, init_method: collections.abc.Callable[[torch.Tensor], None], bias: bool, skip_bias_add: bool, is_expert: bool, tp_comm_buffer_name: str | None, pg_collection: megatron.core.transformer.moe.moe_utils.ProcessGroupCollection | None, name: str | None = None, ) → core.transformer.moe.experts.GroupedLinearFc1Interface#: Builds a linear_fc1 layer for TEGroupedMLP.

class core.transformer.moe.experts.GroupedLinearFc2Interface#

Bases: typing.Protocol

Protocol for linear_fc2 module in TEGroupedMLP.

forward( intermediate_parallel: torch.Tensor, tokens_per_expert: list[int], /, ) → tuple[torch.Tensor, torch.Tensor | None]#: Forward method for linear_fc2 module.

backward_dw() → None#: Backward method for linear_fc2 module.

class core.transformer.moe.experts.GroupedLinearFc2Builder#

Bases: typing.Protocol

Protocol describing how to build a linear_fc2 layer in TEGroupedMLP.

__call__( num_local_experts: int, input_size: int, output_size: int, /, *, num_virtual_experts: int = 0, config: megatron.core.transformer.transformer_config.TransformerConfig, init_method: collections.abc.Callable[[torch.Tensor], None], bias: bool, skip_bias_add: bool, is_expert: bool, tp_comm_buffer_name: str | None, pg_collection: megatron.core.transformer.moe.moe_utils.ProcessGroupCollection | None, name: str | None = None, ) → core.transformer.moe.experts.GroupedLinearFc2Interface#: Builds a linear_fc2 layer for TEGroupedMLP.

class core.transformer.moe.experts.GroupedMLPSubmodules#

The dataclass for ModuleSpecs of TEGroupedMLP submodules including linear fc1, activation function, linear fc2.

linear_fc1: core.transformer.moe.experts.GroupedLinearFc1Builder#: None

linear_fc2: core.transformer.moe.experts.GroupedLinearFc2Builder#: None

activation_func: megatron.core.transformer.mlp.TEActivationFunctionBuilder | None#

None

Builder for an activation function module; only used if config.use_te_activation_func is True.

core.transformer.moe.experts._te_grouped_tensor_supports_sharded_weights() → bool#

Whether TE threads the sharded parameters through its grouped-tensor path.

Added by https://github.com/NVIDIA/TransformerEngine/pull/3517.

TODO: replace with is_te_min_version once that fix ships in a bumped TE release. Fixed and unfixed builds both report 2.20.0.dev0 today, so no version can distinguish them yet.

class core.transformer.moe.experts._VirtualExpertFC2WgradStore(load_balancer)#

Start the FC2 virtual-expert gradient reduction when TE enqueues FC2’s wgrad GEMM.

Installed on the executing FC2 GroupedLinear module or op. TE’s delayed-wgrad protocol hands the wgrad GEMM to put instead of launching it; this store runs the GEMM at once and starts the reduction behind FC1’s dgrad and wgrad GEMMs instead of the dispatch-backward all-to-all. The queue stays empty, so TE’s backward_dw has nothing to run.

Initialization

context#: None

static delay_wgrad_compute() → bool#: TE routes the wgrad GEMM through :meth:put when this is true.

put(tensors, wgrad_gemm) → None#: Run the wgrad GEMM now and start FC2’s virtual-expert reduction behind it.

class core.transformer.moe.experts.TEGroupedMLP( num_local_experts: int, config: megatron.core.transformer.transformer_config.TransformerConfig, submodules: core.transformer.moe.experts.GroupedMLPSubmodules, pg_collection: Optional[megatron.core.transformer.moe.moe_utils.ProcessGroupCollection] = None, name: str | None = None, )#

Bases: megatron.core.transformer.module.MegatronModule

An efficient implementation of the Experts layer using TE’s GroupedLinear.

Executes multiple experts in parallel to maximize computational efficiency.

Initialization

Parameters:: name (str | None) – module instance name passed top-down from its paranet module

bind_virtual_experts(load_balancer) → None#: Bind native and virtual runtime weights after main gradients are initialized.

_apply_packed_bias( intermediate_parallel, packed_bias, tokens_per_expert, permuted_probs, )#: Apply a packed expert bias without reading token counts on the host.

_apply_bias( intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs, )#

_is_fused_impl_supported() → bool#: Check if the TE op fuser supports implementing this module.

_make_fused_ops() → torch.nn.Module#: Construct fused module for FC1, activation, and FC2.

_make_fused_impl_pre_forward_hook() → collections.abc.Callable#

Make function that calls submodule pre-forward callback hooks.

This is intended for compatibility with DistributedDataParallel/FSDP hooks that trigger parameter all-gathers. It does not support general pre-forward hooks since they may manipulate intermediate tensors that are never instantiated by the fused implementation.

static _ensure_main_grad(linear_module: torch.nn.Module) → None#: Expose FSDP main_grad buffers required by TE fused wgrad accumulation.

_ensure_main_grad_for_fused_impl() → None#: Expose wrapper parameter main_grad buffers before TE fused ops run.

_get_paged_stash_scope( permuted_local_hidden_states, tokens_per_expert, )#: Start a paged-stash group and return the input and saved-tensor context.

_fused_forward( permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor, permuted_probs: torch.Tensor, output_buffer: Optional[torch.Tensor] = None, grad_input_buffer: Optional[torch.Tensor] = None, ) → torch.Tensor#: Forward pass using Transformer Engine operation fuser API.

static _remove_glu_interleaving( x: torch.Tensor, interleave_size: int, ) → torch.Tensor#: Reorder interleaved GLU blocks so gate and linear halves are contiguous.

_unfused_forward( permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor | list[int], permuted_probs: torch.Tensor, ) → torch.Tensor#: Run FC1, activation, and FC2 without the TE operation fuser.

forward( permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor, permuted_probs: torch.Tensor, output_buffer: Optional[torch.Tensor] = None, grad_input_buffer: Optional[torch.Tensor] = None, ) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

Forward of TEGroupedMLP

Parameters:

permuted_local_hidden_states (torch.Tensor) – The permuted input hidden states of the
experts. (local)
tokens_per_expert (torch.Tensor) – The number of tokens per expert.
permuted_probs (torch.Tensor) – The permuted probs of each token produced by the router.
output_buffer (torch.Tensor, optional) – Preallocated buffer to write the fc2 output into
it. ((NCCL-EP zero-copy fwd combine); only the fused op-fuser path supports)
grad_input_buffer (torch.Tensor, optional) – Preallocated buffer to write the fc1 dgrad
into (NCCL-EP zero-copy bwd dispatch)

Returns:

The output of the local experts.

Return type:

output (torch.Tensor)

sharded_state_dict( prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None, ) → megatron.core.dist_checkpointing.mapping.ShardedStateDict#: Maps local expert to global experts. The sharded state dict is interchangable with SequentialMLP’s.

backward_dw()#

Performs backward pass for weight gradients in TEGroupedMLP.

This method executes the backward pass for weight gradients by calling backward_dw() on the linear layers in reverse order (fc2 followed by fc1). If an error occurs during execution, it is caught and re-raised with a descriptive message.

class core.transformer.moe.experts.InferenceGroupedMLP( num_local_experts: int, config: megatron.core.transformer.transformer_config.TransformerConfig, submodules: core.transformer.moe.experts.GroupedMLPSubmodules, pg_collection: Optional[megatron.core.transformer.moe.moe_utils.ProcessGroupCollection] = None, name: str | None = None, )#

Bases: core.transformer.moe.experts.TEGroupedMLP

Inference-optimized GroupedMLP with GPU-resident offsets.

Inherits from TEGroupedMLP to reuse weight initialization and checkpoint compatibility. Supports three forward paths:

Training: delegates to parent TEGroupedMLP
Inference + FlashInfer: CUTLASS fused MoE for BF16 and tanh-clamped MXFP8, routed block-scale MoE for other MXFP8 activations and per-token NVFP4
Inference + torch: torch.nn.functional.grouped_mm with GPU-resident cumsum offsets
Inference + vLLM: Triton fused MoE for BF16, MCore scaled grouped GEMM for MXFP8

Initialization

Parameters:: name (str | None) – module instance name passed top-down from its paranet module

_EXPERT_WEIGHT_GROUPS#: ((‘linear_fc1’, ‘_fc1_weight’), (‘linear_fc2’, ‘_fc2_weight’))

_build_flashinfer_clamp_limit(device: torch.device) → None#: Allocate the FP32 limit once during inference warmup, before graph capture.

_resolve_flashinfer_activation_type()#: Map megatron activation config to FlashInfer ActivationType.

_resolve_mcore_activation_type()#: Map megatron activation config to mcore_fused_moe ActivationType.

static _unwrap_mxfp8_weight( weight: object, ) → megatron.core.inference.quantization.mxfp8_tensor.MXFP8Tensor | None#: Return the MCore MXFP8 storage carried by a weight, if any.

static _require_uniform_weight_format( format_flags: collections.abc.Iterable[bool], format_name: str, ) → bool#: Return whether every expert projection uses a format, rejecting mixtures.

_expert_weights_use_mxfp8() → bool#: Return whether all per-expert FC1 and FC2 weights use MCore MXFP8 storage.

_stack_mxfp8_linear_weight( linear_name: str, backend: str, ) → megatron.core.inference.quantization.mxfp8_tensor.MXFP8Tensor#: Stack one linear’s per-expert MXFP8 weights in canonical layout.

_build_concatenated_mxfp8_weights()#

Build contiguous expert stacks after checkpoint loading.

Torch, vLLM and FlashInfer CUTLASS rebind each per-expert MXFP8Tensor to its stacked view. FlashInfer routed derives a shuffled Major-K stack and makes the canonical tensors write through to it, releasing their storage.

refresh_flashinfer_mxfp8_weights() → bool#

Refresh derived FlashInfer expert weights in place after an MXFP8 refit.

Returns whether derived FlashInfer weights were refreshed.

_build_flashinfer_nvfp4_weights()#

Quantize the BF16 expert stacks into routed per-token NVFP4 weights.

The BF16 stacks stay resident as the refit destination; the NVFP4 stacks are a derived representation consumed only by the FlashInfer kernel.

refresh_flashinfer_nvfp4_weights() → bool#

Requantize routed NVFP4 expert weights in place after the BF16 weights change.

Returns whether derived FlashInfer weights were refreshed.

refresh_cache() → None#

Requantize the per-token NVFP4 expert weights after a refit.

Refit writes the BF16 stacks, so without this the NVFP4 kernel keeps serving the weights quantized at the first forward.

_build_concatenated_weights()#

Create big contiguous weight tensors that share storage with TE’s per-expert parameters.

Creates _fc1_weight and _fc2_weight as contiguous tensors of shape [num_experts, out_features, in_features]. Instead of replacing TE’s parameters (which breaks TE’s internal bookkeeping), we redirect each parameter’s .data to be a view into the contiguous buffer. The nn.Parameter objects themselves remain untouched in TE’s module, preserving FP8 scaling state, etc.

This allows:

TE’s forward to work correctly (same Parameter objects, same internal state)
Training updates to flow through (param.data is a view into the big tensor)
torch.nn.functional.grouped_mm / FlashInfer to use the big tensor directly

_flashinfer_forward(hidden_states, routing_map, probs)#: FlashInfer fused MoE kernel for CUDA-graphed inference iterations.

_mcore_fused_moe_forward(hidden_states, probs, routing_map)#: MCore grouped GEMM for Torch and the vLLM-selected MXFP8 fallback.

_vllm_forward(hidden_states, probs, routing_map)#: vLLM Triton fused MoE kernel forward (BF16, CUDA-graph safe).

forward( permuted_local_hidden_states: torch.Tensor, tokens_per_expert: Optional[torch.Tensor], permuted_probs: torch.Tensor, routing_map: Optional[torch.Tensor] = None, ) → Tuple[torch.Tensor, Optional[torch.Tensor]]#

Forward pass with backend-selected inference grouped GEMMs:

Training: delegates to parent TEGroupedMLP.
Inference + FlashInfer: fused BF16 or MXFP8 MoE. tokens_per_expert is not used in this path; the FlashInfer kernels operate directly on routing_map.
Inference + torch: torch.nn.functional.grouped_mm with GPU-resident cumsum offsets.
Inference + vLLM: Triton fused MoE for BF16; MXFP8 layers use MCore’s scaled grouped-GEMM path because the vLLM kernel is BF16-only.

Parameters:

permuted_local_hidden_states – [num_tokens, hidden_size] input hidden states.
tokens_per_expert – [num_experts] number of tokens routed to each expert. None when using the CUDA-graphed FlashInfer path.
permuted_probs – [num_tokens, topk] routing probabilities.
routing_map – [num_tokens, topk] token-to-expert assignment indices. Required for the FlashInfer CUDA-graphed path, None otherwise.

class core.transformer.moe.experts.SequentialMLP( num_local_experts: int, config: megatron.core.transformer.transformer_config.TransformerConfig, submodules: megatron.core.transformer.mlp.MLPSubmodules, pg_collection: Optional[megatron.core.transformer.moe.moe_utils.ProcessGroupCollection] = None, name: str | None = None, )#

Bases: megatron.core.transformer.module.MegatronModule

An implementation of the Experts layer using a sequence of MLP layers.

This class executes each expert sequentially.

Initialization

_pad_tensor_for_quantization(hidden, probs)#: Padding tensor shape to multiples of 16/32.

forward( permuted_local_hidden_states: torch.Tensor, tokens_per_expert: torch.Tensor, permuted_probs: torch.Tensor, )#: Forward step of the SequentialMLP.

backward_dw()#: Backward pass for weight gradients in SequentialMLP.

sharded_state_dict(prefix='', sharded_offsets=(), metadata=None)#: Maps local expert to global experts.

core.transformer.moe.experts#

Module Contents#

Classes#

Functions#

Data#

API#

`core.transformer.moe.experts`#