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_moe`	Returns whether the module is an MOE 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)#

Get the list of linear names for the experts.

Parameters:: module (Module)
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) – type(root_model).__name__.lower() (may change after ModelOpt quantize).

is_layernorm(module)#

Returns whether the module is a layernorm layer.

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

is_moe(module)#

Returns whether the module is an MOE 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)#

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.

Returns:: Number of expert gate/up pairs whose amaxes were synced.
Parameters:: model (Module)
Return type:: int