`core.transformer.transformer_layer`#

Module Contents#

Classes#

`MlpInterface`	Interface for MLP implementations in the transformer layer.
`MlpBuilder`	MLP builder protocol for building MLPs in the transformer layer.
`TransformerLayerSubmodules`	Configuration class for specifying the submodules of a transformer layer.
`BaseTransformerLayer`	A common parent class for `TransformerLayer` like implementations.
`TransformerLayer`	A single transformer layer.
`HyperConnectionTransformerLayer`	A transformer layer with Manifold-Constrained Hyper-Connections (mHC).
`MoETransformerLayer`	A Transformer layer specialized for Mixture-of-Experts (MoE) architectures.

Functions#

`_get_offloading_interface`	Get the offloading interface for fine-grained activation offloading.
`get_transformer_layer_offset`	Get the index offset of current pipeline stage, given the level of pipelining.

Data#

logger

API#

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

core.transformer.transformer_layer._get_offloading_interface()#: Get the offloading interface for fine-grained activation offloading.

core.transformer.transformer_layer.get_transformer_layer_offset( config: megatron.core.transformer.transformer_config.TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None, )#: Get the index offset of current pipeline stage, given the level of pipelining.

class core.transformer.transformer_layer.MlpInterface#

Bases: typing.Protocol

Interface for MLP implementations in the transformer layer.

forward( hidden_states: torch.Tensor, /, *, intermediate_tensors: tuple[torch.Tensor, ...] | None = None, padding_mask: torch.Tensor | None = None, ) → tuple[torch.Tensor, torch.Tensor | None]#: Forward method for the MLP interface.

class core.transformer.transformer_layer.MlpBuilder#

Bases: typing.Protocol

MLP builder protocol for building MLPs in the transformer layer.

__call__( *, config: megatron.core.transformer.transformer_config.TransformerConfig, pg_collection: megatron.core.process_groups_config.ProcessGroupCollection, is_mtp_layer: bool, name: str | None = None, ) → core.transformer.transformer_layer.MlpInterface#

class core.transformer.transformer_layer.TransformerLayerSubmodules#

Configuration class for specifying the submodules of a transformer layer.

This class defines the structure and default implementations for various components of a transformer layer, allowing for flexible customization of the layer’s architecture.

Parameters:

input_layernorm – Specification for the input layer normalization.
self_attention (Union[ModuleSpec, type]) – Specification for the self-attention mechanism.
self_attn_bda (Union[ModuleSpec, type]) – Specification for the bias-dropout-add operation after self-attention.
pre_cross_attn_layernorm – Specification for the layer normalization before cross-attention.
cross_attention (Union[ModuleSpec, type]) – Specification for the cross-attention mechanism.
cross_attn_bda (Union[ModuleSpec, type]) – Specification for the bias-dropout-add operation after cross-attention.
pre_mlp_layernorm – Specification for the layer normalization before the MLP.
mlp (Union[ModuleSpec, type]) – Specification for the MLP in Dense layer.
mlp_bda (Union[ModuleSpec, type]) – Specification for the bias-dropout-add operation after the MLP.
sharded_state_dict_keys_map (Dict[str, str]) – Mapping for sharded tensor keys to be applied in the sharded_state_dict method.

input_layernorm: megatron.core.transformer.torch_norm.LayerNormBuilder#: None

self_attention_hyper_connection: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

self_attention: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

self_attn_bda: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

pre_cross_attn_layernorm: megatron.core.transformer.torch_norm.LayerNormBuilder#: None

cross_attention_hyper_connection: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

cross_attention: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

cross_attn_bda: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

pre_mlp_layernorm: megatron.core.transformer.torch_norm.LayerNormBuilder#: None

mlp_hyper_connection: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

mlp: core.transformer.transformer_layer.MlpBuilder | type[megatron.core.transformer.identity_op.IdentityOp]#: None

mlp_bda: Union[megatron.core.transformer.spec_utils.ModuleSpec, type]#: None

sharded_state_dict_keys_map: Dict[str, str]#: ‘field(…)’

class core.transformer.transformer_layer.BaseTransformerLayer#

Bases: abc.ABC

A common parent class for TransformerLayer like implementations.

