feature: add routed MoE emulator substrate

transactional_emulator/lib/memory/src/lib.rs

@@ -44,12 +44,18 @@ pub trait MemoryModel: Send + Sync {
 
     /// Write 64-bytes of memory.
     async fn write(&self, addr: u64, bytes: [u8; 64]);
+
+    /// Optional utilization statistics for wrappers that collect them.
+    fn statistics(&self) -> Option<Statistics> {
+        None
+    }
 }
 
 #[async_trait::async_trait]
 pub trait ErasedMemoryModel: Send + Sync {
     async fn box_read(&self, addr: u64) -> [u8; 64];
     async fn box_write(&self, addr: u64, bytes: [u8; 64]);
+    fn statistics(&self) -> Option<Statistics>;
 }
 
 #[async_trait::async_trait]
@@ -61,6 +67,10 @@ impl<T: MemoryModel> ErasedMemoryModel for T {
     async fn box_write(&self, addr: u64, bytes: [u8; 64]) {
         self.write(addr, bytes).await
     }
+
+    fn statistics(&self) -> Option<Statistics> {
+        MemoryModel::statistics(self)
+    }
 }
 
 impl MemoryModel for dyn ErasedMemoryModel {
@@ -71,6 +81,10 @@ impl MemoryModel for dyn ErasedMemoryModel {
     async fn write(&self, addr: u64, bytes: [u8; 64]) {
         self.box_write(addr, bytes).await
     }
+
+    fn statistics(&self) -> Option<Statistics> {
+        ErasedMemoryModel::statistics(self)
+    }
 }
 
 /// A memory that discards all written data.
@@ -206,6 +220,10 @@ impl<T: MemoryModel> MemoryModel for WithStats<T> {
         }
         self.model.write(addr, bytes).await
     }
+
+    fn statistics(&self) -> Option<Statistics> {
+        Some(WithStats::statistics(self))
+    }
 }
 
 #[cfg(test)]

transactional_emulator/lib/quantize/src/tensor.rs

@@ -3,6 +3,34 @@ use tch::Tensor;
 
 use crate::dtype::{DataType, FpType, MxDataType};
 
