core.tensor_parallel.inference_layers

Module Contents

Classes

InferenceLinear	Inference optimized version of TELinear.
InferenceLayerNormColumnParallelLinear	Inference optimized version of TELayerNormColumnParallelLinear.
InferenceColumnParallelLinear	Inference optimized version of TEColumnParallelLinear.
InferenceRowParallelLinear	Inference optimized version of TERowParallelLinear.

Functions

_te_rms_norm_kernel
_apply_linear	Helper to apply either MXFP8 or standard GEMM based on the configuration.
is_inference_column_parallel_linear	Whether module is an inference-optimized column-parallel linear.
inference_all_gather_from_tensor_model_parallel_region	NVLS-optimized all-gather along the last dimension, with NCCL fallback.
inference_reduce_scatter_to_sequence_parallel_region	NVLS-optimized reduce-scatter along the first dimension, with NCCL fallback.

API

core.tensor_parallel.inference_layers._te_rms_norm_kernel(x: torch.Tensor, weight: torch.Tensor, eps: float)

core.tensor_parallel.inference_layers._apply_linear(

x: torch.Tensor,

weight: Union[torch.Tensor, megatron.core.inference.quantization.mxfp8_tensor.MXFP8Tensor],

config: megatron.core.transformer.transformer_config.TransformerConfig,

out: Optional[torch.Tensor] = None,

) → torch.Tensor

Helper to apply either MXFP8 or standard GEMM based on the configuration.

class core.tensor_parallel.inference_layers.InferenceLinear(

input_size: int,

output_size: int,

*,

parallel_mode: Optional[str],

config: megatron.core.model_parallel_config.ModelParallelConfig,

init_method: Callable,

bias: bool,

skip_bias_add: bool,

skip_weight_param_allocation: bool,

tp_comm_buffer_name: Optional[str] = None,

is_expert: bool = False,

symmetric_ar_type: Optional[str] = None,

tp_group: Optional[torch.distributed.ProcessGroup] = None,

name: str | None = None,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

)

Bases: megatron.core.extensions.transformer_engine.TELinear

Inference optimized version of TELinear.

Initialization

Parameters:

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

forward(x: torch.Tensor) → Tuple[torch.Tensor, None]

Forward pass.

class core.tensor_parallel.inference_layers.InferenceLayerNormColumnParallelLinear(

input_size: int,

output_size: int,

*,

config: megatron.core.transformer.transformer_config.TransformerConfig,

init_method: Callable,

gather_output: bool,

bias: bool,

skip_bias_add: bool,

is_expert: bool,

stride: int = 1,

skip_weight_param_allocation: bool = False,

tp_comm_buffer_name: Optional[str] = None,

tp_group: Optional[torch.distributed.ProcessGroup] = None,

name: str | None = None,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

)

Bases: megatron.core.extensions.transformer_engine.TELayerNormColumnParallelLinear

Inference optimized version of TELayerNormColumnParallelLinear.

Initialization

Parameters:

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

• pg_collection (ProcessGroupCollection | None) – process groups used by this layer. Falls back to the MPU global process groups when not given.

set_barrier_before_all_gather(value: bool = True) → None

Request a barrier before this layer’s input all-gather reuses the buffer.

Set by callers whose op sequence places another all-gather on the shared symmetric buffer immediately before this layer’s all-gather (e.g. the MTP eh_proj projection), so the kernel synchronizes ranks before overwriting.

_maybe_allocate_symmetric_buffer(x: torch.Tensor)

Attempt to allocate symmetric memory buffer for all-gather.

_all_gather(x: torch.Tensor, symm_mem_buffer: dict) → None

Attempt an NVLS all-gather into symmetric memory. If not possible, revert to torch dist (NCCL) all-gather.

forward(x: torch.Tensor) → Tuple[torch.Tensor, None]

Forward pass.

class core.tensor_parallel.inference_layers.InferenceColumnParallelLinear(

input_size: int,

output_size: int,

*,

config: megatron.core.transformer.transformer_config.TransformerConfig,

init_method: Callable,

gather_output: bool,

bias: bool,

skip_bias_add: bool,

is_expert: bool,

stride: int = 1,

skip_weight_param_allocation: bool = False,

tp_comm_buffer_name: Optional[str] = None,

tp_group: Optional[torch.distributed.ProcessGroup] = None,

name: str | None = None,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

)

