From f6dd91c34b3e4be3e7710debf6296b48be24c6f1 Mon Sep 17 00:00:00 2001 From: Michael C Li <63594253+Happymic@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:53:19 +0100 Subject: [PATCH 1/5] feat(moe): review independent normal cores with calibrated shared HBM DMA --- PLENA_Compiler | 2 +- doc/moe_normal_review.md | 130 ++ transactional_emulator/Cargo.lock | 2 + transactional_emulator/Cargo.toml | 2 + transactional_emulator/lib/memory/src/lib.rs | 61 + .../lib/ramulator/src/model.rs | 203 ++- .../lib/ramulator/src/raw.rs | 37 +- .../pkgs/ramulator2/ramulator_capi.cc | 25 + .../pkgs/ramulator2/ramulator_capi.h | 7 + .../src/bin/moe_dual_normal.rs | 181 +++ .../src/moe_normal/dma_credits.rs | 132 ++ .../src/moe_normal/engine.rs | 1446 +++++++++++++++++ transactional_emulator/src/moe_normal/mod.rs | 15 + .../src/moe_normal/read_cache.rs | 80 + .../src/moe_normal/tests.rs | 741 +++++++++ .../src/moe_normal/types.rs | 241 +++ .../moe_normal_review/large_small.json | 48 + .../moe_normal_review/measured_latency.csv | 9 + .../testbench/moe_normal_review/single.json | 37 + .../testbench/moe_normal_review/smoke.py | 65 + .../moe_timing/replay/compare_moe_normal.py | 385 +++++ .../replay/test_compare_moe_normal.py | 432 +++++ 22 files changed, 4270 insertions(+), 11 deletions(-) create mode 100644 doc/moe_normal_review.md create mode 100644 transactional_emulator/src/bin/moe_dual_normal.rs create mode 100644 transactional_emulator/src/moe_normal/dma_credits.rs create mode 100644 transactional_emulator/src/moe_normal/engine.rs create mode 100644 transactional_emulator/src/moe_normal/mod.rs create mode 100644 transactional_emulator/src/moe_normal/read_cache.rs create mode 100644 transactional_emulator/src/moe_normal/tests.rs create mode 100644 transactional_emulator/src/moe_normal/types.rs create mode 100644 transactional_emulator/testbench/moe_normal_review/large_small.json create mode 100644 transactional_emulator/testbench/moe_normal_review/measured_latency.csv create mode 100644 transactional_emulator/testbench/moe_normal_review/single.json create mode 100644 transactional_emulator/testbench/moe_normal_review/smoke.py create mode 100644 transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py create mode 100644 transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py diff --git a/PLENA_Compiler b/PLENA_Compiler index d89ad594c..510ed3c4d 160000 --- a/PLENA_Compiler +++ b/PLENA_Compiler @@ -1 +1 @@ -Subproject commit d89ad594c798daa54f63f914aebad5317653489a +Subproject commit 510ed3c4deb91494db8e43e6e6cc48defde00bab diff --git a/doc/moe_normal_review.md b/doc/moe_normal_review.md new file mode 100644 index 000000000..5c7cccf2d --- /dev/null +++ b/doc/moe_normal_review.md @@ -0,0 +1,130 @@ +# MoE normal-buffer cores and HBM DMA: review guide + +This review implements the meeting's first experiment: independently configured +normal-buffer cores with private accumulators, shared HBM, and a fair single-core +comparison. It executes actual nonzero MoE arithmetic. The best tested large/small +pair is **slower than the optimized single core**; this is not a speedup claim. + +## Read the implementation in this order + +1. `transactional_emulator/src/moe_normal/types.rs`: workload, per-core resources + and report contract. `testbench/moe_normal_review/{single,large_small}.json` + are the actual fixed configurations selected by the completed search. +2. `src/moe_normal/engine.rs`: validate finite resources, group/dispatch expert + jobs, gather activations, run gate/up → SwiGLU → down, and combine routes in + deterministic token/slot order. Each core owns its normal weight SRAM and + accumulator; the existing ISA runner is not the entry point for this path. +3. `engine.rs` DMA path plus `lib/ramulator/src/model.rs`: bounded prefetch, + in-flight 32-byte sector merging and per-channel request submission to one + Ramulator HBM2 instance. Returned elements/scales are copied and decoded into + reserved BF16 slots before compute. No unbounded cache or free extra bandwidth. +4. `testbench/moe_timing/replay/compare_moe_normal.py`: execute the same workload + on both configurations, enforce fair budgets, compare against the Compiler + oracle, and reject resource/counter/hash/repeat mismatches before reporting time. + +Paths in items 2–4 are relative to `transactional_emulator/`. + +## Current dimensions + +Physical tiles have M=N=BLEN and K=MLEN. A tile is serviced over multiple cycles; +multiplier count is BLEN×MLEN, not M×K×N. Whole-model D/F dimensions are separate. + +| Resource | Single baseline | Large core | Small core | +|---|---:|---:|---:| +| Tile (M,K,N) | (8,512,8) | (4,768,4) | (2,512,2) | +| Multipliers | 4096 | 3072 | 1024 | +| Weight SRAM | 64 KiB | 48 KiB | 16 KiB | +| Weight slots | 4 | 4 | 4 | +| Activation/intermediate SRAM | 4 MiB | 2 MiB | 2 MiB | +| Accumulator storage | 1 MiB | 512 KiB | 512 KiB | + +Both systems share the same 8-channel HBM2 preset, clock, total activation supply, +DMA budget and vector unit. The calibrated native transfer is 32 bytes; the +previous Rust wrapper's 16-byte assumption was corrected, not used as a speedup. +The DRAM command scheduler and HBM timing preset are retained. Logical DMA credits +are capped at 128, native trackers at 256; DMA frontend SRAM is 44 KiB including +staging. The two cores do not receive two independent HBM bandwidth budgets. + +HBM weights are output-major `[N,K]` rows with separate element/scale ranges. +The local E4M3/E8M0 codec uses 8 elements per scale. An aligned MLEN-long row +fragment therefore needs MLEN elements and MLEN/8 scales. SRAM slots reserve both +packed ingress and decoded BF16 data. Both buffers are normal buffers in this PR. + +## How time and correctness are measured + +- Input activations and fixed routes are ready at time zero. Stop when final + combined BF16 outputs are ready. Router execution and input/output HBM DMA are + outside this boundary; weight DMA, decode, compute and combine are inside it. +- HBM delay comes from native Ramulator requests/callbacks. Core, SRAM-port, + lookup/copy and vector service use explicit analytical timing. Time is simulated + elapsed time, not the wall-clock duration of the Rust process. Same-output K + accumulation dependencies are retained; cross-core waits cannot be summed as + total latency. Those dependency/port assumptions remain subjects for review. +- The independent Compiler oracle decodes actual exported weight bytes, performs + ascending-K FP32 arithmetic and applies the same BF16 boundaries. The comparison + checks exact BF16 outputs as well as HBM completed bytes, finite SRAM/queue + limits, task accounting, binary/input identity and repeat determinism. + +## Essential measured result + +The completed bounded search used archived decode routes, synthetic nonzero +weights/activations, D=2048 and F=512 (Qwen) or F=1408 (DeepSeek), with 8/32-token +windows. It selected a fixed winner per architecture class; those fixed choices +were then evaluated on separate route windows. This is a numerical MoE operator +experiment, not full-model inference with trained weights. + +| Large/small pair versus optimized single | Relative elapsed time | +|---|---:| +| Four search workloads, geometric mean | +8.88% (slower) | +| Four held-out workloads, geometric mean | +8.53% (slower) | + +`testbench/moe_normal_review/measured_latency.csv` contains only the eight pairs +needed to audit these percentages: `exp(mean(log(dual_ps/single_ps)))-1`. +These are extracted from the completed, hash-audited campaign, not a newly run +search on the review branch. Its generated weight images, route archives, raw +reports and historical DSE/ablation drivers are deliberately not in this PR. +The standalone runner can validate supplied full-size fixtures with the same +checked-in configurations; the smoke test below does not reproduce this table. + +## Run the paired Compiler → Simulator smoke test + +After checking out this branch, initialize the submodules (the Compiler pin is +the paired review commit). Install Python dependencies with the existing uv setup: + +```sh +git submodule update --init --recursive +uv sync +nix develop --command bash -c \ + 'cd transactional_emulator && cargo build --bin moe_dual_normal' +nix develop --command bash -c \ + 'PYTHONPATH="$PWD/PLENA_Tools" uv run python \ + transactional_emulator/testbench/moe_normal_review/smoke.py \ + --compiler "$PWD/PLENA_Compiler" \ + --binary "$PWD/transactional_emulator/target/debug/moe_dual_normal" \ + --output-dir /tmp/plena-moe-review-smoke' +``` + +Use a new output directory for each invocation. If CARGO_TARGET_DIR is set, pass +the binary in that directory instead. The existing Nix derivation builds the +updated native calibration C API from its pinned Ramulator source. + +The smoke generates D=31/F=47, 12 tokens, uneven expert groups and a shared expert, +then runs each checked-in architecture twice with zero numerical tolerance. It +fails on any correctness/resource/repeat mismatch and prints `all_gates_passed`. +The review checkout passed this end-to-end smoke, exporter tests, Rust library +and MoE tests, Python evidence-gate tests, formatting and strict clippy checks. + +For a supplied full-size fixture, run `compare_moe_normal.py --help`; pass its +`workload.json`, `golden.json`, the two checked-in architecture files (single +first), `--atol 0 --rtol 0`, and an output directory. A full-size Qwen B8 check on +the review checkout passed exact output/resource/repeat checks and reproduced +the archived 0.524575 ms single / 0.564819 ms dual times. This one-case validation +is separate from the archived search result. + +## Review boundary + +No transpose buffer, Attention execution, runtime router or complete dual-core +instruction lowering is claimed. Private buffer allocation does not establish +equal physical area. The immediate review question is whether the data path, +resource accounting and timing assumptions form a sound comparison before any +further scheduling change; extra prefetch helps both the single and dual systems. diff --git a/transactional_emulator/Cargo.lock b/transactional_emulator/Cargo.lock index 5f9465613..f6ba6ae49 100644 --- a/transactional_emulator/Cargo.lock +++ b/transactional_emulator/Cargo.lock @@ -1734,6 +1734,7 @@ version = "0.1.0" dependencies = [ "buddy", "clap", + "futures", "half", "memory", "mimalloc", @@ -1742,6 +1743,7 @@ dependencies = [ "runtime", "serde", "serde_json", + "sha2", "sram", "tch", "tokio", diff --git a/transactional_emulator/Cargo.toml b/transactional_emulator/Cargo.toml index 175430435..915b0480a 100644 --- a/transactional_emulator/Cargo.toml +++ b/transactional_emulator/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition = "2024" [dependencies] +futures = "0.3" runtime = { path = "lib/runtime" } memory = { path = "lib/memory" } ramulator = { path = "lib/ramulator" } @@ -16,6 +17,7 @@ half = "2" clap = { version = "4", features = ["derive"] } serde = { version = "1.0", features = ["derive"] } serde_json = "1.0" +sha2 = "0.10" toml = "0.8" tracing = "0.1" tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/transactional_emulator/lib/memory/src/lib.rs b/transactional_emulator/lib/memory/src/lib.rs index c47719dc1..82d7e345f 100644 --- a/transactional_emulator/lib/memory/src/lib.rs +++ b/transactional_emulator/lib/memory/src/lib.rs @@ -23,6 +23,17 @@ pub trait MemoryTimingModel: Send + Sync { /// We fix to 64-bytes to accomodate memory emulators. async fn read(&self, addr: u64); + /// Optional 32-byte sector reads; the default deliberately fetches a full line. + fn read_mask(&self, addr: u64, mask: u8) -> impl Future + Send { + async move { + assert!((1..=3).contains(&mask)); + self.read(addr).await; + } + } + fn supports_sector_reads(&self) -> bool { + false + } + /// Write 64-bytes of memory. async fn write(&self, addr: u64); } @@ -32,6 +43,12 @@ impl MemoryTimingModel for ManuallyDrop { T::read(self, addr).await } + async fn read_mask(&self, addr: u64, mask: u8) { + T::read_mask(self, addr, mask).await + } + fn supports_sector_reads(&self) -> bool { + T::supports_sector_reads(self) + } async fn write(&self, addr: u64) { T::write(self, addr).await } @@ -42,6 +59,16 @@ pub trait MemoryModel: Send + Sync { /// Read 64-bytes of memory. async fn read(&self, addr: u64) -> [u8; 64]; + fn read_mask(&self, addr: u64, mask: u8) -> impl Future + Send { + async move { + assert!((1..=3).contains(&mask)); + self.read(addr).await + } + } + fn supports_sector_reads(&self) -> bool { + false + } + /// Write 64-bytes of memory. async fn write(&self, addr: u64, bytes: [u8; 64]); @@ -54,12 +81,20 @@ pub trait MemoryModel: Send + Sync { #[async_trait::async_trait] pub trait ErasedMemoryModel: Send + Sync { async fn box_read(&self, addr: u64) -> [u8; 64]; + async fn box_read_mask(&self, addr: u64, mask: u8) -> [u8; 64]; + fn supports_sector_reads(&self) -> bool; async fn box_write(&self, addr: u64, bytes: [u8; 64]); fn statistics(&self) -> Option; } #[async_trait::async_trait] impl ErasedMemoryModel for T { + async fn box_read_mask(&self, addr: u64, mask: u8) -> [u8; 64] { + self.read_mask(addr, mask).await + } + fn supports_sector_reads(&self) -> bool { + MemoryModel::supports_sector_reads(self) + } async fn box_read(&self, addr: u64) -> [u8; 64] { self.read(addr).await } @@ -74,6 +109,12 @@ impl ErasedMemoryModel for T { } impl MemoryModel for dyn ErasedMemoryModel { + async fn read_mask(&self, addr: u64, mask: u8) -> [u8; 64] { + self.box_read_mask(addr, mask).await + } + fn supports_sector_reads(&self) -> bool { + ErasedMemoryModel::supports_sector_reads(self) + } async fn read(&self, addr: u64) -> [u8; 64] { self.box_read(addr).await } @@ -164,6 +205,13 @@ impl WithTiming { } impl MemoryModel for WithTiming { + async fn read_mask(&self, addr: u64, mask: u8) -> [u8; 64] { + self.timing.read_mask(addr, mask).await; + self.data.read(addr).await + } + fn supports_sector_reads(&self) -> bool { + self.timing.supports_sector_reads() + } /// Read 64-bytes of memory. async fn read(&self, addr: u64) -> [u8; 64] { self.timing.read(addr).await; @@ -205,6 +253,19 @@ impl WithStats { } impl MemoryModel for WithStats { + async fn read_mask(&self, addr: u64, mask: u8) -> [u8; 64] { + assert!((1..=3).contains(&mask)); + let bytes = if self.model.supports_sector_reads() { + 32 * u64::from(mask.count_ones()) + } else { + 64 + }; + self.statistics.lock().unwrap().total_bytes_read += bytes; + self.model.read_mask(addr, mask).await + } + fn supports_sector_reads(&self) -> bool { + self.model.supports_sector_reads() + } async fn read(&self, addr: u64) -> [u8; 64] { { let mut guard = self.statistics.lock().unwrap(); diff --git a/transactional_emulator/lib/ramulator/src/model.rs b/transactional_emulator/lib/ramulator/src/model.rs index 178e66490..3ce6c7094 100644 --- a/transactional_emulator/lib/ramulator/src/model.rs +++ b/transactional_emulator/lib/ramulator/src/model.rs @@ -1,4 +1,4 @@ -use core::sync::atomic::AtomicU32; +use core::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use anyhow::Result; @@ -18,6 +18,14 @@ struct State { ticker_running: bool, } +#[derive(Clone, Copy, Debug, Default, serde::Serialize, serde::Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum IssuePolicy { + #[default] + GlobalFifo, + PerChannel, +} + struct Inner { // Use atomic (but only use relaxed ordering) as this can be accessed while holding the mutex. pending_accesses: AtomicU32, @@ -29,13 +37,36 @@ struct Inner { // Size of a single transfer transfer_size: u32, + policy: IssuePolicy, + channel_shift: u32, + channels: usize, + issue_period: Duration, + port_locks: Vec>, + next_issue: Mutex>, + queue_capacity: tokio::sync::Semaphore, + accepted: Vec, + rejected: Vec, + admission_wait_ps: AtomicU64, + native_peak: AtomicU32, } /// A wrapped ramulator that works with the event-based simulation. +#[derive(Clone)] pub struct Ramulator(Arc); impl Ramulator { pub fn new(config: crate::config::Config) -> Result { + let channels = config.controllers.len(); + let shift = match config.channel_mapper { + crate::config::ChannelMapper::CacheLineInterleave { interleave_bits } => { + interleave_bits.unwrap_or(0) + } + _ => anyhow::bail!("calibrated wrapper currently supports CacheLineInterleave only"), + }; + anyhow::ensure!( + channels.is_power_of_two(), + "controller count must be a power of two" + ); let mut ramulator = RawRamulator::new(config)?; let period = Duration::from_picos(ramulator.period() as _); let transfer_size = ramulator.burst_size() * (ramulator.channel_width() / 8); @@ -51,9 +82,54 @@ impl Ramulator { lock: tokio::sync::Mutex::new(()), transfer_size, + policy: IssuePolicy::GlobalFifo, + channel_shift: transfer_size.ilog2() + shift, + channels, + issue_period: period, + port_locks: (0..channels).map(|_| tokio::sync::Mutex::new(())).collect(), + next_issue: Mutex::new(vec![Instant::INIT; channels]), + queue_capacity: tokio::sync::Semaphore::new(256), + accepted: (0..channels).map(|_| AtomicU64::new(0)).collect(), + rejected: (0..channels).map(|_| AtomicU64::new(0)).collect(), + admission_wait_ps: AtomicU64::new(0), + native_peak: AtomicU32::new(0), }))) } + /// Configure only before cloning/submitting. Both policies have the same + /// one-attempt-per-output-port-per-cycle bandwidth and finite queue pool. + pub fn with_issue_policy(mut self, policy: IssuePolicy, period: Duration) -> Self { + assert!(period.as_picos() > 0); + let inner = Arc::get_mut(&mut self.0).expect("configure before sharing"); + inner.policy = policy; + inner.issue_period = period; + self + } + + pub fn channel_for(&self, addr: u64) -> usize { + ((addr >> self.0.channel_shift) as usize) & (self.0.channels - 1) + } + + pub fn telemetry(&self) -> serde_json::Value { + let mut state = self.0.mutable.lock().unwrap(); + serde_json::json!({ + "native_transaction_bytes": self.0.transfer_size, + "channel_shift": self.0.channel_shift, + "channels": self.0.channels, + "mapper": "CacheLineInterleave", + "capi_version": 2, + "issue_policy": self.0.policy, + "issue_period_ps": self.0.issue_period.as_picos(), + "submission_entries": 256, + "accepted_per_channel": self.0.accepted.iter().map(|v| v.load(Ordering::Relaxed)).collect::>(), + "rejected_per_channel": self.0.rejected.iter().map(|v| v.load(Ordering::Relaxed)).collect::>(), + "admission_wait_ps": self.0.admission_wait_ps.load(Ordering::Relaxed), + "native_inflight_peak": self.0.native_peak.load(Ordering::Relaxed), + "native_pending": self.0.pending_accesses.load(Ordering::Relaxed), + "native_stats": state.ramulator.native_stats(), + }) + } + pub fn period(&self) -> Duration { let mut guard = self.0.mutable.lock().unwrap(); Duration::from_picos(guard.ramulator.period().into()) @@ -150,6 +226,10 @@ impl Ramulator { return Err(()); } + self.0.native_peak.fetch_max( + self.0.pending_accesses.load(Ordering::Relaxed), + Ordering::Relaxed, + ); // Drive the model only if something is still outstanding, and only if // nothing is driving it already. Testing `pending_accesses` alone would // start a second loop whenever a synchronous completion dropped the @@ -171,14 +251,41 @@ impl Ramulator { /// Send a request to ramulator. pub async fn access(&self, addr: u64, write: bool) { - let guard = self.0.lock.lock().await; - let mut fut = self.try_access(addr, write); - while fut.is_err() { - Executor::current().resolve_at(self.0.period).await; - fut = self.try_access(addr, write); - } + assert_eq!( + addr % u64::from(self.0.transfer_size), + 0, + "native request must be aligned" + ); + let ex = Executor::current(); + let begin = ex.now().as_picos(); + let _entry = self.0.queue_capacity.acquire().await.unwrap(); + let port = self.channel_for(addr); + let guard = match self.0.policy { + IssuePolicy::GlobalFifo => self.0.lock.lock().await, + IssuePolicy::PerChannel => self.0.port_locks[port].lock().await, + }; + let completion = loop { + let due = self.0.next_issue.lock().unwrap()[port]; + if due > ex.now() { + ex.resolve_at(due).await; + } + self.0.next_issue.lock().unwrap()[port] = ex.now() + self.0.issue_period; + match self.try_access(addr, write) { + Ok(future) => { + self.0.accepted[port].fetch_add(1, Ordering::Relaxed); + break future; + } + Err(()) => { + self.0.rejected[port].fetch_add(1, Ordering::Relaxed); + } + } + }; + self.0 + .admission_wait_ps + .fetch_add(ex.now().as_picos() - begin, Ordering::Relaxed); drop(guard); - fut.unwrap().await + // Entry is retained until the response; tracker storage is not free. + completion.await; } /// Send a read request to ramulator. @@ -193,6 +300,18 @@ impl Ramulator { } impl memory::MemoryTimingModel for Ramulator { + async fn read_mask(&self, addr: u64, mask: u8) { + assert!(self.supports_sector_reads() && addr.is_multiple_of(64) && (1..=3).contains(&mask)); + let transfers: Vec<_> = (0..2u64) + .filter(|s| mask & (1 << s) != 0) + .map(|s| self.read_transfer(addr + s * 32)) + .collect(); + futures::future::join_all(transfers).await; + } + fn supports_sector_reads(&self) -> bool { + self.0.transfer_size == 32 + } + async fn read(&self, addr: u64) { let transfers: Vec<_> = (0..64u64) .step_by(self.0.transfer_size as usize) @@ -363,3 +482,71 @@ mod ticker_tests { ); } } + +#[cfg(test)] +mod calibration_tests { + use super::*; + use memory::MemoryTimingModel; + + #[tokio::test] + async fn native_granularity_mapping_and_read_accounting() { + let ram = Ramulator::hbm2_preset(8).unwrap(); + assert_eq!(ram.transfer_size(), 32); + assert_eq!( + (0..8).map(|p| ram.channel_for(p * 32)).collect::>(), + (0..8).collect::>() + ); + let ex = Executor::new(); + let r = ram.clone(); + ex.spawn(async move { + r.read(0).await; // two native transactions + r.read_mask(64, 1).await; // exactly one sector + r.read_mask(64, 1).await; // repeated read is not a native cache hit + }); + ex.enter(Instant::ETERNITY).await; + let stats = ram.telemetry(); + assert_eq!( + stats["accepted_per_channel"], + serde_json::json!([1, 1, 2, 0, 0, 0, 0, 0]) + ); + assert_eq!(stats["native_pending"], 0); + eprintln!("CALIBRATION {}", stats); + } + + async fn blocked_port_probe(policy: IssuePolicy) -> (u64, serde_json::Value) { + let ram = Ramulator::hbm2_preset(8) + .unwrap() + .with_issue_policy(policy, Duration::from_picos(1000)); + let ex = Executor::new(); + // More independent misses than the default 32-entry channel queue. + for i in 0..96 { + let r = ram.clone(); + ex.spawn(async move { + r.read_transfer(i * 65536).await; + }); + } + let completed = Arc::new(AtomicU64::new(0)); + let observed = completed.clone(); + let r = ram.clone(); + ex.spawn(async move { + r.read_transfer(32).await; + observed.store(Executor::current().now().as_picos(), Ordering::Relaxed); + }); + ex.enter(Instant::ETERNITY).await; + (completed.load(Ordering::Relaxed), ram.telemetry()) + } + + #[tokio::test] + async fn blocked_channel_does_not_hold_an_independent_channel() { + let (global, a) = blocked_port_probe(IssuePolicy::GlobalFifo).await; + let (ported, b) = blocked_port_probe(IssuePolicy::PerChannel).await; + assert!( + ported < global, + "per-channel queue must remove cross-channel HOL: {ported} >= {global}" + ); + assert_eq!(a["accepted_per_channel"], b["accepted_per_channel"]); + assert_eq!(a["native_pending"], 0); + assert_eq!(b["native_pending"], 0); + eprintln!("HOL global_ps={global} ported_ps={ported}"); + } +} diff --git a/transactional_emulator/lib/ramulator/src/raw.rs b/transactional_emulator/lib/ramulator/src/raw.rs index 32ef23574..aaf9dddc4 100644 --- a/transactional_emulator/lib/ramulator/src/raw.rs +++ b/transactional_emulator/lib/ramulator/src/raw.rs @@ -1,7 +1,7 @@ use anyhow::Result; use core::ffi::c_void; use core::mem::ManuallyDrop; -use std::ffi::CString; +use std::ffi::{CStr, CString}; use crate::config::*; @@ -27,6 +27,10 @@ mod bindings { ) -> bool; pub fn ramulator_period(val: *mut ramulator) -> f32; pub fn ramulator_tick(val: *mut ramulator); + pub fn ramulator_capi_version() -> u32; + pub fn ramulator_tx_bytes(val: *mut ramulator) -> u32; + pub fn ramulator_stats(val: *mut ramulator, buffer: *mut c_char, capacity: u64) -> u64; + pub fn ramulator_library_path() -> *const c_char; } } @@ -79,6 +83,10 @@ impl Drop for Ramulator { impl Ramulator { pub fn new(config: crate::config::Config) -> Result { + anyhow::ensure!( + !config.controllers.is_empty(), + "Ramulator needs at least one controller" + ); let dram_impl = match &config.controllers[0] { Controller::GenericDDR(c) => &c.dram, Controller::HBM12(c) => &c.dram, @@ -94,7 +102,7 @@ impl Ramulator { // "LPDDR5" => 8, // "GDDR6" => 8, // "HBM" => 2, - DRAM::HBM2(_) => 2, + DRAM::HBM2(_) => 4, // "HBM3" => 2, }; @@ -132,6 +140,12 @@ impl Ramulator { if raw.is_null() { anyhow::bail!("Ramulator failed to initialize"); } + let native_tx = unsafe { bindings::ramulator_tx_bytes(raw) }; + let expected_tx = internal_prefetch_size * default_channel_width / 8; + if unsafe { bindings::ramulator_capi_version() } != 2 || native_tx != expected_tx { + unsafe { bindings::ramulator_finalize(raw) }; + anyhow::bail!("native transaction mismatch: native={native_tx}, wrapper={expected_tx}"); + } Ok(Ramulator { raw, burst_size: internal_prefetch_size, @@ -141,6 +155,25 @@ impl Ramulator { }) } + pub fn native_stats(&mut self) -> serde_json::Value { + let size = unsafe { bindings::ramulator_stats(self.raw, std::ptr::null_mut(), 0) }; + assert!((1..=16 * 1024 * 1024).contains(&size)); + let mut buffer = vec![0u8; size as usize]; + let written = + unsafe { bindings::ramulator_stats(self.raw, buffer.as_mut_ptr().cast(), size) }; + assert_eq!(written, size); + serde_yaml_ng::from_slice(&buffer[..buffer.len() - 1]) + .expect("native statistics must be valid YAML") + } + + pub fn library_path() -> String { + let path = unsafe { bindings::ramulator_library_path() }; + assert!(!path.is_null(), "native library identity unavailable"); + unsafe { CStr::from_ptr(path) } + .to_string_lossy() + .into_owned() + } + pub fn period(&mut self) -> u32 { let period_in_ns = unsafe { bindings::ramulator_period(self.raw) }; (period_in_ns * 1000.).round() as u32 diff --git a/transactional_emulator/pkgs/ramulator2/ramulator_capi.cc b/transactional_emulator/pkgs/ramulator2/ramulator_capi.cc index f1bdc3bfe..797f2cae3 100644 --- a/transactional_emulator/pkgs/ramulator2/ramulator_capi.cc +++ b/transactional_emulator/pkgs/ramulator2/ramulator_capi.cc @@ -6,6 +6,9 @@ #include #include +#include +#include +#include #include "ramulator_capi.h" @@ -51,3 +54,25 @@ float ramulator_period(ramulator *val) { void ramulator_tick(ramulator *val) { val->memory_system->tick(); } + +uint32_t ramulator_capi_version() { return 2; } + +uint32_t ramulator_tx_bytes(ramulator *val) { + return val->memory_system->get_tx_bytes(); +} + +uint64_t ramulator_stats(ramulator *val, char *buffer, uint64_t capacity) { + val->memory_system->update_stats_recursive(); + std::ostringstream stream; + val->memory_system->print_stats(stream); + const auto text = stream.str(); + const auto required = text.size() + 1; + if (buffer && capacity >= required) std::memcpy(buffer, text.c_str(), required); + return required; +} + +const char *ramulator_library_path() { + Dl_info info{}; + if (!dladdr(reinterpret_cast(&ramulator_capi_version), &info)) return nullptr; + return info.dli_fname; +} diff --git a/transactional_emulator/pkgs/ramulator2/ramulator_capi.h b/transactional_emulator/pkgs/ramulator2/ramulator_capi.h index 6ec7e99fe..6731f216e 100644 --- a/transactional_emulator/pkgs/ramulator2/ramulator_capi.h +++ b/transactional_emulator/pkgs/ramulator2/ramulator_capi.h @@ -17,6 +17,13 @@ float ramulator_period(ramulator *val); void ramulator_tick(ramulator *val); +// Versioned calibration interface. All getters run under the Rust model lock. +uint32_t ramulator_capi_version(); +uint32_t ramulator_tx_bytes(ramulator *val); +// YAML snapshot; returns required bytes including NUL, without partial writes. +uint64_t ramulator_stats(ramulator *val, char *buffer, uint64_t capacity); +const char *ramulator_library_path(); + #ifdef __cplusplus } #endif diff --git a/transactional_emulator/src/bin/moe_dual_normal.rs b/transactional_emulator/src/bin/moe_dual_normal.rs new file mode 100644 index 000000000..6ffe63623 --- /dev/null +++ b/transactional_emulator/src/bin/moe_dual_normal.rs @@ -0,0 +1,181 @@ +//! Fixed-route MoE functional execution on independently configured normal cores. +//! +//! This entry point consumes the compiler's versioned job manifest. It keeps +//! the legacy ISA runner intact while sharing the simulator's memory/runtime +//! libraries and one real Ramulator HBM instance across all cores. + +#[path = "../moe_normal/mod.rs"] +mod moe_normal; + +use std::io::{Read, Write}; +use std::mem::ManuallyDrop; +use std::path::{Path, PathBuf}; +use std::sync::Arc; + +use clap::Parser; +use serde_json::json; +use sha2::{Digest, Sha256}; + +#[derive(Parser)] +#[command(about = "Execute grouped MoE on normal-buffer cores with shared HBM")] +struct Opts { + #[arg(long)] + workload: PathBuf, + #[arg(long)] + architecture: PathBuf, + #[arg(long)] + output: PathBuf, + /// Shared HBM2 channel count, identical for baseline and candidate. + #[arg(long, default_value_t = 8)] + hbm_channels: u32, + /// Bound the small experiment's host allocation before loading any image. + #[arg(long, default_value_t = 536_870_912)] + max_hbm_bytes: u64, +} + +fn sha256_file(path: &Path) -> Result> { + let mut stream = std::fs::File::open(path)?; + let mut digest = Sha256::new(); + let mut buffer = [0u8; 64 * 1024]; + loop { + let count = stream.read(&mut buffer)?; + if count == 0 { + break; + } + digest.update(&buffer[..count]); + } + Ok(format!("{:x}", digest.finalize())) +} + +async fn execute(opts: Opts) -> Result<(), Box> { + if !opts.hbm_channels.is_power_of_two() || opts.hbm_channels > 32 { + return Err("hbm-channels must be a power of two between 1 and 32".into()); + } + let workload_path = opts.workload.canonicalize()?; + let architecture_path = opts.architecture.canonicalize()?; + let workload_bytes = std::fs::read(&workload_path)?; + let architecture_bytes = std::fs::read(&architecture_path)?; + let workload_json: serde_json::Value = serde_json::from_slice(&workload_bytes)?; + let architecture_json: serde_json::Value = serde_json::from_slice(&architecture_bytes)?; + let image_name = workload_json["hbm_file"] + .as_str() + .ok_or("workload.hbm_file must name the encoded weight image")?; + let image_path = workload_path + .parent() + .unwrap() + .join(image_name) + .canonicalize()?; + let image_len = std::fs::metadata(&image_path)?.len(); + if image_len == 0 || image_len % 64 != 0 || image_len > opts.max_hbm_bytes { + return Err(format!( + "HBM image must be nonempty, 64-byte aligned and at most {} bytes; got {}", + opts.max_hbm_bytes, image_len + ) + .into()); + } + let output_parent = opts + .output + .parent() + .filter(|p| !p.as_os_str().is_empty()) + .unwrap_or(Path::new(".")); + let output_name = opts + .output + .file_name() + .ok_or("output requires a filename")?; + let output_path = output_parent.canonicalize()?.join(output_name); + let existing_output = output_path + .canonicalize() + .unwrap_or_else(|_| output_path.clone()); + if [&workload_path, &architecture_path, &image_path].contains(&&existing_output) { + return Err("output may not overwrite workload, architecture or HBM input".into()); + } + let workload: moe_normal::Workload = serde_json::from_slice(&workload_bytes)?; + let architecture: moe_normal::Architecture = serde_json::from_slice(&architecture_bytes)?; + moe_normal::validate(&workload, &architecture, image_len).map_err(std::io::Error::other)?; + let image = std::fs::read(&image_path)?; + let image_sha256 = format!("{:x}", Sha256::digest(&image)); + if let Some(expected) = workload_json.pointer("/metadata/hbm_sha256") + && expected.as_str() != Some(image_sha256.as_str()) + { + return Err("encoded HBM image does not match manifest metadata.hbm_sha256".into()); + } + let backing = memory::MemoryBacked::with_capacity(usize::try_from(image_len)?); + backing.with_data(|destination| destination.copy_from_slice(&image)); + drop(image); + // Match the existing runner's native-model lifetime: the process owns this + // one HBM model until exit, after the executor drains all memory requests. + let native = ramulator::Ramulator::hbm2_preset(opts.hbm_channels as usize)?.with_issue_policy( + architecture + .dma + .as_ref() + .map(|d| d.issue_policy) + .unwrap_or_default(), + runtime::Duration::from_picos(architecture.clock_period_ps), + ); + let native_observer = native.clone(); + let library_path = PathBuf::from(ramulator::raw::Ramulator::library_path()).canonicalize()?; + let library_sha256 = sha256_file(&library_path)?; + let hbm: Arc = Arc::new(memory::WithStats::new( + memory::WithTiming::new(ManuallyDrop::new(native), backing), + )); + let report = moe_normal::run(workload, architecture, hbm, image_len) + .await + .map_err(std::io::Error::other)?; + let native_telemetry = native_observer.telemetry(); + if native_telemetry["native_pending"] != 0 { + return Err("native requests not drained".into()); + } + if sha256_file(&library_path)? != library_sha256 { + return Err("native library changed during execution".into()); + } + let envelope = json!({ + "schema_version": 1, + "evidence_level": "fixed_route_numerical_moe_with_ramulator_and_analytical_core_timing", + "provenance": { + "workload_path": workload_path, + "workload_sha256": format!("{:x}", Sha256::digest(&workload_bytes)), + "architecture_path": architecture_path, + "architecture_sha256": format!("{:x}", Sha256::digest(&architecture_bytes)), + "hbm_path": image_path, + "hbm_sha256": image_sha256, + "hbm_image_bytes": image_len, + "native_library_path": library_path, + "native_library_sha256": library_sha256, + "executable_sha256": sha256_file(&std::env::current_exe()?)?, + }, + "memory_model": {"name": "Ramulator HBM2 preset", "channels": opts.hbm_channels, "upper_burst_bytes": 64, "calibration": native_telemetry}, + "workload_manifest": workload_json, + "architecture_manifest": architecture_json, + "result": report, + }); + let payload = serde_json::to_vec_pretty(&envelope)?; + let temporary = output_path.with_file_name(format!( + ".{}.{}.tmp", + output_name.to_string_lossy(), + std::process::id() + )); + let mut file = std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&temporary)?; + let write_result = (|| -> std::io::Result<()> { + file.write_all(&payload)?; + file.write_all(b"\n")?; + file.sync_all()?; + std::fs::rename(&temporary, &output_path) + })(); + if write_result.is_err() { + let _ = std::fs::remove_file(&temporary); + } + write_result?; + println!("MoE result: {}", output_path.display()); + Ok(()) +} + +#[tokio::main(flavor = "current_thread")] +async fn main() { + if let Err(error) = execute(Opts::parse()).await { + eprintln!("error: {error}"); + std::process::exit(1); + } +} diff --git a/transactional_emulator/src/moe_normal/dma_credits.rs b/transactional_emulator/src/moe_normal/dma_credits.rs new file mode 100644 index 000000000..e3b8b73e6 --- /dev/null +++ b/transactional_emulator/src/moe_normal/dma_credits.rs @@ -0,0 +1,132 @@ +//! Borrowable per-core minimum credits. Reservations apply only while another +//! core is waiting, so an idle core cannot strand shared response storage. +use std::sync::{Arc, Mutex}; +use tokio::sync::Notify; + +struct State { + used: Vec, + waiting: Vec, +} +pub(super) struct CreditPool { + total: usize, + reserve: usize, + state: Mutex, + changed: Notify, +} + +pub(super) struct Credit { + pool: Arc, + core: usize, +} +struct Waiting { + pool: Arc, + core: usize, + active: bool, +} + +impl Drop for Waiting { + fn drop(&mut self) { + if self.active { + self.pool.state.lock().unwrap().waiting[self.core] -= 1; + self.pool.changed.notify_waiters(); + } + } +} +impl Drop for Credit { + fn drop(&mut self) { + self.pool.state.lock().unwrap().used[self.core] -= 1; + self.pool.changed.notify_waiters(); + } +} +impl CreditPool { + pub(super) fn new(total: usize, cores: usize) -> Self { + assert!(cores > 0 && total >= cores); + Self { + total, + reserve: total / 8, + state: Mutex::new(State { + used: vec![0; cores], + waiting: vec![0; cores], + }), + changed: Notify::new(), + } + } + pub(super) async fn acquire(self: &Arc, core: usize) -> Credit { + self.state.lock().unwrap().waiting[core] += 1; + let mut waiting = Waiting { + pool: self.clone(), + core, + active: true, + }; + loop { + let notification = self.changed.notified(); + tokio::pin!(notification); + // Register before checking state; a release between the check and + // await must not be lost. Cancellation removes the waiter as well. + notification.as_mut().enable(); + { + let mut s = self.state.lock().unwrap(); + let free = self.total - s.used.iter().sum::(); + let protected: usize = s + .used + .iter() + .enumerate() + .filter(|(i, _)| *i != core && s.waiting[*i] > 0) + .map(|(_, used)| self.reserve.saturating_sub(*used)) + .sum(); + if free > protected { + s.used[core] += 1; + s.waiting[core] -= 1; + waiting.active = false; + return Credit { + pool: self.clone(), + core, + }; + } + } + notification.await; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use futures::FutureExt; + + #[tokio::test] + async fn idle_reservation_is_borrowed_and_waiting_core_recovers_first_credit() { + let pool = Arc::new(CreditPool::new(16, 2)); + let mut held = Vec::new(); + for _ in 0..16 { + held.push(pool.acquire(0).await); + } + let waiting_small = pool.acquire(1); + tokio::pin!(waiting_small); + assert!(waiting_small.as_mut().now_or_never().is_none()); + drop(held.pop()); + let waiting_large = pool.acquire(0); + tokio::pin!(waiting_large); + assert!(waiting_large.as_mut().now_or_never().is_none()); + let small = waiting_small.await; + assert_eq!(pool.state.lock().unwrap().used, vec![15, 1]); + drop(small); + let large = waiting_large.await; + drop(large); + drop(held); + assert_eq!(pool.state.lock().unwrap().used, vec![0, 0]); + } + + #[tokio::test] + async fn cancelled_waiter_does_not_reserve_forever() { + let pool = Arc::new(CreditPool::new(8, 2)); + let held: Vec<_> = futures::future::join_all((0..8).map(|_| pool.acquire(0))).await; + { + let mut waiting = Box::pin(pool.acquire(1)); + assert!(waiting.as_mut().now_or_never().is_none()); + } + assert_eq!(pool.state.lock().unwrap().waiting, vec![0, 0]); + drop(held); + drop(pool.acquire(0).await); + } +} diff --git a/transactional_emulator/src/moe_normal/engine.rs b/transactional_emulator/src/moe_normal/engine.rs new file mode 100644 index 000000000..82df55589 --- /dev/null +++ b/transactional_emulator/src/moe_normal/engine.rs @@ -0,0 +1,1446 @@ +use futures::{ + FutureExt, + future::{BoxFuture, Shared as SharedFuture}, +}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +use half::bf16; +use memory::ErasedMemoryModel; +use quantize::{DataType, FpType}; +use runtime::{Duration, Executor, Instant}; +use tokio::sync::{Semaphore, oneshot}; + +use super::dma_credits::CreditPool; +use super::read_cache::{ENTRY_BYTES, ReadCache}; +use super::types::*; + +const BLOCK: usize = 8; + +fn checked(values: &[usize], what: &str) -> Result { + values.iter().try_fold(1usize, |a, b| { + a.checked_mul(*b) + .ok_or_else(|| format!("{what} overflows usize")) + }) +} + +fn extra_cycles(a: &Architecture, c: &CoreConfig) -> u64 { + if a.matrix_timing == MatrixTiming::LegacySerialized { + return 0; + } + (usize::BITS - (c.mlen - 1).leading_zeros()) as u64 + a.mac_pipeline_cycles +} + +fn service_cycles(a: &Architecture, c: &CoreConfig) -> u64 { + match a.matrix_timing { + MatrixTiming::Pipelined => { + let supply = c.activation_elements_per_cycle.unwrap_or(c.mlen); + (c.blen as u64).max((c.blen as u64 * c.mlen as u64).div_ceil(supply as u64)) + } + MatrixTiming::LegacySerialized => c.mlen as u64 + a.mac_pipeline_cycles, + } +} + +fn pipeline_bytes(a: &Architecture, c: &CoreConfig) -> Result { + checked( + &[c.blen, extra_cycles(a, c) as usize, 4], + "pipeline storage", + ) +} + +fn weight_slot_bytes(c: &CoreConfig) -> Result { + // BF16 decoded tile, packed E4M3 elements and packed E8M0 scales. + checked(&[c.blen, c.mlen / BLOCK, 25], "weight slot") +} + +fn vector_bytes(w: &Workload, rows: usize) -> Result { + let per_row = checked(&[2, w.input_dim], "X and output")? + .checked_add(checked(&[3, w.expert_hidden_dim], "gate/up/Z")?) + .ok_or("vector storage overflow")?; + checked(&[rows, per_row, 2], "BF16 vector storage") +} + +fn accumulator_bytes( + w: &Workload, + rows: usize, + a: &Architecture, + c: &CoreConfig, +) -> Result { + checked( + &[rows, w.input_dim.max(w.expert_hidden_dim), 4], + "accumulator", + )? + .checked_add(pipeline_bytes(a, c)?) + .ok_or_else(|| "accumulator plus pipeline storage overflow".into()) +} + +fn region_valid(r: &MatrixRegion, n: usize, k: usize, len: u64) -> Result<(), String> { + if r.rows != n || r.cols != k { + return Err(format!( + "matrix shape is [{},{}], expected [{n},{k}]", + r.rows, r.cols + )); + } + let blocks = k.div_ceil(BLOCK) as u64; + let padded_k = blocks + .checked_mul(BLOCK as u64) + .ok_or("padded K overflow")?; + if r.element_row_stride < padded_k || r.scale_row_stride < blocks { + return Err("MX row strides must cover K padded to block 8 and all row scales".into()); + } + for (name, base, stride, width) in [ + ("element", r.element_base, r.element_row_stride, padded_k), + ("scale", r.scale_base, r.scale_row_stride, blocks), + ] { + let end = (n as u64 - 1) + .checked_mul(stride) + .and_then(|v| v.checked_add(base)) + .and_then(|v| v.checked_add(width)) + .ok_or_else(|| format!("{name} address overflow"))?; + if end > len { + return Err(format!( + "{name} matrix ends at {end}, beyond HBM image length {len}" + )); + } + } + Ok(()) +} + +#[derive(Clone)] +struct Row { + token: usize, + output: usize, +} + +#[derive(Clone)] +struct Job { + id: usize, + expert: usize, + shared: bool, + rows: Vec, +} + +struct Plan { + queues: Vec>, + routes: Vec, + output_rows: usize, + shared_base: usize, + global_storage: usize, +} + +fn job_fits(w: &Workload, a: &Architecture, c: &CoreConfig, job: &Job) -> bool { + vector_bytes(w, job.rows.len()).is_ok_and(|v| v <= c.vector_sram_bytes) + && accumulator_bytes(w, job.rows.len(), a, c).is_ok_and(|v| v <= c.accumulator_bytes) +} + +fn plan(w: &Workload, a: &Architecture) -> Result { + let mut routes = w.routes.clone(); + routes.sort_by_key(|r| (r.token, r.slot)); + let mut grouped: BTreeMap> = BTreeMap::new(); + for (i, r) in routes.iter().enumerate() { + grouped.entry(r.expert).or_default().push(Row { + token: r.token, + output: i, + }); + } + let mut jobs: Vec = grouped + .into_iter() + .enumerate() + .map(|(id, (expert, rows))| Job { + id, + expert, + shared: false, + rows, + }) + .collect(); + let shared_base = routes.len(); + let output_rows = routes + .len() + .checked_add(if w.shared_expert.is_some() { + w.inputs_bf16.len() + } else { + 0 + }) + .ok_or("route count overflow")?; + if let Some(shared) = &w.shared_expert { + jobs.push(Job { + id: jobs.len(), + expert: shared.expert, + shared: true, + rows: (0..w.inputs_bf16.len()) + .map(|t| Row { + token: t, + output: shared_base + t, + }) + .collect(), + }); + } + // Inputs and final BF16 output each need 2 bytes, final FP32 sum needs 4. + // Every route result is BF16, held in a finite reorder region until combine. + let global_storage = checked( + &[w.inputs_bf16.len(), w.input_dim, 8], + "shared inputs/output/sum", + )? + .checked_add(checked( + &[output_rows, w.input_dim, 2], + "route reorder storage", + )?) + .ok_or("shared storage overflow")?; + if global_storage > a.combine_sram_bytes { + return Err(format!( + "shared input/combine SRAM needs {global_storage} bytes, capacity {}", + a.combine_sram_bytes + )); + } + if checked(&[jobs.len(), 64], "ready-job descriptors")? > a.dispatch_queue_bytes { + return Err("ready-job descriptors exceed dispatch_queue_bytes".into()); + } + let mut queues = vec![Vec::new(); a.cores.len()]; + for job in jobs { + let mut target = if job.rows.len() >= a.dispatch_threshold { + a.large_core + } else { + a.small_core + }; + if a.dispatch_policy == DispatchPolicy::WorkConserving + && !job_fits(w, a, &a.cores[target], &job) + { + target = a + .cores + .iter() + .position(|c| job_fits(w, a, c, &job)) + .ok_or("expert group does not fit any core; group spilling is not enabled")?; + } + let c = &a.cores[target]; + let v = vector_bytes(w, job.rows.len())?; + let acc = accumulator_bytes(w, job.rows.len(), a, c)?; + if v > c.vector_sram_bytes { + return Err(format!( + "core {} expert {} M={} requires {v} vector SRAM bytes, capacity {}; V0 does not spill or tile groups", + c.id, + job.expert, + job.rows.len(), + c.vector_sram_bytes + )); + } + if acc > c.accumulator_bytes { + return Err(format!( + "core {} expert {} M={} requires {acc} accumulator/pipeline bytes, capacity {}", + c.id, + job.expert, + job.rows.len(), + c.accumulator_bytes + )); + } + queues[target].push(job); + } + Ok(Plan { + queues, + routes, + output_rows, + shared_base, + global_storage, + }) +} + +fn validate_model_bounds(w: &Workload, a: &Architecture, p: &Plan) -> Result<(), String> { + // Reject impossible counter/timer ranges before entering the executor. + // This bounds our own computation and vector work; backend HBM timing is + // still the backend's contract. u128 keeps the validation itself checked. + let lanes = a.vector_elements_per_cycle; + let mut cycles = + checked(&[p.output_rows, w.input_dim], "combine elements")?.div_ceil(lanes) as u128; + cycles += + checked(&[w.inputs_bf16.len(), w.input_dim], "output elements")?.div_ceil(lanes) as u128; + let mut issued = 0u128; + let mut multipliers = 0u128; + for (c, queue) in a.cores.iter().zip(&p.queues) { + multipliers += checked(&[c.blen, c.mlen], "multipliers")? as u128; + let per_tile_macs = checked(&[c.blen, c.blen, c.mlen], "tile MACs")? as u128; + let tile_latency = service_cycles(a, c) as u128 + extra_cycles(a, c) as u128; + let decode_cycles = checked(&[c.blen, c.mlen], "decode elements")?.div_ceil(lanes) as u128; + let candidates: Vec<&Job> = if a.dispatch_policy == DispatchPolicy::WorkConserving { + p.queues + .iter() + .flatten() + .filter(|j| job_fits(w, a, c, j)) + .collect() + } else { + queue.iter().collect() + }; + for job in candidates { + cycles = cycles + .checked_add(a.dispatch_cycles as u128) + .ok_or("dispatch timing overflow")?; + let m = job.rows.len(); + cycles += + 2 * checked(&[m, w.input_dim], "gather/copy elements")?.div_ceil(lanes) as u128; + cycles += + checked(&[m, w.expert_hidden_dim], "SwiGLU elements")?.div_ceil(lanes) as u128; + for (n, k) in [ + (w.expert_hidden_dim, w.input_dim), + (w.expert_hidden_dim, w.input_dim), + (w.input_dim, w.expert_hidden_dim), + ] { + let weight_tiles = checked( + &[n.div_ceil(c.blen), k.div_ceil(c.mlen)], + "weight tile count", + )? as u128; + let macro_tiles = weight_tiles + .checked_mul(m.div_ceil(c.blen) as u128) + .ok_or("macro tile count overflow")?; + issued = issued + .checked_add( + macro_tiles + .checked_mul(per_tile_macs) + .ok_or("issued MAC overflow")?, + ) + .ok_or("issued MAC overflow")?; + let upper = macro_tiles + .checked_mul(tile_latency) + .and_then(|v| { + weight_tiles + .checked_mul(decode_cycles) + .and_then(|decode| v.checked_add(decode)) + }) + .ok_or("timing bound overflow")?; + cycles = cycles.checked_add(upper).ok_or("timing bound overflow")?; + // Conservative: at most one line request per payload byte, + // one lookup and one fill cycle per request. + if c.read_cache_bytes > 0 { + let cache_cycles = weight_tiles + .checked_mul(per_tile_macs) + .and_then(|v| v.checked_mul(4)) + .ok_or("cache timing overflow")?; + cycles = cycles + .checked_add(cache_cycles) + .ok_or("cache timing overflow")?; + } + } + } + } + if issued > u64::MAX as u128 || multipliers > u64::MAX as u128 { + return Err("workload/architecture exceeds u64 MAC counter range".into()); + } + if cycles + .checked_mul(a.clock_period_ps as u128) + .is_none_or(|t| t > u64::MAX as u128) + { + return Err( + "workload/architecture computation or vector duration exceeds u64 picoseconds".into(), + ); + } + Ok(()) +} + +/// Validate all capacities, shape/route contracts and complete HBM ranges before +/// spawning work. Out-of-range MemoryBacked reads must never silently turn to 0. +pub fn validate(w: &Workload, a: &Architecture, hbm_len: u64) -> Result<(), String> { + if w.schema_version != 1 || a.schema_version != 1 { + return Err("only workload/architecture schema_version 1 is supported".into()); + } + if w.input_dim == 0 || w.expert_hidden_dim == 0 || w.inputs_bf16.is_empty() { + return Err("input_dim, expert_hidden_dim and token count must be positive".into()); + } + if !hbm_len.is_multiple_of(64) || hbm_len == 0 { + return Err("HBM image length must be a positive multiple of 64 bytes".into()); + } + if !(1..=2).contains(&a.cores.len()) + || a.large_core >= a.cores.len() + || a.small_core >= a.cores.len() + { + return Err("V0 supports one or two cores, with in-range dispatch core indices".into()); + } + if a.cores.len() == 2 && a.large_core == a.small_core { + return Err("a two-core architecture must select distinct large/small cores".into()); + } + if a.clock_period_ps == 0 + || a.vector_elements_per_cycle == 0 + || a.dispatch_threshold == 0 + || a.global_dma_credits == 0 + { + return Err("clock, vector throughput, threshold and DMA credits must be positive".into()); + } + if a.global_dma_credits > Semaphore::MAX_PERMITS + || checked(&[a.global_dma_credits, 64], "DMA staging")? > a.global_dma_staging_bytes + { + return Err("global DMA staging must provide 64 bytes per credit".into()); + } + if let Some(dma) = &a.dma { + if !(1..=2).contains(&dma.lookup_ii_cycles) { + return Err("lookup_ii_cycles must be 1 or 2".into()); + } + if dma.fair_credits && a.global_dma_credits < 8 { + return Err("fair DMA credits require at least 8 credits".into()); + } + if a.global_dma_credits > 128 { + return Err("DMA V1 supports at most 128 line/waiter credits".into()); + } + if a.cores.iter().any(|c| c.read_cache_bytes != 0) { + return Err("DMA V1 uses element bypass; legacy read cache must be disabled".into()); + } + if frontend_bytes(a)? > dma.frontend_sram_bytes { + return Err("DMA frontend metadata/response storage exceeds capacity".into()); + } + } + let mut ids = BTreeSet::new(); + for c in &a.cores { + if c.activation_elements_per_cycle == Some(0) { + return Err("activation supply must be positive".into()); + } + if c.read_cache_bytes > 0 && c.read_cache_bytes < ENTRY_BYTES { + return Err("read cache needs at least one 80-byte data/tag entry".into()); + } + if c.id.is_empty() + || !ids.insert(c.id.clone()) + || c.blen == 0 + || c.mlen == 0 + || !c.mlen.is_multiple_of(BLOCK) + { + return Err( + "core ids must be unique/nonempty; BLEN positive and MLEN a positive multiple of 8" + .into(), + ); + } + // Keep the legacy square MatrixMachine's complete BLEN column slices + // inside an MLEN row. The local normal-buffer loop itself could handle + // other ratios; those would be a separate mapping/ISA extension. + if !c.mlen.is_multiple_of(c.blen) { + return Err(format!( + "core {} MLEN must be divisible by BLEN for the square mapping", + c.id + )); + } + checked(&[c.blen, c.mlen], "multiplier count")?; + checked(&[c.blen, c.blen, c.mlen], "issued MACs per tile")?; + if a.mac_pipeline_cycles > u32::MAX as u64 { + return Err("mac_pipeline_cycles exceeds supported range".into()); + } + if !(2..=4).contains(&c.weight_slots) { + return Err("weight_slots must be 2, 3 or 4".into()); + } + let need = weight_slot_bytes(c)? + .checked_mul(c.weight_slots) + .ok_or("weight SRAM overflow")?; + if need > c.weight_sram_bytes { + return Err(format!( + "core {} needs {need} bytes for configured MX/BF16 weight slots, capacity {}", + c.id, c.weight_sram_bytes + )); + } + let cycles = service_cycles(a, c) + .checked_add(extra_cycles(a, c)) + .ok_or("tile timing overflow")?; + cycles + .checked_mul(a.clock_period_ps) + .ok_or("tile duration overflow")?; + pipeline_bytes(a, c)?; + } + for row in &w.inputs_bf16 { + if row.len() != w.input_dim + || row + .iter() + .any(|v| !bf16::from_bits(*v).to_f32().is_finite()) + { + return Err("inputs must contain input_dim finite BF16 values per token".into()); + } + } + let mut experts = BTreeSet::new(); + for e in &w.experts { + if !experts.insert(e.id) { + return Err(format!("duplicate expert id {}", e.id)); + } + region_valid(&e.gate, w.expert_hidden_dim, w.input_dim, hbm_len)?; + region_valid(&e.up, w.expert_hidden_dim, w.input_dim, hbm_len)?; + region_valid(&e.down, w.input_dim, w.expert_hidden_dim, hbm_len)?; + } + let mut slots = BTreeSet::new(); + for r in &w.routes { + if r.token >= w.inputs_bf16.len() || !experts.contains(&r.expert) || !r.weight.is_finite() { + return Err("route references an invalid token/expert or non-finite weight".into()); + } + if !slots.insert((r.token, r.slot)) { + return Err(format!("duplicate route slot ({},{})", r.token, r.slot)); + } + } + if let Some(s) = &w.shared_expert + && (!experts.contains(&s.expert) || !s.weight.is_finite()) + { + return Err("shared expert reference or weight is invalid".into()); + } + let p = plan(w, a)?; + validate_model_bounds(w, a, &p)?; + Ok(()) +} + +struct CoreState { + cache: ReadCache, + config: CoreConfig, + report: Mutex, + slots: AtomicUsize, + slots_peak: AtomicUsize, + weight_bytes: AtomicUsize, + weight_peak: AtomicUsize, +} + +impl CoreState { + fn new(c: &CoreConfig, a: &Architecture) -> Self { + Self { + cache: ReadCache::new(c.read_cache_bytes, a.clock_period_ps), + config: c.clone(), + report: Mutex::new(CoreReport { + id: c.id.clone(), + blen: c.blen, + mlen: c.mlen, + multipliers: (c.blen * c.mlen) as u64, + jobs: 0, + useful_macs: 0, + issued_macs: 0, + compute_busy_ps: 0, + accumulator_dependency_stall_ps: 0, + pipeline_drain_ps: 0, + pipeline_register_bytes: pipeline_bytes(a, c).unwrap(), + weight_ready_wait_ps: 0, + vector_wait_ps: 0, + vector_sram_peak_bytes: 0, + accumulator_peak_bytes: 0, + weight_sram_peak_bytes: 0, + weight_slots_peak: 0, + hbm_read_bytes: 0, + compute_busy_fraction: 0.0, + mac_utilization: 0.0, + cache_requests: 0, + cache_hits: 0, + cache_port_busy_ps: 0, + cache_peak_bytes: 0, + }), + slots: AtomicUsize::new(0), + slots_peak: AtomicUsize::new(0), + weight_bytes: AtomicUsize::new(0), + weight_peak: AtomicUsize::new(0), + } + } +} + +struct SlotReservation { + core: Arc, + bytes: usize, +} + +impl SlotReservation { + fn new(core: Arc) -> Self { + let bytes = weight_slot_bytes(&core.config).unwrap(); + let slots = core.slots.fetch_add(1, Ordering::SeqCst) + 1; + let occupied = core.weight_bytes.fetch_add(bytes, Ordering::SeqCst) + bytes; + assert!( + slots <= core.config.weight_slots && occupied <= core.config.weight_sram_bytes, + "weight slot ownership violated" + ); + core.slots_peak.fetch_max(slots, Ordering::SeqCst); + core.weight_peak.fetch_max(occupied, Ordering::SeqCst); + Self { core, bytes } + } +} + +impl Drop for SlotReservation { + fn drop(&mut self) { + self.core.slots.fetch_sub(1, Ordering::SeqCst); + self.core + .weight_bytes + .fetch_sub(self.bytes, Ordering::SeqCst); + } +} + +struct Dma { + hbm: Arc, + credits: Semaphore, + fair_pool: Option>, + core_indices: BTreeMap, + inflight: AtomicUsize, + peak: AtomicUsize, + bytes: AtomicU64, + config: Option, + clock: u64, + lookup: Vec, + copy: Vec, + lines: Mutex>, + report: Mutex, +} + +// Conservative physical reservation, independent of host future/Vec sizes. +// Per line: 64B response + 32B MSHR + 24B waiter. Native pool: 256 * 16B. +// Tile copy fragments cost 24B each; 128B state per resident tile. 4KiB +// covers queue heads, lookup/copy pipelines, and bounded arbitration state. +fn frontend_bytes(a: &Architecture) -> Result { + let mut bytes = checked(&[a.global_dma_credits, 120], "DMA line metadata")? + .checked_add(8192) + .ok_or("DMA control storage overflow")?; + for c in &a.cores { + let fragments = c + .mlen + .div_ceil(64) + .checked_add(c.mlen.div_ceil(BLOCK).div_ceil(64)) + .and_then(|v| v.checked_add(2)) + .ok_or("DMA fragment count overflow")?; + let tile = checked(&[c.blen, fragments, 24], "DMA fragment descriptors")? + .checked_add(128) + .ok_or("DMA tile state overflow")?; + bytes = bytes + .checked_add(checked( + &[c.weight_slots, tile], + "DMA resident descriptors", + )?) + .ok_or("DMA frontend storage overflow")?; + } + Ok(bytes) +} + +type SectorFuture = SharedFuture>; +struct PendingLine { + users: usize, + sectors: [Option; 2], +} + +impl Dma { + fn sector(self: &Arc, address: u64, sector: usize, core: Arc) -> SectorFuture { + let dma = self.clone(); + async move { + let data = dma.hbm.box_read_mask(address, 1 << sector).await; + dma.bytes.fetch_add(32, Ordering::SeqCst); + core.report.lock().unwrap().hbm_read_bytes += 32; + let mut bytes = [0; 32]; + bytes.copy_from_slice(&data[sector * 32..sector * 32 + 32]); + bytes + } + .boxed() + .shared() + } + + async fn read( + self: &Arc, + address: u64, + requested: u8, + useful: usize, + core: Arc, + ) -> [u8; 64] { + let config = self.config.as_ref().unwrap(); + let mask = if config.sector_reads { requested } else { 3 }; + let ex = Executor::current(); + { + // Four banks, conservative 2-cycle occupied lookup, common to all + // DMA V1 policies including global-FIFO and coalescing-off controls. + let _port = self.lookup[(address as usize / 64) % 4] + .acquire() + .await + .unwrap(); + // Two-cycle latency with an independently pipelined lookup port. + ex.resolve_at(Duration::from_picos( + config.lookup_ii_cycles as u64 * self.clock, + )) + .await; + drop(_port); + if config.lookup_ii_cycles == 1 { + ex.resolve_at(Duration::from_picos(self.clock)).await; + } + let mut r = self.report.lock().unwrap(); + r.lookup_busy_ps += config.lookup_ii_cycles as u64 * self.clock; + r.line_requests += 1; + r.sector_requests += u64::from(mask.count_ones()); + r.useful_copy_bytes += useful as u64; + } + let sectors = if config.coalesce { + let mut lines = self.lines.lock().unwrap(); + let entry = lines.entry(address).or_insert_with(|| PendingLine { + users: 0, + sectors: [None, None], + }); + entry.users += 1; + let mut needed = Vec::new(); + for sector in 0..2 { + if mask & (1 << sector) == 0 { + continue; + } + if entry.sectors[sector].is_some() { + self.report.lock().unwrap().merged_sectors += 1; + } else { + entry.sectors[sector] = Some(self.sector(address, sector, core.clone())); + } + needed.push((sector, entry.sectors[sector].as_ref().unwrap().clone())); + } + let mut report = self.report.lock().unwrap(); + report.mshr_peak = report.mshr_peak.max(lines.len()); + needed + } else { + (0..2) + .filter(|s| mask & (1 << s) != 0) + .map(|s| (s, self.sector(address, s, core.clone()))) + .collect() + }; + let responses = futures::future::join_all( + sectors + .into_iter() + .map(|(index, future)| async move { (index, future.await) }), + ) + .await; + let mut bytes = [0; 64]; + for (index, data) in responses { + bytes[index * 32..index * 32 + 32].copy_from_slice(&data); + } + { + // Eight return/copy banks, 32 B per cycle each. Duplicate consumers + // each pay copy service; multicast and SRAM writes are not free. + let _port = self.copy[(address as usize / 64) % 8] + .acquire() + .await + .unwrap(); + let duration = useful.div_ceil(32) as u64 * self.clock; + ex.resolve_at(Duration::from_picos(duration)).await; + self.report.lock().unwrap().copy_busy_ps += duration; + } + if config.coalesce { + let mut lines = self.lines.lock().unwrap(); + let entry = lines.get_mut(&address).unwrap(); + entry.users -= 1; + if entry.users == 0 { + lines.remove(&address); + } + } + bytes + } +} + +struct VectorUnit { + permit: Semaphore, + lanes: usize, + clock: u64, + busy_ps: AtomicU64, +} + +impl VectorUnit { + async fn work(&self, elements: usize) -> u64 { + if elements == 0 { + return 0; + } + let ex = Executor::current(); + let begin = ex.now().as_picos(); + let _permit = self.permit.acquire().await.unwrap(); + let waiting = ex.now().as_picos() - begin; + let duration = elements.div_ceil(self.lanes) as u64 * self.clock; + self.busy_ps.fetch_add(duration, Ordering::SeqCst); + ex.resolve_at(Duration::from_picos(duration)).await; + waiting + } +} + +struct Shared { + ready: Mutex>, + dispatcher: Semaphore, + dispatcher_busy: AtomicU64, + dma: Arc, + vector: Arc, + reorder: Mutex>, + completions: Mutex>, +} + +#[derive(Clone, Copy)] +struct TileSpec { + n: usize, + k: usize, +} + +struct WeightTile { + values: Vec, + _reservation: SlotReservation, +} + +#[derive(Clone, Copy)] +struct CopySpan { + src: usize, + dst: usize, + len: usize, +} + +fn add_reads(reads: &mut BTreeMap>, base: u64, len: usize, dst: usize) { + let mut offset = 0; + while offset < len { + let addr = base + offset as u64; + let aligned = addr / 64 * 64; + let src = (addr % 64) as usize; + let count = (64 - src).min(len - offset); + reads.entry(aligned).or_default().push(CopySpan { + src, + dst: dst + offset, + len: count, + }); + offset += count; + } +} + +fn spawn_load( + core: Arc, + shared: Arc, + region: MatrixRegion, + tile: TileSpec, +) -> oneshot::Receiver> { + let reservation = SlotReservation::new(core.clone()); + let (tx, rx) = oneshot::channel(); + Executor::current().spawn(async move { + let result = load_tile(core, shared, region, tile, reservation).await; + let _ = tx.send(result); + }); + rx +} + +async fn load_tile( + core: Arc, + shared: Arc, + region: MatrixRegion, + tile: TileSpec, + reservation: SlotReservation, +) -> Result { + let c = &core.config; + let nr = c.blen.min(region.rows - tile.n); + let kr = c.mlen.min(region.cols - tile.k); + let elements = c.blen * c.mlen; + let scale_count = elements / BLOCK; + let packed = Arc::new(Mutex::new(vec![0u8; elements + scale_count])); + let mut reads = BTreeMap::new(); + for row in 0..nr { + add_reads( + &mut reads, + region.element_base + (tile.n + row) as u64 * region.element_row_stride + tile.k as u64, + kr, + row * c.mlen, + ); + add_reads( + &mut reads, + region.scale_base + + (tile.n + row) as u64 * region.scale_row_stride + + (tile.k / BLOCK) as u64, + kr.div_ceil(BLOCK), + elements + row * (c.mlen / BLOCK), + ); + } + let mut pending = Vec::with_capacity(reads.len()); + for (address, spans) in reads { + let dma = shared.dma.clone(); + let packed = packed.clone(); + let core = core.clone(); + let (tx, rx) = oneshot::channel(); + pending.push(rx); + Executor::current().spawn(async move { + let (credit, fair_credit) = if let Some(pool) = &dma.fair_pool { + let begin = Executor::current().now().as_picos(); + let credit = pool.acquire(dma.core_indices[&core.config.id]).await; + dma.report.lock().unwrap().fair_credit_wait_ps += + Executor::current().now().as_picos() - begin; + (None, Some(credit)) + } else { + (Some(dma.credits.acquire().await.unwrap()), None) + }; + let current = dma.inflight.fetch_add(1, Ordering::SeqCst) + 1; + dma.peak.fetch_max(current, Ordering::SeqCst); + let bytes = if dma.config.is_some() { + let mut mask = 0u8; + for span in &spans { + mask |= 1 << (span.src / 32); + mask |= 1 << ((span.src + span.len - 1) / 32); + } + dma.read( + address, + mask, + spans.iter().map(|s| s.len).sum(), + core.clone(), + ) + .await + } else if let Some(bytes) = core.cache.lookup(address).await { + bytes + } else { + let bytes = dma.hbm.box_read(address).await; + dma.bytes.fetch_add(64, Ordering::SeqCst); + core.report.lock().unwrap().hbm_read_bytes += 64; + core.cache.insert(address, bytes).await; + bytes + }; + { + let mut buffer = packed.lock().unwrap(); + for s in spans { + buffer[s.dst..s.dst + s.len].copy_from_slice(&bytes[s.src..s.src + s.len]); + } + } + dma.inflight.fetch_sub(1, Ordering::SeqCst); + // Credit covers both the HBM request and its 64-byte staging data. + drop(credit); + drop(fair_credit); + let _ = tx.send(()); + }); + } + for rx in pending { + rx.await.map_err(|_| "DMA request task failed")?; + } + // MX decode and BF16 placement use the same finite shared vector actor; + // decoded data cannot become available for free at the final HBM response. + let wait = shared.vector.work(nr * kr).await; + core.report.lock().unwrap().vector_wait_ps += wait; + let packed = Arc::try_unwrap(packed) + .map_err(|_| "DMA buffer retained after completion")? + .into_inner() + .unwrap(); + // Reuse the actual repository codec, including its local E4M3 subnormal rule. + // Decode one scalar at a time, avoiding an unaccounted full FP32 tile. + let element_type = DataType::Fp(FpType { + sign: true, + exponent: 4, + mantissa: 3, + }); + let scale_type = DataType::Fp(FpType::E8M0); + let mut values = vec![bf16::ZERO; elements]; + for row in 0..nr { + for k in 0..kr { + let element = element_type.convert_bits_to_f32(packed[row * c.mlen + k] as u32); + let scale = scale_type + .convert_bits_to_f32(packed[elements + row * (c.mlen / BLOCK) + k / BLOCK] as u32); + let v = element * scale; + let value = bf16::from_f32(v); + if !value.is_finite() { + return Err("decoded weight is not finite BF16".into()); + } + values[row * c.mlen + k] = value; + } + } + Ok(WeightTile { + values, + _reservation: reservation, + }) +} + +// Keep the three separately owned data buffers explicit at the execution boundary. +#[allow(clippy::too_many_arguments)] +async fn gemm( + input: &[bf16], + output: &mut [bf16], + accumulator: &mut [f32], + m: usize, + region: &MatrixRegion, + core: &Arc, + a: &Architecture, + shared: &Arc, +) -> Result<(), String> { + let n = region.rows; + let k = region.cols; + let c = &core.config; + let ex = Executor::current(); + let acc = &mut accumulator[..m * n]; + acc.fill(0.0); + let mb = m.div_ceil(c.blen); + let nb = n.div_ceil(c.blen); + // One readiness time per physical M,N output block. The complete output + // accumulator and fixed in-flight result registers were already reserved. + let mut ready = vec![ex.now(); mb * nb]; + let mut specs = Vec::new(); + for n0 in (0..n).step_by(c.blen) { + for k0 in (0..k).step_by(c.mlen) { + specs.push(TileSpec { n: n0, k: k0 }); + } + } + let mut pending = VecDeque::new(); + let mut issued = 0; + // Initial window reserves every slot before DMA. At each retirement exactly + // one slot becomes reusable; READY tiles keep their ownership meanwhile. + for spec in specs.iter().take(c.weight_slots) { + pending.push_back(spawn_load( + core.clone(), + shared.clone(), + region.clone(), + *spec, + )); + issued += 1; + } + for spec in specs.iter().copied() { + let wait_start = ex.now().as_picos(); + let tile = pending + .pop_front() + .unwrap() + .await + .map_err(|_| "weight loader task failed")??; + core.report.lock().unwrap().weight_ready_wait_ps += ex.now().as_picos() - wait_start; + for m0 in (0..m).step_by(c.blen) { + let block_index = (spec.n / c.blen) * mb + m0 / c.blen; + if ready[block_index] > ex.now() { + let stall = (ready[block_index] - ex.now()).as_picos(); + core.report.lock().unwrap().accumulator_dependency_stall_ps += stall; + ex.resolve_at(ready[block_index]).await; + } + let mr = c.blen.min(m - m0); + let nr = c.blen.min(n - spec.n); + let kr = c.mlen.min(k - spec.k); + for row in 0..mr { + for col in 0..nr { + let target = (m0 + row) * n + spec.n + col; + for kk in 0..kr { + // Separate FP32 multiply/add, in ascending global K. + let product = input[(m0 + row) * k + spec.k + kk].to_f32() + * tile.values[col * c.mlen + kk].to_f32(); + acc[target] += product; + } + } + } + let service = service_cycles(a, c) * a.clock_period_ps; + ready[block_index] = + ex.now() + Duration::from_picos(service + extra_cycles(a, c) * a.clock_period_ps); + { + let mut report = core.report.lock().unwrap(); + report.useful_macs += (mr * nr * kr) as u64; + report.issued_macs += (c.blen * c.blen * c.mlen) as u64; + report.compute_busy_ps += service; + } + ex.resolve_at(Duration::from_picos(service)).await; + } + drop(tile); + if issued < specs.len() { + pending.push_back(spawn_load( + core.clone(), + shared.clone(), + region.clone(), + specs[issued], + )); + issued += 1; + } + } + let last = ready.into_iter().max().unwrap(); + if last > ex.now() { + core.report.lock().unwrap().pipeline_drain_ps += (last - ex.now()).as_picos(); + ex.resolve_at(last).await; + } + for (value, sum) in output.iter_mut().zip(acc.iter().copied()) { + *value = bf16::from_f32(sum); + if !value.is_finite() { + return Err("GEMM produced non-finite BF16".into()); + } + } + Ok(()) +} + +async fn run_core( + core: Arc, + queue: Vec, + w: Arc, + a: Arc, + shared: Arc, +) -> Result<(), String> { + let mut private = queue.into_iter(); + loop { + let job = { + let _dispatch = shared.dispatcher.acquire().await.unwrap(); + let selected = if a.dispatch_policy == DispatchPolicy::Threshold { + private.next() + } else { + let mut ready = shared.ready.lock().unwrap(); + // Prefer the core's own M class, then take any fitting group. + // A free large core can drain remaining small jobs (and vice versa). + let core_index = a.cores.iter().position(|c| c.id == core.config.id).unwrap(); + let index = ready + .iter() + .enumerate() + .filter(|(_, j)| job_fits(&w, &a, &core.config, j)) + .min_by_key(|(_, j)| { + let preferred = if j.rows.len() >= a.dispatch_threshold { + a.large_core + } else { + a.small_core + }; + (preferred != core_index, j.id) + }) + .map(|(index, _)| index); + index.map(|index| ready.remove(index)) + }; + if selected.is_some() { + let duration = a.dispatch_cycles * a.clock_period_ps; + shared.dispatcher_busy.fetch_add(duration, Ordering::SeqCst); + Executor::current() + .resolve_at(Duration::from_picos(duration)) + .await; + } + selected + }; + let Some(job) = job else { + break; + }; + let start_ps = Executor::current().now().as_picos(); + let m = job.rows.len(); + let d = w.input_dim; + let e = w.expert_hidden_dim; + { + let mut report = core.report.lock().unwrap(); + report.jobs += 1; + report.vector_sram_peak_bytes = report.vector_sram_peak_bytes.max(vector_bytes(&w, m)?); + report.accumulator_peak_bytes = + report + .accumulator_peak_bytes + .max(accumulator_bytes(&w, m, &a, &core.config)?); + } + // These five BF16 regions are explicitly reserved, and retain their + // contents until the output copy completes. No cross-core aliases. + let mut x = vec![bf16::ZERO; m * d]; + let mut gate = vec![bf16::ZERO; m * e]; + let mut up = vec![bf16::ZERO; m * e]; + let mut z = vec![bf16::ZERO; m * e]; + let mut output = vec![bf16::ZERO; m * d]; + let mut accumulator = vec![0f32; m * d.max(e)]; + let wait = shared.vector.work(m * d).await; + core.report.lock().unwrap().vector_wait_ps += wait; + for (row, reference) in job.rows.iter().enumerate() { + for col in 0..d { + x[row * d + col] = bf16::from_bits(w.inputs_bf16[reference.token][col]); + } + } + let expert = w.experts.iter().find(|x| x.id == job.expert).unwrap(); + gemm( + &x, + &mut gate, + &mut accumulator, + m, + &expert.gate, + &core, + &a, + &shared, + ) + .await?; + gemm( + &x, + &mut up, + &mut accumulator, + m, + &expert.up, + &core, + &a, + &shared, + ) + .await?; + let wait = shared.vector.work(m * e).await; + core.report.lock().unwrap().vector_wait_ps += wait; + for index in 0..m * e { + let g = gate[index].to_f32(); + z[index] = bf16::from_f32((g / (1.0 + (-g).exp())) * up[index].to_f32()); + if !z[index].is_finite() { + return Err("SwiGLU produced non-finite BF16".into()); + } + } + gemm( + &z, + &mut output, + &mut accumulator, + m, + &expert.down, + &core, + &a, + &shared, + ) + .await?; + let compute_done_ps = Executor::current().now().as_picos(); + let wait = shared.vector.work(m * d).await; + core.report.lock().unwrap().vector_wait_ps += wait; + { + let mut reorder = shared.reorder.lock().unwrap(); + for (row, reference) in job.rows.iter().enumerate() { + reorder[reference.output * d..(reference.output + 1) * d] + .copy_from_slice(&output[row * d..(row + 1) * d]); + } + } + shared.completions.lock().unwrap().push(JobCompletion { + job: job.id, + expert: job.expert, + shared: job.shared, + core: core.config.id.clone(), + rows: m, + start_ps, + compute_done_ps, + output_copied_ps: Executor::current().now().as_picos(), + }); + // All per-job buffers are dropped here, before the next job is admitted. + } + Ok(()) +} + +/// Execute inside an already-entered runtime::Executor. The public `run` +/// convenience wrapper owns one executor; callers integrating Ramulator may +/// instead spawn this future on their existing executor. +pub async fn execute( + w: Workload, + a: Architecture, + hbm: Arc, + hbm_len: u64, +) -> Result { + validate(&w, &a, hbm_len)?; + if a.dma.is_some() && !hbm.supports_sector_reads() { + return Err( + "DMA sector/coalescing mode requires an explicitly capable memory backend".into(), + ); + } + let plan = plan(&w, &a)?; + let ex = Executor::current(); + let begin = ex.now().as_picos(); + let w = Arc::new(w); + let a = Arc::new(a); + let queue_peak = plan.queues.iter().map(Vec::len).sum::() * 64; + let shared = Arc::new(Shared { + ready: Mutex::new(if a.dispatch_policy == DispatchPolicy::WorkConserving { + plan.queues.iter().flatten().cloned().collect() + } else { + Vec::new() + }), + dispatcher: Semaphore::new(1), + dispatcher_busy: AtomicU64::new(0), + dma: Arc::new(Dma { + hbm, + credits: Semaphore::new(a.global_dma_credits), + fair_pool: a + .dma + .as_ref() + .filter(|d| d.fair_credits) + .map(|_| Arc::new(CreditPool::new(a.global_dma_credits, a.cores.len()))), + core_indices: a + .cores + .iter() + .enumerate() + .map(|(i, c)| (c.id.clone(), i)) + .collect(), + inflight: AtomicUsize::new(0), + peak: AtomicUsize::new(0), + bytes: AtomicU64::new(0), + config: a.dma.clone(), + clock: a.clock_period_ps, + lookup: (0..4).map(|_| Semaphore::new(1)).collect(), + copy: (0..8).map(|_| Semaphore::new(1)).collect(), + lines: Mutex::new(BTreeMap::new()), + report: Mutex::new(DmaReport { + fair_credit_reserve_per_core: if a.dma.as_ref().is_some_and(|d| d.fair_credits) { + a.global_dma_credits / 8 + } else { + 0 + }, + reserved_bytes: if a.dma.is_some() { + frontend_bytes(&a)? + } else { + 0 + }, + ..Default::default() + }), + }), + vector: Arc::new(VectorUnit { + permit: Semaphore::new(1), + lanes: a.vector_elements_per_cycle, + clock: a.clock_period_ps, + busy_ps: AtomicU64::new(0), + }), + reorder: Mutex::new(vec![bf16::ZERO; plan.output_rows * w.input_dim]), + completions: Mutex::new(Vec::new()), + }); + let cores: Vec> = a + .cores + .iter() + .map(|c| Arc::new(CoreState::new(c, &a))) + .collect(); + let mut pending = Vec::new(); + for (core, queue) in cores.iter().cloned().zip(plan.queues) { + let w = w.clone(); + let a = a.clone(); + let shared = shared.clone(); + let (tx, rx) = oneshot::channel(); + pending.push(rx); + ex.spawn(async move { + let _ = tx.send(run_core(core, queue, w, a, shared).await); + }); + } + // All queues drain even if one returns an error; no orphaned tasks or credits. + let mut failure = None; + for rx in pending { + match rx.await { + Ok(Ok(())) => (), + Ok(Err(e)) => { + failure.get_or_insert(e); + } + Err(_) => { + failure.get_or_insert("core task terminated without result".into()); + } + } + } + if let Some(error) = failure { + return Err(error); + } + let d = w.input_dim; + let tokens = w.inputs_bf16.len(); + let mut sums = vec![vec![0f32; d]; tokens]; + // One shared unit, one deterministic accumulation stream. A per-route + // BF16 result is retained in bounded reorder SRAM until this point. + shared.vector.work(plan.output_rows * d).await; + { + let reorder = shared.reorder.lock().unwrap(); + for (index, route) in plan.routes.iter().enumerate() { + for col in 0..d { + sums[route.token][col] += route.weight * reorder[index * d + col].to_f32(); + } + } + if let Some(s) = &w.shared_expert { + for token in 0..tokens { + for col in 0..d { + sums[token][col] += + s.weight * reorder[(plan.shared_base + token) * d + col].to_f32(); + } + } + } + } + shared.vector.work(tokens * d).await; + let output_bf16: Vec> = sums + .iter() + .map(|row| row.iter().map(|v| bf16::from_f32(*v).to_bits()).collect()) + .collect(); + let output_f32: Vec> = output_bf16 + .iter() + .map(|row| row.iter().map(|v| bf16::from_bits(*v).to_f32()).collect()) + .collect(); + if output_f32.iter().flatten().any(|v| !v.is_finite()) { + return Err("combined output is non-finite BF16".into()); + } + let total_ps = ex.now().as_picos() - begin; + let reports: Vec = cores + .iter() + .map(|c| { + let mut r = c.report.lock().unwrap().clone(); + r.weight_slots_peak = c.slots_peak.load(Ordering::SeqCst); + r.weight_sram_peak_bytes = c.weight_peak.load(Ordering::SeqCst); + r.cache_requests = c.cache.requests.load(Ordering::SeqCst); + r.cache_hits = c.cache.hits.load(Ordering::SeqCst); + r.cache_port_busy_ps = c.cache.busy.load(Ordering::SeqCst); + r.cache_peak_bytes = c.cache.peak.load(Ordering::SeqCst); + if total_ps > 0 { + r.compute_busy_fraction = r.compute_busy_ps as f64 / total_ps as f64; + r.mac_utilization = r.useful_macs as f64 * a.clock_period_ps as f64 + / (r.multipliers as f64 * total_ps as f64); + } + r + }) + .collect(); + let inflight_peak = shared.dma.peak.load(Ordering::SeqCst); + let mut completions = shared.completions.lock().unwrap().clone(); + for j in &mut completions { + j.start_ps -= begin; + j.compute_done_ps -= begin; + j.output_copied_ps -= begin; + } + Ok(RunReport { + schema_version: 1, workload: w.name.clone(), architecture: a.name.clone(), + timing_model: format!("{:?}: pipelined max(BLEN, ceil(BLEN*MLEN/activation_supply)) service with log2(MLEN)+pipeline readiness, or legacy serialized MLEN+overhead per instruction; shared ErasedMemoryModel; not RTL calibrated", a.matrix_timing), + timing_boundary: "ready BF16 inputs/routes and resident MX weights -> core gathers, weight reads/decode, three numerical GEMMs, SwiGLU, route reorder, deterministic weighted combine -> ready BF16 output; excludes router, initial input/weight placement, output HBM store".into(), + weight_format: "PLENA local E4M3/E8M0 block8, output-major [N,K], separate element/scale streams, decoded BF16 normal SRAM; not an OCP MX conformance claim".into(), + total_ps, multipliers: reports.iter().map(|r| r.multipliers).sum(), + useful_macs: reports.iter().map(|r| r.useful_macs).sum(), issued_macs: reports.iter().map(|r| r.issued_macs).sum(), + hbm_read_bytes: shared.dma.bytes.load(Ordering::SeqCst), hbm_write_bytes: 0, + global_dma_inflight_peak: inflight_peak, global_dma_staging_peak_bytes: inflight_peak * 64, + combine_sram_peak_bytes: plan.global_storage, shared_vector_busy_ps: shared.vector.busy_ps.load(Ordering::SeqCst), + dispatch_queue_peak_bytes: queue_peak, dispatcher_busy_ps: shared.dispatcher_busy.load(Ordering::SeqCst), + dma_frontend: a.dma.as_ref().map(|_| shared.dma.report.lock().unwrap().clone()), + cores: reports, job_completions: completions, output_bf16, output_f32, pre_round_output_f32: sums, + }) +} + +/// Run a fresh deterministic simulation. HBM backing and timing are provided by +/// the caller and shared by all cores. The HBM image length is mandatory. +pub async fn run( + w: Workload, + a: Architecture, + hbm: Arc, + hbm_len: u64, +) -> Result { + validate(&w, &a, hbm_len)?; + let ex = Executor::new(); + let (tx, mut rx) = oneshot::channel(); + ex.spawn(async move { + let _ = tx.send(execute(w, a, hbm, hbm_len).await); + }); + ex.enter(Instant::ETERNITY).await; + rx.try_recv().map_err(|_| { + "simulation stopped before completion (deadlock or task failure)".to_string() + })? +} + +#[cfg(test)] +mod dma_lifetime_tests { + use super::*; + + #[tokio::test] + async fn late_join_fetches_missing_sector_and_completed_entries_are_released() { + let mut a = super::super::tests::architecture(); + a.dma = Some(DmaConfig { + issue_policy: ramulator::model::IssuePolicy::PerChannel, + sector_reads: true, + coalesce: true, + fair_credits: false, + lookup_ii_cycles: 2, + frontend_sram_bytes: 45056, + }); + let backing = memory::MemoryBacked::with_capacity(64); + backing.with_data(|bytes| { + for (i, b) in bytes.iter_mut().enumerate() { + *b = i as u8; + } + }); + let native = ramulator::Ramulator::hbm2_preset(8).unwrap(); + let hbm = Arc::new(memory::WithStats::new(memory::WithTiming::new( + native, backing, + ))); + let dma = Arc::new(Dma { + hbm: hbm.clone(), + credits: Semaphore::new(2), + fair_pool: None, + core_indices: BTreeMap::new(), + inflight: AtomicUsize::new(0), + peak: AtomicUsize::new(0), + bytes: AtomicU64::new(0), + config: a.dma.clone(), + clock: 1000, + lookup: (0..4).map(|_| Semaphore::new(1)).collect(), + copy: (0..8).map(|_| Semaphore::new(1)).collect(), + lines: Mutex::new(BTreeMap::new()), + report: Mutex::new(DmaReport::default()), + }); + let core = Arc::new(CoreState::new(&a.cores[0], &a)); + let ex = Executor::new(); + let d = dma.clone(); + ex.spawn(async move { + let first = d.read(0, 1, 32, core.clone()); + let late = async { + Executor::current() + .resolve_at(Duration::from_picos(3000)) + .await; + d.read(0, 3, 64, core.clone()).await + }; + let (a, b) = futures::join!(first, late); + assert_eq!(&a[..32], &(0..32u8).collect::>()); + assert!( + a[32..].iter().all(|b| *b == 0), + "unrequested sector leaked into result" + ); + assert_eq!(b, std::array::from_fn(|i| i as u8)); + assert!( + d.lines.lock().unwrap().is_empty(), + "MSHR retained after final copy" + ); + d.read(0, 1, 32, core).await; + }); + ex.enter(Instant::ETERNITY).await; + let r = dma.report.lock().unwrap(); + assert_eq!( + ( + r.line_requests, + r.sector_requests, + r.merged_sectors, + r.mshr_peak + ), + (3, 4, 1, 1) + ); + assert_eq!(hbm.statistics().total_bytes_read, 96); + assert_eq!(dma.bytes.load(Ordering::SeqCst), 96); + } +} diff --git a/transactional_emulator/src/moe_normal/mod.rs b/transactional_emulator/src/moe_normal/mod.rs new file mode 100644 index 000000000..68591b7b7 --- /dev/null +++ b/transactional_emulator/src/moe_normal/mod.rs @@ -0,0 +1,15 @@ +//! Executable V0: expert grouping, independent normal SRAM/accumulators, +//! finite double-buffered MX ingress, and one shared HBM/vector path. +//! This is a numerical architecture experiment, not the ISA or RTL runner. +mod dma_credits; +mod engine; +mod read_cache; +mod types; + +// `execute` is retained for an embedding caller with an existing executor. +#[allow(unused_imports)] +pub use engine::{execute, run, validate}; +pub use types::*; + +#[cfg(test)] +mod tests; diff --git a/transactional_emulator/src/moe_normal/read_cache.rs b/transactional_emulator/src/moe_normal/read_cache.rs new file mode 100644 index 000000000..e0dcf2c64 --- /dev/null +++ b/transactional_emulator/src/moe_normal/read_cache.rs @@ -0,0 +1,80 @@ +//! One finite FIFO line cache per core. Port service is one cycle per lookup +//! and insertion. A loader credit retains returned data until its tile copy. +use runtime::{Duration, Executor}; +use std::collections::{BTreeMap, VecDeque}; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use tokio::sync::Semaphore; + +pub const ENTRY_BYTES: usize = 80; + +#[derive(Default)] +struct State { + lines: BTreeMap, + fifo: VecDeque, +} + +pub struct ReadCache { + entries: usize, + clock: u64, + port: Semaphore, + state: Mutex, + pub requests: AtomicU64, + pub hits: AtomicU64, + pub busy: AtomicU64, + pub peak: AtomicUsize, +} + +impl ReadCache { + pub fn new(bytes: usize, clock: u64) -> Self { + Self { + entries: bytes / ENTRY_BYTES, + clock, + port: Semaphore::new(1), + state: Mutex::new(State::default()), + requests: AtomicU64::new(0), + hits: AtomicU64::new(0), + busy: AtomicU64::new(0), + peak: AtomicUsize::new(0), + } + } + + pub async fn lookup(&self, address: u64) -> Option<[u8; 64]> { + self.requests.fetch_add(1, Ordering::SeqCst); + if self.entries == 0 { + return None; + } + let _port = self.port.acquire().await.unwrap(); + self.busy.fetch_add(self.clock, Ordering::SeqCst); + Executor::current() + .resolve_at(Duration::from_picos(self.clock)) + .await; + let value = self.state.lock().unwrap().lines.get(&address).copied(); + if value.is_some() { + self.hits.fetch_add(1, Ordering::SeqCst); + } + value + } + + pub async fn insert(&self, address: u64, bytes: [u8; 64]) { + if self.entries == 0 { + return; + } + let _port = self.port.acquire().await.unwrap(); + self.busy.fetch_add(self.clock, Ordering::SeqCst); + Executor::current() + .resolve_at(Duration::from_picos(self.clock)) + .await; + let mut state = self.state.lock().unwrap(); + if !state.lines.contains_key(&address) { + if state.lines.len() == self.entries { + let oldest = state.fifo.pop_front().unwrap(); + state.lines.remove(&oldest); + } + state.fifo.push_back(address); + } + state.lines.insert(address, bytes); + self.peak + .fetch_max(state.lines.len() * ENTRY_BYTES, Ordering::SeqCst); + } +} diff --git a/transactional_emulator/src/moe_normal/tests.rs b/transactional_emulator/src/moe_normal/tests.rs new file mode 100644 index 000000000..e91af7798 --- /dev/null +++ b/transactional_emulator/src/moe_normal/tests.rs @@ -0,0 +1,741 @@ +use std::sync::Arc; + +use half::bf16; +use memory::{MemoryBacked, MemoryTimingModel, WithStats, WithTiming}; +use runtime::{Duration, Executor}; + +use super::*; + +/// One serialized shared transfer resource, with real bytes behind it. Unlike +/// a request-count-only test, wrong addresses or missing reads change results. +struct TestTiming { + bus: tokio::sync::Semaphore, +} +impl MemoryTimingModel for TestTiming { + async fn read(&self, _: u64) { + let _permit = self.bus.acquire().await.unwrap(); + Executor::current() + .resolve_at(Duration::from_picos(3_000)) + .await; + } + async fn write(&self, _: u64) { + let _permit = self.bus.acquire().await.unwrap(); + Executor::current() + .resolve_at(Duration::from_picos(3_000)) + .await; + } +} + +#[derive(Clone)] +struct DenseExpert { + gate: Vec, + up: Vec, + down: Vec, +} + +fn append_matrix( + bytes: &mut Vec, + rows: usize, + cols: usize, + expert: usize, + projection: usize, +) -> (MatrixRegion, Vec) { + bytes.resize(bytes.len().div_ceil(64) * 64, 0xa5); + let element_base = bytes.len() as u64; + let stride = cols.div_ceil(8) * 8; + // Padding deliberately contains nonzero values; masked tails must ignore it. + bytes.resize(bytes.len() + rows * stride, 0x38); + let mut values = vec![0f32; rows * cols]; + for row in 0..rows { + for col in 0..cols { + let byte = if row % cols == col { + 0x38 + } else if (row + col + expert + projection).is_multiple_of(7) { + 0xb0 + } else { + 0x00 + }; + bytes[element_base as usize + row * stride + col] = byte; + // Alternate explicit shared scales by row. All values exact BF16. + let scale = if row % 2 == 0 { 1.0 } else { 0.5 }; + values[row * cols + col] = scale + * match byte { + 0x38 => 1.0, + 0xb0 => -0.5, + _ => 0.0, + }; + } + } + bytes.resize(bytes.len().div_ceil(64) * 64, 0xa5); + let scale_base = bytes.len() as u64; + let scale_stride = cols.div_ceil(8); + for row in 0..rows { + bytes.extend(std::iter::repeat_n( + if row % 2 == 0 { 127 } else { 126 }, + scale_stride, + )); + } + bytes.resize(bytes.len().div_ceil(64) * 64, 0xa5); + ( + MatrixRegion { + rows, + cols, + element_base, + scale_base, + element_row_stride: stride as u64, + scale_row_stride: scale_stride as u64, + }, + values, + ) +} + +fn fixture() -> (Workload, Vec, Vec) { + let d = 9; + let e = 11; + let mut bytes = Vec::new(); + let mut experts = Vec::new(); + let mut dense = Vec::new(); + for expert in 0..4 { + let (gate, g) = append_matrix(&mut bytes, e, d, expert, 0); + let (up, u) = append_matrix(&mut bytes, e, d, expert, 1); + let (down, out) = append_matrix(&mut bytes, d, e, expert, 2); + experts.push(Expert { + id: expert, + gate, + up, + down, + }); + dense.push(DenseExpert { + gate: g, + up: u, + down: out, + }); + } + let input: Vec> = (0..7) + .map(|t| { + (0..d) + .map(|k| bf16::from_f32((t as f32 - k as f32 + 2.0) / 8.0).to_bits()) + .collect() + }) + .collect(); + let routes = (0..7) + .map(|token| Route { + token, + slot: 0, + expert: if token < 4 { + 0 + } else if token < 6 { + 1 + } else { + 2 + }, + weight: 0.75, + }) + .collect(); + ( + Workload { + schema_version: 1, + name: "unit_nonzero_tails".into(), + hbm_file: "unused.bin".into(), + input_dim: d, + expert_hidden_dim: e, + inputs_bf16: input, + routes, + experts, + shared_expert: None, + metadata: None, + grouped_routes: None, + }, + bytes, + dense, + ) +} + +pub(super) fn architecture() -> Architecture { + Architecture { + schema_version: 1, + name: "dual_normal".into(), + cores: vec![ + CoreConfig { + id: "large".into(), + blen: 4, + mlen: 16, + vector_sram_bytes: 4096, + accumulator_bytes: 4096, + weight_sram_bytes: 4096, + read_cache_bytes: 0, + weight_slots: 2, + activation_elements_per_cycle: None, + }, + CoreConfig { + id: "small".into(), + blen: 2, + mlen: 8, + vector_sram_bytes: 4096, + accumulator_bytes: 4096, + weight_sram_bytes: 4096, + read_cache_bytes: 0, + weight_slots: 2, + activation_elements_per_cycle: None, + }, + ], + dispatch_threshold: 4, + large_core: 0, + small_core: 1, + global_dma_credits: 2, + global_dma_staging_bytes: 128, + combine_sram_bytes: 16384, + clock_period_ps: 1000, + mac_pipeline_cycles: 2, + vector_elements_per_cycle: 16, + dispatch_policy: DispatchPolicy::Threshold, + dispatch_queue_bytes: 262144, + dispatch_cycles: 1, + matrix_timing: MatrixTiming::Pipelined, + dma: None, + } +} + +#[tokio::test] +async fn serialized_instruction_sensitivity_preserves_values_and_counts() { + let (w, bytes, _) = fixture(); + let a = architecture(); + let pipeline = simulate(w.clone(), a.clone(), &bytes).await; + let mut serial = a.clone(); + serial.matrix_timing = MatrixTiming::LegacySerialized; + let report = simulate(w, serial, &bytes).await; + assert_eq!(report.output_bf16, pipeline.output_bf16); + assert_eq!(report.useful_macs, pipeline.useful_macs); + for (c, config) in report.cores.iter().zip(&a.cores) { + let tile_macs = (config.blen * config.blen * config.mlen) as u64; + assert_eq!( + c.compute_busy_ps, + c.issued_macs / tile_macs + * (config.mlen as u64 + a.mac_pipeline_cycles) + * a.clock_period_ps + ); + assert_eq!(c.pipeline_drain_ps, 0); + } +} + +#[tokio::test] +async fn finite_cache_reuses_real_bytes_and_charges_one_serial_port() { + let (w, bytes, _) = fixture(); + let a = architecture(); + let baseline = simulate(w.clone(), a.clone(), &bytes).await; + for capacity in [80, 4096] { + let mut cached = a.clone(); + for core in &mut cached.cores { + core.read_cache_bytes = capacity; + } + let r = simulate(w.clone(), cached, &bytes).await; + assert_eq!(r.output_bf16, baseline.output_bf16); + assert!(r.hbm_read_bytes <= baseline.hbm_read_bytes); + for c in &r.cores { + assert!(c.cache_peak_bytes <= capacity); + assert_eq!(c.cache_requests, c.cache_hits + c.hbm_read_bytes / 64); + assert_eq!( + c.cache_port_busy_ps, + (c.cache_requests + c.hbm_read_bytes / 64) * a.clock_period_ps + ); + assert!(c.cache_port_busy_ps <= r.total_ps); + } + if capacity == 4096 { + assert!(r.hbm_read_bytes < baseline.hbm_read_bytes); + assert!(r.cores.iter().any(|c| c.cache_hits > 0)); + } + } +} + +#[tokio::test] +async fn ready_dispatch_steals_small_jobs_and_preserves_numerics() { + let (mut w, bytes, _) = fixture(); + for id in 10..30 { + let mut expert = w.experts[1].clone(); + expert.id = id; + w.experts.push(expert); + w.routes.push(Route { + token: 0, + slot: id, + expert: id, + weight: 0.03125, + }); + } + let fixed = simulate(w.clone(), architecture(), &bytes).await; + let mut a = architecture(); + a.dispatch_policy = DispatchPolicy::WorkConserving; + let ready = simulate(w.clone(), a.clone(), &bytes).await; + assert_eq!(ready.output_bf16, fixed.output_bf16); + assert!( + ready + .job_completions + .iter() + .any(|j| j.rows < a.dispatch_threshold && j.core == "large") + ); + assert_eq!(ready.job_completions.len(), fixed.job_completions.len()); + assert_eq!( + ready.dispatcher_busy_ps, + ready.job_completions.len() as u64 * a.dispatch_cycles * a.clock_period_ps + ); + assert!(ready.dispatch_queue_peak_bytes <= a.dispatch_queue_bytes); + a.dispatch_queue_bytes = 64; + assert!(validate(&w, &a, bytes.len() as u64).is_err()); +} + +async fn simulate(w: Workload, a: Architecture, bytes: &[u8]) -> RunReport { + let backing = MemoryBacked::with_capacity(bytes.len()); + backing.with_data(|dst| dst.copy_from_slice(bytes)); + let memory = Arc::new(WithStats::new(WithTiming::new( + TestTiming { + bus: tokio::sync::Semaphore::new(1), + }, + backing, + ))); + let report = run(w, a, memory.clone(), bytes.len() as u64).await.unwrap(); + assert_eq!(report.hbm_read_bytes, memory.statistics().total_bytes_read); + assert_eq!(report.hbm_write_bytes, 0); + report +} + +fn gemm_reference(x: &[f32], w: &[f32], n: usize, k: usize) -> Vec { + (0..n) + .map(|col| { + let mut sum = 0f32; + for kk in 0..k { + sum += x[kk] * w[col * k + kk]; + } + bf16::from_f32(sum).to_f32() + }) + .collect() +} + +fn expert_reference(x: &[f32], w: &DenseExpert, d: usize, e: usize) -> Vec { + let gate = gemm_reference(x, &w.gate, e, d); + let up = gemm_reference(x, &w.up, e, d); + let z: Vec = gate + .into_iter() + .zip(up) + .map(|(g, u)| bf16::from_f32((g / (1.0 + (-g).exp())) * u).to_f32()) + .collect(); + gemm_reference(&z, &w.down, d, e) +} + +fn reference(w: &Workload, dense: &[DenseExpert]) -> Vec> { + let mut routes = w.routes.clone(); + routes.sort_by_key(|r| (r.token, r.slot)); + (0..w.inputs_bf16.len()) + .map(|token| { + let x: Vec = w.inputs_bf16[token] + .iter() + .map(|v| bf16::from_bits(*v).to_f32()) + .collect(); + let mut sums = vec![0f32; w.input_dim]; + for r in routes.iter().filter(|r| r.token == token) { + let y = expert_reference(&x, &dense[r.expert], w.input_dim, w.expert_hidden_dim); + for (s, v) in sums.iter_mut().zip(y) { + *s += r.weight * v; + } + } + if let Some(r) = &w.shared_expert { + let y = expert_reference(&x, &dense[r.expert], w.input_dim, w.expert_hidden_dim); + for (s, v) in sums.iter_mut().zip(y) { + *s += r.weight * v; + } + } + sums.into_iter() + .map(|v| bf16::from_f32(v).to_bits()) + .collect() + }) + .collect() +} + +#[tokio::test] +async fn heterogeneous_nonzero_full_moe_and_tail_masks_match_scalar_reference() { + let (w, bytes, dense) = fixture(); + let expected = reference(&w, &dense); + let report = simulate(w.clone(), architecture(), &bytes).await; + assert_eq!(report.output_bf16, expected); + assert!(report.output_f32.iter().flatten().any(|v| *v != 0.0)); + assert_eq!( + report.useful_macs, + (w.routes.len() * 3 * w.input_dim * w.expert_hidden_dim) as u64 + ); + assert!(report.issued_macs > report.useful_macs); + assert!(report.cores.iter().all(|r| r.jobs > 0)); + assert!(report.cores.iter().all(|r| r.weight_slots_peak == 2)); + assert!( + report + .cores + .iter() + .all(|r| r.accumulator_peak_bytes > 0 && r.vector_sram_peak_bytes > 0) + ); + for c in &report.cores { + assert_eq!(c.compute_busy_ps, c.issued_macs / c.multipliers * 1000); + assert!(c.mac_utilization > 0.0 && c.mac_utilization <= 1.0); + } +} + +#[tokio::test] +async fn baseline_has_same_numerical_path_and_multiplies_and_streaming_service() { + let (w, bytes, _) = fixture(); + let dual = architecture(); + let mut single = dual.clone(); + single.name = "single".into(); + // 4*16 + 2*8 = 80 = 2*40, with a non-power-of-two MLEN. + single.cores.truncate(1); + single.cores[0].blen = 2; + single.cores[0].mlen = 40; + single.large_core = 0; + single.small_core = 0; + let a = simulate(w.clone(), dual, &bytes).await; + let b = simulate(w, single, &bytes).await; + assert_eq!(a.multipliers, b.multipliers); + assert_eq!(a.output_bf16, b.output_bf16); + assert_eq!(a.useful_macs, b.useful_macs); + assert_eq!(b.cores[0].weight_slots_peak, 2); + assert_eq!( + b.cores[0].compute_busy_ps, + b.issued_macs / b.multipliers * 1000 + ); + assert!(b.cores[0].pipeline_drain_ps > 0); +} + +#[tokio::test] +async fn duplicate_expert_routes_shared_expert_and_input_order_are_deterministic() { + let (mut w, bytes, dense) = fixture(); + w.routes.push(Route { + token: 0, + slot: 3, + expert: 0, + weight: -0.25, + }); + w.routes.push(Route { + token: 0, + slot: 1, + expert: 2, + weight: 0.5, + }); + w.shared_expert = Some(SharedExpert { + expert: 3, + weight: 0.125, + }); + let expected = reference(&w, &dense); + let first = simulate(w.clone(), architecture(), &bytes).await; + w.routes.reverse(); + let reversed = simulate(w, architecture(), &bytes).await; + assert_eq!(first.output_bf16, expected); + assert_eq!(first.output_bf16, reversed.output_bf16); + assert_eq!(first.total_ps, reversed.total_ps); + assert_eq!(first.job_completions.len(), 4); + assert!(first.job_completions.iter().any(|j| j.shared)); +} + +#[tokio::test] +async fn global_credit_limit_is_shared_and_hbm_byte_counters_are_real() { + let (w, bytes, _) = fixture(); + for credits in [1, 2, 3] { + let mut a = architecture(); + a.global_dma_credits = credits; + a.global_dma_staging_bytes = 64 * credits; + let report = simulate(w.clone(), a, &bytes).await; + assert_eq!(report.global_dma_inflight_peak, credits); + assert_eq!(report.global_dma_staging_peak_bytes, credits * 64); + assert_eq!( + report.hbm_read_bytes, + report.cores.iter().map(|r| r.hbm_read_bytes).sum::() + ); + assert!(report.hbm_read_bytes > 0); + } +} + +#[test] +fn finite_capacity_shortages_fail_before_execution() { + let (w, bytes, _) = fixture(); + let valid = architecture(); + validate(&w, &valid, bytes.len() as u64).unwrap(); + for resource in 0..5 { + let mut a = valid.clone(); + match resource { + 0 => a.cores[0].vector_sram_bytes = 1, + 1 => a.cores[0].accumulator_bytes = 1, + 2 => a.cores[0].weight_sram_bytes = 1, + 3 => a.combine_sram_bytes = 1, + 4 => a.global_dma_staging_bytes = 1, + _ => unreachable!(), + } + assert!( + validate(&w, &a, bytes.len() as u64).is_err(), + "resource {resource}" + ); + } +} + +#[test] +fn malformed_routes_geometry_and_hbm_ranges_are_rejected() { + let (w, bytes, _) = fixture(); + let a = architecture(); + let mut bad = w.clone(); + bad.routes.push(bad.routes[0].clone()); + assert!( + validate(&bad, &a, bytes.len() as u64) + .unwrap_err() + .contains("duplicate route slot") + ); + let mut bad = w.clone(); + bad.experts[0].gate.scale_base = bytes.len() as u64; + assert!( + validate(&bad, &a, bytes.len() as u64) + .unwrap_err() + .contains("beyond HBM") + ); + let mut bad = w.clone(); + bad.experts[0].gate.element_row_stride = 9; + assert!( + validate(&bad, &a, bytes.len() as u64) + .unwrap_err() + .contains("strides") + ); + let mut bad = w.clone(); + bad.inputs_bf16[0][0] = bf16::NAN.to_bits(); + assert!(validate(&bad, &a, bytes.len() as u64).is_err()); + let mut bad_a = a; + bad_a.cores[0].mlen = 7; + assert!(validate(&w, &bad_a, bytes.len() as u64).is_err()); +} + +#[tokio::test] +async fn unselected_experts_do_not_read_hbm_and_empty_routes_return_zero() { + let (mut w, bytes, _) = fixture(); + w.routes.clear(); + let report = simulate(w, architecture(), &bytes).await; + assert_eq!(report.hbm_read_bytes, 0); + assert_eq!(report.useful_macs, 0); + assert!(report.output_bf16.iter().flatten().all(|v| *v == 0)); + assert!(report.job_completions.is_empty()); + assert!(report.shared_vector_busy_ps > 0); +} + +#[tokio::test] +async fn out_of_order_completion_preserves_route_reduction_order() { + let (mut w, bytes, dense) = fixture(); + // Large expert is assigned first, but receives many extra rows. Smaller + // expert must complete while it is still running; combine cannot reorder. + for slot in 1..100 { + w.routes.push(Route { + token: 0, + slot, + expert: 0, + weight: 0.0625, + }); + } + let expected = reference(&w, &dense); + let mut a = architecture(); + a.cores[0].vector_sram_bytes = 16384; + a.cores[0].accumulator_bytes = 16384; + let report = simulate(w, a, &bytes).await; + assert_ne!(report.job_completions[0].job, 0); + assert_eq!(report.output_bf16, expected); + let large = report.job_completions.iter().find(|j| j.job == 0).unwrap(); + assert!(report.job_completions[0].output_copied_ps < large.output_copied_ps); +} + +#[test] +fn serde_contract_rejects_misspelled_fields() { + let (w, _, _) = fixture(); + let mut value = serde_json::to_value(w).unwrap(); + value["hidden_dimm"] = serde_json::json!(9); + assert!(serde_json::from_value::(value).is_err()); +} + +#[test] +fn extreme_clock_and_invalid_square_mapping_fail_before_timers() { + let (w, bytes, _) = fixture(); + let mut a = architecture(); + a.cores.truncate(1); + a.large_core = 0; + a.small_core = 0; + a.cores[0].blen = 1; + a.cores[0].mlen = 8; + a.mac_pipeline_cycles = 0; + a.clock_period_ps = u64::MAX / 4; + a.vector_elements_per_cycle = 1; + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("duration exceeds") + ); + a.clock_period_ps = 1000; + a.cores[0].blen = 3; + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("divisible by BLEN") + ); +} + +async fn simulate_native(w: Workload, a: Architecture, bytes: &[u8]) -> RunReport { + let backing = memory::MemoryBacked::with_capacity(bytes.len()); + backing.with_data(|data| data.copy_from_slice(bytes)); + let native = ramulator::Ramulator::hbm2_preset(8) + .unwrap() + .with_issue_policy( + a.dma.as_ref().unwrap().issue_policy, + Duration::from_picos(a.clock_period_ps), + ); + let observer = native.clone(); + let memory = Arc::new(memory::WithStats::new(memory::WithTiming::new( + native, backing, + ))); + let result = run(w, a, memory.clone(), bytes.len() as u64).await.unwrap(); + assert_eq!(result.hbm_read_bytes, memory.statistics().total_bytes_read); + let stats = observer.telemetry(); + let controllers = stats["native_stats"]["memory_system"]["controller"] + .as_array() + .unwrap(); + let served: u64 = controllers + .iter() + .map(|c| c["num_read_reqs_served"].as_u64().unwrap()) + .sum(); + assert_eq!(served * 32, result.hbm_read_bytes); + assert_eq!(stats["native_pending"], 0); + result +} + +#[tokio::test] +async fn native_dma_variants_preserve_moe_and_shared_expert_with_tail_masks() { + let (mut w, bytes, dense) = fixture(); + w.shared_expert = Some(SharedExpert { + expert: 0, + weight: 0.5, + }); + let expected = reference(&w, &dense); + for slots in [2, 3, 4] { + for (policy, sector, coalesce, fair) in [ + ( + ramulator::model::IssuePolicy::GlobalFifo, + false, + false, + false, + ), + ( + ramulator::model::IssuePolicy::PerChannel, + true, + false, + false, + ), + (ramulator::model::IssuePolicy::PerChannel, true, true, false), + (ramulator::model::IssuePolicy::PerChannel, true, true, true), + ] { + let mut a = architecture(); + a.global_dma_credits = 16; + a.global_dma_staging_bytes = 1024; + for core in &mut a.cores { + core.weight_slots = slots; + } + a.dma = Some(DmaConfig { + issue_policy: policy, + sector_reads: sector, + coalesce, + fair_credits: fair, + lookup_ii_cycles: 2, + frontend_sram_bytes: 45056, + }); + let result = simulate_native(w.clone(), a, &bytes).await; + assert_eq!(result.output_bf16, expected); + let dma = result.dma_frontend.unwrap(); + assert_eq!( + (dma.sector_requests - dma.merged_sectors) * 32, + result.hbm_read_bytes + ); + assert!(dma.mshr_peak <= result.global_dma_inflight_peak); + assert!(result.cores.iter().all(|c| c.weight_slots_peak <= slots)); + } + } +} + +#[test] +fn dma_metadata_and_prefetch_slots_fail_closed() { + let (w, bytes, _) = fixture(); + let mut a = architecture(); + a.dma = Some(DmaConfig { + issue_policy: ramulator::model::IssuePolicy::PerChannel, + sector_reads: true, + coalesce: true, + fair_credits: false, + lookup_ii_cycles: 2, + frontend_sram_bytes: 1, + }); + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("metadata") + ); + a.dma.as_mut().unwrap().frontend_sram_bytes = 45056; + a.cores[0].weight_slots = 4; + a.cores[0].weight_sram_bytes = 600; + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("weight slots") + ); + a.cores[0].weight_sram_bytes = 4096; + a.cores[0].weight_slots = 5; + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("weight_slots") + ); +} + +#[tokio::test] +async fn activation_supply_bounds_compute_without_changing_values_or_work() { + let (w, bytes, _) = fixture(); + let mut a = architecture(); + let reference = simulate(w.clone(), a.clone(), &bytes).await; + for c in &mut a.cores { + c.activation_elements_per_cycle = Some(c.mlen / 2); + } + let limited = simulate(w.clone(), a.clone(), &bytes).await; + assert_eq!(reference.output_bf16, limited.output_bf16); + assert_eq!(reference.useful_macs, limited.useful_macs); + assert_eq!(reference.issued_macs, limited.issued_macs); + for (before, after) in reference.cores.iter().zip(&limited.cores) { + assert_eq!(after.compute_busy_ps, 2 * before.compute_busy_ps); + } + a.cores[0].activation_elements_per_cycle = Some(0); + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("activation supply") + ); +} + +#[tokio::test] +async fn pipelined_dma_lookup_accounts_port_occupancy_and_preserves_numerics() { + let (w, bytes, dense) = fixture(); + for ii in [1, 2] { + let mut a = architecture(); + a.dma = Some(DmaConfig { + issue_policy: ramulator::model::IssuePolicy::PerChannel, + sector_reads: true, + coalesce: true, + fair_credits: false, + lookup_ii_cycles: ii, + frontend_sram_bytes: 45056, + }); + let result = simulate_native(w.clone(), a.clone(), &bytes).await; + assert_eq!(result.output_bf16, reference(&w, &dense)); + let d = result.dma_frontend.unwrap(); + assert_eq!( + d.lookup_busy_ps, + ii as u64 * d.line_requests * a.clock_period_ps + ); + assert!(d.lookup_busy_ps <= 4 * result.total_ps); + a.dma.as_mut().unwrap().lookup_ii_cycles = 0; + assert!( + validate(&w, &a, bytes.len() as u64) + .unwrap_err() + .contains("lookup_ii_cycles") + ); + } +} diff --git a/transactional_emulator/src/moe_normal/types.rs b/transactional_emulator/src/moe_normal/types.rs new file mode 100644 index 000000000..8d2136f55 --- /dev/null +++ b/transactional_emulator/src/moe_normal/types.rs @@ -0,0 +1,241 @@ +//! Versioned input and output contract for the normal-buffer MoE experiment. +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct MatrixRegion { + pub rows: usize, + pub cols: usize, + pub element_base: u64, + pub scale_base: u64, + pub element_row_stride: u64, + pub scale_row_stride: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct Expert { + pub id: usize, + pub gate: MatrixRegion, + pub up: MatrixRegion, + pub down: MatrixRegion, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct Route { + pub token: usize, + pub slot: usize, + pub expert: usize, + pub weight: f32, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct SharedExpert { + pub expert: usize, + pub weight: f32, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct Workload { + pub schema_version: u32, + pub name: String, + pub hbm_file: String, + pub input_dim: usize, + pub expert_hidden_dim: usize, + pub inputs_bf16: Vec>, + pub routes: Vec, + pub experts: Vec, + #[serde(default)] + pub shared_expert: Option, + #[serde(default)] + pub metadata: Option, + #[serde(default)] + pub grouped_routes: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct CoreConfig { + pub id: String, + /// Physical M and N tile dimensions. V0 retains square tiles. + pub blen: usize, + /// Physical K tile dimension; a multiple of local MX block size 8. + pub mlen: usize, + /// BF16 X, gate, up, Z and output, all reserved for one expert group. + pub vector_sram_bytes: usize, + /// FP32 output accumulator, reused across the three projections. + pub accumulator_bytes: usize, + /// Configurable slots, each holding packed MX ingress and decoded BF16 tile. + pub weight_sram_bytes: usize, + /// FIFO read cache: each entry reserves 64 data + 16 tag/control bytes. + #[serde(default)] + pub read_cache_bytes: usize, + #[serde(default = "default_weight_slots")] + pub weight_slots: usize, + /// Explicit BF16 activation supply per cycle. None retains the legacy + /// full-width analytical assumption; this is not a physical area estimate. + #[serde(default)] + pub activation_elements_per_cycle: Option, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum DispatchPolicy { + #[default] + Threshold, + WorkConserving, +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum MatrixTiming { + #[default] + Pipelined, + /// Legacy MatrixMachine-like per-instruction K + overhead serialization. + LegacySerialized, +} + +fn default_weight_slots() -> usize { + 2 +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct DmaConfig { + pub issue_policy: ramulator::model::IssuePolicy, + pub sector_reads: bool, + pub coalesce: bool, + #[serde(default)] + pub fair_credits: bool, + /// Lookup latency is two cycles; initiation interval is independently 1/2. + #[serde(default = "default_weight_slots")] + pub lookup_ii_cycles: usize, + /// Includes response data, MSHRs, waiters, native trackers, and tile descriptors. + pub frontend_sram_bytes: usize, +} + +fn default_queue_bytes() -> usize { + 262144 +} +fn default_dispatch_cycles() -> u64 { + 1 +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(deny_unknown_fields)] +pub struct Architecture { + pub schema_version: u32, + pub name: String, + pub cores: Vec, + /// Groups with M >= threshold select large_core, otherwise small_core. + pub dispatch_threshold: usize, + pub large_core: usize, + pub small_core: usize, + pub global_dma_credits: usize, + /// Each in-flight 64-byte HBM request retains its credit until copied. + pub global_dma_staging_bytes: usize, + /// Ready input BF16 + all BF16 route outputs + FP32 combine + BF16 final. + pub combine_sram_bytes: usize, + pub clock_period_ps: u64, + /// Analytical fill/drain overhead, in addition to BLEN + ceil(log2 MLEN). + pub mac_pipeline_cycles: u64, + /// One shared vector actor serves gathers, SwiGLU, result copies/combine. + pub vector_elements_per_cycle: usize, + #[serde(default)] + pub dispatch_policy: DispatchPolicy, + /// Fixed 64-byte descriptor per ready expert job. + #[serde(default = "default_queue_bytes")] + pub dispatch_queue_bytes: usize, + #[serde(default = "default_dispatch_cycles")] + pub dispatch_cycles: u64, + #[serde(default)] + pub matrix_timing: MatrixTiming, + #[serde(default)] + pub dma: Option, +} + +#[derive(Clone, Debug, Serialize)] +pub struct CoreReport { + pub id: String, + pub blen: usize, + pub mlen: usize, + pub multipliers: u64, + pub jobs: usize, + pub useful_macs: u64, + pub issued_macs: u64, + pub compute_busy_ps: u64, + pub accumulator_dependency_stall_ps: u64, + pub pipeline_drain_ps: u64, + pub pipeline_register_bytes: usize, + pub weight_ready_wait_ps: u64, + pub vector_wait_ps: u64, + pub vector_sram_peak_bytes: usize, + pub accumulator_peak_bytes: usize, + pub weight_sram_peak_bytes: usize, + pub weight_slots_peak: usize, + pub hbm_read_bytes: u64, + pub compute_busy_fraction: f64, + pub mac_utilization: f64, + pub cache_requests: u64, + pub cache_hits: u64, + pub cache_port_busy_ps: u64, + pub cache_peak_bytes: usize, +} + +#[derive(Clone, Debug, Serialize)] +pub struct JobCompletion { + pub job: usize, + pub expert: usize, + pub shared: bool, + pub core: String, + pub rows: usize, + pub start_ps: u64, + pub compute_done_ps: u64, + pub output_copied_ps: u64, +} + +#[derive(Clone, Debug, Serialize)] +pub struct RunReport { + pub schema_version: u32, + pub workload: String, + pub architecture: String, + pub timing_model: String, + pub timing_boundary: String, + pub weight_format: String, + pub total_ps: u64, + pub multipliers: u64, + pub useful_macs: u64, + pub issued_macs: u64, + pub hbm_read_bytes: u64, + pub hbm_write_bytes: u64, + pub global_dma_inflight_peak: usize, + pub global_dma_staging_peak_bytes: usize, + pub combine_sram_peak_bytes: usize, + pub shared_vector_busy_ps: u64, + pub dispatch_queue_peak_bytes: usize, + pub dispatcher_busy_ps: u64, + pub dma_frontend: Option, + pub cores: Vec, + /// Completion order, which may differ from deterministic reduction order. + pub job_completions: Vec, + pub output_bf16: Vec>, + pub output_f32: Vec>, + pub pre_round_output_f32: Vec>, +} + +#[derive(Clone, Debug, Default, Serialize)] +pub struct DmaReport { + pub reserved_bytes: usize, + pub line_requests: u64, + pub sector_requests: u64, + pub merged_sectors: u64, + pub useful_copy_bytes: u64, + pub lookup_busy_ps: u64, + pub copy_busy_ps: u64, + pub mshr_peak: usize, + pub fair_credit_reserve_per_core: usize, + pub fair_credit_wait_ps: u64, +} diff --git a/transactional_emulator/testbench/moe_normal_review/large_small.json b/transactional_emulator/testbench/moe_normal_review/large_small.json new file mode 100644 index 000000000..1aa42cef0 --- /dev/null +++ b/transactional_emulator/testbench/moe_normal_review/large_small.json @@ -0,0 +1,48 @@ +{ + "clock_period_ps": 1000, + "combine_sram_bytes": 4194304, + "cores": [ + { + "accumulator_bytes": 524288, + "activation_elements_per_cycle": 768, + "blen": 4, + "id": "core0", + "mlen": 768, + "read_cache_bytes": 0, + "vector_sram_bytes": 2097152, + "weight_slots": 4, + "weight_sram_bytes": 49152 + }, + { + "accumulator_bytes": 524288, + "activation_elements_per_cycle": 256, + "blen": 2, + "id": "core1", + "mlen": 512, + "read_cache_bytes": 0, + "vector_sram_bytes": 2097152, + "weight_slots": 4, + "weight_sram_bytes": 16384 + } + ], + "dispatch_cycles": 1, + "dispatch_policy": "work_conserving", + "dispatch_queue_bytes": 16384, + "dispatch_threshold": 8, + "dma": { + "coalesce": true, + "fair_credits": false, + "frontend_sram_bytes": 45056, + "issue_policy": "per_channel", + "lookup_ii_cycles": 2, + "sector_reads": true + }, + "global_dma_credits": 128, + "global_dma_staging_bytes": 8192, + "large_core": 0, + "mac_pipeline_cycles": 16, + "name": "p0353", + "schema_version": 1, + "small_core": 1, + "vector_elements_per_cycle": 512 +} diff --git a/transactional_emulator/testbench/moe_normal_review/measured_latency.csv b/transactional_emulator/testbench/moe_normal_review/measured_latency.csv new file mode 100644 index 000000000..7663b3d6c --- /dev/null +++ b/transactional_emulator/testbench/moe_normal_review/measured_latency.csv @@ -0,0 +1,9 @@ +set,case,D,F,tokens,single_ps,large_small_ps +search,qwen_full_decode_b8,2048,512,8,524575000,564819000 +search,qwen_full_decode_b32,2048,512,32,713878000,792445000 +search,deepseek_full_decode_b8,2048,1408,8,1484505000,1645167000 +search,deepseek_full_decode_b32,2048,1408,32,1504962000,1596841000 +holdout,qwen_holdout_b8,2048,512,8,801578000,869781000 +holdout,qwen_holdout_b32,2048,512,32,1268039000,1377461000 +holdout,deepseek_holdout_b8,2048,1408,8,1483785000,1647373000 +holdout,deepseek_holdout_b32,2048,1408,32,1506083000,1596643000 diff --git a/transactional_emulator/testbench/moe_normal_review/single.json b/transactional_emulator/testbench/moe_normal_review/single.json new file mode 100644 index 000000000..c11c54209 --- /dev/null +++ b/transactional_emulator/testbench/moe_normal_review/single.json @@ -0,0 +1,37 @@ +{ + "clock_period_ps": 1000, + "combine_sram_bytes": 4194304, + "cores": [ + { + "accumulator_bytes": 1048576, + "activation_elements_per_cycle": 1024, + "blen": 8, + "id": "core0", + "mlen": 512, + "read_cache_bytes": 0, + "vector_sram_bytes": 4194304, + "weight_slots": 4, + "weight_sram_bytes": 65536 + } + ], + "dispatch_cycles": 1, + "dispatch_policy": "work_conserving", + "dispatch_queue_bytes": 16384, + "dispatch_threshold": 8, + "dma": { + "coalesce": true, + "fair_credits": false, + "frontend_sram_bytes": 45056, + "issue_policy": "per_channel", + "lookup_ii_cycles": 2, + "sector_reads": true + }, + "global_dma_credits": 128, + "global_dma_staging_bytes": 8192, + "large_core": 0, + "mac_pipeline_cycles": 16, + "name": "p0034", + "schema_version": 1, + "small_core": 0, + "vector_elements_per_cycle": 512 +} diff --git a/transactional_emulator/testbench/moe_normal_review/smoke.py b/transactional_emulator/testbench/moe_normal_review/smoke.py new file mode 100644 index 000000000..f0b01ca1b --- /dev/null +++ b/transactional_emulator/testbench/moe_normal_review/smoke.py @@ -0,0 +1,65 @@ +#!/usr/bin/env python3 +"""Export a small nonzero workload and validate both checked-in architectures. + +This checks the compiler-to-Rust path, not the full-dimension performance claim. +Requires Python 3.10+, NumPy, torch and PLENA_Tools on PYTHONPATH. +""" +import argparse +import importlib.util +import json +from pathlib import Path +import sys + + +def matrix(rows, cols, seed): + return [[(((r * 17 + c * 13 + seed * 7) % 41) - 20) / 128 + for c in range(cols)] for r in range(rows)] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--compiler", type=Path, required=True) + parser.add_argument("--binary", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + args = parser.parse_args() + root = args.output_dir.resolve() + root.mkdir(parents=True, exist_ok=False) + exporter_path = args.compiler.resolve() / "aten/plena/moe_normal_export.py" + spec = importlib.util.spec_from_file_location("moe_normal_export", exporter_path) + exporter = importlib.util.module_from_spec(spec) + spec.loader.exec_module(exporter) + here = Path(__file__).resolve().parent + sys.path.insert(0, str(here.parent / "moe_timing/replay")) + from compare_moe_normal import require, run_comparison + + # Uneven expert groups exercise both core choices; odd D/F exercise tails. + d, f, tokens = 31, 47, 12 + experts = [dict(id=i, gate=matrix(f, d, 3*i+1), + up=matrix(f, d, 3*i+2), down=matrix(d, f, 3*i+3)) + for i in range(3)] + exporter.export_workload( + root / "fixture", name="review_nonzero_tails_and_shared", + inputs=matrix(tokens, d, 91), experts=experts, + routes=[dict(token=t, slot=0, expert=0 if t < 10 else 1, weight=1.0) + for t in range(tokens)], + shared_expert=dict(expert=2, weight=0.25), + provenance={"scope": "synthetic smoke test; not a performance benchmark"}, + ) + summary = run_comparison( + args.binary.resolve(), root / "fixture/workload.json", + root / "fixture/golden.json", + [here / "single.json", here / "large_small.json"], + root / "comparison", repeats=2, atol=0, rtol=0, + ) + dual = summary["comparisons"][1]["result"] + require(all(core["jobs"] > 0 for core in dual["cores"]), + "smoke must exercise both cores") + require(any(value & 0x7FFF for row in dual["output_bf16"] for value in row), + "smoke must produce nonzero output") + print(json.dumps({"status": summary["status"], + "all_gates_passed": summary["all_gates_passed"], + "result": str(root / "comparison/comparison.json")}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py b/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py new file mode 100644 index 000000000..47269d9f1 --- /dev/null +++ b/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py @@ -0,0 +1,385 @@ +#!/usr/bin/env python3 +"""Run an equal-multiplier fixed-route MoE comparison with mandatory gates. + +No cycle estimate or speedup is produced unless every architecture and repeat +passes numerical, identity, resource and deterministic-repeat checks. This is +an operator experiment; it does not measure router execution or full requests. +""" + +import argparse +import hashlib +import json +import math +from pathlib import Path +import subprocess +import struct +import sys +import uuid +from concurrent.futures import ThreadPoolExecutor + + +def read_json(path): + return json.loads(Path(path).read_text(encoding="utf-8")) + + +def digest(path): + value = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + value.update(block) + return value.hexdigest() + + +def require(condition, message): + if not condition: + raise ValueError(message) + + +def numerical_gate(actual, expected, atol, rtol): + require(len(actual) == len(expected), "output token count differs from golden") + maximum_error = 0.0 + for token, (row, reference) in enumerate(zip(actual, expected)): + require(len(row) == len(reference), "output width differs from golden") + for column, (value, target) in enumerate(zip(row, reference)): + require(math.isfinite(value) and math.isfinite(target), "nonfinite output") + error = abs(value - target) + maximum_error = max(maximum_error, error) + require(error <= atol + rtol * abs(target), + "golden mismatch at token {}, column {}: {} != {}".format( + token, column, value, target)) + return maximum_error + + +def validate_output_pair(bits, values): + require(len(bits) == len(values), "BF16/FP32 token count mismatch") + for bit_row, value_row in zip(bits, values): + require(len(bit_row) == len(value_row), "BF16/FP32 width mismatch") + for bit, value in zip(bit_row, value_row): + require(isinstance(bit, int) and not isinstance(bit, bool) and 0 <= bit <= 65535, + "output BF16 must be uint16") + decoded = struct.unpack("= 0 for v in cal[key]), "invalid native channel counter") + native = cal["native_stats"]["memory_system"] + controllers = native["controller"] + require(len(controllers) == channels, "native controller count differs") + for port, controller in enumerate(controllers): + require(controller["id"] == "Channel " + str(port), "native channel identity mismatch") + require(controller["num_read_reqs"] == controller["num_read_reqs_served"] + == cal["accepted_per_channel"][port], "accepted/read-command counts differ") + require(controller["num_read_reqs_forwarded"] == controller["num_write_reqs"] == 0, + "unexpected forwarding or writes in read-only weight experiment") + require(sum(cal["accepted_per_channel"]) == native["total_num_read_requests"] + and native["total_num_read_requests"] * 32 == result["hbm_read_bytes"], "native byte accounting mismatch") + identity = envelope["provenance"] + require(digest(identity["native_library_path"]) == identity["native_library_sha256"], "native library hash differs") + if architecture.get("dma"): + dma = result["dma_frontend"] + require(all(type(v) is int and v >= 0 for v in dma.values()), "invalid DMA metric") + require(dma["reserved_bytes"] <= architecture["dma"]["frontend_sram_bytes"], "DMA metadata SRAM exceeded") + require(dma["mshr_peak"] <= result["global_dma_inflight_peak"], "MSHR entries exceed admitted waiters") + require(dma["sector_requests"] >= dma["merged_sectors"] + and (dma["sector_requests"] - dma["merged_sectors"]) * 32 == result["hbm_read_bytes"], "MSHR sectors do not reconcile") + require(dma["lookup_busy_ps"] == architecture["dma"].get("lookup_ii_cycles", 2) * dma["line_requests"] * architecture["clock_period_ps"], "lookup service unaccounted") + require(dma["copy_busy_ps"] * 32 >= dma["useful_copy_bytes"] * architecture["clock_period_ps"], "copy bandwidth exceeded") + require(dma["lookup_busy_ps"] <= 4 * result["total_ps"] + and dma["copy_busy_ps"] <= 8 * result["total_ps"], "DMA port occupancy exceeds elapsed time") + + +def validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_channels=8): + result = envelope["result"] + require(result.get("numerical_execution", True) is True, "timing-only output is not numerical evidence") + require(envelope.get("evidence_level") != "fixed_route_timing_only_no_numerical_validation", "timing-only report rejected") + require(envelope["memory_model"]["channels"] == hbm_channels + and envelope["memory_model"]["upper_burst_bytes"] == 64 + and envelope["memory_model"]["name"] == "Ramulator HBM2 preset", + "unexpected memory model or bandwidth configuration") + for field in ("total_ps", "multipliers", "useful_macs", "issued_macs", "hbm_read_bytes", + "hbm_write_bytes", "global_dma_inflight_peak", "global_dma_staging_peak_bytes", + "combine_sram_peak_bytes", "shared_vector_busy_ps"): + value = result[field] + require(isinstance(value, int) and not isinstance(value, bool) and value >= 0, + "invalid nonnegative integer metric: " + field) + require(envelope["workload_manifest"] == workload, "workload was changed") + require(envelope["architecture_manifest"] == architecture, "architecture was changed") + expected_multipliers = sum(core["blen"] * core["mlen"] for core in architecture["cores"]) + require(result["multipliers"] == expected_multipliers, "incorrect multiplier count") + token_count = len(workload["inputs_bf16"]) + rows = len(workload["routes"]) + if workload.get("shared_expert") is not None: + rows += token_count + expected_macs = 3 * rows * workload["input_dim"] * workload["expert_hidden_dim"] + require(result["useful_macs"] == expected_macs, "useful MAC count differs from MoE work") + require(result["issued_macs"] >= expected_macs, "issued MAC count is too small") + require(result["global_dma_inflight_peak"] <= architecture["global_dma_credits"], + "global DMA credits exceeded") + require(result["global_dma_staging_peak_bytes"] <= architecture["global_dma_staging_bytes"], + "DMA staging capacity exceeded") + require(result["combine_sram_peak_bytes"] <= architecture["combine_sram_bytes"], + "combine SRAM capacity exceeded") + require(len(result["cores"]) == len(architecture["cores"]), "core count mismatch") + for observed, core in zip(result["cores"], architecture["cores"]): + require(observed["id"] == core["id"], "core identity mismatch") + require(observed["multipliers"] == core["blen"] * core["mlen"], "core multiplier mismatch") + for metric in ("useful_macs", "issued_macs", "jobs", "hbm_read_bytes", "compute_busy_ps", + "vector_sram_peak_bytes", "accumulator_peak_bytes", "weight_sram_peak_bytes", + "weight_slots_peak"): + require(isinstance(observed[metric], int) and not isinstance(observed[metric], bool) + and observed[metric] >= 0, "invalid per-core metric: " + metric) + require(observed["issued_macs"] >= observed["useful_macs"], "per-core issued MAC deficit") + require(observed["compute_busy_ps"] <= result["total_ps"], "core busy time exceeds run") + for peak, capacity in (("vector_sram_peak_bytes", "vector_sram_bytes"), + ("accumulator_peak_bytes", "accumulator_bytes"), + ("weight_sram_peak_bytes", "weight_sram_bytes")): + require(observed[peak] <= core[capacity], core["id"] + " exceeded " + capacity) + require(observed["weight_slots_peak"] <= core.get("weight_slots", 2), "weight slot count exceeded") + if "read_cache_bytes" in core and not architecture.get("dma"): + for metric in ("cache_requests", "cache_hits", "cache_port_busy_ps", "cache_peak_bytes"): + require(type(observed[metric]) is int and observed[metric] >= 0, "invalid cache metric: " + metric) + require(observed["cache_requests"] == observed["cache_hits"] + observed["hbm_read_bytes"] // 64, + "cache requests do not reconcile with real HBM reads") + require(observed["cache_peak_bytes"] <= core["read_cache_bytes"] + and observed["cache_peak_bytes"] % 80 == 0, "cache data/tag capacity exceeded") + expected_cache_cycles = (observed["cache_requests"] + observed["hbm_read_bytes"] // 64) if core["read_cache_bytes"] else 0 + require(observed["cache_port_busy_ps"] == expected_cache_cycles * architecture["clock_period_ps"], + "cache port service is not accounted") + require(observed["cache_port_busy_ps"] <= result["total_ps"], "cache port busy time exceeds run") + for metric in ("compute_busy_fraction", "mac_utilization"): + require(math.isfinite(observed[metric]) and 0 <= observed[metric] <= 1.000000001, + "invalid " + metric) + jobs = result["job_completions"] + if "dispatch_queue_bytes" in architecture: + require(result["dispatch_queue_peak_bytes"] == len(jobs) * 64 + and result["dispatch_queue_peak_bytes"] <= architecture["dispatch_queue_bytes"], + "ready-job descriptor capacity mismatch") + require(result["dispatcher_busy_ps"] == len(jobs) * architecture["dispatch_cycles"] * architecture["clock_period_ps"] + and result["dispatcher_busy_ps"] <= result["total_ps"], "dispatcher service mismatch") + identities = [(job["expert"], job["shared"]) for job in jobs] + expected_rows = {} + for route in workload["routes"]: + identity = (route["expert"], False) + expected_rows[identity] = expected_rows.get(identity, 0) + 1 + if workload.get("shared_expert") is not None and token_count: + expected_rows[(workload["shared_expert"]["expert"], True)] = token_count + require(len(identities) == len(set(identities)) and set(identities) == set(expected_rows), + "missing or duplicate expert completion") + require(sum(core["jobs"] for core in result["cores"]) == len(jobs), "job count mismatch") + for field in ("useful_macs", "issued_macs", "hbm_read_bytes"): + require(sum(core[field] for core in result["cores"]) == result[field], + "per-core totals do not reconcile: " + field) + jobs_per_core = {core["id"]: 0 for core in architecture["cores"]} + rows_per_core = {core["id"]: 0 for core in architecture["cores"]} + for job in jobs: + require(job["core"] in jobs_per_core, "job completion names an unknown core") + require(type(job["rows"]) is int + and job["rows"] == expected_rows[(job["expert"], job["shared"])], + "expert completion row count differs from workload") + jobs_per_core[job["core"]] += 1 + rows_per_core[job["core"]] += job["rows"] + require(all(type(job[field]) is int for field in + ("start_ps", "compute_done_ps", "output_copied_ps")), + "job completion timestamps must be integers") + require(0 <= job["start_ps"] <= job["compute_done_ps"] <= job["output_copied_ps"] + <= result["total_ps"], "invalid job completion timestamps") + for core in result["cores"]: + require(core["jobs"] == jobs_per_core[core["id"]], "per-core job count mismatch") + require(core["useful_macs"] == rows_per_core[core["id"]] * 3 + * workload["input_dim"] * workload["expert_hidden_dim"], + "per-core useful MACs differ from completed jobs") + require(result["hbm_read_bytes"] % (32 if architecture.get("dma") else 64) == 0, "HBM request bytes must match transfer granularity") + require(result["hbm_write_bytes"] == 0, "unexpected HBM output/intermediate write") + require(result["shared_vector_busy_ps"] <= result["total_ps"], "vector busy time exceeds run") + for source in (result, golden): + for field in ("output_bf16", "output_f32"): + validate_output_shape(source[field], token_count, workload["input_dim"]) + validate_output_pair(result["output_bf16"], result["output_f32"]) + validate_output_pair(golden["output_bf16"], golden["output_f32"]) + if architecture.get("dma") or "calibration" in envelope["memory_model"]: + validate_native(envelope, architecture, hbm_channels) + reference = golden["output_f32"] + error = numerical_gate(result["output_f32"], reference, atol, rtol) + return {"passed": True, "max_absolute_error": error, + "output_bit_exact": result["output_bf16"] == golden["output_bf16"]} + + +def _run_comparison(binary, workload_path, golden_path, architecture_paths, output_dir, + repeats=2, hbm_channels=8, atol=1e-5, rtol=0.01, timeout=180, workers=1): + require(type(workers) is int and 1 <= workers <= 8, "workers must be an integer between 1 and 8") + require(type(repeats) is int and repeats >= 2, "at least two integer repeats are required") + require(type(hbm_channels) is int and 1 <= hbm_channels <= 32 + and hbm_channels & (hbm_channels - 1) == 0, + "HBM channels must be a power of two between 1 and 32") + require(type(timeout) in (int, float) and math.isfinite(timeout) and timeout > 0, + "timeout must be finite and positive") + require(len(architecture_paths) >= 2, "provide baseline and at least one candidate") + require(all(type(x) in (int, float) and math.isfinite(x) and x >= 0 for x in (atol, rtol)), + "invalid tolerance") + binary = Path(binary).resolve() + workload_path = Path(workload_path).resolve() + golden_path = Path(golden_path).resolve() + output_dir = Path(output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + # Bind decoded manifests and the numerical reference to the same bytes + # whose hashes will be published, even if another process edits a file. + workload_bytes, golden_bytes = workload_path.read_bytes(), golden_path.read_bytes() + workload, golden = json.loads(workload_bytes), json.loads(golden_bytes) + architecture_bytes = [Path(path).read_bytes() for path in architecture_paths] + architectures = [json.loads(payload) for payload in architecture_bytes] + architecture_hashes = [hashlib.sha256(payload).hexdigest() for payload in architecture_bytes] + golden_hash = hashlib.sha256(golden_bytes).hexdigest() + names = [arch["name"] for arch in architectures] + require(len(names) == len(set(names)), "architecture names must be unique") + total_pes = [sum(c["blen"] * c["mlen"] for c in a["cores"]) for a in architectures] + require(len(set(total_pes)) == 1, "comparison requires equal multiplier totals") + for field in ("clock_period_ps", "mac_pipeline_cycles", "vector_elements_per_cycle", + "global_dma_credits", "global_dma_staging_bytes", "combine_sram_bytes"): + require(len({a[field] for a in architectures}) == 1, + "comparison requires same shared resource: " + field) + for field, default in (("dispatch_cycles", 1), ("dispatch_queue_bytes", 262144), + ("matrix_timing", "pipelined")): + require(len({a.get(field, default) for a in architectures}) == 1, + "comparison requires same shared resource/timing: " + field) + cache_fields = ["read_cache_bytes" in c for a in architectures for c in a["cores"]] + require(not any(cache_fields) or all(cache_fields), "cache budgets must be explicit for every core") + if all(cache_fields): + for field in ("vector_sram_bytes", "accumulator_bytes", "weight_sram_bytes", "read_cache_bytes"): + require(len({sum(c[field] for c in a["cores"]) for a in architectures}) == 1, + "full-shape comparison requires equal total configured SRAM: " + field) + dma_configs = [a.get("dma") for a in architectures] + require(all(d is None for d in dma_configs) or all(d is not None for d in dma_configs), "DMA configuration must be explicit for all architectures") + if all(d is not None for d in dma_configs): + require(len({json.dumps(d, sort_keys=True) for d in dma_configs}) == 1, "comparison requires identical DMA policy and budget") + hbm_path = workload_path.parent / workload["hbm_file"] + expected_hashes = {"workload_sha256": hashlib.sha256(workload_bytes).hexdigest(), + "hbm_sha256": digest(hbm_path), "executable_sha256": digest(binary)} + for field in ("workload_sha256", "hbm_sha256"): + require(golden[field] == expected_hashes[field], "golden identity mismatch: " + field) + run_dir = output_dir / ("run_" + uuid.uuid4().hex) + run_dir.mkdir() + def execute_architecture(index): + path, architecture = architecture_paths[index], architectures[index] + repeat_results, gates, native_repeats, library_hashes = [], [], [], [] + architecture_hash = architecture_hashes[index] + for repeat in range(repeats): + report_path = run_dir / "arch{:02d}_repeat{:02d}.json".format(index, repeat) + log_path = report_path.with_suffix(".log") + command = [str(binary), "--workload", str(workload_path), "--architecture", + str(Path(path).resolve()), "--output", str(report_path), + "--hbm-channels", str(hbm_channels)] + with log_path.open("wb") as log: + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, + check=True, timeout=timeout) + envelope = read_json(report_path) + for field, value in dict(expected_hashes, architecture_sha256=architecture_hash).items(): + require(envelope["provenance"][field] == value, "input/binary identity changed: " + field) + gates.append(validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_channels)) + repeat_results.append(envelope["result"]) + native_repeats.append(envelope["memory_model"].get("calibration")) + library_hashes.append(envelope["provenance"].get("native_library_sha256")) + require(all(result == repeat_results[0] for result in repeat_results[1:]), + "repeat mismatch for " + architecture["name"]) + require(all(n == native_repeats[0] for n in native_repeats), "native counters changed between repeats") + require(len(set(library_hashes)) == 1, "native library changed between repeats") + result = repeat_results[0] + require(result["total_ps"] > 0, "nonempty comparison must advance time") + return {"architecture": architecture, "result": result, "gates": gates, "native": native_repeats[0], "native_library_sha256": library_hashes[0]} + # Each process owns an independent Ramulator/executor. Preserve manifest + # order and propagate every exception before publishing any speedup. + with ThreadPoolExecutor(max_workers=workers) as pool: + results = list(pool.map(execute_architecture, range(len(architectures)))) + for path, expected, label in ( + (workload_path, expected_hashes["workload_sha256"], "workload"), + (golden_path, golden_hash, "golden"), + (hbm_path, expected_hashes["hbm_sha256"], "HBM"), + (binary, expected_hashes["executable_sha256"], "executable")): + require(digest(path) == expected, label + " changed during comparison") + for path, expected in zip(architecture_paths, architecture_hashes): + require(digest(path) == expected, "architecture changed during comparison") + require(len({r["native_library_sha256"] for r in results}) == 1, "architectures used different native libraries") + baseline_ps = results[0]["result"]["total_ps"] + for row in results: + row["speedup_vs_baseline"] = baseline_ps / row["result"]["total_ps"] + summary = { + "schema_version": 1, "all_gates_passed": True, + "status": "passed", "run_directory": str(run_dir), + "scope": "Fixed-route numerical MoE operator; analytical core timing and shared Ramulator. " + "Not RTL-calibrated or whole-model/request latency.", + "source_metadata": workload.get("metadata", {}), + "workload_sha256": expected_hashes["workload_sha256"], + "golden_sha256": golden_hash, + "hbm_sha256": expected_hashes["hbm_sha256"], + "executable_sha256": expected_hashes["executable_sha256"], + "repeats": repeats, "hbm_channels": hbm_channels, + "numerical_tolerance": {"atol": atol, "rtol": rtol}, "comparisons": results, + } + (output_dir / "comparison.json").write_text( + json.dumps(summary, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8") + return summary + + +def run_comparison(binary, workload_path, golden_path, architecture_paths, output_dir, + repeats=2, hbm_channels=8, atol=1e-5, rtol=0.01, timeout=180, workers=1): + output_dir = Path(output_dir).resolve() + output_dir.mkdir(parents=True, exist_ok=True) + summary_path = output_dir / "comparison.json" + summary_path.write_text(json.dumps({"schema_version": 1, "status": "running", + "all_gates_passed": False}) + "\n", encoding="utf-8") + try: + return _run_comparison(binary, workload_path, golden_path, architecture_paths, + output_dir, repeats, hbm_channels, atol, rtol, timeout, workers) + except Exception as error: + summary_path.write_text(json.dumps({"schema_version": 1, "status": "failed", + "all_gates_passed": False, "error": str(error)}) + "\n", + encoding="utf-8") + raise + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--binary", required=True) + parser.add_argument("--workload", required=True) + parser.add_argument("--golden", required=True) + parser.add_argument("--architecture", action="append", required=True, + help="Baseline first, then candidate(s)") + parser.add_argument("--output-dir", required=True) + parser.add_argument("--repeats", type=int, default=2) + parser.add_argument("--hbm-channels", type=int, default=8) + parser.add_argument("--atol", type=float, default=1e-5) + parser.add_argument("--rtol", type=float, default=0.01) + parser.add_argument("--timeout", type=float, default=180) + parser.add_argument("--workers", type=int, default=1) + args = parser.parse_args() + try: + summary = run_comparison(args.binary, args.workload, args.golden, args.architecture, + args.output_dir, args.repeats, args.hbm_channels, + args.atol, args.rtol, args.timeout, args.workers) + except (ValueError, OSError, KeyError, subprocess.SubprocessError) as error: + parser.error(str(error)) + for row in summary["comparisons"]: + print("{}: {} ps, {:.4f}x, gates passed".format( + row["architecture"]["name"], row["result"]["total_ps"], row["speedup_vs_baseline"])) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py b/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py new file mode 100644 index 000000000..b0af606e1 --- /dev/null +++ b/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py @@ -0,0 +1,432 @@ +"""Evidence gate tests through a real fallible executable; Python 3.6+ stdlib.""" +import copy +import importlib.util +import json +from pathlib import Path +import sys +import tempfile +import unittest + +SPEC = importlib.util.spec_from_file_location( + "compare_moe_normal", str(Path(__file__).with_name("compare_moe_normal.py"))) +compare = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(compare) + +FAKE_RUNNER = r''' +import argparse +import fcntl +import hashlib +import json +from pathlib import Path +import sys +parser = argparse.ArgumentParser() +for name in ("workload", "architecture", "output", "hbm-channels"): + parser.add_argument("--" + name, required=True) +args = parser.parse_args() +root = Path(__file__).parent +scenario = json.loads((root / "scenario.json").read_text()) +workload = json.loads(Path(args.workload).read_text()) +architecture = json.loads(Path(args.architecture).read_text()) +name = architecture["name"] +state_path = root / "calls.json" +state_lock = (root / "calls.lock").open('w') +fcntl.flock(state_lock, fcntl.LOCK_EX) +state = json.loads(state_path.read_text()) if state_path.exists() else {} +repeat = state.get(name, 0) +state[name] = repeat + 1 +state_path.write_text(json.dumps(state)) +state_lock.close() +if scenario.get("exit_code"): + sys.exit(scenario["exit_code"]) +if scenario.get("skip_write"): + sys.exit(0) +if scenario.get("rewrite_golden"): + golden_path = root / "golden.json" + golden = json.loads(golden_path.read_text()) + golden["output_f32"] = [[7.0] * 8, [7.0] * 8] + golden["output_bf16"] = [[16608] * 8, [16608] * 8] + golden_path.write_text(json.dumps(golden)) +envelope = json.loads((root / "reports.json").read_text())[name] +def sha(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() +envelope["workload_manifest"] = workload +envelope["architecture_manifest"] = architecture +envelope["provenance"] = { + "workload_sha256": sha(args.workload), "architecture_sha256": sha(args.architecture), + "hbm_sha256": sha(Path(args.workload).parent / workload["hbm_file"]), + "executable_sha256": sha(__file__), +} +envelope["memory_model"]["channels"] = int(args.hbm_channels) +for mutation in scenario.get("mutations", []): + if mutation.get("architecture", name) != name or mutation.get("repeat", repeat) != repeat: + continue + target = envelope + for key in mutation["path"][:-1]: + target = target[key] + target[mutation["path"][-1]] = mutation["value"] +Path(args.output).write_text(json.dumps(envelope)) +''' + + +def core_config(name, blen, mlen): + return dict(id=name, blen=blen, mlen=mlen, vector_sram_bytes=4096, + accumulator_bytes=1024, weight_sram_bytes=2048) + + +def architecture(name, cores): + return dict(schema_version=1, name=name, cores=cores, dispatch_threshold=2, + large_core=0, small_core=len(cores) - 1, global_dma_credits=2, + global_dma_staging_bytes=128, combine_sram_bytes=4096, + clock_period_ps=10, mac_pipeline_cycles=8, vector_elements_per_cycle=8) + + +def core_report(core, jobs, rows, total_ps, busy_ps): + useful = rows * 3 * 8 * 8 + multipliers = core["blen"] * core["mlen"] + return dict(id=core["id"], blen=core["blen"], mlen=core["mlen"], + multipliers=multipliers, jobs=jobs, useful_macs=useful, + issued_macs=useful, compute_busy_ps=busy_ps, + accumulator_dependency_stall_ps=0, pipeline_drain_ps=10, + pipeline_register_bytes=64, weight_ready_wait_ps=0, + vector_wait_ps=0, vector_sram_peak_bytes=256, + accumulator_peak_bytes=128, weight_sram_peak_bytes=256, + weight_slots_peak=2, hbm_read_bytes=192 * jobs, + compute_busy_fraction=busy_ps / total_ps, + mac_utilization=useful / (multipliers * total_ps / 10)) + + +def envelope(arch): + single = len(arch["cores"]) == 1 + total = 1000 if single else 800 + jobs = [ + dict(job=0, expert=0, shared=False, core=arch["cores"][0]["id"], rows=2, + start_ps=0, compute_done_ps=200, output_copied_ps=220), + dict(job=1, expert=2, shared=True, core=arch["cores"][0]["id"], rows=2, + start_ps=220, compute_done_ps=420, output_copied_ps=440), + dict(job=2, expert=1, shared=False, core=arch["cores"][-1]["id"], rows=1, + start_ps=440 if single else 0, compute_done_ps=600 if single else 180, + output_copied_ps=620 if single else 200), + ] + cores = ([core_report(arch["cores"][0], 3, 5, total, 600)] if single else + [core_report(arch["cores"][0], 2, 4, total, 600), + core_report(arch["cores"][1], 1, 1, total, 200)]) + result = dict(schema_version=1, workload="fixture", architecture=arch["name"], + timing_model="analytical", timing_boundary="operator", + weight_format="local_mx", total_ps=total, multipliers=32, + useful_macs=960, issued_macs=960, hbm_read_bytes=576, + hbm_write_bytes=0, global_dma_inflight_peak=2, + global_dma_staging_peak_bytes=128, combine_sram_peak_bytes=256, + shared_vector_busy_ps=100, cores=cores, + job_completions=jobs if single else [jobs[2], jobs[0], jobs[1]], + output_f32=[[1.0] * 8, [-1.0] * 8], + output_bf16=[[16256] * 8, [49024] * 8], + pre_round_output_f32=[[1.0] * 8, [-1.0] * 8]) + return dict(schema_version=1, evidence_level="fixture_only", provenance={}, + memory_model=dict(name="Ramulator HBM2 preset", channels=8, + upper_burst_bytes=64), result=result) + + +class ComparisonEvidenceTests(unittest.TestCase): + def setUp(self): + self.temporary = tempfile.TemporaryDirectory(prefix="plena-moe-gates-", dir="/tmp") + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.binary = self.root / "fake_moe_runner" + self.binary.write_text("#!" + sys.executable + "\n" + FAKE_RUNNER) + self.binary.chmod(0o700) + self.hbm = self.root / "weights.bin" + self.hbm.write_bytes(bytes(1152)) + self.workload = dict(schema_version=1, name="fixture", hbm_file="weights.bin", + input_dim=8, expert_hidden_dim=8, + inputs_bf16=[[16256] * 8, [16256] * 8], + routes=[dict(token=0, slot=0, expert=0, weight=0.5), + dict(token=0, slot=1, expert=1, weight=0.5), + dict(token=1, slot=0, expert=0, weight=1.0)], + shared_expert=dict(expert=2, weight=1.0), metadata={}) + self.workload["experts"] = [] + for expert in range(3): + record = dict(id=expert) + for projection, name in enumerate(("gate", "up", "down")): + base = (expert * 3 + projection) * 128 + record[name] = dict(rows=8, cols=8, element_base=base, + scale_base=base + 64, element_row_stride=8, + scale_row_stride=1) + self.workload["experts"].append(record) + self.workload_path = self.root / "workload.json" + self.write(self.workload_path, self.workload) + self.architectures = [architecture("single", [core_config("single", 4, 8)]), + architecture("dual", [core_config("large", 2, 8), + core_config("small", 1, 16)])] + self.arch_paths = [self.root / "single.json", self.root / "dual.json"] + for path, arch in zip(self.arch_paths, self.architectures): + self.write(path, arch) + self.reports = {arch["name"]: envelope(arch) for arch in self.architectures} + self.write(self.root / "reports.json", self.reports) + self.golden = dict(schema_version=1, + workload_sha256=compare.digest(self.workload_path), + hbm_sha256=compare.digest(self.hbm), + output_f32=[[1.0] * 8, [-1.0] * 8], + output_bf16=[[16256] * 8, [49024] * 8]) + self.golden_path = self.root / "golden.json" + self.write(self.golden_path, self.golden) + self.scenario = {} + self.write(self.root / "scenario.json", self.scenario) + self.output = self.root / "comparison" + + @staticmethod + def write(path, data): + path.write_text(json.dumps(data), encoding="utf-8") + + def run_comparison(self, **kwargs): + self.write(self.root / "scenario.json", self.scenario) + return compare.run_comparison(str(self.binary), str(self.workload_path), + str(self.golden_path), self.arch_paths, + self.output, **kwargs) + + def mutate(self, path, value, **filters): + self.scenario.setdefault("mutations", []).append(dict(path=path, value=value, **filters)) + + def rejected(self, **kwargs): + with self.assertRaises((ValueError, KeyError)): + self.run_comparison(**kwargs) + summary = self.output / "comparison.json" + if summary.exists(): + self.assertIsNot(compare.read_json(summary).get("all_gates_passed"), True) + + def test_complete_deterministic_reports_publish_speedup_and_hashes(self): + summary = self.run_comparison() + self.assertTrue(summary["all_gates_passed"]) + self.assertEqual(1.25, summary["comparisons"][1]["speedup_vs_baseline"]) + self.assertEqual(compare.digest(self.golden_path), summary["golden_sha256"]) + self.assertEqual(compare.digest(self.hbm), summary["hbm_sha256"]) + self.assertEqual({"single": 2, "dual": 2}, compare.read_json(self.root / "calls.json")) + + def test_wrong_numeric_result_is_rejected(self): + self.mutate(["result", "output_f32", 0, 0], 7.0) + self.rejected() + + def test_parallel_architectures_keep_order_and_repeat_gates(self): + result = self.run_comparison(workers=2) + self.assertEqual([r['architecture']['name'] for r in result['comparisons']], ['single','dual']) + self.assertEqual(compare.read_json(self.root/'calls.json'),{'single':2,'dual':2}) + self.mutate(['result','total_ps'],801,architecture='dual',repeat=3) + self.rejected(workers=2) + + def test_invalid_parallelism_is_rejected(self): + for workers in [0,9,True,1.5]: + with self.subTest(workers=workers): self.rejected(workers=workers) + + def enable_full_shape_resources(self): + for path, arch in zip(self.arch_paths, self.architectures): + arch.update(dispatch_queue_bytes=256, dispatch_cycles=1) + parts = len(arch['cores']) + for core in arch['cores']: + core.update(read_cache_bytes=160 // parts, vector_sram_bytes=4096 // parts, + accumulator_bytes=1024 // parts, weight_sram_bytes=2048 // parts) + report = self.reports[arch['name']]['result'] + report.update(dispatch_queue_peak_bytes=192, dispatcher_busy_ps=30) + for core in report['cores']: + misses = core['hbm_read_bytes'] // 64 + core.update(cache_requests=misses+2, cache_hits=2, + cache_port_busy_ps=(2*misses+2)*10, cache_peak_bytes=80) + self.write(path, arch) + self.write(self.root / 'reports.json', self.reports) + + def test_full_shape_resource_accounting_passes(self): + self.enable_full_shape_resources() + self.assertTrue(self.run_comparison()['all_gates_passed']) + + def test_cache_capacity_traffic_and_port_undercharging_are_rejected(self): + self.enable_full_shape_resources() + for metric,value in [('cache_peak_bytes',240),('cache_requests',0), + ('cache_port_busy_ps',0),('cache_hits',-1)]: + with self.subTest(metric=metric): + self.scenario = {} + self.mutate(['result','cores',0,metric],value) + self.rejected() + + def test_dispatch_capacity_and_service_are_required(self): + self.enable_full_shape_resources() + for metric,value in [('dispatch_queue_peak_bytes',0),('dispatcher_busy_ps',0)]: + with self.subTest(metric=metric): + self.scenario = {} + self.mutate(['result',metric],value) + self.rejected() + + def test_unequal_or_partly_missing_cache_budgets_fail_before_execution(self): + self.enable_full_shape_resources() + self.architectures[0]['cores'][0]['read_cache_bytes'] += 80 + self.write(self.arch_paths[0],self.architectures[0]) + self.rejected() + self.assertFalse((self.root/'calls.json').exists()) + del self.architectures[0]['cores'][0]['read_cache_bytes'] + self.write(self.arch_paths[0],self.architectures[0]) + self.rejected() + self.assertFalse((self.root/'calls.json').exists()) + + def test_unequal_timing_modes_fail_before_execution(self): + self.architectures[0]['matrix_timing'] = 'legacy_serialized' + self.write(self.arch_paths[0],self.architectures[0]) + self.rejected() + self.assertFalse((self.root/'calls.json').exists()) + + def test_nonfinite_output_is_rejected(self): + for value in (float("nan"), float("inf"), -float("inf")): + with self.subTest(value=value): + self.scenario = {} + self.mutate(["result", "output_f32", 0, 0], value) + self.rejected() + + def test_bf16_output_cannot_disagree_with_decoded_float_output(self): + self.mutate(["result", "output_bf16", 0, 0], 0) + self.rejected() + + def test_output_shape_must_match_workload_even_if_golden_is_empty(self): + self.golden.update(output_f32=[], output_bf16=[]) + self.write(self.golden_path, self.golden) + self.mutate(["result", "output_f32"], []) + self.mutate(["result", "output_bf16"], []) + self.rejected() + + def test_every_finite_storage_bound_is_enforced(self): + cases = [(["global_dma_inflight_peak"], 3), + (["global_dma_staging_peak_bytes"], 129), + (["combine_sram_peak_bytes"], 4097), + (["cores", 0, "vector_sram_peak_bytes"], 4097), + (["cores", 0, "accumulator_peak_bytes"], 1025), + (["cores", 0, "weight_sram_peak_bytes"], 2049), + (["cores", 0, "weight_slots_peak"], 3)] + for path, value in cases: + with self.subTest(metric=path): + self.scenario = {} + self.mutate(["result"] + path, value) + self.rejected() + + def test_negative_resource_counter_is_rejected(self): + self.mutate(["result", "global_dma_staging_peak_bytes"], -64) + self.rejected() + + def test_useful_and_issued_macs_cannot_omit_work(self): + for metric in ("useful_macs", "issued_macs"): + with self.subTest(metric=metric): + self.scenario = {} + self.mutate(["result", metric], 959) + self.rejected() + + def test_per_core_macs_must_close_to_total(self): + for metric in ("useful_macs", "issued_macs"): + with self.subTest(metric=metric): + self.scenario = {} + self.mutate(["result", "cores", 0, metric], 959) + self.rejected() + + def test_per_core_hbm_bytes_must_close_to_total(self): + self.mutate(["result", "cores", 0, "hbm_read_bytes"], 512) + self.rejected() + + def test_missing_or_duplicate_expert_completions_are_rejected(self): + original = self.reports["single"]["result"]["job_completions"] + for jobs in (original[:-1], original + [copy.deepcopy(original[0])]): + with self.subTest(count=len(jobs)): + self.scenario = {} + self.mutate(["result", "job_completions"], jobs, architecture="single") + self.rejected() + + def test_completed_expert_must_report_all_routed_rows(self): + self.mutate(["result", "job_completions", 0, "rows"], 1, architecture="single") + self.rejected() + + def test_completed_job_must_name_a_real_core(self): + self.mutate(["result", "job_completions", 0, "core"], "missing") + self.rejected() + + def test_completion_after_total_time_is_rejected(self): + self.mutate(["result", "job_completions", 0, "output_copied_ps"], 1001) + self.rejected() + + def test_nonfinite_or_negative_timings_are_rejected(self): + for metric, value in (("total_ps", float("inf")), ("total_ps", -1), + ("shared_vector_busy_ps", float("nan")), + ("shared_vector_busy_ps", -1)): + with self.subTest(metric=metric, value=value): + self.scenario = {} + self.mutate(["result", metric], value) + self.rejected() + + def test_actual_memory_configuration_must_match_requested_hbm(self): + for field, value in (("channels", 16), ("upper_burst_bytes", 32)): + with self.subTest(field=field): + self.scenario = {} + self.mutate(["memory_model", field], value) + self.rejected() + + def test_input_or_executable_hash_mismatch_is_rejected(self): + for field in ("workload_sha256", "architecture_sha256", "hbm_sha256", "executable_sha256"): + with self.subTest(field=field): + self.scenario = {} + self.mutate(["provenance", field], "0" * 64) + self.rejected() + + def test_golden_must_be_bound_to_exact_workload_and_hbm(self): + for field in ("workload_sha256", "hbm_sha256"): + with self.subTest(field=field): + bad = dict(self.golden) + bad[field] = "0" * 64 + self.write(self.golden_path, bad) + self.rejected() + self.write(self.golden_path, self.golden) + + def test_golden_without_identity_is_rejected(self): + del self.golden["workload_sha256"] + self.write(self.golden_path, self.golden) + self.rejected() + + def test_golden_changed_after_loading_cannot_be_published_as_reference(self): + self.scenario = dict(rewrite_golden=True) + self.rejected() + + def test_unequal_multiplier_architectures_are_rejected_before_execution(self): + self.architectures[1]["cores"][0]["mlen"] = 16 + self.write(self.arch_paths[1], self.architectures[1]) + self.rejected() + self.assertFalse((self.root / "calls.json").exists()) + + def test_different_shared_resources_are_rejected_before_execution(self): + original = copy.deepcopy(self.architectures[1]) + for field in ("clock_period_ps", "mac_pipeline_cycles", "vector_elements_per_cycle", + "global_dma_credits", "global_dma_staging_bytes", "combine_sram_bytes"): + with self.subTest(field=field): + altered = copy.deepcopy(original) + altered[field] += 1 + self.write(self.arch_paths[1], altered) + self.rejected() + self.assertFalse((self.root / "calls.json").exists()) + + def test_repeat_difference_prevents_speedup_publication(self): + self.mutate(["result", "total_ps"], 1001, architecture="single", repeat=1) + self.rejected() + + def test_failed_rerun_cannot_leave_previous_success_summary(self): + self.run_comparison() + self.mutate(["result", "output_f32", 0, 0], 7.0) + self.rejected() + + def test_success_exit_without_new_report_cannot_reuse_previous_reports(self): + self.run_comparison() + self.scenario = dict(skip_write=True) + with self.assertRaises((ValueError, OSError, KeyError)): + self.run_comparison() + summary = self.output / "comparison.json" + if summary.exists(): + self.assertIsNot(compare.read_json(summary).get("all_gates_passed"), True) + + def test_invalid_run_controls_are_rejected(self): + for kwargs in (dict(repeats=1), dict(repeats=2.5), dict(atol=-1.0), + dict(rtol=float("nan")), dict(hbm_channels=0), + dict(hbm_channels=3), dict(timeout=0), dict(timeout=float("inf"))): + with self.subTest(kwargs=kwargs): + self.rejected(**kwargs) + + +if __name__ == "__main__": + unittest.main() From cc0d26181f6f20d00333d1368aeaf87d00ddbfa3 Mon Sep 17 00:00:00 2001 From: Michael C Li <63594253+Happymic@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:57:03 +0100 Subject: [PATCH 2/5] ci: wire numerical MoE checks and format review tooling --- .github/workflows/transactional_emulator.yml | 9 + PLENA_Compiler | 2 +- .../testbench/moe_normal_review/smoke.py | 47 +- .../moe_timing/replay/compare_moe_normal.py | 463 ++++++++++++------ .../replay/test_compare_moe_normal.py | 328 ++++++++----- 5 files changed, 576 insertions(+), 273 deletions(-) diff --git a/.github/workflows/transactional_emulator.yml b/.github/workflows/transactional_emulator.yml index e63788de2..ce7abb1f2 100644 --- a/.github/workflows/transactional_emulator.yml +++ b/.github/workflows/transactional_emulator.yml @@ -86,6 +86,7 @@ jobs: python3 -m pytest \ transactional_emulator/testbench/test_testbench_stage_attribution.py \ transactional_emulator/testbench/moe_timing/replay/test_branch_split.py \ + transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py \ transactional_emulator/testbench/moe_timing/campaign/test_run_subset_gates.py \ -v @@ -128,6 +129,14 @@ jobs: uv python install 3.12 uv sync + - name: Run Compiler to Rust MoE normal-core smoke + run: | + nix develop --command bash -c 'PYTHONPATH="$PWD/PLENA_Tools" uv run python \ + transactional_emulator/testbench/moe_normal_review/smoke.py \ + --compiler "$PWD/PLENA_Compiler" \ + --binary "$PWD/transactional_emulator/target/release/moe_dual_normal" \ + --output-dir /tmp/plena-moe-review-smoke' + - name: Run ATen linear integration test run: | nix develop --command bash -c "just test-aten-linear" diff --git a/PLENA_Compiler b/PLENA_Compiler index 510ed3c4d..5c968f897 160000 --- a/PLENA_Compiler +++ b/PLENA_Compiler @@ -1 +1 @@ -Subproject commit 510ed3c4deb91494db8e43e6e6cc48defde00bab +Subproject commit 5c968f89749ebcc36378956721078c4b1fc46e96 diff --git a/transactional_emulator/testbench/moe_normal_review/smoke.py b/transactional_emulator/testbench/moe_normal_review/smoke.py index f0b01ca1b..c9306b3ce 100644 --- a/transactional_emulator/testbench/moe_normal_review/smoke.py +++ b/transactional_emulator/testbench/moe_normal_review/smoke.py @@ -4,6 +4,7 @@ This checks the compiler-to-Rust path, not the full-dimension performance claim. Requires Python 3.10+, NumPy, torch and PLENA_Tools on PYTHONPATH. """ + import argparse import importlib.util import json @@ -12,8 +13,7 @@ def matrix(rows, cols, seed): - return [[(((r * 17 + c * 13 + seed * 7) % 41) - 20) / 128 - for c in range(cols)] for r in range(rows)] + return [[(((r * 17 + c * 13 + seed * 7) % 41) - 20) / 128 for c in range(cols)] for r in range(rows)] def main(): @@ -34,31 +34,42 @@ def main(): # Uneven expert groups exercise both core choices; odd D/F exercise tails. d, f, tokens = 31, 47, 12 - experts = [dict(id=i, gate=matrix(f, d, 3*i+1), - up=matrix(f, d, 3*i+2), down=matrix(d, f, 3*i+3)) - for i in range(3)] + experts = [ + dict(id=i, gate=matrix(f, d, 3 * i + 1), up=matrix(f, d, 3 * i + 2), down=matrix(d, f, 3 * i + 3)) + for i in range(3) + ] exporter.export_workload( - root / "fixture", name="review_nonzero_tails_and_shared", - inputs=matrix(tokens, d, 91), experts=experts, - routes=[dict(token=t, slot=0, expert=0 if t < 10 else 1, weight=1.0) - for t in range(tokens)], + root / "fixture", + name="review_nonzero_tails_and_shared", + inputs=matrix(tokens, d, 91), + experts=experts, + routes=[dict(token=t, slot=0, expert=0 if t < 10 else 1, weight=1.0) for t in range(tokens)], shared_expert=dict(expert=2, weight=0.25), provenance={"scope": "synthetic smoke test; not a performance benchmark"}, ) summary = run_comparison( - args.binary.resolve(), root / "fixture/workload.json", + args.binary.resolve(), + root / "fixture/workload.json", root / "fixture/golden.json", [here / "single.json", here / "large_small.json"], - root / "comparison", repeats=2, atol=0, rtol=0, + root / "comparison", + repeats=2, + atol=0, + rtol=0, ) dual = summary["comparisons"][1]["result"] - require(all(core["jobs"] > 0 for core in dual["cores"]), - "smoke must exercise both cores") - require(any(value & 0x7FFF for row in dual["output_bf16"] for value in row), - "smoke must produce nonzero output") - print(json.dumps({"status": summary["status"], - "all_gates_passed": summary["all_gates_passed"], - "result": str(root / "comparison/comparison.json")}, indent=2)) + require(all(core["jobs"] > 0 for core in dual["cores"]), "smoke must exercise both cores") + require(any(value & 0x7FFF for row in dual["output_bf16"] for value in row), "smoke must produce nonzero output") + print( + json.dumps( + { + "status": summary["status"], + "all_gates_passed": summary["all_gates_passed"], + "result": str(root / "comparison/comparison.json"), + }, + indent=2, + ) + ) if __name__ == "__main__": diff --git a/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py b/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py index 47269d9f1..932ca200a 100644 --- a/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py +++ b/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py @@ -44,9 +44,10 @@ def numerical_gate(actual, expected, atol, rtol): require(math.isfinite(value) and math.isfinite(target), "nonfinite output") error = abs(value - target) maximum_error = max(maximum_error, error) - require(error <= atol + rtol * abs(target), - "golden mismatch at token {}, column {}: {} != {}".format( - token, column, value, target)) + require( + error <= atol + rtol * abs(target), + f"golden mismatch at token {token}, column {column}: {value} != {target}", + ) return maximum_error @@ -55,42 +56,61 @@ def validate_output_pair(bits, values): for bit_row, value_row in zip(bits, values): require(len(bit_row) == len(value_row), "BF16/FP32 width mismatch") for bit, value in zip(bit_row, value_row): - require(isinstance(bit, int) and not isinstance(bit, bool) and 0 <= bit <= 65535, - "output BF16 must be uint16") + require( + isinstance(bit, int) and not isinstance(bit, bool) and 0 <= bit <= 65535, "output BF16 must be uint16" + ) decoded = struct.unpack("= 0 for v in cal[key]), "invalid native channel counter") + require( + len(cal[key]) == channels and all(type(v) is int and v >= 0 for v in cal[key]), + "invalid native channel counter", + ) native = cal["native_stats"]["memory_system"] controllers = native["controller"] require(len(controllers) == channels, "native controller count differs") for port, controller in enumerate(controllers): require(controller["id"] == "Channel " + str(port), "native channel identity mismatch") - require(controller["num_read_reqs"] == controller["num_read_reqs_served"] - == cal["accepted_per_channel"][port], "accepted/read-command counts differ") - require(controller["num_read_reqs_forwarded"] == controller["num_write_reqs"] == 0, - "unexpected forwarding or writes in read-only weight experiment") - require(sum(cal["accepted_per_channel"]) == native["total_num_read_requests"] - and native["total_num_read_requests"] * 32 == result["hbm_read_bytes"], "native byte accounting mismatch") + require( + controller["num_read_reqs"] == controller["num_read_reqs_served"] == cal["accepted_per_channel"][port], + "accepted/read-command counts differ", + ) + require( + controller["num_read_reqs_forwarded"] == controller["num_write_reqs"] == 0, + "unexpected forwarding or writes in read-only weight experiment", + ) + require( + sum(cal["accepted_per_channel"]) == native["total_num_read_requests"] + and native["total_num_read_requests"] * 32 == result["hbm_read_bytes"], + "native byte accounting mismatch", + ) identity = envelope["provenance"] require(digest(identity["native_library_path"]) == identity["native_library_sha256"], "native library hash differs") if architecture.get("dma"): @@ -98,28 +118,56 @@ def validate_native(envelope, architecture, channels): require(all(type(v) is int and v >= 0 for v in dma.values()), "invalid DMA metric") require(dma["reserved_bytes"] <= architecture["dma"]["frontend_sram_bytes"], "DMA metadata SRAM exceeded") require(dma["mshr_peak"] <= result["global_dma_inflight_peak"], "MSHR entries exceed admitted waiters") - require(dma["sector_requests"] >= dma["merged_sectors"] - and (dma["sector_requests"] - dma["merged_sectors"]) * 32 == result["hbm_read_bytes"], "MSHR sectors do not reconcile") - require(dma["lookup_busy_ps"] == architecture["dma"].get("lookup_ii_cycles", 2) * dma["line_requests"] * architecture["clock_period_ps"], "lookup service unaccounted") - require(dma["copy_busy_ps"] * 32 >= dma["useful_copy_bytes"] * architecture["clock_period_ps"], "copy bandwidth exceeded") - require(dma["lookup_busy_ps"] <= 4 * result["total_ps"] - and dma["copy_busy_ps"] <= 8 * result["total_ps"], "DMA port occupancy exceeds elapsed time") + require( + dma["sector_requests"] >= dma["merged_sectors"] + and (dma["sector_requests"] - dma["merged_sectors"]) * 32 == result["hbm_read_bytes"], + "MSHR sectors do not reconcile", + ) + require( + dma["lookup_busy_ps"] + == architecture["dma"].get("lookup_ii_cycles", 2) * dma["line_requests"] * architecture["clock_period_ps"], + "lookup service unaccounted", + ) + require( + dma["copy_busy_ps"] * 32 >= dma["useful_copy_bytes"] * architecture["clock_period_ps"], + "copy bandwidth exceeded", + ) + require( + dma["lookup_busy_ps"] <= 4 * result["total_ps"] and dma["copy_busy_ps"] <= 8 * result["total_ps"], + "DMA port occupancy exceeds elapsed time", + ) def validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_channels=8): result = envelope["result"] require(result.get("numerical_execution", True) is True, "timing-only output is not numerical evidence") - require(envelope.get("evidence_level") != "fixed_route_timing_only_no_numerical_validation", "timing-only report rejected") - require(envelope["memory_model"]["channels"] == hbm_channels - and envelope["memory_model"]["upper_burst_bytes"] == 64 - and envelope["memory_model"]["name"] == "Ramulator HBM2 preset", - "unexpected memory model or bandwidth configuration") - for field in ("total_ps", "multipliers", "useful_macs", "issued_macs", "hbm_read_bytes", - "hbm_write_bytes", "global_dma_inflight_peak", "global_dma_staging_peak_bytes", - "combine_sram_peak_bytes", "shared_vector_busy_ps"): + require( + envelope.get("evidence_level") != "fixed_route_timing_only_no_numerical_validation", + "timing-only report rejected", + ) + require( + envelope["memory_model"]["channels"] == hbm_channels + and envelope["memory_model"]["upper_burst_bytes"] == 64 + and envelope["memory_model"]["name"] == "Ramulator HBM2 preset", + "unexpected memory model or bandwidth configuration", + ) + for field in ( + "total_ps", + "multipliers", + "useful_macs", + "issued_macs", + "hbm_read_bytes", + "hbm_write_bytes", + "global_dma_inflight_peak", + "global_dma_staging_peak_bytes", + "combine_sram_peak_bytes", + "shared_vector_busy_ps", + ): value = result[field] - require(isinstance(value, int) and not isinstance(value, bool) and value >= 0, - "invalid nonnegative integer metric: " + field) + require( + isinstance(value, int) and not isinstance(value, bool) and value >= 0, + "invalid nonnegative integer metric: " + field, + ) require(envelope["workload_manifest"] == workload, "workload was changed") require(envelope["architecture_manifest"] == architecture, "architecture was changed") expected_multipliers = sum(core["blen"] * core["mlen"] for core in architecture["cores"]) @@ -131,49 +179,74 @@ def validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_chann expected_macs = 3 * rows * workload["input_dim"] * workload["expert_hidden_dim"] require(result["useful_macs"] == expected_macs, "useful MAC count differs from MoE work") require(result["issued_macs"] >= expected_macs, "issued MAC count is too small") - require(result["global_dma_inflight_peak"] <= architecture["global_dma_credits"], - "global DMA credits exceeded") - require(result["global_dma_staging_peak_bytes"] <= architecture["global_dma_staging_bytes"], - "DMA staging capacity exceeded") - require(result["combine_sram_peak_bytes"] <= architecture["combine_sram_bytes"], - "combine SRAM capacity exceeded") + require(result["global_dma_inflight_peak"] <= architecture["global_dma_credits"], "global DMA credits exceeded") + require( + result["global_dma_staging_peak_bytes"] <= architecture["global_dma_staging_bytes"], + "DMA staging capacity exceeded", + ) + require(result["combine_sram_peak_bytes"] <= architecture["combine_sram_bytes"], "combine SRAM capacity exceeded") require(len(result["cores"]) == len(architecture["cores"]), "core count mismatch") for observed, core in zip(result["cores"], architecture["cores"]): require(observed["id"] == core["id"], "core identity mismatch") require(observed["multipliers"] == core["blen"] * core["mlen"], "core multiplier mismatch") - for metric in ("useful_macs", "issued_macs", "jobs", "hbm_read_bytes", "compute_busy_ps", - "vector_sram_peak_bytes", "accumulator_peak_bytes", "weight_sram_peak_bytes", - "weight_slots_peak"): - require(isinstance(observed[metric], int) and not isinstance(observed[metric], bool) - and observed[metric] >= 0, "invalid per-core metric: " + metric) + for metric in ( + "useful_macs", + "issued_macs", + "jobs", + "hbm_read_bytes", + "compute_busy_ps", + "vector_sram_peak_bytes", + "accumulator_peak_bytes", + "weight_sram_peak_bytes", + "weight_slots_peak", + ): + require( + isinstance(observed[metric], int) and not isinstance(observed[metric], bool) and observed[metric] >= 0, + "invalid per-core metric: " + metric, + ) require(observed["issued_macs"] >= observed["useful_macs"], "per-core issued MAC deficit") require(observed["compute_busy_ps"] <= result["total_ps"], "core busy time exceeds run") - for peak, capacity in (("vector_sram_peak_bytes", "vector_sram_bytes"), - ("accumulator_peak_bytes", "accumulator_bytes"), - ("weight_sram_peak_bytes", "weight_sram_bytes")): + for peak, capacity in ( + ("vector_sram_peak_bytes", "vector_sram_bytes"), + ("accumulator_peak_bytes", "accumulator_bytes"), + ("weight_sram_peak_bytes", "weight_sram_bytes"), + ): require(observed[peak] <= core[capacity], core["id"] + " exceeded " + capacity) require(observed["weight_slots_peak"] <= core.get("weight_slots", 2), "weight slot count exceeded") if "read_cache_bytes" in core and not architecture.get("dma"): for metric in ("cache_requests", "cache_hits", "cache_port_busy_ps", "cache_peak_bytes"): require(type(observed[metric]) is int and observed[metric] >= 0, "invalid cache metric: " + metric) - require(observed["cache_requests"] == observed["cache_hits"] + observed["hbm_read_bytes"] // 64, - "cache requests do not reconcile with real HBM reads") - require(observed["cache_peak_bytes"] <= core["read_cache_bytes"] - and observed["cache_peak_bytes"] % 80 == 0, "cache data/tag capacity exceeded") - expected_cache_cycles = (observed["cache_requests"] + observed["hbm_read_bytes"] // 64) if core["read_cache_bytes"] else 0 - require(observed["cache_port_busy_ps"] == expected_cache_cycles * architecture["clock_period_ps"], - "cache port service is not accounted") + require( + observed["cache_requests"] == observed["cache_hits"] + observed["hbm_read_bytes"] // 64, + "cache requests do not reconcile with real HBM reads", + ) + require( + observed["cache_peak_bytes"] <= core["read_cache_bytes"] and observed["cache_peak_bytes"] % 80 == 0, + "cache data/tag capacity exceeded", + ) + expected_cache_cycles = ( + (observed["cache_requests"] + observed["hbm_read_bytes"] // 64) if core["read_cache_bytes"] else 0 + ) + require( + observed["cache_port_busy_ps"] == expected_cache_cycles * architecture["clock_period_ps"], + "cache port service is not accounted", + ) require(observed["cache_port_busy_ps"] <= result["total_ps"], "cache port busy time exceeds run") for metric in ("compute_busy_fraction", "mac_utilization"): - require(math.isfinite(observed[metric]) and 0 <= observed[metric] <= 1.000000001, - "invalid " + metric) + require(math.isfinite(observed[metric]) and 0 <= observed[metric] <= 1.000000001, "invalid " + metric) jobs = result["job_completions"] if "dispatch_queue_bytes" in architecture: - require(result["dispatch_queue_peak_bytes"] == len(jobs) * 64 - and result["dispatch_queue_peak_bytes"] <= architecture["dispatch_queue_bytes"], - "ready-job descriptor capacity mismatch") - require(result["dispatcher_busy_ps"] == len(jobs) * architecture["dispatch_cycles"] * architecture["clock_period_ps"] - and result["dispatcher_busy_ps"] <= result["total_ps"], "dispatcher service mismatch") + require( + result["dispatch_queue_peak_bytes"] == len(jobs) * 64 + and result["dispatch_queue_peak_bytes"] <= architecture["dispatch_queue_bytes"], + "ready-job descriptor capacity mismatch", + ) + require( + result["dispatcher_busy_ps"] + == len(jobs) * architecture["dispatch_cycles"] * architecture["clock_period_ps"] + and result["dispatcher_busy_ps"] <= result["total_ps"], + "dispatcher service mismatch", + ) identities = [(job["expert"], job["shared"]) for job in jobs] expected_rows = {} for route in workload["routes"]: @@ -181,32 +254,44 @@ def validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_chann expected_rows[identity] = expected_rows.get(identity, 0) + 1 if workload.get("shared_expert") is not None and token_count: expected_rows[(workload["shared_expert"]["expert"], True)] = token_count - require(len(identities) == len(set(identities)) and set(identities) == set(expected_rows), - "missing or duplicate expert completion") + require( + len(identities) == len(set(identities)) and set(identities) == set(expected_rows), + "missing or duplicate expert completion", + ) require(sum(core["jobs"] for core in result["cores"]) == len(jobs), "job count mismatch") for field in ("useful_macs", "issued_macs", "hbm_read_bytes"): - require(sum(core[field] for core in result["cores"]) == result[field], - "per-core totals do not reconcile: " + field) + require( + sum(core[field] for core in result["cores"]) == result[field], "per-core totals do not reconcile: " + field + ) jobs_per_core = {core["id"]: 0 for core in architecture["cores"]} rows_per_core = {core["id"]: 0 for core in architecture["cores"]} for job in jobs: require(job["core"] in jobs_per_core, "job completion names an unknown core") - require(type(job["rows"]) is int - and job["rows"] == expected_rows[(job["expert"], job["shared"])], - "expert completion row count differs from workload") + require( + type(job["rows"]) is int and job["rows"] == expected_rows[(job["expert"], job["shared"])], + "expert completion row count differs from workload", + ) jobs_per_core[job["core"]] += 1 rows_per_core[job["core"]] += job["rows"] - require(all(type(job[field]) is int for field in - ("start_ps", "compute_done_ps", "output_copied_ps")), - "job completion timestamps must be integers") - require(0 <= job["start_ps"] <= job["compute_done_ps"] <= job["output_copied_ps"] - <= result["total_ps"], "invalid job completion timestamps") + require( + all(type(job[field]) is int for field in ("start_ps", "compute_done_ps", "output_copied_ps")), + "job completion timestamps must be integers", + ) + require( + 0 <= job["start_ps"] <= job["compute_done_ps"] <= job["output_copied_ps"] <= result["total_ps"], + "invalid job completion timestamps", + ) for core in result["cores"]: require(core["jobs"] == jobs_per_core[core["id"]], "per-core job count mismatch") - require(core["useful_macs"] == rows_per_core[core["id"]] * 3 - * workload["input_dim"] * workload["expert_hidden_dim"], - "per-core useful MACs differ from completed jobs") - require(result["hbm_read_bytes"] % (32 if architecture.get("dma") else 64) == 0, "HBM request bytes must match transfer granularity") + require( + core["useful_macs"] + == rows_per_core[core["id"]] * 3 * workload["input_dim"] * workload["expert_hidden_dim"], + "per-core useful MACs differ from completed jobs", + ) + require( + result["hbm_read_bytes"] % (32 if architecture.get("dma") else 64) == 0, + "HBM request bytes must match transfer granularity", + ) require(result["hbm_write_bytes"] == 0, "unexpected HBM output/intermediate write") require(result["shared_vector_busy_ps"] <= result["total_ps"], "vector busy time exceeds run") for source in (result, golden): @@ -218,22 +303,37 @@ def validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_chann validate_native(envelope, architecture, hbm_channels) reference = golden["output_f32"] error = numerical_gate(result["output_f32"], reference, atol, rtol) - return {"passed": True, "max_absolute_error": error, - "output_bit_exact": result["output_bf16"] == golden["output_bf16"]} + return { + "passed": True, + "max_absolute_error": error, + "output_bit_exact": result["output_bf16"] == golden["output_bf16"], + } -def _run_comparison(binary, workload_path, golden_path, architecture_paths, output_dir, - repeats=2, hbm_channels=8, atol=1e-5, rtol=0.01, timeout=180, workers=1): +def _run_comparison( + binary, + workload_path, + golden_path, + architecture_paths, + output_dir, + repeats=2, + hbm_channels=8, + atol=1e-5, + rtol=0.01, + timeout=180, + workers=1, +): require(type(workers) is int and 1 <= workers <= 8, "workers must be an integer between 1 and 8") require(type(repeats) is int and repeats >= 2, "at least two integer repeats are required") - require(type(hbm_channels) is int and 1 <= hbm_channels <= 32 - and hbm_channels & (hbm_channels - 1) == 0, - "HBM channels must be a power of two between 1 and 32") - require(type(timeout) in (int, float) and math.isfinite(timeout) and timeout > 0, - "timeout must be finite and positive") + require( + type(hbm_channels) is int and 1 <= hbm_channels <= 32 and hbm_channels & (hbm_channels - 1) == 0, + "HBM channels must be a power of two between 1 and 32", + ) + require( + type(timeout) in (int, float) and math.isfinite(timeout) and timeout > 0, "timeout must be finite and positive" + ) require(len(architecture_paths) >= 2, "provide baseline and at least one candidate") - require(all(type(x) in (int, float) and math.isfinite(x) and x >= 0 for x in (atol, rtol)), - "invalid tolerance") + require(all(type(x) in (int, float) and math.isfinite(x) and x >= 0 for x in (atol, rtol)), "invalid tolerance") binary = Path(binary).resolve() workload_path = Path(workload_path).resolve() golden_path = Path(golden_path).resolve() @@ -251,44 +351,69 @@ def _run_comparison(binary, workload_path, golden_path, architecture_paths, outp require(len(names) == len(set(names)), "architecture names must be unique") total_pes = [sum(c["blen"] * c["mlen"] for c in a["cores"]) for a in architectures] require(len(set(total_pes)) == 1, "comparison requires equal multiplier totals") - for field in ("clock_period_ps", "mac_pipeline_cycles", "vector_elements_per_cycle", - "global_dma_credits", "global_dma_staging_bytes", "combine_sram_bytes"): - require(len({a[field] for a in architectures}) == 1, - "comparison requires same shared resource: " + field) - for field, default in (("dispatch_cycles", 1), ("dispatch_queue_bytes", 262144), - ("matrix_timing", "pipelined")): - require(len({a.get(field, default) for a in architectures}) == 1, - "comparison requires same shared resource/timing: " + field) + for field in ( + "clock_period_ps", + "mac_pipeline_cycles", + "vector_elements_per_cycle", + "global_dma_credits", + "global_dma_staging_bytes", + "combine_sram_bytes", + ): + require(len({a[field] for a in architectures}) == 1, "comparison requires same shared resource: " + field) + for field, default in (("dispatch_cycles", 1), ("dispatch_queue_bytes", 262144), ("matrix_timing", "pipelined")): + require( + len({a.get(field, default) for a in architectures}) == 1, + "comparison requires same shared resource/timing: " + field, + ) cache_fields = ["read_cache_bytes" in c for a in architectures for c in a["cores"]] require(not any(cache_fields) or all(cache_fields), "cache budgets must be explicit for every core") if all(cache_fields): for field in ("vector_sram_bytes", "accumulator_bytes", "weight_sram_bytes", "read_cache_bytes"): - require(len({sum(c[field] for c in a["cores"]) for a in architectures}) == 1, - "full-shape comparison requires equal total configured SRAM: " + field) + require( + len({sum(c[field] for c in a["cores"]) for a in architectures}) == 1, + "full-shape comparison requires equal total configured SRAM: " + field, + ) dma_configs = [a.get("dma") for a in architectures] - require(all(d is None for d in dma_configs) or all(d is not None for d in dma_configs), "DMA configuration must be explicit for all architectures") + require( + all(d is None for d in dma_configs) or all(d is not None for d in dma_configs), + "DMA configuration must be explicit for all architectures", + ) if all(d is not None for d in dma_configs): - require(len({json.dumps(d, sort_keys=True) for d in dma_configs}) == 1, "comparison requires identical DMA policy and budget") + require( + len({json.dumps(d, sort_keys=True) for d in dma_configs}) == 1, + "comparison requires identical DMA policy and budget", + ) hbm_path = workload_path.parent / workload["hbm_file"] - expected_hashes = {"workload_sha256": hashlib.sha256(workload_bytes).hexdigest(), - "hbm_sha256": digest(hbm_path), "executable_sha256": digest(binary)} + expected_hashes = { + "workload_sha256": hashlib.sha256(workload_bytes).hexdigest(), + "hbm_sha256": digest(hbm_path), + "executable_sha256": digest(binary), + } for field in ("workload_sha256", "hbm_sha256"): require(golden[field] == expected_hashes[field], "golden identity mismatch: " + field) run_dir = output_dir / ("run_" + uuid.uuid4().hex) run_dir.mkdir() + def execute_architecture(index): path, architecture = architecture_paths[index], architectures[index] repeat_results, gates, native_repeats, library_hashes = [], [], [], [] architecture_hash = architecture_hashes[index] for repeat in range(repeats): - report_path = run_dir / "arch{:02d}_repeat{:02d}.json".format(index, repeat) + report_path = run_dir / f"arch{index:02d}_repeat{repeat:02d}.json" log_path = report_path.with_suffix(".log") - command = [str(binary), "--workload", str(workload_path), "--architecture", - str(Path(path).resolve()), "--output", str(report_path), - "--hbm-channels", str(hbm_channels)] + command = [ + str(binary), + "--workload", + str(workload_path), + "--architecture", + str(Path(path).resolve()), + "--output", + str(report_path), + "--hbm-channels", + str(hbm_channels), + ] with log_path.open("wb") as log: - subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, - check=True, timeout=timeout) + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True, timeout=timeout) envelope = read_json(report_path) for field, value in dict(expected_hashes, architecture_sha256=architecture_hash).items(): require(envelope["provenance"][field] == value, "input/binary identity changed: " + field) @@ -296,22 +421,32 @@ def execute_architecture(index): repeat_results.append(envelope["result"]) native_repeats.append(envelope["memory_model"].get("calibration")) library_hashes.append(envelope["provenance"].get("native_library_sha256")) - require(all(result == repeat_results[0] for result in repeat_results[1:]), - "repeat mismatch for " + architecture["name"]) + require( + all(result == repeat_results[0] for result in repeat_results[1:]), + "repeat mismatch for " + architecture["name"], + ) require(all(n == native_repeats[0] for n in native_repeats), "native counters changed between repeats") require(len(set(library_hashes)) == 1, "native library changed between repeats") result = repeat_results[0] require(result["total_ps"] > 0, "nonempty comparison must advance time") - return {"architecture": architecture, "result": result, "gates": gates, "native": native_repeats[0], "native_library_sha256": library_hashes[0]} + return { + "architecture": architecture, + "result": result, + "gates": gates, + "native": native_repeats[0], + "native_library_sha256": library_hashes[0], + } + # Each process owns an independent Ramulator/executor. Preserve manifest # order and propagate every exception before publishing any speedup. with ThreadPoolExecutor(max_workers=workers) as pool: results = list(pool.map(execute_architecture, range(len(architectures)))) for path, expected, label in ( - (workload_path, expected_hashes["workload_sha256"], "workload"), - (golden_path, golden_hash, "golden"), - (hbm_path, expected_hashes["hbm_sha256"], "HBM"), - (binary, expected_hashes["executable_sha256"], "executable")): + (workload_path, expected_hashes["workload_sha256"], "workload"), + (golden_path, golden_hash, "golden"), + (hbm_path, expected_hashes["hbm_sha256"], "HBM"), + (binary, expected_hashes["executable_sha256"], "executable"), + ): require(digest(path) == expected, label + " changed during comparison") for path, expected in zip(architecture_paths, architecture_hashes): require(digest(path) == expected, "architecture changed during comparison") @@ -320,37 +455,67 @@ def execute_architecture(index): for row in results: row["speedup_vs_baseline"] = baseline_ps / row["result"]["total_ps"] summary = { - "schema_version": 1, "all_gates_passed": True, - "status": "passed", "run_directory": str(run_dir), + "schema_version": 1, + "all_gates_passed": True, + "status": "passed", + "run_directory": str(run_dir), "scope": "Fixed-route numerical MoE operator; analytical core timing and shared Ramulator. " - "Not RTL-calibrated or whole-model/request latency.", + "Not RTL-calibrated or whole-model/request latency.", "source_metadata": workload.get("metadata", {}), "workload_sha256": expected_hashes["workload_sha256"], "golden_sha256": golden_hash, "hbm_sha256": expected_hashes["hbm_sha256"], "executable_sha256": expected_hashes["executable_sha256"], - "repeats": repeats, "hbm_channels": hbm_channels, - "numerical_tolerance": {"atol": atol, "rtol": rtol}, "comparisons": results, + "repeats": repeats, + "hbm_channels": hbm_channels, + "numerical_tolerance": {"atol": atol, "rtol": rtol}, + "comparisons": results, } (output_dir / "comparison.json").write_text( - json.dumps(summary, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8") + json.dumps(summary, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8" + ) return summary -def run_comparison(binary, workload_path, golden_path, architecture_paths, output_dir, - repeats=2, hbm_channels=8, atol=1e-5, rtol=0.01, timeout=180, workers=1): +def run_comparison( + binary, + workload_path, + golden_path, + architecture_paths, + output_dir, + repeats=2, + hbm_channels=8, + atol=1e-5, + rtol=0.01, + timeout=180, + workers=1, +): output_dir = Path(output_dir).resolve() output_dir.mkdir(parents=True, exist_ok=True) summary_path = output_dir / "comparison.json" - summary_path.write_text(json.dumps({"schema_version": 1, "status": "running", - "all_gates_passed": False}) + "\n", encoding="utf-8") + summary_path.write_text( + json.dumps({"schema_version": 1, "status": "running", "all_gates_passed": False}) + "\n", encoding="utf-8" + ) try: - return _run_comparison(binary, workload_path, golden_path, architecture_paths, - output_dir, repeats, hbm_channels, atol, rtol, timeout, workers) + return _run_comparison( + binary, + workload_path, + golden_path, + architecture_paths, + output_dir, + repeats, + hbm_channels, + atol, + rtol, + timeout, + workers, + ) except Exception as error: - summary_path.write_text(json.dumps({"schema_version": 1, "status": "failed", - "all_gates_passed": False, "error": str(error)}) + "\n", - encoding="utf-8") + summary_path.write_text( + json.dumps({"schema_version": 1, "status": "failed", "all_gates_passed": False, "error": str(error)}) + + "\n", + encoding="utf-8", + ) raise @@ -359,8 +524,7 @@ def main(): parser.add_argument("--binary", required=True) parser.add_argument("--workload", required=True) parser.add_argument("--golden", required=True) - parser.add_argument("--architecture", action="append", required=True, - help="Baseline first, then candidate(s)") + parser.add_argument("--architecture", action="append", required=True, help="Baseline first, then candidate(s)") parser.add_argument("--output-dir", required=True) parser.add_argument("--repeats", type=int, default=2) parser.add_argument("--hbm-channels", type=int, default=8) @@ -370,14 +534,27 @@ def main(): parser.add_argument("--workers", type=int, default=1) args = parser.parse_args() try: - summary = run_comparison(args.binary, args.workload, args.golden, args.architecture, - args.output_dir, args.repeats, args.hbm_channels, - args.atol, args.rtol, args.timeout, args.workers) + summary = run_comparison( + args.binary, + args.workload, + args.golden, + args.architecture, + args.output_dir, + args.repeats, + args.hbm_channels, + args.atol, + args.rtol, + args.timeout, + args.workers, + ) except (ValueError, OSError, KeyError, subprocess.SubprocessError) as error: parser.error(str(error)) for row in summary["comparisons"]: - print("{}: {} ps, {:.4f}x, gates passed".format( - row["architecture"]["name"], row["result"]["total_ps"], row["speedup_vs_baseline"])) + print( + "{}: {} ps, {:.4f}x, gates passed".format( + row["architecture"]["name"], row["result"]["total_ps"], row["speedup_vs_baseline"] + ) + ) return 0 diff --git a/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py b/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py index b0af606e1..9d0d91fa3 100644 --- a/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py +++ b/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py @@ -1,4 +1,5 @@ """Evidence gate tests through a real fallible executable; Python 3.6+ stdlib.""" + import copy import importlib.util import json @@ -8,11 +9,12 @@ import unittest SPEC = importlib.util.spec_from_file_location( - "compare_moe_normal", str(Path(__file__).with_name("compare_moe_normal.py"))) + "compare_moe_normal", str(Path(__file__).with_name("compare_moe_normal.py")) +) compare = importlib.util.module_from_spec(SPEC) SPEC.loader.exec_module(compare) -FAKE_RUNNER = r''' +FAKE_RUNNER = r""" import argparse import fcntl import hashlib @@ -65,65 +67,127 @@ def sha(path): target = target[key] target[mutation["path"][-1]] = mutation["value"] Path(args.output).write_text(json.dumps(envelope)) -''' +""" def core_config(name, blen, mlen): - return dict(id=name, blen=blen, mlen=mlen, vector_sram_bytes=4096, - accumulator_bytes=1024, weight_sram_bytes=2048) + return dict(id=name, blen=blen, mlen=mlen, vector_sram_bytes=4096, accumulator_bytes=1024, weight_sram_bytes=2048) def architecture(name, cores): - return dict(schema_version=1, name=name, cores=cores, dispatch_threshold=2, - large_core=0, small_core=len(cores) - 1, global_dma_credits=2, - global_dma_staging_bytes=128, combine_sram_bytes=4096, - clock_period_ps=10, mac_pipeline_cycles=8, vector_elements_per_cycle=8) + return dict( + schema_version=1, + name=name, + cores=cores, + dispatch_threshold=2, + large_core=0, + small_core=len(cores) - 1, + global_dma_credits=2, + global_dma_staging_bytes=128, + combine_sram_bytes=4096, + clock_period_ps=10, + mac_pipeline_cycles=8, + vector_elements_per_cycle=8, + ) def core_report(core, jobs, rows, total_ps, busy_ps): useful = rows * 3 * 8 * 8 multipliers = core["blen"] * core["mlen"] - return dict(id=core["id"], blen=core["blen"], mlen=core["mlen"], - multipliers=multipliers, jobs=jobs, useful_macs=useful, - issued_macs=useful, compute_busy_ps=busy_ps, - accumulator_dependency_stall_ps=0, pipeline_drain_ps=10, - pipeline_register_bytes=64, weight_ready_wait_ps=0, - vector_wait_ps=0, vector_sram_peak_bytes=256, - accumulator_peak_bytes=128, weight_sram_peak_bytes=256, - weight_slots_peak=2, hbm_read_bytes=192 * jobs, - compute_busy_fraction=busy_ps / total_ps, - mac_utilization=useful / (multipliers * total_ps / 10)) + return dict( + id=core["id"], + blen=core["blen"], + mlen=core["mlen"], + multipliers=multipliers, + jobs=jobs, + useful_macs=useful, + issued_macs=useful, + compute_busy_ps=busy_ps, + accumulator_dependency_stall_ps=0, + pipeline_drain_ps=10, + pipeline_register_bytes=64, + weight_ready_wait_ps=0, + vector_wait_ps=0, + vector_sram_peak_bytes=256, + accumulator_peak_bytes=128, + weight_sram_peak_bytes=256, + weight_slots_peak=2, + hbm_read_bytes=192 * jobs, + compute_busy_fraction=busy_ps / total_ps, + mac_utilization=useful / (multipliers * total_ps / 10), + ) def envelope(arch): single = len(arch["cores"]) == 1 total = 1000 if single else 800 jobs = [ - dict(job=0, expert=0, shared=False, core=arch["cores"][0]["id"], rows=2, - start_ps=0, compute_done_ps=200, output_copied_ps=220), - dict(job=1, expert=2, shared=True, core=arch["cores"][0]["id"], rows=2, - start_ps=220, compute_done_ps=420, output_copied_ps=440), - dict(job=2, expert=1, shared=False, core=arch["cores"][-1]["id"], rows=1, - start_ps=440 if single else 0, compute_done_ps=600 if single else 180, - output_copied_ps=620 if single else 200), + dict( + job=0, + expert=0, + shared=False, + core=arch["cores"][0]["id"], + rows=2, + start_ps=0, + compute_done_ps=200, + output_copied_ps=220, + ), + dict( + job=1, + expert=2, + shared=True, + core=arch["cores"][0]["id"], + rows=2, + start_ps=220, + compute_done_ps=420, + output_copied_ps=440, + ), + dict( + job=2, + expert=1, + shared=False, + core=arch["cores"][-1]["id"], + rows=1, + start_ps=440 if single else 0, + compute_done_ps=600 if single else 180, + output_copied_ps=620 if single else 200, + ), ] - cores = ([core_report(arch["cores"][0], 3, 5, total, 600)] if single else - [core_report(arch["cores"][0], 2, 4, total, 600), - core_report(arch["cores"][1], 1, 1, total, 200)]) - result = dict(schema_version=1, workload="fixture", architecture=arch["name"], - timing_model="analytical", timing_boundary="operator", - weight_format="local_mx", total_ps=total, multipliers=32, - useful_macs=960, issued_macs=960, hbm_read_bytes=576, - hbm_write_bytes=0, global_dma_inflight_peak=2, - global_dma_staging_peak_bytes=128, combine_sram_peak_bytes=256, - shared_vector_busy_ps=100, cores=cores, - job_completions=jobs if single else [jobs[2], jobs[0], jobs[1]], - output_f32=[[1.0] * 8, [-1.0] * 8], - output_bf16=[[16256] * 8, [49024] * 8], - pre_round_output_f32=[[1.0] * 8, [-1.0] * 8]) - return dict(schema_version=1, evidence_level="fixture_only", provenance={}, - memory_model=dict(name="Ramulator HBM2 preset", channels=8, - upper_burst_bytes=64), result=result) + cores = ( + [core_report(arch["cores"][0], 3, 5, total, 600)] + if single + else [core_report(arch["cores"][0], 2, 4, total, 600), core_report(arch["cores"][1], 1, 1, total, 200)] + ) + result = dict( + schema_version=1, + workload="fixture", + architecture=arch["name"], + timing_model="analytical", + timing_boundary="operator", + weight_format="local_mx", + total_ps=total, + multipliers=32, + useful_macs=960, + issued_macs=960, + hbm_read_bytes=576, + hbm_write_bytes=0, + global_dma_inflight_peak=2, + global_dma_staging_peak_bytes=128, + combine_sram_peak_bytes=256, + shared_vector_busy_ps=100, + cores=cores, + job_completions=jobs if single else [jobs[2], jobs[0], jobs[1]], + output_f32=[[1.0] * 8, [-1.0] * 8], + output_bf16=[[16256] * 8, [49024] * 8], + pre_round_output_f32=[[1.0] * 8, [-1.0] * 8], + ) + return dict( + schema_version=1, + evidence_level="fixture_only", + provenance={}, + memory_model=dict(name="Ramulator HBM2 preset", channels=8, upper_burst_bytes=64), + result=result, + ) class ComparisonEvidenceTests(unittest.TestCase): @@ -136,37 +200,48 @@ def setUp(self): self.binary.chmod(0o700) self.hbm = self.root / "weights.bin" self.hbm.write_bytes(bytes(1152)) - self.workload = dict(schema_version=1, name="fixture", hbm_file="weights.bin", - input_dim=8, expert_hidden_dim=8, - inputs_bf16=[[16256] * 8, [16256] * 8], - routes=[dict(token=0, slot=0, expert=0, weight=0.5), - dict(token=0, slot=1, expert=1, weight=0.5), - dict(token=1, slot=0, expert=0, weight=1.0)], - shared_expert=dict(expert=2, weight=1.0), metadata={}) + self.workload = dict( + schema_version=1, + name="fixture", + hbm_file="weights.bin", + input_dim=8, + expert_hidden_dim=8, + inputs_bf16=[[16256] * 8, [16256] * 8], + routes=[ + dict(token=0, slot=0, expert=0, weight=0.5), + dict(token=0, slot=1, expert=1, weight=0.5), + dict(token=1, slot=0, expert=0, weight=1.0), + ], + shared_expert=dict(expert=2, weight=1.0), + metadata={}, + ) self.workload["experts"] = [] for expert in range(3): record = dict(id=expert) for projection, name in enumerate(("gate", "up", "down")): base = (expert * 3 + projection) * 128 - record[name] = dict(rows=8, cols=8, element_base=base, - scale_base=base + 64, element_row_stride=8, - scale_row_stride=1) + record[name] = dict( + rows=8, cols=8, element_base=base, scale_base=base + 64, element_row_stride=8, scale_row_stride=1 + ) self.workload["experts"].append(record) self.workload_path = self.root / "workload.json" self.write(self.workload_path, self.workload) - self.architectures = [architecture("single", [core_config("single", 4, 8)]), - architecture("dual", [core_config("large", 2, 8), - core_config("small", 1, 16)])] + self.architectures = [ + architecture("single", [core_config("single", 4, 8)]), + architecture("dual", [core_config("large", 2, 8), core_config("small", 1, 16)]), + ] self.arch_paths = [self.root / "single.json", self.root / "dual.json"] for path, arch in zip(self.arch_paths, self.architectures): self.write(path, arch) self.reports = {arch["name"]: envelope(arch) for arch in self.architectures} self.write(self.root / "reports.json", self.reports) - self.golden = dict(schema_version=1, - workload_sha256=compare.digest(self.workload_path), - hbm_sha256=compare.digest(self.hbm), - output_f32=[[1.0] * 8, [-1.0] * 8], - output_bf16=[[16256] * 8, [49024] * 8]) + self.golden = dict( + schema_version=1, + workload_sha256=compare.digest(self.workload_path), + hbm_sha256=compare.digest(self.hbm), + output_f32=[[1.0] * 8, [-1.0] * 8], + output_bf16=[[16256] * 8, [49024] * 8], + ) self.golden_path = self.root / "golden.json" self.write(self.golden_path, self.golden) self.scenario = {} @@ -179,9 +254,9 @@ def write(path, data): def run_comparison(self, **kwargs): self.write(self.root / "scenario.json", self.scenario) - return compare.run_comparison(str(self.binary), str(self.workload_path), - str(self.golden_path), self.arch_paths, - self.output, **kwargs) + return compare.run_comparison( + str(self.binary), str(self.workload_path), str(self.golden_path), self.arch_paths, self.output, **kwargs + ) def mutate(self, path, value, **filters): self.scenario.setdefault("mutations", []).append(dict(path=path, value=value, **filters)) @@ -207,68 +282,81 @@ def test_wrong_numeric_result_is_rejected(self): def test_parallel_architectures_keep_order_and_repeat_gates(self): result = self.run_comparison(workers=2) - self.assertEqual([r['architecture']['name'] for r in result['comparisons']], ['single','dual']) - self.assertEqual(compare.read_json(self.root/'calls.json'),{'single':2,'dual':2}) - self.mutate(['result','total_ps'],801,architecture='dual',repeat=3) + self.assertEqual([r["architecture"]["name"] for r in result["comparisons"]], ["single", "dual"]) + self.assertEqual(compare.read_json(self.root / "calls.json"), {"single": 2, "dual": 2}) + self.mutate(["result", "total_ps"], 801, architecture="dual", repeat=3) self.rejected(workers=2) def test_invalid_parallelism_is_rejected(self): - for workers in [0,9,True,1.5]: - with self.subTest(workers=workers): self.rejected(workers=workers) + for workers in [0, 9, True, 1.5]: + with self.subTest(workers=workers): + self.rejected(workers=workers) def enable_full_shape_resources(self): for path, arch in zip(self.arch_paths, self.architectures): arch.update(dispatch_queue_bytes=256, dispatch_cycles=1) - parts = len(arch['cores']) - for core in arch['cores']: - core.update(read_cache_bytes=160 // parts, vector_sram_bytes=4096 // parts, - accumulator_bytes=1024 // parts, weight_sram_bytes=2048 // parts) - report = self.reports[arch['name']]['result'] + parts = len(arch["cores"]) + for core in arch["cores"]: + core.update( + read_cache_bytes=160 // parts, + vector_sram_bytes=4096 // parts, + accumulator_bytes=1024 // parts, + weight_sram_bytes=2048 // parts, + ) + report = self.reports[arch["name"]]["result"] report.update(dispatch_queue_peak_bytes=192, dispatcher_busy_ps=30) - for core in report['cores']: - misses = core['hbm_read_bytes'] // 64 - core.update(cache_requests=misses+2, cache_hits=2, - cache_port_busy_ps=(2*misses+2)*10, cache_peak_bytes=80) + for core in report["cores"]: + misses = core["hbm_read_bytes"] // 64 + core.update( + cache_requests=misses + 2, + cache_hits=2, + cache_port_busy_ps=(2 * misses + 2) * 10, + cache_peak_bytes=80, + ) self.write(path, arch) - self.write(self.root / 'reports.json', self.reports) + self.write(self.root / "reports.json", self.reports) def test_full_shape_resource_accounting_passes(self): self.enable_full_shape_resources() - self.assertTrue(self.run_comparison()['all_gates_passed']) + self.assertTrue(self.run_comparison()["all_gates_passed"]) def test_cache_capacity_traffic_and_port_undercharging_are_rejected(self): self.enable_full_shape_resources() - for metric,value in [('cache_peak_bytes',240),('cache_requests',0), - ('cache_port_busy_ps',0),('cache_hits',-1)]: + for metric, value in [ + ("cache_peak_bytes", 240), + ("cache_requests", 0), + ("cache_port_busy_ps", 0), + ("cache_hits", -1), + ]: with self.subTest(metric=metric): self.scenario = {} - self.mutate(['result','cores',0,metric],value) + self.mutate(["result", "cores", 0, metric], value) self.rejected() def test_dispatch_capacity_and_service_are_required(self): self.enable_full_shape_resources() - for metric,value in [('dispatch_queue_peak_bytes',0),('dispatcher_busy_ps',0)]: + for metric, value in [("dispatch_queue_peak_bytes", 0), ("dispatcher_busy_ps", 0)]: with self.subTest(metric=metric): self.scenario = {} - self.mutate(['result',metric],value) + self.mutate(["result", metric], value) self.rejected() def test_unequal_or_partly_missing_cache_budgets_fail_before_execution(self): self.enable_full_shape_resources() - self.architectures[0]['cores'][0]['read_cache_bytes'] += 80 - self.write(self.arch_paths[0],self.architectures[0]) + self.architectures[0]["cores"][0]["read_cache_bytes"] += 80 + self.write(self.arch_paths[0], self.architectures[0]) self.rejected() - self.assertFalse((self.root/'calls.json').exists()) - del self.architectures[0]['cores'][0]['read_cache_bytes'] - self.write(self.arch_paths[0],self.architectures[0]) + self.assertFalse((self.root / "calls.json").exists()) + del self.architectures[0]["cores"][0]["read_cache_bytes"] + self.write(self.arch_paths[0], self.architectures[0]) self.rejected() - self.assertFalse((self.root/'calls.json').exists()) + self.assertFalse((self.root / "calls.json").exists()) def test_unequal_timing_modes_fail_before_execution(self): - self.architectures[0]['matrix_timing'] = 'legacy_serialized' - self.write(self.arch_paths[0],self.architectures[0]) + self.architectures[0]["matrix_timing"] = "legacy_serialized" + self.write(self.arch_paths[0], self.architectures[0]) self.rejected() - self.assertFalse((self.root/'calls.json').exists()) + self.assertFalse((self.root / "calls.json").exists()) def test_nonfinite_output_is_rejected(self): for value in (float("nan"), float("inf"), -float("inf")): @@ -289,17 +377,19 @@ def test_output_shape_must_match_workload_even_if_golden_is_empty(self): self.rejected() def test_every_finite_storage_bound_is_enforced(self): - cases = [(["global_dma_inflight_peak"], 3), - (["global_dma_staging_peak_bytes"], 129), - (["combine_sram_peak_bytes"], 4097), - (["cores", 0, "vector_sram_peak_bytes"], 4097), - (["cores", 0, "accumulator_peak_bytes"], 1025), - (["cores", 0, "weight_sram_peak_bytes"], 2049), - (["cores", 0, "weight_slots_peak"], 3)] + cases = [ + (["global_dma_inflight_peak"], 3), + (["global_dma_staging_peak_bytes"], 129), + (["combine_sram_peak_bytes"], 4097), + (["cores", 0, "vector_sram_peak_bytes"], 4097), + (["cores", 0, "accumulator_peak_bytes"], 1025), + (["cores", 0, "weight_sram_peak_bytes"], 2049), + (["cores", 0, "weight_slots_peak"], 3), + ] for path, value in cases: with self.subTest(metric=path): self.scenario = {} - self.mutate(["result"] + path, value) + self.mutate(["result", *path], value) self.rejected() def test_negative_resource_counter_is_rejected(self): @@ -326,7 +416,7 @@ def test_per_core_hbm_bytes_must_close_to_total(self): def test_missing_or_duplicate_expert_completions_are_rejected(self): original = self.reports["single"]["result"]["job_completions"] - for jobs in (original[:-1], original + [copy.deepcopy(original[0])]): + for jobs in (original[:-1], [*original, copy.deepcopy(original[0])]): with self.subTest(count=len(jobs)): self.scenario = {} self.mutate(["result", "job_completions"], jobs, architecture="single") @@ -345,9 +435,12 @@ def test_completion_after_total_time_is_rejected(self): self.rejected() def test_nonfinite_or_negative_timings_are_rejected(self): - for metric, value in (("total_ps", float("inf")), ("total_ps", -1), - ("shared_vector_busy_ps", float("nan")), - ("shared_vector_busy_ps", -1)): + for metric, value in ( + ("total_ps", float("inf")), + ("total_ps", -1), + ("shared_vector_busy_ps", float("nan")), + ("shared_vector_busy_ps", -1), + ): with self.subTest(metric=metric, value=value): self.scenario = {} self.mutate(["result", metric], value) @@ -393,8 +486,14 @@ def test_unequal_multiplier_architectures_are_rejected_before_execution(self): def test_different_shared_resources_are_rejected_before_execution(self): original = copy.deepcopy(self.architectures[1]) - for field in ("clock_period_ps", "mac_pipeline_cycles", "vector_elements_per_cycle", - "global_dma_credits", "global_dma_staging_bytes", "combine_sram_bytes"): + for field in ( + "clock_period_ps", + "mac_pipeline_cycles", + "vector_elements_per_cycle", + "global_dma_credits", + "global_dma_staging_bytes", + "combine_sram_bytes", + ): with self.subTest(field=field): altered = copy.deepcopy(original) altered[field] += 1 @@ -421,9 +520,16 @@ def test_success_exit_without_new_report_cannot_reuse_previous_reports(self): self.assertIsNot(compare.read_json(summary).get("all_gates_passed"), True) def test_invalid_run_controls_are_rejected(self): - for kwargs in (dict(repeats=1), dict(repeats=2.5), dict(atol=-1.0), - dict(rtol=float("nan")), dict(hbm_channels=0), - dict(hbm_channels=3), dict(timeout=0), dict(timeout=float("inf"))): + for kwargs in ( + dict(repeats=1), + dict(repeats=2.5), + dict(atol=-1.0), + dict(rtol=float("nan")), + dict(hbm_channels=0), + dict(hbm_channels=3), + dict(timeout=0), + dict(timeout=float("inf")), + ): with self.subTest(kwargs=kwargs): self.rejected(**kwargs) From ddcfae5d0f2fca8296f64157379e05460c0d80c7 Mon Sep 17 00:00:00 2001 From: Michael C Li <63594253+Happymic@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:57:55 +0100 Subject: [PATCH 3/5] chore: pin the paired Compiler review with complete CI dependencies --- PLENA_Compiler | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PLENA_Compiler b/PLENA_Compiler index 5c968f897..a6730b5d6 160000 --- a/PLENA_Compiler +++ b/PLENA_Compiler @@ -1 +1 @@ -Subproject commit 5c968f89749ebcc36378956721078c4b1fc46e96 +Subproject commit a6730b5d6cfd3190d571236bb02a4a2babaaabba From bead42b2effde6ea57ab2c62d7ae12ee39f632c3 Mon Sep 17 00:00:00 2001 From: Michael C Li <63594253+Happymic@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:51:44 +0100 Subject: [PATCH 4/5] fix(moe): enforce exact outputs and complete comparison resource parity --- PLENA_Compiler | 2 +- doc/moe_normal_review.md | 11 ++++ .../moe_timing/replay/compare_moe_normal.py | 31 +++++++---- .../replay/test_compare_moe_normal.py | 54 +++++++++++++++++++ 4 files changed, 88 insertions(+), 10 deletions(-) diff --git a/PLENA_Compiler b/PLENA_Compiler index a6730b5d6..c45cf7b06 160000 --- a/PLENA_Compiler +++ b/PLENA_Compiler @@ -1 +1 @@ -Subproject commit a6730b5d6cfd3190d571236bb02a4a2babaaabba +Subproject commit c45cf7b069c9f16cfe2d5a2a454ee4e78d35f0be diff --git a/doc/moe_normal_review.md b/doc/moe_normal_review.md index 5c7cccf2d..44923ca33 100644 --- a/doc/moe_normal_review.md +++ b/doc/moe_normal_review.md @@ -37,6 +37,7 @@ multiplier count is BLEN×MLEN, not M×K×N. Whole-model D/F dimensions are sepa | Weight slots | 4 | 4 | 4 | | Activation/intermediate SRAM | 4 MiB | 2 MiB | 2 MiB | | Accumulator storage | 1 MiB | 512 KiB | 512 KiB | +| BF16 activation elements per cycle | 1024 | 768 | 256 | Both systems share the same 8-channel HBM2 preset, clock, total activation supply, DMA budget and vector unit. The calibrated native transfer is 32 bytes; the @@ -45,6 +46,12 @@ The DRAM command scheduler and HBM timing preset are retained. Logical DMA credi are capped at 128, native trackers at 256; DMA frontend SRAM is 44 KiB including staging. The two cores do not receive two independent HBM bandwidth budgets. +The activation split is fixed, not dynamically borrowed. Under +`service_cycles=max(BLEN,ceil(BLEN*MLEN/activation_supply))`, the small core +needs 4 service cycles per tile rather than the 2 cycles possible with sufficient +operand supply. This is a constraint of the current on-chip port allocation, +not a DRAM delay. The search did not vary this split independently of core size. + HBM weights are output-major `[N,K]` rows with separate element/scale ranges. The local E4M3/E8M0 codec uses 8 elements per scale. An aligned MLEN-long row fragment therefore needs MLEN elements and MLEN/8 scales. SRAM slots reserve both @@ -60,6 +67,10 @@ packed ingress and decoded BF16 data. Both buffers are normal buffers in this PR elapsed time, not the wall-clock duration of the Rust process. Same-output K accumulation dependencies are retained; cross-core waits cannot be summed as total latency. Those dependency/port assumptions remain subjects for review. + The kernel visits N tiles, then K tiles, then token blocks; it does not yet + interleave independent N tiles to hide same-output accumulation waits. A full + result-ready delay is used for the next K update, so these timings do not prove + that a better feedback path or instruction schedule would have the same result. - The independent Compiler oracle decodes actual exported weight bytes, performs ascending-K FP32 arithmetic and applies the same BF16 boundaries. The comparison checks exact BF16 outputs as well as HBM completed bytes, finite SRAM/queue diff --git a/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py b/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py index 932ca200a..5af8eb64c 100644 --- a/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py +++ b/transactional_emulator/testbench/moe_timing/replay/compare_moe_normal.py @@ -303,10 +303,13 @@ def validate_run(envelope, golden, workload, architecture, atol, rtol, hbm_chann validate_native(envelope, architecture, hbm_channels) reference = golden["output_f32"] error = numerical_gate(result["output_f32"], reference, atol, rtol) + bit_exact = result["output_bf16"] == golden["output_bf16"] + if atol == 0 and rtol == 0: + require(bit_exact, "exact comparison requires identical BF16 bits") return { "passed": True, "max_absolute_error": error, - "output_bit_exact": result["output_bf16"] == golden["output_bf16"], + "output_bit_exact": bit_exact, } @@ -351,6 +354,15 @@ def _run_comparison( require(len(names) == len(set(names)), "architecture names must be unique") total_pes = [sum(c["blen"] * c["mlen"] for c in a["cores"]) for a in architectures] require(len(set(total_pes)) == 1, "comparison requires equal multiplier totals") + # Match Rust's Option: missing/null activation supply defaults to MLEN. + activation_totals = [ + sum( + c["mlen"] if c.get("activation_elements_per_cycle") is None else c["activation_elements_per_cycle"] + for c in a["cores"] + ) + for a in architectures + ] + require(len(set(activation_totals)) == 1, "comparison requires equal total activation supply") for field in ( "clock_period_ps", "mac_pipeline_cycles", @@ -365,14 +377,15 @@ def _run_comparison( len({a.get(field, default) for a in architectures}) == 1, "comparison requires same shared resource/timing: " + field, ) - cache_fields = ["read_cache_bytes" in c for a in architectures for c in a["cores"]] - require(not any(cache_fields) or all(cache_fields), "cache budgets must be explicit for every core") - if all(cache_fields): - for field in ("vector_sram_bytes", "accumulator_bytes", "weight_sram_bytes", "read_cache_bytes"): - require( - len({sum(c[field] for c in a["cores"]) for a in architectures}) == 1, - "full-shape comparison requires equal total configured SRAM: " + field, - ) + for field in ("vector_sram_bytes", "accumulator_bytes", "weight_sram_bytes"): + require( + len({sum(c[field] for c in a["cores"]) for a in architectures}) == 1, + "comparison requires equal total configured SRAM: " + field, + ) + require( + len({sum(c.get("read_cache_bytes", 0) for c in a["cores"]) for a in architectures}) == 1, + "comparison requires equal total configured SRAM: read_cache_bytes", + ) dma_configs = [a.get("dma") for a in architectures] require( all(d is None for d in dma_configs) or all(d is not None for d in dma_configs), diff --git a/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py b/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py index 9d0d91fa3..1cf4c1561 100644 --- a/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py +++ b/transactional_emulator/testbench/moe_timing/replay/test_compare_moe_normal.py @@ -232,6 +232,10 @@ def setUp(self): ] self.arch_paths = [self.root / "single.json", self.root / "dual.json"] for path, arch in zip(self.arch_paths, self.architectures): + for core in arch["cores"]: + core["activation_elements_per_cycle"] = 24 // len(arch["cores"]) + for field in ("vector_sram_bytes", "accumulator_bytes", "weight_sram_bytes"): + core[field] //= len(arch["cores"]) self.write(path, arch) self.reports = {arch["name"]: envelope(arch) for arch in self.architectures} self.write(self.root / "reports.json", self.reports) @@ -280,6 +284,32 @@ def test_wrong_numeric_result_is_rejected(self): self.mutate(["result", "output_f32", 0, 0], 7.0) self.rejected() + def test_exact_mode_accepts_identical_bf16_bits(self): + result = self.run_comparison(atol=0, rtol=0) + self.assertTrue(result["all_gates_passed"]) + self.assertTrue(all(gate["output_bit_exact"] for row in result["comparisons"] for gate in row["gates"])) + + def use_opposite_signed_zeros(self): + self.golden["output_f32"][0][0] = 0.0 + self.golden["output_bf16"][0][0] = 0 + self.write(self.golden_path, self.golden) + for report in self.reports.values(): + report["result"]["output_f32"][0][0] = -0.0 + report["result"]["output_bf16"][0][0] = 32768 + self.write(self.root / "reports.json", self.reports) + + def test_exact_mode_rejects_opposite_signed_zero_bits(self): + self.use_opposite_signed_zeros() + with self.assertRaisesRegex(ValueError, "identical BF16 bits"): + self.run_comparison(atol=0, rtol=0) + self.assertFalse(compare.read_json(self.output / "comparison.json")["all_gates_passed"]) + + def test_nonzero_tolerance_keeps_bit_difference_informational(self): + self.use_opposite_signed_zeros() + result = self.run_comparison(atol=1e-5, rtol=0) + self.assertTrue(result["all_gates_passed"]) + self.assertTrue(all(not gate["output_bit_exact"] for row in result["comparisons"] for gate in row["gates"])) + def test_parallel_architectures_keep_order_and_repeat_gates(self): result = self.run_comparison(workers=2) self.assertEqual([r["architecture"]["name"] for r in result["comparisons"]], ["single", "dual"]) @@ -358,6 +388,30 @@ def test_unequal_timing_modes_fail_before_execution(self): self.rejected() self.assertFalse((self.root / "calls.json").exists()) + def test_unequal_activation_supply_fails_before_execution(self): + self.architectures[1]["cores"][0]["activation_elements_per_cycle"] += 1 + self.write(self.arch_paths[1], self.architectures[1]) + with self.assertRaisesRegex(ValueError, "equal total activation supply"): + self.run_comparison() + self.assertFalse((self.root / "calls.json").exists()) + + def test_missing_and_null_activation_supply_use_mlen_default(self): + del self.architectures[1]["cores"][0]["activation_elements_per_cycle"] + self.architectures[1]["cores"][1]["activation_elements_per_cycle"] = None + self.write(self.arch_paths[1], self.architectures[1]) + self.assertTrue(self.run_comparison()["all_gates_passed"]) + + def test_unequal_sram_fails_even_when_all_cache_fields_are_absent(self): + original = copy.deepcopy(self.architectures[1]) + for field in ("vector_sram_bytes", "accumulator_bytes", "weight_sram_bytes"): + with self.subTest(field=field): + changed = copy.deepcopy(original) + changed["cores"][0][field] += 1 + self.write(self.arch_paths[1], changed) + with self.assertRaisesRegex(ValueError, "equal total configured SRAM: " + field): + self.run_comparison() + self.assertFalse((self.root / "calls.json").exists()) + def test_nonfinite_output_is_rejected(self): for value in (float("nan"), float("inf"), -float("inf")): with self.subTest(value=value): From fe6768aac2e427ae431d431b3c99adb164c5ae40 Mon Sep 17 00:00:00 2001 From: Michael C Li <63594253+Happymic@users.noreply.github.com> Date: Mon, 7 Sep 2026 02:13:48 +0100 Subject: [PATCH 5/5] docs(moe): specify architecture refinement and validation contracts --- doc/moe_architecture_refinement.md | 228 +++++++++++++++++++++++++++++ doc/moe_normal_review.md | 6 +- 2 files changed, 233 insertions(+), 1 deletion(-) create mode 100644 doc/moe_architecture_refinement.md diff --git a/doc/moe_architecture_refinement.md b/doc/moe_architecture_refinement.md new file mode 100644 index 000000000..6048d9963 --- /dev/null +++ b/doc/moe_architecture_refinement.md @@ -0,0 +1,228 @@ +# MoE architecture refinement: compute, SRAM, DMA and dispatch + +Status: **design proposal, not implemented or performance-validated** (2026-09-07). +The runnable reference is Simulator `bead42b2` and Compiler `c45cf7b`. +The [existing review guide](moe_normal_review.md) describes that implementation. +This document defines the missing contracts and the order for implementing them +in Compiler and Simulator. Passing a numerical smoke test does not freeze an +architecture, its port structure, or its arithmetic pipeline. + +## Goal and evidence boundary + +Determine whether heterogeneous compute shapes improve fixed-route MoE execution +when task partitioning, operand delivery, accumulation and shared memory service +are designed together. Hold workload, numerical contract and declared resources +constant; measure complete operator latency, not only occupied multipliers. + +The existing negative result applies to **whole-expert dispatch, square tiles, +fixed activation allocation and the current analytical timing rules**. It does +not test arbitrary splitting of one expert across cores. The performance cases +cover routed-expert gate/up/down; the independent shared-expert path is exercised +by functional tests. Initial input HBM transfers, runtime router execution and +output HBM stores are excluded from the present timing boundary and must remain +labelled as such. Input gather and activation supply are included. + +## 1. Separate the three meanings of matrix size + +| Meaning | Contract | +|---|---| +| Workload | `X[Mw,Kw] * W[Nw,Kw]^T -> Y[Mw,Nw]`; Mw is the number of routed token rows for this expert. | +| Scheduled tile | Explicit ranges `(m0, mt), (n0, nt), (k0, kt)`, with valid tails. | +| Compute organization | Modeled output lanes `P` and reduction width `R`, with `P*R` multipliers. Token rows are serviced over time. | + +Today nominal tile capacities are `Mt=Nt=P=BLEN` and `Kt=R=MLEN`; valid extents +`mt,nt,kt` may be smaller at tails. These equalities are a mapping restriction, +not the definition of MoE or proof that a three-dimensional multiplier array +exists. Gate/up use `(Mw,Nw,Kw)=(Me,F,D)`; down uses `(Me,D,F)`. +Compiler descriptors must retain full dimensions and ranges independently of +which core is eventually chosen. + +## 2. Make task partitioning an explicit experiment axis + +| Mode | Ownership and additional costs | +|---|---| +| Whole expert — implemented | One core owns an expert through gather, gate/up, SwiGLU, down and output copy. Dispatch may choose a different core for the next expert. | +| M partition — proposed later | Different token rows go to different cores. Outputs remain independent, but repeated weight reads, activation copies and lost weight reuse must be counted. | +| N partition — proposed later | Different output columns of a projection go to different cores. Gate/up partitions must align for SwiGLU; down needs the complete required Z features, so gathering/replication and barriers must be represented. | +| K partition — outside the next experiment | Requires partial-sum communication and a separately specified floating-point reduction order. It cannot reuse the current exact oracle without justification. | + +First refine the whole-expert reference so memory/scheduling changes can be +isolated. A claim about intra-expert splitting requires a separate executable +graph with the transfers above. A task is not migrated after execution starts +in the next implementation; pre-dispatch placement is a different decision. + +## 3. Give Compiler and runtime different responsibilities + +Compiler supplies validated tensor regions, token/route order, GEMM ranges, +element/scale format and logical producer-consumer dependencies. For a tile: + +``` +element address = element_base + output_row * element_stride + k0 +scale address = scale_base + output_row * scale_stride + k0 / block_size +``` + +The next contract keeps `k0` block-aligned and charges tail sectors explicitly. +Runtime selects a core, reserves physical storage, creates the concrete request +fragments and assigns destination addresses. The HBM command scheduler sees +memory requests; the task dispatcher decides which core executes an expert. +HBM request priority must not silently change that ownership. + +## 4. Specify storage lifetime before scheduling it + +A resident weight tile has the following proposed observable states: + +``` +FREE -> RESERVED -> FETCHING -> PACKED_READY -> DECODE_QUEUED + -> READY -> CONSUMING -> FREE +``` + +- Reserve the complete slot before issuing a request. Track element and scale + completion separately; both must cover the same valid reduction range. +- Decode only after the required packed bytes have been copied. Compute requires + decoded readiness, activation readiness and an available output context. +- A READY tile still owns its slot. Release it after the last admitted consumer + finishes reading/capturing the weight operands. Independently reserved result + and accumulator state may remain in flight until result completion. No silent + eviction, free reload or persistent weight cache. +- A destination is `(core, slot, packed-or-BF16, row, column)`, not just an HBM + source address. Copies to two destinations are charged twice. + +Existing weight SRAM budgets are 64 KiB single, or 48/16 KiB large/small. +Four slots actually occupy 50 KiB single, or 37.5/12.5 KiB large/small: +`slot_bytes = BLEN * (MLEN/8) * 25` includes packed E4M3, E8M0 and decoded BF16. +Configured capacity and occupied bytes must be reported separately. + +The complete expert accumulator data is already reserved. Introduce a bounded +`OutputContext(job, projection, m0, n0, next_k, ready_time, state)` table for active +outputs; explicitly charge its encoded bytes against the accumulator/control +budget. The current host `ready[Mb*Nb]` vector is not a complete metadata budget. +The first interleave experiment admits at most two N tiles per core; the context +count is bounded by those tiles and the admitted token blocks. Reject an expert +if data, pipeline storage and metadata together exceed capacity. + +## 5. Close the operand-port and accumulation contracts + +These are blocking architecture decisions, not parameters to hide behind a +successful functional test: + +| Path | Existing model | Required refinement | +|---|---|---| +| Activation -> compute | Aggregate per-core supply; single 1024 versus 768+256 BF16 elements/cycle. | Declare local bank/read limits and the destination interface. Fixed allocations remain a control; dynamic borrowing needs a real shared arbiter and transport budget. | +| Decoded weight -> compute | Values are directly consumed from a host vector. | Choose SRAM streaming or stationary operand latches; declare preload/read bandwidth, latch storage if present, and contention with decode writes. | +| Returned bytes -> SRAM | Copy bank selected by `(HBM source address / 64) % 8`. | Define destination bank mapping and ports; source-address striping is not a destination-SRAM conflict model. | +| Accumulator | Full data capacity plus analytical result-ready timestamps. | Explicit read/write service, result-queue capacity, forwarding rules and same-output feedback latency. | + +Do not add operand latches or crossbars without charging their storage and service. +If an access model cannot feed the claimed multiplier rate, report its sustainable +rate instead. The current small core takes 4 service cycles rather than its ideal +2 under `max(BLEN,ceil(BLEN*MLEN/activation_supply))`; changing its supply is a +named design change, not an HBM optimization. + +The current oracle performs ascending-global-K FP32 multiply/add, whereas the +timing model includes a `ceil(log2 MLEN)` term. That does not define the numerical +behavior of a parallel reduction tree. Preserve current arithmetic and feedback +latency for the first scheduling comparison. A chosen reduction structure must +subsequently specify rounding/order and get its own independent numerical check; +never shorten feedback latency or relax accuracy solely to obtain a speedup. + +## 6. Schedule independent outputs without breaking dependencies + +First compare the current N->K->M order with a two-N-tile active window. Only +issue a context when its weight/activation data are ready, its next K range is +correct, its accumulator feedback is ready, and all operand/result resources +have been reserved. Select the oldest legal context with deterministic ties. + +The four weight slots can hold independent N tiles rather than only successive K +tiles of one output. Preserve every scalar output's ascending K order and the +current full feedback delay in this comparison. When no context is legal, wait +for a real readiness event. Do not manufacture overlap by subtracting stalls. +Reserve/admit each active context's earliest unfinished K fragment before +speculative later-K fragments. Never fill every slot with future fragments while +the predecessor needed for progress has neither a resident nor a reserved slot. +Ready-context selection alone does not guarantee this progress invariant. + +## 7. Make the DMA frontend aware of tile demand + +The existing per-channel mutex removes cross-channel blocking, but holds its +position across native rejection retries. Element/scale fragments are initially +created in address order. This is not progress-aware arbitration. + +Proposed frontend rules: + +1. Keep non-issued fragments in already-budgeted resident-tile descriptors. + Each has core/tile identity, element-or-scale tag, destination and age. +2. Distinguish current demand from future prefetch. A fragment that completes a + demanded tile is useful whether it contains elements or scales; all scales + do not receive unconditional priority. An unissued coalesced line/sector + inherits the strongest current consumer demand and relevant age, so demand + merged into prefetch cannot remain at prefetch priority. Destination copies + remain separately charged; accepted requests are not preempted. +3. At each legal channel issue opportunity, choose among eligible fragments. + Demand wins over ordinary prefetch; use per-core byte fairness and bounded + age promotion to prevent starvation. Expose the age threshold as a parameter. +4. A rejected attempt remains queued and can be reconsidered at the next legal + opportunity. Accepted native requests are neither cancelled nor preempted. +5. Reserve response/destination space before submission. Drain returns without + requiring a new request credit, so request pressure cannot deadlock response + completion. Release each waiter/line/slot at its own defined lifetime boundary. + +Retain 128 logical line/waiter credits, 256 native 32-byte trackers and 44 KiB +total frontend storage, including staging. Coalescing remains in-flight only. +Current accounting reserves 32,512 B single or 32,832 B dual within that frontend +budget. Allocate and validate new queue heads, age fields and completion bitmaps +within the remaining capacity; do not use unbounded host queues as free hardware. +Keep the native HBM configuration and command scheduler fixed for this control. + +## 8. Refine core assignment only after the estimator is observable + +Keep the M-threshold/work-conserving policy as a baseline. A proposed alternative +compares estimated finish times using known shapes, tile counts/tails, current +core work, resident data, channel queues and shared vector service. It must use +only current state and completed-service statistics, never future native return +times or the measured answer of that same candidate run. + +Bound candidate scanning and account for scheduler service and descriptor space. +Record both candidate estimates, the choice and actual completion. Freeze any +fitted parameters before held-out windows. Treat this as a heuristic whose error +is measured, not an optimal scheduler. Do not change dispatch, DMA arbitration +and memory size simultaneously and attribute the result to one component. + +## 9. Instrument the evidence needed to distinguish causes + +Record tile milestones: slot reservation, last element copy, last scale copy, +decode queue/start/end, first compute, last consumer and release. Record request +credit wait, lookup, native admission, response and destination-copy service. +Separate vector service into gather/decode/SwiGLU/combine. Track per-core/channel +demand and prefetch bytes, queue peaks and oldest request age. + +Trace collection may stream to a bounded host sink; state used to make scheduling +decisions must remain inside declared modeled storage. Use per-core timelines +and final completion time: overlapping waits cannot be summed into latency. +The current `weight_ready_wait_ps` alone cannot identify a bad HBM controller. + +## 10. Implementation and acceptance order + +1. Add observation-only milestones and audit storage accounting; reproduce the + present configurations without changing issue order or latency rules. Retain + this version as the legacy timing control. +2. Instantiate the operand/accumulator ports and bounded result contexts from + section 5 as a separately labelled model. Check insufficient-port cases; + added service constraints may change latency. Do not mix that correction with + a scheduling speedup or require it to reproduce the legacy timing. +3. Implement bounded independent-output scheduling, preserving existing feedback + and arithmetic. Test one-context and multiple-context cases to distinguish + unavoidable dependencies from waiting that can legally overlap. +4. Implement demand-aware DMA separately; test delayed scale, same-channel + prefetch blocking, asymmetric consumers, full queues and starvation bounds. +5. Compare dispatch policies and activation allocations only after their inputs + and service constraints are explicit. Then consider M/N task partitions. +6. Search SRAM sizes and core shapes under the resulting contract, with the same + policy options available to single and dual cores. Transpose/Attention uses a + later access contract; adding it does not resolve the normal-path gaps above. + +For every stage: exact current-oracle checks, tails, deterministic repetition, +byte conservation and finite-capacity/backpressure assertions precede latency +ranking. Include all-small-M, mixed/hot experts, one dominant expert, demand +blocked by prefetch, and insufficient-port cases. Publish both wins and losses. +Make compact routing inputs and the fixture-generation recipe available before +claiming external reproduction of the complete performance table. diff --git a/doc/moe_normal_review.md b/doc/moe_normal_review.md index 44923ca33..baf11a84a 100644 --- a/doc/moe_normal_review.md +++ b/doc/moe_normal_review.md @@ -5,6 +5,10 @@ normal-buffer cores with private accumulators, shared HBM, and a fair single-cor comparison. It executes actual nonzero MoE arithmetic. The best tested large/small pair is **slower than the optimized single core**; this is not a speedup claim. +For the missing port, arithmetic, DMA and dispatch contracts, see the +[architecture refinement proposal](moe_architecture_refinement.md). Proposed +behavior there is not included in the measured results below. + ## Read the implementation in this order 1. `transactional_emulator/src/moe_normal/types.rs`: workload, per-core resources @@ -26,7 +30,7 @@ Paths in items 2–4 are relative to `transactional_emulator/`. ## Current dimensions -Physical tiles have M=N=BLEN and K=MLEN. A tile is serviced over multiple cycles; +Scheduled tiles have M=N=BLEN and K=MLEN. A tile is serviced over multiple cycles; multiplier count is BLEN×MLEN, not M×K×N. Whole-model D/F dimensions are separate. | Resource | Single baseline | Large core | Small core |