A dummy class that is subclassed by similar TransformerLayers e.g. the TransformerLayer in this file and possibly other TransformerLayer implementations that aim to use TransformerBlock as the base module. The main purpose is to check if any layer (or module) provided in the spec is a subclass of this class to allow fanning-out of that spec for all the layers in the TransformerBlock. See _get_block_submodules method implementation in transformer_block.py file for more details.

Initialization

supports_mhc_connections: bool#: False

class core.transformer.transformer_layer.TransformerLayer( config: megatron.core.transformer.transformer_config.TransformerConfig, submodules: core.transformer.transformer_layer.TransformerLayerSubmodules, layer_number: int = 1, hidden_dropout: Optional[float] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, vp_stage: Optional[int] = None, is_mtp_layer: bool = False, add_layer_offset: bool = True, pp_layer_offset: Optional[int] = None, name: str | None = None, )#

Bases: megatron.core.transformer.module.GraphableMegatronModule, core.transformer.transformer_layer.BaseTransformerLayer, megatron.core.transformer.module.TwoStageAttentionLayer

A single transformer layer.

Transformer layer takes input with size [s, b, h] and returns an output of the same size.

Initialization

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

get_inner_quantization_context() → contextlib.AbstractContextManager#: Return the quantization context for fine-grained layer execution.

create_mcore_cudagraph_manager(config)#: Register the transformer layer for cudagraphs.

static _get_layer_offset( config: megatron.core.transformer.transformer_config.TransformerConfig, )#

Get the layer offset for the current pipeline stage.

Deprecated: please use get_transformer_layer_offset instead.

static _group_offload_output_with_bias( output_with_bias, offload_manager, forced_released_tensors: Optional[list[torch.Tensor]] = None, )#: Commit a fine-grained offload group for a raw branch output tuple.

_forward_self_attention_output_with_bias( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, rotary_pos_emb: Optional[torch.Tensor] = None, rotary_pos_cos: Optional[torch.Tensor] = None, rotary_pos_sin: Optional[torch.Tensor] = None, rotary_pos_cos_sin: Optional[torch.Tensor] = None, attention_bias: Optional[torch.Tensor] = None, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, sequence_len_offset: Optional[torch.Tensor] = None, *, inference_params: Optional[Any] = None, )#: Run input norm and self-attention, returning the raw output before BDA.

supports_two_stage_attention() → bool#: Return whether this is an attention-only layer that supports two-stage execution.

attention_bda_and_cross_attention( attention_output_with_bias, residual: torch.Tensor, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, attn_state=(), )#: Apply checkpoint bookkeeping, self-attention BDA, and cross-attention.

_forward_attention( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, rotary_pos_emb: Optional[torch.Tensor] = None, rotary_pos_cos: Optional[torch.Tensor] = None, rotary_pos_sin: Optional[torch.Tensor] = None, rotary_pos_cos_sin: Optional[torch.Tensor] = None, attention_bias: Optional[torch.Tensor] = None, inference_context: Optional[megatron.core.inference.contexts.BaseInferenceContext] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, sequence_len_offset: Optional[torch.Tensor] = None, padding_mask: Optional[torch.Tensor] = None, *, inference_params: Optional[Any] = None, )#

Perform a forward pass through the attention layer and the layernorms before and after the attention operations.

Parameters:

hidden_states (Tensor) – Input tensor of shape [s, b, h] where s is sequence length, b is batch size, and h is hidden size.
attention_mask (Tensor) – Mask tensor for self-attention.
context (Tensor, optional) – Context tensor for cross-attention.
context_mask (Tensor, optional) – Mask tensor for cross-attention.
rotary_pos_emb (Tensor, optional) – Rotary positional embeddings.
rotary_pos_cos (Optional[Tensor]) – Rotary embedding cosine.
rotary_pos_sin (Optional[Tensor]) – Rotary embedding sine.
rotary_pos_cos_sin (Optional[Tensor]) – Combined rotary embedding cosine and sine.
RoPE. (Currently used exclusively for inference with dynamic batching and flashinfer)
attention_bias (Tensor, optional) – Bias tensor for Q * K.T.
inference_context (object, optional) – Parameters for inference-time optimizations.
packed_seq_params (object, optional) – Parameters for packed sequence processing.
sequence_len_offset (Tensor, optional) – Offset along sequence dimension during inference.

Returns:

A tuple containing: hidden_states (Tensor): Transformed hidden states before the MLP layernorm. context (Tensor): Updated context tensor if cross-attention is used, otherwise None.

