layer_utils#

Module-shape predicates and MoE quantizer helpers shared by every export backend.

The TensorRT-LLM build_*_config builders that turn these modules into checkpoint configs live in modelopt.torch.export.trtllm.layer_utils.

Functions

`get_expert_linear_names`	Get the list of linear names for the experts.
`get_experts_list`	Returns list of grouped experts by linear name for given module.
`is_layernorm`	Returns whether the module is a layernorm layer.
`is_quantlinear`	Returns whether the module is a quantized linear layer.
`set_expert_quantizer_amax`	Set amax values for expert quantizers using smart fallback logic.
`sync_moe_gate_up_amax`	Take element-wise max of gate and up weight quantizer amaxes per expert.

get_expert_linear_names(module, model_type)#

Get the list of linear names for the experts.

The model’s own spec always wins where it has an answer for this module’s layout. That qualifier matters: a layout’s names describe one layout, and the same model type materializes per-expert on transformers 4 and fused on 5, so naming from the wrong layout would be wrong rather than merely generic. The structural fallbacks below run only where the spec declines.

Raises NotImplementedError when nothing resolves, so a new MoE model fails loudly instead of silently inheriting another model’s naming.

Parameters:

module (Module) – the MoE block.
model_type (str | None) – the model’s HF model type (model.config.model_type).

Return type:

list[str]

get_experts_list(module, model_type)#

Returns list of grouped experts by linear name for given module.

Parameters:

module (Module) – MoE block (e.g. MixtralSparseMoeBlock, NemotronHMOE).
model_type (str | None) – the model’s HF model type (model.config.model_type), used to resolve the model’s own spec.

is_layernorm(module)#

Returns whether the module is a layernorm layer.

Parameters:: module (Module)
Return type:: bool

is_quantlinear(module)#

Returns whether the module is a quantized linear layer.

Parameters:: module (Module)
Return type:: bool

set_expert_quantizer_amax(modules, quantizer_attrs=None, fallback_value=0.5, device=None)#

Set amax values for expert quantizers using smart fallback logic.

Uses smart fallback logic:

Use max from existing quantizers in current batch (best - direct from calibration)
If no existing values found, then: - For weight quantizers: calculate from weight statistics - For input quantizers: use max from other experts, fallback if none found
Use fallback value as last resort

This ensures we always have semantically appropriate amax values for export.

Parameters:

modules (Module | list[Module]) – Single module or list of modules containing quantizers
quantizer_attrs (str | list[str] | None) – Specific quantizer attributes to handle. If None, defaults to [“input_quantizer”] for backward compatibility.
fallback_value (float) – Final fallback value when other methods fail (default: 0.5)
device (device | None) – Target device for tensors (auto-detected if None)

Returns:

a list of uncalibrated experts

Return type:

uncalibrated_modules

sync_moe_gate_up_amax(model, model_type=None)#

Take element-wise max of gate and up weight quantizer amaxes per expert.

Serving engines fuse gate_proj and up_proj into a single gate_up_proj and require a single weight_scale_2. Since weight_scale_2 = amax / (6 * m_fp8) (m_fp8=448 normally, 256 for NVFP4 4/6 mode), syncing amaxes before quantization ensures the per-block weight_scale values are computed against a consistent global scale.

Only affects standard MoE models with separate gate/up linear layers (e.g. Qwen MoE, DeepSeek). Models with already-fused gate_up_proj (e.g. Llama4, GptOss) are unaffected.

model_type is the root model’s HF model type; callers passing a sub-tree (layerwise export passes one decoder layer) must supply it, since it cannot be resolved from a decoder layer.

Returns:

Number of expert gate/up pairs whose amaxes were synced.

Parameters:

model (Module)
model_type (str | None)

Return type:

int