feat: add tract_moe_ffn operator for Mixture-of-Experts FFN

harness/nnef-test-cases/moe-ffn/qwen3-tiny/graph.nnef

@@ -0,0 +1,39 @@
+version 1.0;
+
+extension tract_registry tract_core;
+extension tract_registry tract_transformers;
+
+fragment tract_core_properties(
+) -> (properties: (string, tensor<scalar>)[])
+{
+  properties = [
+    ("tract_target_version", "0.22.0"),
+    ("torch_to_nnef_version", "0.21.0"),
+    ("torch_version", "2.6.0"),
+    ("transformers_version", "5.5.0"),
+    ("os", "Darwin SNS009332 24.6.0 Darwin Kernel Version 24.6.0: Mon Jul 14 11:30:40 PDT 2025; root:xnu-11417.140.69~1/RELEASE_ARM64_T6041 arm64 Darwin"),
+    ("hostname", "SNS009332"),
+    ("user", "julien.balian"),
+    ("py_version", "3.12.10 (main, Apr  9 2025, 03:49:38) [Clang 20.1.0 ] (64-bit runtime)"),
+    ("export_date", "2026-04-03 11:27:03.768935"),
+    ("exported_py_class", "Qwen3TinyMoE"),
+    ("export_cmd", "scripts/export_moe_test_asset.py /Users/julien.balian/SONOS/src/tract/harness/nnef-test-cases/moe-ffn/qwen3-tiny")
+  ];
+}
+
+
+
+
+
+
+
+graph network(input_0) -> (output_0)
+{
+    input_0 = tract_core_external(shape = [1, 3, 16], datum_type = 'f32');
+    moe_gate_weight = variable<scalar>(label = 'moe.gate.weight', shape = [4, 16]);
+    output_0_w1 = variable<scalar>(label = 'output_0_w1', shape = [4, 16, 32]);
+    output_0_w2 = variable<scalar>(label = 'output_0_w2', shape = [4, 32, 16]);
+    output_0_w3 = variable<scalar>(label = 'output_0_w3', shape = [4, 16, 32]);
+    moe_gate_weight_aligned_rank_expanded = unsqueeze(moe_gate_weight, axes = [0]);
+    output_0, _router_logits = tract_moe_ffn(input_0, moe_gate_weight_aligned_rank_expanded, output_0_w1, output_0_w2, output_0_w3, k = 2, activation = 'swiglu', normalize_gates = true);
+}

harness/nnef-test-cases/moe-ffn/qwen3-tiny/runme.sh

@@ -0,0 +1,8 @@
+#!/bin/sh
+
+cd `dirname $0`
+set -ex
+
+: ${TRACT_RUN:=cargo run -p tract $CARGO_OPTS --}
+
+$TRACT_RUN --nnef-tract-core --nnef-tract-transformers . run --input-from-bundle io.npz --assert-output-bundle io.npz

transformers/Cargo.toml

@@ -23,3 +23,6 @@ tract-nnef.workspace = true
 # the op stays single-threaded.
 [target.'cfg(not(target_family = "wasm"))'.dependencies]
 rayon.workspace = true
+
+[dev-dependencies]
+ndarray-npy.workspace = true

transformers/src/lib.rs

@@ -22,6 +22,7 @@ pub fn register(registry: &mut Registry) {
     ops::sdpa::register(registry);
     ops::dyn_kv_cache::register(registry);
     ops::window_kv_cache::register(registry);
+    ops::moe_ffn::register(registry);
 }
 
 pub trait WithTractTransformers {

transformers/src/ops/mod.rs

@@ -2,6 +2,7 @@ pub mod apply_rope;
 pub mod diag_gather;
 pub mod dyn_kv_cache;
 pub mod flash_sdpa;
+pub mod moe_ffn;
 pub mod scaled_masked_softmax;
 pub mod sdpa;
 pub mod streamed_sdpa;