Return type:

Tuple[Tensor, Tensor]

forward_pre_attn_and_core_attn( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, rotary_pos_emb: Optional[torch.Tensor] = None, attention_bias: Optional[torch.Tensor] = None, packed_seq_params: Optional[megatron.core.packed_seq_params.PackedSeqParams] = None, *, packed_sequence_cp_metadata=None, )#: Run the training path through pre-attention and core attention.

forward_post_core_attn( attention_intermediate: torch.Tensor, residual: torch.Tensor, context: Optional[torch.Tensor] = None, attn_state=(), context_mask: Optional[torch.Tensor] = None, )#: Run the training path after core attention.

_run_input_layernorm(hidden_states)#

Run input layernorm with optional output-discarding checkpoint and fine-grained activation offloading.

Sets self._input_layernorm_checkpoint_active so the caller can gate the post-attention discard-and-register hook on the same condition. The flag is consumed by the next self._apply_self_attn_bda_step step.

Returns:: Tuple (input_layernorm_output, residual, attn_state) where attn_state is an opaque payload subclasses can use to thread extra intermediates (e.g. mHC h_res/h_post) through to _apply_self_attn_bda_step. Base returns ().

_apply_self_attn_bda_step( attention_output_with_bias, residual, attn_state=(), )#

bias-dropout-add for self-attention output + post-step offload commit.

Subclasses override this to swap in a fused kernel that consumes extra intermediates threaded via attn_state (the third element returned by _run_input_layernorm). Base ignores attn_state.

_run_cross_attention( hidden_states, context, context_mask, inference_context, )#: Optional pre-cross-attn layernorm + cross-attention + bda block.

forward(*args, **kwargs)#

Perform a forward pass through the transformer layer.

This method calls the core computation of a transformer layer, including self-attention, cross-attention (if applicable), and feed-forward operations.

_forward_pre_mlp_layernorm(hidden_states: torch.Tensor)#

_maybe_unflatten_for_moe( hidden_states, padding_mask, packed_seq_params, )#

Un-flatten packed sequences to restore the batch dimension for MoE.

When inter-document masking flattens MBS > 1 into [mbs*S, 1, H], the MoE router sees bsz=1 and computes seq_aux_loss over the entire flattened sequence instead of per sample. Un-flattening to [S, mbs, H] before the MoE layer restores the correct per-sample structure.

Returns:: (hidden_states, padding_mask, mbs) where mbs is None if no un-flattening was applied.

_maybe_reflatten_from_moe(output, packed_seq_params, mbs)#: Re-flatten MoE output back to [mbs*S, 1, H] for the residual add.

_pre_mlp_layernorm_and_residual(hidden_states)#

Run pre-MLP layernorm (with optional recompute and offload), unpack a tuple-output layernorm, and apply the fp32-residual cast.

Returns:: Tuple (pre_mlp_layernorm_output, residual, mlp_state) where mlp_state is an opaque payload subclasses can use to thread extra intermediates (e.g. mHC mlp_h_res / mlp_hc_h_post) through to _apply_mlp_bda_step. Base returns ().

_forward_mlp_output_with_bias( hidden_states: torch.Tensor, inference_context: megatron.core.inference.contexts.BaseInferenceContext | None = None, padding_mask: torch.Tensor | None = None, packed_seq_params=None, ) → tuple[tuple[torch.Tensor, torch.Tensor | None], torch.Tensor]#: Run pre-MLP norm and MLP/MoE, returning the raw output before BDA.

_forward_mlp( hidden_states: torch.Tensor, inference_context: megatron.core.inference.contexts.BaseInferenceContext | None = None, padding_mask: torch.Tensor | None = None, packed_seq_params=None, ) → torch.Tensor | list[torch.Tensor | None]#

Perform a forward pass through the feed-forward layer.

Parameters:

hidden_states (Tensor) – Transformed hidden states before the MLP layernorm. Shape [seq_length, batch_size, hidden_size].
inference_context – Inference context for optimizations.
padding_mask (Tensor, optional) – Padding mask for MoE routing. Shape [bsz, seq_length]. True = padding (exclude), False = valid (include). Only used for MoE layers to exclude padding tokens from aux loss computations. The MoELayer will internally transform this to [seq_length, bsz] format.
packed_seq_params – Packed sequence parameters, used to detect flattened batches that need reshaping for MoE sequence load balancing.

