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

api/rs/Cargo.toml

@@ -28,6 +28,7 @@ tract-extra.workspace = true
 tract-pulse.workspace = true
 tract-libcli.workspace = true
 tract-transformers.workspace = true
+tract-gpu.workspace = true
 serde_json.workspace = true
 
 [target.'cfg(any(target_vendor = "apple"))'.dependencies]

api/rs/src/lib.rs

@@ -496,6 +496,31 @@ impl TensorInterface for Tensor {
     }
 }
 
+impl Tensor {
+    /// Slice `[start, end)` along `axis`. Plain host tensors are sliced by
+    /// copy; tensors backed by device storage (e.g. in-place KV-cache views)
+    /// are sliced as metadata only when possible, without synchronizing or
+    /// copying, keeping the backing device buffer alive.
+    pub fn sliced(&self, axis: usize, start: usize, end: usize) -> Result<Tensor> {
+        use tract_gpu::tensor::{DeviceTensor, DeviceTensorExt};
+        if let Some(dev) = self.0.as_device_tensor() {
+            let sliced = match dev {
+                DeviceTensor::ArenaView(view) => {
+                    DeviceTensor::ArenaView(view.sliced(axis, start, end)?)
+                }
+                DeviceTensor::Owned(_) => {
+                    // Rare: owned device tensors have no shareable arena, so
+                    // go through the host. Correct, but synchronizes.
+                    let host = dev.to_host()?;
+                    return Ok(Tensor(host.slice(axis, start, end)?.into_arc_tensor()));
+                }
+            };
+            return Ok(Tensor(sliced.into_tensor().into_arc_tensor()));
+        }
+        Ok(Tensor(self.0.slice(axis, start, end)?.into_arc_tensor()))
+    }
+}
+
 impl PartialEq for Tensor {
     fn eq(&self, other: &Self) -> bool {
         let Ok((me_dt, me_shape, me_data)) = self.as_bytes() else { return false };

gpu/src/sync.rs

@@ -48,6 +48,12 @@ impl EvalOp for DeviceSync {
                 Ok(tvec![tensor.into_tvalue()])
             }
             DeviceSyncKind::ToDevice => {
+                // Already-device tensors (e.g. cache views fed back by the
+                // caller) pass through untouched: uploading them would read
+                // opaque storage as host bytes and panic.
+                if input.to_device_tensor().is_ok() {
+                    return Ok(tvec![input]);
+                }
                 let device_input = if let Some(t) = input.as_arc_tensor() {
                     Arc::clone(t).into_device()?
                 } else {

gpu/src/tensor/arena_view.rs

@@ -21,6 +21,54 @@ pub struct DeviceArenaView {
 }
 
 impl DeviceArenaView {
+    /// Build a view over any owned device tensor. `shape`/`strides` are in
+    /// elements of `dt`; `offset_bytes` from the buffer start. The backing
+    /// tensor stays alive as long as any view of it does.
+    pub fn from_owned(
+        arena: Arc<Box<dyn OwnedDeviceTensor>>,
+        dt: DatumType,
+        shape: TVec<usize>,
+        strides: TVec<isize>,
+        offset_bytes: usize,
+    ) -> TractResult<Self> {
+        // Unlike arena slots, these views may be non-dense (e.g. the valid
+        // region of a capacity buffer); validate bounds, not density.
+        ensure!(shape.len() == strides.len());
+        ensure!(strides.iter().all(|&s| s >= 0), "negative strides unsupported");
+        let max_index: usize = shape
+            .iter()
+            .zip(strides.iter())
+            .map(|(&d, &s)| d.saturating_sub(1) * s as usize)
+            .sum();
+        let needed = offset_bytes + (max_index + 1) * dt.size_of();
+        let arena_bytes = arena.len() * arena.datum_type().size_of();
+        ensure!(
+            shape.iter().product::<usize>() == 0 || needed <= arena_bytes,
+            "view out of bounds: needs {needed} bytes, arena has {arena_bytes}"
+        );
+        let len = shape.iter().product();
+        Ok(DeviceArenaView { arena, dt, len, shape, strides, offset_bytes, exotic_fact: None })
+    }
+
+    /// Metadata-only slice keeping `[start, end)` along `axis`: same arena,
+    /// same strides, adjusted shape and byte offset. No bytes move, and the
+    /// backing buffer stays alive through the new view's Arc, so other views
+    /// of the arena (e.g. longer KV-cache snapshots) remain valid.
+    pub fn sliced(&self, axis: usize, start: usize, end: usize) -> TractResult<Self> {
+        ensure!(self.exotic_fact.is_none(), "cannot slice a view with an exotic fact");
+        ensure!(axis < self.shape.len(), "axis {axis} out of rank {}", self.shape.len());
+        ensure!(
+            start <= end && end <= self.shape[axis],
+            "invalid slice [{start}, {end}) on axis {axis} of len {}",
+            self.shape[axis]
+        );
+        let mut shape = self.shape.clone();
+        shape[axis] = end - start;
+        let offset_bytes =
+            self.offset_bytes + start * self.strides[axis] as usize * self.dt.size_of();
+        Self::from_owned(self.arena.clone(), self.dt, shape, self.strides.clone(), offset_bytes)
+    }
+
     #[inline]
     pub fn shape(&self) -> &[usize] {
         self.shape.as_slice()
@@ -71,7 +119,49 @@ impl DeviceArenaView {
         } else {
             self.len() * self.dt.size_of()
         };
-        self.arena.get_bytes_slice(self.offset_bytes, len)
+        if self.is_dense() {
+            return self.arena.get_bytes_slice(self.offset_bytes, len);
+        }
+        // Non-dense view: gather row by row (contiguous rows in one slice,
+        // element-wise when the last axis is strided too, e.g. transposed
+        // KV-cache layouts).
+        let esize = self.dt.size_of();
+        let rank = self.shape.len();
+        let row = self.shape[rank - 1];
+        let last_stride = self.strides[rank - 1] as usize;
+        let outer: usize = self.shape[..rank - 1].iter().product();
+        let mut out = Vec::with_capacity(len);
+        for r in 0..outer {
+            let mut rem = r;
+            let mut offset = self.offset_bytes;
+            for ax in (0..rank - 1).rev() {
+                let ix = rem % self.shape[ax];
+                rem /= self.shape[ax];
+                offset += ix * self.strides[ax] as usize * esize;
+            }
+            if last_stride == 1 {
+                out.extend_from_slice(&self.arena.get_bytes_slice(offset, row * esize));
+            } else {
+                for i in 0..row {
+                    out.extend_from_slice(
+                        &self.arena.get_bytes_slice(offset + i * last_stride * esize, esize),
+                    );
+                }
+            }
+        }
+        out
+    }
+
+    /// True when the view's strides are the natural (packed) strides.
+    pub fn is_dense(&self) -> bool {
+        let mut expect = 1isize;
+        for (d, s) in self.shape.iter().zip(self.strides.iter()).rev() {
+            if *d != 1 && *s != expect {
+                return false;
+            }
+            expect *= *d as isize;
+        }
+        true
     }
 
     /// Reshaped tensor with given shape.

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

@@ -0,0 +1,38 @@
+version 1.0;
+
+extension tract_registry tract_transformers;
+extension tract_registry tract_core;
+
+fragment tract_core_properties(
+) -> (properties: (string, tensor<scalar>)[])
+{
+  properties = [
+    ("tract_target_version", "0.23.2-pre"),
+    ("torch_to_nnef_version", "0.24.0"),
+    ("torch_version", "2.9.1"),
+    ("transformers_version", "5.8.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.11.13 (main, Sep 18 2025, 19:53:58) [Clang 20.1.4 ] (64-bit runtime)"),
+    ("export_date", "2026-06-23 14:24:11.387210"),
+    ("exported_py_class", "Qwen3TinyMoE"),
+    ("export_cmd", "/tmp/regen_moe_asset.py")
+  ];
+}
+
+
+
+
+
+
+
+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]);
+    output_0 = tract_moe_ffn(input_0, moe_gate_weight, output_0_w1, output_0_w2, output_0_w3, k = 2, activation = 'swiglu', gate = 'softmax_topk');
+}

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