+fn round_ties_to_even_nonnegative(x: f32) -> u32 {
+    debug_assert!(x >= 0.0);
+    let floor = x.floor();
+    let frac = x - floor;
+    if frac > 0.5 {
+        (floor as u32) + 1
+    } else if frac < 0.5 {
+        floor as u32
+    } else {
+        let floor_int = floor as u32;
+        if floor_int & 1 == 0 {
+            floor_int
+        } else {
+            floor_int + 1
+        }
+    }
+}
+
+fn hardware_round_nonnegative(x: f32) -> u32 {
+    debug_assert!(x >= 0.0);
+    // Match PLENA_Tools/plena_quant/common/hardware_utils.py:
+    //   x = floor(x * 2**round_bits) / 2**round_bits
+    //   return x.round()
+    // with round_bits=2 and PyTorch round-to-even semantics.
+    let quarter_truncated = (x * 4.0).floor() / 4.0;
+    round_ties_to_even_nonnegative(quarter_truncated)
+}
+
 /// Quantize a single FP32 value to minifloat format using IEEE hardware quantization
 /// This matches the Python _minifloat_ieee_quantize_hardware function
 fn minifloat_ieee_quantize_hardware(value: f32, fp_type: FpType) -> u32 {
@@ -47,12 +75,12 @@ fn minifloat_ieee_quantize_hardware(value: f32, fp_type: FpType) -> u32 {
     // Quantize mantissa
     let shift = 1u32 << mantissa_bits;
     let shifted_mantissa = if is_normal {
-        // Normal: (mantissa - 1) * shift, round, clamp
-        let shifted = ((mantissa - 1.0) * shift as f32).round() as u32;
+        // Normal: (mantissa - 1) * shift, hardware_round, clamp.
+        let shifted = hardware_round_nonnegative((mantissa - 1.0) * shift as f32);
         shifted.max(shifted_mantissa_min).min(shifted_mantissa_max)
     } else {
-        // Subnormal: mantissa * shift, round, clamp
-        let shifted = (mantissa * shift as f32).round() as u32;
+        // Subnormal: mantissa * shift, hardware_round, clamp
+        let shifted = hardware_round_nonnegative(mantissa * shift as f32);
         shifted.max(shifted_mantissa_min).min(shifted_mantissa_max)
     };
 
@@ -295,6 +323,21 @@ mod tests {
         assert_eq!(minifloat_ieee_quantize_hardware(-1.0, ty), 0xB8); // sign | 0x38
     }
 
+    #[test]
+    fn test_minifloat_quantize_hardware_rounds_fractional_mantissa() {
+        let ty = e4m3();
+        // E4M3 has 0.125 spacing in [1,2).  PLENA's MX hardware quantizer
+        // first truncates the shifted mantissa to 2 fractional bits, then
+        // applies round-to-even.  This is neither plain nearest rounding nor
+        // plain truncation.
+        assert_eq!(minifloat_ieee_quantize_hardware(1.0625, ty), 0x38);
+        assert_eq!(minifloat_ieee_quantize_hardware(1.0859375, ty), 0x38);
+        assert_eq!(minifloat_ieee_quantize_hardware(1.1875, ty), 0x3A);
+        assert_eq!(minifloat_ieee_quantize_hardware(1.328125, ty), 0x3A);
+        assert_eq!(minifloat_ieee_quantize_hardware(-1.0625, ty), 0xB8);
+        assert_eq!(minifloat_ieee_quantize_hardware(-1.1875, ty), 0xBA);
+    }
+
     #[test]
     fn test_into_bytes_plain_packs_elements_and_no_scale_stream() {
         let elem = DataType::Fp(e4m3());

transactional_emulator/lib/sram/src/matrix.rs

@@ -14,6 +14,19 @@ pub struct MatrixSram {
 }
 
 impl MatrixSram {
+    fn tensor_to_f32_vec(tensor: &tch::Tensor) -> Vec<f32> {
+        let len = tensor.size1().unwrap() as usize;
+        let tensor_f32 = tensor.to_kind(tch::Kind::Float).contiguous();
+        let mut data = vec![0.0f32; len];
+        if tensor_f32.f_copy_data(&mut data, len).is_ok() {
+            return data;
+        }
+        for (idx, value) in data.iter_mut().enumerate() {
+            *value = tensor_f32.double_value(&[idx as i64]) as f32;
+        }
+        data
+    }
+
     /// Create a matrix SRAM with given tile size and depth.
     pub fn new(tile_size: u32, depth: usize, ty: MxDataType) -> Self {
         let tiles = (0..(depth / tile_size as usize))
@@ -116,14 +129,13 @@ impl MatrixSram {
             let mut guard = tile_mutex.lock().await;
             let tensor = guard.resolve().await;
             let tensor_data = tensor.as_tensor();
-            let len = tensor_data.size1().unwrap() as usize;
-            let f32_slice =
-                unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) };
+            let f32_vec = Self::tensor_to_f32_vec(tensor_data);
+            let len = f32_vec.len();
             // Calculate bytes needed for THIS tile's actual size
             let total_bits = len * element_ty.size_in_bits() as usize;
             let bytes_needed = (total_bits + 7) / 8;
             let mut tile_bytes = vec![0u8; bytes_needed];
-            element_ty.bytes_from_f32(f32_slice, &mut tile_bytes);
+            element_ty.bytes_from_f32(&f32_vec, &mut tile_bytes);
             result.extend_from_slice(&tile_bytes);
         }
 

transactional_emulator/lib/sram/src/vector.rs

@@ -28,6 +28,35 @@ pub struct VectorSram {
 }
 
 impl VectorSram {
+    fn tensor_to_f32_vec(tensor: &Tensor) -> Vec<f32> {
+        let len = tensor.size1().unwrap() as usize;
+        let tensor_f32 = tensor.to_kind(tch::Kind::Float).contiguous();
+        let mut data = vec![0.0f32; len];
+        if tensor_f32.f_copy_data(&mut data, len).is_ok() {
+            return data;
+        }
+        for (idx, value) in data.iter_mut().enumerate() {
+            *value = tensor_f32.double_value(&[idx as i64]) as f32;
+        }
+        data
+    }
+
+    fn tensor_from_f32_slice(data: &[f32]) -> Tensor {
+        if data.is_empty() {
+            return Tensor::zeros([0], (tch::Kind::Float, tch::Device::Cpu));
+        }
+        unsafe {
+            Tensor::from_blob(
+                data.as_ptr() as *const u8,
+                &[data.len() as i64],
+                &[],
+                tch::Kind::Float,
+                tch::Device::Cpu,
+            )
+            .internal_to_copy((tch::Kind::Float, tch::Device::Cpu), false)
+        }
+    }
+
     fn row_width_bytes(vlen: u32, fp_type: DataType) -> usize {
         (vlen as usize * fp_type.size_in_bits() as usize).div_ceil(8)
     }
@@ -163,11 +192,7 @@ impl VectorSram {
         let tensor_data = tensor.as_tensor();
         let total_elements = tensor_data.size1().unwrap() as usize;
 
-        // Extract f32 data from tensor to make it Send-safe
-        let len = total_elements;
-        let f32_slice =
-            unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) };
-        let data_vec: Vec<f32> = f32_slice.to_vec();
+        let data_vec = Self::tensor_to_f32_vec(tensor_data);
 
         let chunk_size = self.vlen as usize;
         let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32);
@@ -188,7 +213,7 @@ impl VectorSram {
             padded_data[..chunk_len].copy_from_slice(chunk_data);
 
             // Create tensor from padded data and convert to bytes
-            let padded_tensor = Tensor::from_slice(&padded_data);
+            let padded_tensor = Self::tensor_from_f32_slice(&padded_data);
             let chunk_qt = QuantTensor::quantize(padded_tensor, MxDataType::Plain(self.fp_type));
             let row_bytes = self.quant_tensor_to_bytes(&chunk_qt);
 
@@ -262,7 +287,7 @@ impl VectorSram {
             }
 
             // Create QuantTensor and convert to bytes
-            let tensor = Tensor::from_slice(&vec);
+            let tensor = Self::tensor_from_f32_slice(&vec);
             let quant_tensor = QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type));
             let row_bytes = self.quant_tensor_to_bytes(&quant_tensor);
             *self.rows[row_idx].lock().await = Cell::Ready(row_bytes);
@@ -305,14 +330,13 @@ impl VectorSram {
     /// Convert QuantTensor to bytes (FP format)
     fn quant_tensor_to_bytes(&self, tensor: &QuantTensor) -> Vec<u8> {
         let tensor_data = tensor.as_tensor();
-        let len = tensor_data.size1().unwrap() as usize;
-        let f32_slice =
-            unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) };
+        let f32_vec = Self::tensor_to_f32_vec(tensor_data);
+        let len = f32_vec.len();
 
         let total_bits = len * self.fp_type.size_in_bits() as usize;
         let bytes_needed = (total_bits + 7) / 8;
         let mut bytes = vec![0u8; bytes_needed];
-        self.fp_type.bytes_from_f32(f32_slice, &mut bytes);
+        self.fp_type.bytes_from_f32(&f32_vec, &mut bytes);
         bytes
     }
 
@@ -330,7 +354,7 @@ impl VectorSram {
             vec.resize(expected_len as usize, 0.0f32);
         }
 
-        let tensor = Tensor::from_slice(&vec);
+        let tensor = Self::tensor_from_f32_slice(&vec);
         QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type))
     }