Returns:

Transformed hidden states of shape [s, b, h].

Return type:

output (Tensor)

_run_mlp( pre_mlp_layernorm_output: torch.Tensor, residual: torch.Tensor, padding_mask: torch.Tensor | None, inference_context: megatron.core.inference.contexts.BaseInferenceContext | None, )#

Execute the MLP submodule with the appropriate variant.

Picks between the recompute (te_checkpoint / tensor_parallel.checkpoint), chunked-prefill, and direct-call paths. Shared by both

Class:: TransformerLayer and :class:HyperConnectionTransformerLayer so the MLP-call branching stays in one place.
Returns:: tuple of (mlp_output, mlp_bias).
Return type:: mlp_output_with_bias

_apply_mlp_bda_step( mlp_output_with_bias: tuple[torch.Tensor, torch.Tensor | None], residual: torch.Tensor, mlp_state: tuple = (), ) → torch.Tensor#

Perform operations after the MLP computation: bias-dropout-add for the MLP output + post-step offload commit + viewless-tensor wrap.

Subclasses override this to swap in a fused kernel that consumes extra intermediates threaded via mlp_state (the third element returned by _pre_mlp_layernorm_and_residual). Base ignores mlp_state.

Parameters:

mlp_output_with_bias (Tensor) – Output tensor of the MLP layer with bias.
residual (Tensor) – Residual tensor.
mlp_state – Opaque payload from _pre_mlp_layernorm_and_residual. Default ().

Returns:

Transformed hidden states of shape [s, b, h].

Return type:

output (Tensor)

sharded_state_dict( prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None, ) → megatron.core.dist_checkpointing.mapping.ShardedStateDict#

Generate a sharded state dictionary for the transformer layer.

Parameters:

prefix (str, optional) – Prefix to be added to all keys in the state dict.
sharded_offsets (tuple, optional) – Tuple of sharding offsets.
metadata (Optional[dict], optional) – Additional metadata for sharding.

Returns:

A dictionary containing the sharded state of the transformer layer.

Return type:

ShardedStateDict

configure_fused_tp_inference( skip_qkv_norm_and_all_gather: bool = False, fc2_next_layer_norm_weights: Optional[torch.Tensor] = None, )#

Configure settings for fused TP communication in inference mode.

Parameters:

skip_qkv_norm (bool) – Whether to skip norm and all-gather for linear_qkv.
fc2_next_layer_norm_weights (Optional[Tensor]) – Next layer’s QKV norm weights for current layer’s MLP FC2.

_set_proj_next_layer_norm_weights(weights: torch.Tensor)#: Set next layer norm weights for attention/mixer’s linear_proj.

_set_fc2_next_layer_norm_weights( weights: Optional[torch.Tensor], )#: Set next layer norm weights for MLP FC2.

_set_proj_residual(residual: torch.Tensor)#: Set residual for attention’s/mixer’s out_proj (linear_proj).

_set_fc2_residual(residual: torch.Tensor)#: Set residual for MLP FC2.

get_mlp_layer_norm_weights() → torch.Tensor#

Get the MLP FC1 layer norm weights.

Returns:: The layer norm weight data.
Return type:: Tensor

get_qkv_layer_norm_weights() → torch.Tensor#

Get the QKV layer norm weights.

Returns:: The layer norm weight data.
Return type:: Tensor

get_layer_static_inputs(seq_length, micro_batch_size)#

Get the static inputs for the transformer layer. Besides the hidden_states that is generated in GraphableMegatronModule, we also add the attention_mask.

Returns:: A dictionary containing the static inputs for the layer.
Return type:: Dict[str, torch.Tensor]

_get_submodules_under_cudagraphs()#: Get the submodules that are covered by cudagraphs.

_te_cuda_graph_capture(*args, **kwargs)#

CUDA Graph capture for this layer using TE interface. There are some differences from the normal pass:

In some conditions CUDA graph cannot cover the entire layer. The cuda_graph_modules attribute can be set to control the scope of the CUDA graph.
If context is None, it cannot be returned as output.

_te_cuda_graph_replay(*args, **kwargs)#: CUDA graph replay for this layer and microbatch self.current_microbatch using TE interface. TransformerEngine versions>=1.10 allow keyword arguments with CUDA graph. However, CUDA graph accepts only Tensor inputs. Hence, inference_context and packed_seq_params are excluded from input list.

