Board-aware weight-streaming + capacity-regime memory models + weight-manifest generator

transactional_emulator/Cargo.toml

@@ -15,6 +15,7 @@ tch = { version = "0.20.0", features = ["download-libtorch"] }
 half = "2"
 clap = { version = "4", features = ["derive"] }
 serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
 toml = "0.8"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

transactional_emulator/lib/memory/Cargo.toml

@@ -6,11 +6,13 @@ edition = "2024"
 [dependencies]
 async-trait = "0.1.88"
 trait-variant = "0.1.2"
+serde = { version = "1", features = ["derive"] }
 zerocopy = "0.8"
 runtime = { path = "../runtime" }
 tokio = { version = "1.45.1", features = ["sync"] }
 futures-util = "0.3.31"
 rand = "0.9.2"
+tracing = "0.1"
 
 [dev-dependencies]
 tokio = { version = "1.45.1", features = ["sync", "rt-multi-thread", "macros"] }

transactional_emulator/lib/memory/src/lib.rs

@@ -2,6 +2,7 @@ pub mod chunked;
 mod frfcfs;
 mod naive;
 mod simple;
+pub mod streaming;
 pub mod testutils;
 
 use std::mem::ManuallyDrop;
@@ -44,12 +45,24 @@ pub trait MemoryModel: Send + Sync {
 
     /// Write 64-bytes of memory.
     async fn write(&self, addr: u64, bytes: [u8; 64]);
+
+    /// Render a human-readable, per-model statistics summary, or `None` if the
+    /// model does not track statistics. `elapsed_secs` is the simulated wall
+    /// time used for bandwidth/utilization figures.
+    ///
+    /// This is intentionally a non-async default so that `trait_variant` (which
+    /// only rewrites async methods) passes it through unchanged, letting it be
+    /// forwarded uniformly via `ErasedMemoryModel`.
+    fn statistics_summary(&self, _elapsed_secs: f64) -> Option<String> {
+        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 box_statistics_summary(&self, elapsed_secs: f64) -> Option<String>;
 }
 
 #[async_trait::async_trait]
@@ -61,6 +74,10 @@ impl<T: MemoryModel> ErasedMemoryModel for T {
     async fn box_write(&self, addr: u64, bytes: [u8; 64]) {
         self.write(addr, bytes).await
     }
+
+    fn box_statistics_summary(&self, elapsed_secs: f64) -> Option<String> {
+        self.statistics_summary(elapsed_secs)
+    }
 }
 
 impl MemoryModel for dyn ErasedMemoryModel {
@@ -71,6 +88,10 @@ impl MemoryModel for dyn ErasedMemoryModel {
     async fn write(&self, addr: u64, bytes: [u8; 64]) {
         self.box_write(addr, bytes).await
     }
+
+    fn statistics_summary(&self, elapsed_secs: f64) -> Option<String> {
+        self.box_statistics_summary(elapsed_secs)
+    }
 }
 
 /// A memory that discards all written data.
@@ -206,6 +227,24 @@ impl<T: MemoryModel> MemoryModel for WithStats<T> {
         }
         self.model.write(addr, bytes).await
     }
+
+    fn statistics_summary(&self, elapsed_secs: f64) -> Option<String> {
+        let s = self.statistics();
+        let hbm_line = format!(
+            "HBM Statistics - Bytes read: {:?} | Bytes written: {:?} | Utilization: {:.2e} bytes/sec",
+            s.total_bytes_read,
+            s.total_bytes_written,
+            (s.total_bytes_read + s.total_bytes_written) as f64 / elapsed_secs
+        );
+        // `WithStats` wraps the streaming models (LayerSwapping/HostStream) in the
+        // runner, so forward the inner model's summary too instead of shadowing it.
+        // The default HBM path's inner is `WithTiming`, whose summary is `None`, so
+        // this returns `hbm_line` unchanged and the HBM output stays byte-for-byte.
+        match self.model.statistics_summary(elapsed_secs) {
+            Some(inner) => Some(format!("{hbm_line}\n{inner}")),
+            None => Some(hbm_line),
+        }
+    }
 }
 
 #[cfg(test)]