Bases: megatron.core.extensions.transformer_engine.TEColumnParallelLinear

Inference optimized version of TEColumnParallelLinear.

Initialization

Parameters:

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

• pg_collection (ProcessGroupCollection | None) – process groups used by this layer. Falls back to the MPU global process groups when not given.

set_barrier_before_all_gather(value: bool = True) → None

Request a barrier before this layer’s input all-gather reuses the buffer.

Set by callers whose op sequence places another all-gather on the shared symmetric buffer immediately before this layer’s all-gather (e.g. the MTP eh_proj projection), so the kernel synchronizes ranks before overwriting.

_maybe_allocate_symmetric_buffer(x: torch.Tensor)

Attempt to allocate symmetric memory buffer for all-gather.

_all_gather(x: torch.Tensor, symm_mem_buffer: dict) → None

Attempt an NVLS all-gather into symmetric memory. If not possible, revert to torch dist (NCCL) all-gather.

forward(x: torch.Tensor) → Tuple[torch.Tensor, None]

Forward pass.

class core.tensor_parallel.inference_layers.InferenceRowParallelLinear(

input_size: int,

output_size: int,

*,

config: megatron.core.model_parallel_config.ModelParallelConfig,

init_method: Callable,

bias: bool,

input_is_parallel: bool,

skip_bias_add: bool,

is_expert: bool,

tp_comm_buffer_name: Optional[str] = None,

tp_group: Optional[torch.distributed.ProcessGroup] = None,

name: str | None = None,

pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None,

)

Bases: megatron.core.extensions.transformer_engine.TERowParallelLinear

Inference optimized version of TERowParallelLinear.

Initialization

Parameters:

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

• pg_collection (ProcessGroupCollection | None) – process groups used by this layer. Falls back to the MPU global process groups when not given.

_matmul_reduce_scatter(x, residual=None)

Multiplies x by the weight matrix and performs a reduce-scatter. It will first try to write the matmul output to symmetric memory and perform an NVLS multicast reduce-scatter. If that is not possible, it will revert to torch.dist (NCCL) reduce-scatter.

_set_next_layer_norm_weights(weights: torch.Tensor)

Set next layer norm weights for fused reduce-scatter + add + rms-norm + all-gather.

_set_residual(residual: torch.Tensor)

Set residual for fused reduce-scatter + add + rms-norm + all-gather.

forward(

x: torch.Tensor,

residual: Optional[torch.Tensor] = None,

) → tuple[torch.Tensor, None]

Forward pass.

core.tensor_parallel.inference_layers.is_inference_column_parallel_linear(module) → bool

Whether module is an inference-optimized column-parallel linear.

These are the layers that perform a symmetric-memory all-gather and therefore expose set_barrier_before_all_gather. Returns False for anything else (including None and non-inference linear implementations).

core.tensor_parallel.inference_layers.inference_all_gather_from_tensor_model_parallel_region(

x: torch.Tensor,

tp_group: torch.distributed.ProcessGroup,

config: megatron.core.transformer.transformer_config.TransformerConfig,

barrier_before: bool = False,

) → torch.Tensor

NVLS-optimized all-gather along the last dimension, with NCCL fallback.

Replaces gather_from_tensor_model_parallel_region in inference paths where autograd is not needed and NVLS symmetric-memory is available.

The NVLS path performs a flat all-gather into symmetric memory (concatenating along dim-0), then rearranges the result to the last dimension — the same semantics as _gather_along_last_dim but using hardware multicast when possible.

barrier_before is forwarded to multimem_all_gather: pass True when this all-gather directly follows another all-gather on the shared symmetric buffer so it barriers before overwriting the previous contents.

core.tensor_parallel.inference_layers.inference_reduce_scatter_to_sequence_parallel_region(

x: torch.Tensor,

tp_group: torch.distributed.ProcessGroup,

config: megatron.core.transformer.transformer_config.TransformerConfig,

) → torch.Tensor

NVLS-optimized reduce-scatter along the first dimension, with NCCL fallback.

Replaces reduce_scatter_to_sequence_parallel_region in inference paths where autograd is not needed and NVLS symmetric-memory is available.