_te_cuda_graph_replay_impl(args, kwargs, context)#: Implementation of _te_cuda_graph_replay, separated for replay mode cleanup.

_get_te_cuda_graph_replay_args(*args, **kwargs)#: Helper function to get tensor arguments for TE CUDA graph.

_should_call_local_cudagraph(*args, **kwargs)#: Check if we should call the local cudagraph path.

_set_offload_modules()#: Set the offload modules for the transformer layer.

get_layer_norm_weights()#

Get the weights of all layernorms (attention and MLP) in the transformer layer.

Returns:: A list of layernorm weight tensors.
Return type:: List[Tensor]

class core.transformer.transformer_layer.HyperConnectionTransformerLayer( config: megatron.core.transformer.transformer_config.TransformerConfig, submodules: core.transformer.transformer_layer.TransformerLayerSubmodules, layer_number: int = 1, hidden_dropout: Optional[float] = None, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, vp_stage: Optional[int] = None, )#

Bases: core.transformer.transformer_layer.TransformerLayer

A transformer layer with Manifold-Constrained Hyper-Connections (mHC).

Extends TransformerLayer by adding hyper connection modules around self-attention and MLP. The n-stream hidden states are aggregated before each sub-layer and expanded back afterwards using learned mappings (H_pre, H_post, H_res).

Cross-attention hyper connection is not supported.

Initialization

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

supports_mhc_connections: bool#: True

__call__(*args, **kwargs)#

get_layer_static_inputs(seq_length, micro_batch_size)#

Override to produce n-stream hidden_states of shape [s, b, n*C].

CUDA graph capture creates static buffers whose shapes are determined by this method. The base class returns [s, b, C], but mHC layers operate on n-stream hidden states of shape [s, b, n*C].

create_mcore_cudagraph_manager(config)#: Create only CUDA graph managers compatible with this mHC layer.

_get_submodules_under_cudagraphs()#

Override to include hyper connection modules.

The base TransformerLayer.get_submodules_under_cudagraphs does not include self_attention_hyper_connection / mlp_hyper_connection. Their learnable parameters (mapping_proj, alpha*, bias) need manual pre-forward hooks during CUDA graph replay so that parameter all-gathers are triggered.

_run_input_layernorm(hidden_states)#

HC input layernorm: hyper-connection pre-wrap + mHC-aware checkpoint.

Threads h_res and h_post (produced by the hyper-connection pre-wrap) to _apply_self_attn_bda_step via the attn_state slot in the return tuple. Also sets self._input_layernorm_checkpoint_active for the post-self-attn discard hook.

Returns (input_layernorm_output, residual, (h_res, h_post)) where residual is the n-stream hidden state captured before aggregation — it flows to _apply_self_attn_bda_step via the base skeleton’s residual argument, and (h_res, h_post) flows via attn_state.

static _reject_residual_returning_norm(layernorm_output, norm_name)#

Reject layernorms that also return a residual.

Base TransformerLayer accepts a (output, residual) tuple from a layernorm and uses the returned residual. mHC cannot: its residual must be the n-stream tensor captured before the hyper-connection aggregates n-stream -> single-stream, because fused_h_res_h_post_bda mixes it with H_res. A norm-supplied single-stream residual would silently produce wrong shapes, so fail fast instead.

_apply_self_attn_bda_step( attention_output_with_bias, residual, attn_state=(), )#

HC fused bias-dropout-add: combines apply_h_res + apply_h_post + bda.

Unpacks h_res and h_post from attn_state (threaded by _run_input_layernorm via the base skeleton). Keeps the base class’s default so base-class call sites that pass only two arguments fail with a readable message rather than a TypeError.

_pre_mlp_layernorm_and_residual(hidden_states)#

HC pre-mlp layernorm: hyper-connection pre-wrap + mHC-aware checkpoint.

Threads mlp_h_res and mlp_hc_h_post (produced by the hyper-connection pre-wrap) to _apply_mlp_bda_step via the mlp_state slot in the return tuple.

Returns (pre_mlp_layernorm_output, residual, (mlp_h_res, mlp_hc_h_post)) where residual is the n-stream hidden state captured before aggregation — it flows to _apply_mlp_bda_step via the base skeleton’s residual argument, and (mlp_h_res, mlp_hc_h_post) flows via mlp_state.

