torch_to_nnef.op.custom_extractors.moe

MoE FFN export to tract_moe_ffn operator (tract_transformers extension).

Supports (transformers): - MoEFFN (reference wrapper) - MixtralSparseMoeBlock / MistralSparseMoeBlock - Mistral4MoE (single-group routing, ungated shared expert) - Gemma4RoutedExperts (gated GELU, folded router/input/output scales) - GptOssMLP (router + expert biases, interleaved gate/up, clamped SwiGLU) - Qwen2/Qwen3/Qwen3.5/Qwen3-Next MoE (shared expert decomposed outside) - OlmoeSparseMoeBlock

All variants are normalized to the same tract_moe_ffn signature

inputs: x [T,D], wg [E,D], w1 [E,D,H], w2 [E,H,D], w3 [E,D,H], optional biases (wg_bias, w1_bias, w3_bias, w2_bias) attrs: k (int), activation (str), gate (softmax_topk | softmax_all | sigmoid | raw), optional act_alpha / act_limit (clamped SwiGLU), optional expert_layout ("canonical" | "linear") output: y [T,D]

The default canonical expert layout matches the signature above. The optional linear layout stores expert weights in their native linear-filter orientation: w1/w3 [E,H,D], w2 [E,D,H] Layout selection is independent from whether those tensors are quantized. A shared expert (Qwen2 / Qwen3.5) is emitted as a standard NNEF subgraph added on top of the routed output, not baked into the op.

Classes:

Name	Description
Gemma4RoutedExperts	Gemma4 routing and experts on an unweighted RMS-normalized input.
Gemma4RoutedExpertsExtractor	Export the Gemma4 boundary installed by the decoder handler.
MoEFFN	Reference MoE FFN block for export testing.
MoEFFNExtractor	Extractor for the reference MoEFFN wrapper.

Functions:

Name	Description
mark_moe_experts_for_q40	Mark every MoE block of model for export-time expert handling.

Gemma4RoutedExperts

Gemma4RoutedExperts(router, experts, expert_norm)

Bases: Module

Gemma4 routing and experts on an unweighted RMS-normalized input.

The decoder handler groups the original FFN into this export boundary. Keep the PyTorch reference unfused; the adapter folds the router scale, expert input norm weight, and per-expert output scale during conversion.

Gemma4RoutedExpertsExtractor

Gemma4RoutedExpertsExtractor()

Bases: MoEFFNExtractor

Export the Gemma4 boundary installed by the decoder handler.

MoEFFN

MoEFFN(num_experts: int, d_model: int, d_hidden: int, k: int = 2, activation: str = 'silu', normalize_gates: bool = True, bias: bool = False)

Bases: Module

Reference MoE FFN block for export testing.

For production models use the transformers extractors directly. Expert biases are supported in eager execution but rejected on export: the operator's positional bias inputs require a SwiGLU expert layout.

Methods:

Name	Description
forward	Forward: [T, D] -> [T, D].

forward

forward(x: Tensor) -> torch.Tensor

Forward: [T, D] -> [T, D].

MoEFFNExtractor

MoEFFNExtractor()

Bases: ModuleInfoExtractor

Extractor for the reference MoEFFN wrapper.

mark_moe_experts_for_q40

mark_moe_experts_for_q40(model: Module, *, quantize_q40: bool = True, layout: Optional[str] = None, module_class_names: Optional[Collection[str]] = None) -> int

Mark every MoE block of model for export-time expert handling.

Sets the module markers the tract_moe_ffn extractor consumes: _t2n_quantize_moe_experts_q40 (quantize expert w1/w2/w3 to tract q4_0 min-max AT EXTRACTION, so a large MoE never materializes dense float experts) and _t2n_moe_expert_layout.

By default any module whose class the MoE extractor itself dispatches on (i.e. with a registered weight adapter) is marked; pass module_class_names to restrict marking to explicit class names.

Returns the number of modules marked; raises if nothing matched (silently exporting dense f16 experts is never what the caller wants).