core.transformer.moe.megakernel.mok.weights

MOK-specific adaptation of MCore-owned expert parameters.

Module Contents

Functions

prepare_shared_expert_config	Construct MOK shared experts as native BF16 MCore modules.
_storage_view	Return a zero-copy dense view over a TE grouped backing tensor.
_single_grouped_mxfp8_scale_view	Expose one single-grouped parameter’s expert-major MXFP8 scale storage.
_swizzle_mxfp8_scale	Convert TE’s logical E8M0 matrix to MOK’s tcgen05 scale layout.
_native_single_grouped_weight_view	Build or refresh one MOK view over a native TE single-grouped parameter.
_parameter_storage_attr	Read a TE storage attribute without materializing the logical parameter.
_native_split_weight_view	Build or refresh one MOK view over independent per-expert parameters.

Data

_SHARED_EXPERT_BF16_WARNING_EMITTED

API

core.transformer.moe.megakernel.mok.weights._SHARED_EXPERT_BF16_WARNING_EMITTED

False

core.transformer.moe.megakernel.mok.weights.prepare_shared_expert_config(

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

) → megatron.core.transformer.transformer_config.TransformerConfig

Construct MOK shared experts as native BF16 MCore modules.

Routed experts continue to use the original model configuration. The shared-only copy expresses the BF16 module configuration; the factory also disables any enclosing TE FP8 model-init context during construction.

core.transformer.moe.megakernel.mok.weights._storage_view(

storage: torch.Tensor,

shape: tuple[int, ...],

*,

dtype: torch.dtype,

name: str,

) → torch.Tensor

Return a zero-copy dense view over a TE grouped backing tensor.

core.transformer.moe.megakernel.mok.weights._single_grouped_mxfp8_scale_view(

param: torch.nn.Parameter,

member_attr: str,

shape: tuple[int, ...],

*,

name: str,

) → torch.Tensor

Expose one single-grouped parameter’s expert-major MXFP8 scale storage.

core.transformer.moe.megakernel.mok.weights._swizzle_mxfp8_scale(

logical_scale: torch.Tensor,

*,

rows: int,

columns: int,

out: torch.Tensor | None = None,

) → torch.Tensor

Convert TE’s logical E8M0 matrix to MOK’s tcgen05 scale layout.

TE stores rowwise scales as [E, M, K / 32]. MOK consumes the tcgen05 1x scale-factor layout [E * M / 128, K / 128, 32, 16]. Only the scale bytes are copied; the much larger FP8 payload stays in native TE storage.

core.transformer.moe.megakernel.mok.weights._native_single_grouped_weight_view(

weight: torch.nn.Parameter,

*,

num_experts: int,

rows: int,

columns: int,

use_mxfp8: bool,

cached_view=None,

)

Build or refresh one MOK view over a native TE single-grouped parameter.

core.transformer.moe.megakernel.mok.weights._parameter_storage_attr(

param: torch.nn.Parameter,

name: str,

) → torch.Tensor | None

Read a TE storage attribute without materializing the logical parameter.

TE releases/wrappers have exposed native storage either directly on the Parameter or on param.data, using public or private attribute names.

core.transformer.moe.megakernel.mok.weights._native_split_weight_view(

params: tuple[torch.nn.Parameter, ...],

*,

rows: int,

columns: int,

use_mxfp8: bool,

cached_view=None,

)

Build or refresh one MOK view over independent per-expert parameters.