_apply_mlp_bda_step(mlp_output_with_bias, residual, mlp_state=())#

HC fused bias-dropout-add for MLP: combines apply_h_res + apply_h_post + bda.

Unpacks mlp_h_res and mlp_hc_h_post from mlp_state (threaded by _pre_mlp_layernorm_and_residual via the base skeleton). Computes the per-call mhc_mlp_bda_manager from self._mhc_recompute_manager: the last layer of a recompute block does NOT pass the manager into the fused-bda checkpoint — the block-end finalize hook handles its output discard.

Keeps the base class’s mlp_state default so base-class call sites that pass only two arguments (_te_cuda_graph_replay_impl, MoETransformerLayer._forward_mlp_postprocess) fail with a readable message rather than a TypeError.

class core.transformer.transformer_layer.MoETransformerLayer(*args, **kwargs)#

Bases: core.transformer.transformer_layer.TransformerLayer

A Transformer layer specialized for Mixture-of-Experts (MoE) architectures.

Implements specific functionality to support CUDA graph capture for MoE layers. Due to the dynamic nature of MoE, capturing the entire layer in a single CUDA graph can be challenging. This class supports “partial” CUDA graphs by decomposing the MLP forward pass into router, expert-compute, and post-process stages.

Initialization

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

_should_call_local_cudagraph(*args, **kwargs)#

Controls whether the full-layer cudagraph_manager captures the entire forward call as a single graph. Returns False to skip full-layer capture and route through _forward_mlp.

MoE layers have two cudagraph modes:

Full-layer (use_partial_cudagraphs=False): the full-layer cudagraph_manager captures the forward pass as one graph. This is used during inference.
Partial (use_partial_cudagraphs=True): the full-layer manager is bypassed (returns False), and _forward_mlp routes through cudagraph_manager_router and cudagraph_manager_postprocess, which are monkey-patched onto _forward_mlp_router and _forward_mlp_postprocess by CudaGraphManager.init. The expert dispatch in between runs eagerly. This is used during training.

transition_cudagraph_scope(mode)#

Transition between full-layer and partial CUDA graph capture.

Parameters:

mode – ‘full’ for inference (full-layer capture) or ‘partial’ for training
captured ((router + postprocess)
eagerly). (expert dispatch runs)

create_mcore_cudagraph_manager(config)#

Initializes the CUDA graph manager(s) for the MoE layer.

Unlike the standard layer which typically uses a single manager, this method can configure multiple graph managers if partial CUDA graphs are enabled via cuda_graph_modules. This allows capturing the static parts of the MoE pass while leaving the expert computation to execute eagerly.

_resolve_token_dispatcher_attr( attr_name: str, ) → tuple[Any, str]#

_restore_token_dispatcher_attrs(attr_outputs)#

_get_token_dispatcher_attrs()#

_synchronize_router_host_outputs(attr_outputs)#: Wait for partial-router graph outputs only when they reside on the host.

_forward_mlp_router(hidden_states, padding_mask=None)#

Executes the router phase of the MoE block.

This includes the pre-MLP layernorm and the routing logic. This method is isolated so it can be captured by cudagraph_manager_router.

_forward_mlp_expert_compute( hidden_states, probs, token_dispatcher_attr_outputs, )#

Executes the actual computation of the experts.

This phase takes the routing information and inputs, dispatches them to the appropriate experts, and computes the results. In partial graph modes, this step runs eagerly between the router and postprocess graph replays.

_forward_mlp_postprocess( residual, output, shared_expert_output, mlp_bias, )#

Executes the post-processing phase of the MoE block.

Handles combining the expert outputs, applying biases, re-registering activation recomputation hooks if necessary, and performing the final Bias-Dropout-Add. This method is isolated so it can be captured by cudagraphs.

_forward_mlp( hidden_states, inference_context=None, padding_mask=None, packed_seq_params=None, )#

Orchestrates the MLP forward pass, handling partial CUDA graph execution logic.

If use_partial_cudagraphs is True, this method stitches together the router, expert_compute, and postprocess calls.

core.transformer.transformer_layer#

Module Contents#

Classes#

Functions#

Data#

API#

`core.transformer.transformer_layer`#