diff --git a/PLENA_Compiler b/PLENA_Compiler index ebdba9ef..54a68e69 160000 --- a/PLENA_Compiler +++ b/PLENA_Compiler @@ -1 +1 @@ -Subproject commit ebdba9efe60e2213e61556bfe83797720767c530 +Subproject commit 54a68e6969651eb5c1c3397ac59bed6585c08806 diff --git a/transactional_emulator/lib/memory/src/lib.rs b/transactional_emulator/lib/memory/src/lib.rs index 4b26e866..a44c7c55 100644 --- a/transactional_emulator/lib/memory/src/lib.rs +++ b/transactional_emulator/lib/memory/src/lib.rs @@ -44,12 +44,18 @@ pub trait MemoryModel: Send + Sync { /// Write 64-bytes of memory. async fn write(&self, addr: u64, bytes: [u8; 64]); + + /// Optional utilization statistics for wrappers that collect them. + fn statistics(&self) -> Option { + None + } } #[async_trait::async_trait] pub trait ErasedMemoryModel: Send + Sync { async fn box_read(&self, addr: u64) -> [u8; 64]; async fn box_write(&self, addr: u64, bytes: [u8; 64]); + fn statistics(&self) -> Option; } #[async_trait::async_trait] @@ -61,6 +67,10 @@ impl ErasedMemoryModel for T { async fn box_write(&self, addr: u64, bytes: [u8; 64]) { self.write(addr, bytes).await } + + fn statistics(&self) -> Option { + MemoryModel::statistics(self) + } } impl MemoryModel for dyn ErasedMemoryModel { @@ -71,6 +81,10 @@ impl MemoryModel for dyn ErasedMemoryModel { async fn write(&self, addr: u64, bytes: [u8; 64]) { self.box_write(addr, bytes).await } + + fn statistics(&self) -> Option { + ErasedMemoryModel::statistics(self) + } } /// A memory that discards all written data. @@ -206,6 +220,10 @@ impl MemoryModel for WithStats { } self.model.write(addr, bytes).await } + + fn statistics(&self) -> Option { + Some(WithStats::statistics(self)) + } } #[cfg(test)] diff --git a/transactional_emulator/lib/quantize/src/tensor.rs b/transactional_emulator/lib/quantize/src/tensor.rs index 072bdbda..45ee0a20 100644 --- a/transactional_emulator/lib/quantize/src/tensor.rs +++ b/transactional_emulator/lib/quantize/src/tensor.rs @@ -3,6 +3,34 @@ use tch::Tensor; use crate::dtype::{DataType, FpType, MxDataType}; +fn round_ties_to_even_nonnegative(x: f32) -> u32 { + debug_assert!(x >= 0.0); + let floor = x.floor(); + let frac = x - floor; + if frac > 0.5 { + (floor as u32) + 1 + } else if frac < 0.5 { + floor as u32 + } else { + let floor_int = floor as u32; + if floor_int & 1 == 0 { + floor_int + } else { + floor_int + 1 + } + } +} + +fn hardware_round_nonnegative(x: f32) -> u32 { + debug_assert!(x >= 0.0); + // Match PLENA_Tools/plena_quant/common/hardware_utils.py: + // x = floor(x * 2**round_bits) / 2**round_bits + // return x.round() + // with round_bits=2 and PyTorch round-to-even semantics. + let quarter_truncated = (x * 4.0).floor() / 4.0; + round_ties_to_even_nonnegative(quarter_truncated) +} + /// Quantize a single FP32 value to minifloat format using IEEE hardware quantization /// This matches the Python _minifloat_ieee_quantize_hardware function fn minifloat_ieee_quantize_hardware(value: f32, fp_type: FpType) -> u32 { @@ -47,12 +75,12 @@ fn minifloat_ieee_quantize_hardware(value: f32, fp_type: FpType) -> u32 { // Quantize mantissa let shift = 1u32 << mantissa_bits; let shifted_mantissa = if is_normal { - // Normal: (mantissa - 1) * shift, round, clamp - let shifted = ((mantissa - 1.0) * shift as f32).round() as u32; + // Normal: (mantissa - 1) * shift, hardware_round, clamp. + let shifted = hardware_round_nonnegative((mantissa - 1.0) * shift as f32); shifted.max(shifted_mantissa_min).min(shifted_mantissa_max) } else { - // Subnormal: mantissa * shift, round, clamp - let shifted = (mantissa * shift as f32).round() as u32; + // Subnormal: mantissa * shift, hardware_round, clamp + let shifted = hardware_round_nonnegative(mantissa * shift as f32); shifted.max(shifted_mantissa_min).min(shifted_mantissa_max) }; @@ -295,6 +323,21 @@ mod tests { assert_eq!(minifloat_ieee_quantize_hardware(-1.0, ty), 0xB8); // sign | 0x38 } + #[test] + fn test_minifloat_quantize_hardware_rounds_fractional_mantissa() { + let ty = e4m3(); + // E4M3 has 0.125 spacing in [1,2). PLENA's MX hardware quantizer + // first truncates the shifted mantissa to 2 fractional bits, then + // applies round-to-even. This is neither plain nearest rounding nor + // plain truncation. + assert_eq!(minifloat_ieee_quantize_hardware(1.0625, ty), 0x38); + assert_eq!(minifloat_ieee_quantize_hardware(1.0859375, ty), 0x38); + assert_eq!(minifloat_ieee_quantize_hardware(1.1875, ty), 0x3A); + assert_eq!(minifloat_ieee_quantize_hardware(1.328125, ty), 0x3A); + assert_eq!(minifloat_ieee_quantize_hardware(-1.0625, ty), 0xB8); + assert_eq!(minifloat_ieee_quantize_hardware(-1.1875, ty), 0xBA); + } + #[test] fn test_into_bytes_plain_packs_elements_and_no_scale_stream() { let elem = DataType::Fp(e4m3()); diff --git a/transactional_emulator/lib/sram/src/matrix.rs b/transactional_emulator/lib/sram/src/matrix.rs index 3cd98396..fca7dab4 100644 --- a/transactional_emulator/lib/sram/src/matrix.rs +++ b/transactional_emulator/lib/sram/src/matrix.rs @@ -14,6 +14,19 @@ pub struct MatrixSram { } impl MatrixSram { + fn tensor_to_f32_vec(tensor: &tch::Tensor) -> Vec { + let len = tensor.size1().unwrap() as usize; + let tensor_f32 = tensor.to_kind(tch::Kind::Float).contiguous(); + let mut data = vec![0.0f32; len]; + if tensor_f32.f_copy_data(&mut data, len).is_ok() { + return data; + } + for (idx, value) in data.iter_mut().enumerate() { + *value = tensor_f32.double_value(&[idx as i64]) as f32; + } + data + } + /// Create a matrix SRAM with given tile size and depth. pub fn new(tile_size: u32, depth: usize, ty: MxDataType) -> Self { let tiles = (0..(depth / tile_size as usize)) @@ -116,14 +129,13 @@ impl MatrixSram { let mut guard = tile_mutex.lock().await; let tensor = guard.resolve().await; let tensor_data = tensor.as_tensor(); - let len = tensor_data.size1().unwrap() as usize; - let f32_slice = - unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) }; + let f32_vec = Self::tensor_to_f32_vec(tensor_data); + let len = f32_vec.len(); // Calculate bytes needed for THIS tile's actual size let total_bits = len * element_ty.size_in_bits() as usize; let bytes_needed = (total_bits + 7) / 8; let mut tile_bytes = vec![0u8; bytes_needed]; - element_ty.bytes_from_f32(f32_slice, &mut tile_bytes); + element_ty.bytes_from_f32(&f32_vec, &mut tile_bytes); result.extend_from_slice(&tile_bytes); } diff --git a/transactional_emulator/lib/sram/src/vector.rs b/transactional_emulator/lib/sram/src/vector.rs index ba6a0ee7..591b4da7 100644 --- a/transactional_emulator/lib/sram/src/vector.rs +++ b/transactional_emulator/lib/sram/src/vector.rs @@ -28,6 +28,35 @@ pub struct VectorSram { } impl VectorSram { + fn tensor_to_f32_vec(tensor: &Tensor) -> Vec { + let len = tensor.size1().unwrap() as usize; + let tensor_f32 = tensor.to_kind(tch::Kind::Float).contiguous(); + let mut data = vec![0.0f32; len]; + if tensor_f32.f_copy_data(&mut data, len).is_ok() { + return data; + } + for (idx, value) in data.iter_mut().enumerate() { + *value = tensor_f32.double_value(&[idx as i64]) as f32; + } + data + } + + fn tensor_from_f32_slice(data: &[f32]) -> Tensor { + if data.is_empty() { + return Tensor::zeros([0], (tch::Kind::Float, tch::Device::Cpu)); + } + unsafe { + Tensor::from_blob( + data.as_ptr() as *const u8, + &[data.len() as i64], + &[], + tch::Kind::Float, + tch::Device::Cpu, + ) + .internal_to_copy((tch::Kind::Float, tch::Device::Cpu), false) + } + } + fn row_width_bytes(vlen: u32, fp_type: DataType) -> usize { (vlen as usize * fp_type.size_in_bits() as usize).div_ceil(8) } @@ -163,11 +192,7 @@ impl VectorSram { let tensor_data = tensor.as_tensor(); let total_elements = tensor_data.size1().unwrap() as usize; - // Extract f32 data from tensor to make it Send-safe - let len = total_elements; - let f32_slice = - unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) }; - let data_vec: Vec = f32_slice.to_vec(); + let data_vec = Self::tensor_to_f32_vec(tensor_data); let chunk_size = self.vlen as usize; let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); @@ -188,7 +213,7 @@ impl VectorSram { padded_data[..chunk_len].copy_from_slice(chunk_data); // Create tensor from padded data and convert to bytes - let padded_tensor = Tensor::from_slice(&padded_data); + let padded_tensor = Self::tensor_from_f32_slice(&padded_data); let chunk_qt = QuantTensor::quantize(padded_tensor, MxDataType::Plain(self.fp_type)); let row_bytes = self.quant_tensor_to_bytes(&chunk_qt); @@ -262,7 +287,7 @@ impl VectorSram { } // Create QuantTensor and convert to bytes - let tensor = Tensor::from_slice(&vec); + let tensor = Self::tensor_from_f32_slice(&vec); let quant_tensor = QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type)); let row_bytes = self.quant_tensor_to_bytes(&quant_tensor); *self.rows[row_idx].lock().await = Cell::Ready(row_bytes); @@ -305,14 +330,13 @@ impl VectorSram { /// Convert QuantTensor to bytes (FP format) fn quant_tensor_to_bytes(&self, tensor: &QuantTensor) -> Vec { let tensor_data = tensor.as_tensor(); - let len = tensor_data.size1().unwrap() as usize; - let f32_slice = - unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) }; + let f32_vec = Self::tensor_to_f32_vec(tensor_data); + let len = f32_vec.len(); let total_bits = len * self.fp_type.size_in_bits() as usize; let bytes_needed = (total_bits + 7) / 8; let mut bytes = vec![0u8; bytes_needed]; - self.fp_type.bytes_from_f32(f32_slice, &mut bytes); + self.fp_type.bytes_from_f32(&f32_vec, &mut bytes); bytes } @@ -330,7 +354,7 @@ impl VectorSram { vec.resize(expected_len as usize, 0.0f32); } - let tensor = Tensor::from_slice(&vec); + let tensor = Self::tensor_from_f32_slice(&vec); QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type)) } diff --git a/transactional_emulator/src/accelerator/dispatch.rs b/transactional_emulator/src/accelerator/dispatch.rs index 6bb4cbb7..b5401947 100644 --- a/transactional_emulator/src/accelerator/dispatch.rs +++ b/transactional_emulator/src/accelerator/dispatch.rs @@ -11,7 +11,9 @@ use crate::runtime_config::{ SCALAR_FP_BASIC_CYCLES, SCALAR_FP_EXP_CYCLES, SCALAR_FP_RECI_CYCLES, SCALAR_FP_SQRT_CYCLES, SCALAR_INT_BASIC_CYCLES, STORE_V_AMOUNT, VECTOR_ACTIVATION_TYPE, VECTOR_KV_TYPE, VLEN, }; +use crate::stage_profile::StageProfiler; use crate::{cycle, dma, op}; +use runtime::Executor; use super::Accelerator; use super::loop_state::LoopDecision; @@ -48,11 +50,26 @@ impl Accelerator { } } - pub(crate) async fn do_ops(&mut self, ops: &[op::Opcode]) { + pub(crate) async fn do_ops( + &mut self, + ops: &[op::Opcode], + mut stage_profiler: Option<&mut StageProfiler>, + ) { let mut pc: usize = 0; // Program counter while pc < ops.len() { + let executed_pc = pc; let op = &ops[pc]; + let profile_start_secs = if stage_profiler.is_some() { + Some(Executor::current().now().to_secs()) + } else { + None + }; + let profile_start_hbm = if stage_profiler.is_some() { + self.hbm.statistics() + } else { + None + }; self.loop_state.record_instruction(); @@ -252,6 +269,64 @@ impl Accelerator { ) .await; } + op::Opcode::V_MAX_VF { + rd, + rs1, + rs2, + rmask, + } => { + let mask = self.resolve_v_mask(*rmask); + self.v_machine + .max_scalar( + self.reg_file.read_gp(*rd), + self.reg_file.read_gp(*rs1), + self.reg_file.read_fp(*rs2).into(), + *rmask, + mask, + ) + .await; + } + op::Opcode::V_MIN_VF { + rd, + rs1, + rs2, + rmask, + } => { + let mask = self.resolve_v_mask(*rmask); + self.v_machine + .min_scalar( + self.reg_file.read_gp(*rd), + self.reg_file.read_gp(*rs1), + self.reg_file.read_fp(*rs2).into(), + *rmask, + mask, + ) + .await; + } + op::Opcode::V_TOPK { + rd, + rs1, + rs2, + rmask, + } => { + let (expert_count, topk) = match *rmask { + 0 => (32, 4), + 1 => (128, 8), + other => panic!( + "unsupported V_TOPK rmask policy {other}; expected 0=32/top4 or 1=128/top8" + ), + }; + let fp_base = self.reg_file.read_gp(*rd) as usize; + let int_base = self.reg_file.read_gp(*rs2) as usize; + let (indices, weights) = self + .v_machine + .topk_softmax(self.reg_file.read_gp(*rs1), expert_count, topk) + .await; + for (offset, (idx, weight)) in indices.iter().zip(weights.iter()).enumerate() { + self.scalar_sram.write_int(int_base + offset, *idx); + self.scalar_sram.write_fp(fp_base + offset, *weight); + } + } op::Opcode::V_EXP_V { rd, rs1, rmask } => { let mask = self.resolve_v_mask(*rmask); self.v_machine @@ -274,7 +349,7 @@ impl Accelerator { ) .await; } - op::Opcode::V_SHIFT_V { rd, rs1, rs2 } => { + op::Opcode::V_SHFT_V { rd, rs1, rs2 } => { self.v_machine .shift_scalar( self.reg_file.read_gp(*rd), @@ -549,6 +624,27 @@ impl Accelerator { } else { pc += 1; } + + if let (Some(start_secs), Some(profiler)) = + (profile_start_secs, stage_profiler.as_deref_mut()) + { + let elapsed_secs = Executor::current().now().to_secs() - start_secs; + let (hbm_bytes_read, hbm_bytes_written) = if let (Some(before), Some(after)) = + (profile_start_hbm, self.hbm.statistics()) + { + ( + after + .total_bytes_read + .saturating_sub(before.total_bytes_read), + after + .total_bytes_written + .saturating_sub(before.total_bytes_written), + ) + } else { + (0, 0) + }; + profiler.record(executed_pc, elapsed_secs, hbm_bytes_read, hbm_bytes_written); + } } } } diff --git a/transactional_emulator/src/accelerator/mod.rs b/transactional_emulator/src/accelerator/mod.rs index 6de4d755..f920f92c 100644 --- a/transactional_emulator/src/accelerator/mod.rs +++ b/transactional_emulator/src/accelerator/mod.rs @@ -83,4 +83,8 @@ impl Accelerator { pub(crate) fn fpsram_dump_bytes(&self) -> Vec { self.scalar_sram.fpsram_to_le_bytes() } + + pub(crate) fn intsram_dump_bytes(&self) -> Vec { + self.scalar_sram.intsram_to_le_bytes() + } } diff --git a/transactional_emulator/src/accelerator/scalar_sram.rs b/transactional_emulator/src/accelerator/scalar_sram.rs index 8e959e8e..4c4e1ae5 100644 --- a/transactional_emulator/src/accelerator/scalar_sram.rs +++ b/transactional_emulator/src/accelerator/scalar_sram.rs @@ -53,6 +53,10 @@ impl ScalarSram { pub(super) fn fpsram_to_le_bytes(&self) -> Vec { self.fpsram.iter().flat_map(|f| f.to_le_bytes()).collect() } + + pub(super) fn intsram_to_le_bytes(&self) -> Vec { + self.intsram.iter().flat_map(|v| v.to_le_bytes()).collect() + } } fn decode_fpsram_f16_bytes(bytes: &[u8]) -> Vec { diff --git a/transactional_emulator/src/cli.rs b/transactional_emulator/src/cli.rs index 74774565..6c2eabe7 100644 --- a/transactional_emulator/src/cli.rs +++ b/transactional_emulator/src/cli.rs @@ -164,4 +164,13 @@ pub(crate) struct Opts { /// Path to plena_settings.toml. Overrides PLENA_SETTINGS_TOML env var and /// the default ../plena_settings.toml lookup. pub(crate) settings: Option, + + #[arg(long)] + /// Optional generated ASM source used to derive PC-to-stage labels for a + /// runtime stage profile. This is diagnostic only; normal runs omit it. + pub(crate) stage_profile_asm: Option, + + #[arg(long)] + /// Optional JSON output path for runtime stage profile results. + pub(crate) stage_profile_out: Option, } diff --git a/transactional_emulator/src/dma.rs b/transactional_emulator/src/dma.rs index 445197c3..c150f8d7 100644 --- a/transactional_emulator/src/dma.rs +++ b/transactional_emulator/src/dma.rs @@ -23,6 +23,22 @@ use runtime::Executor; use sram::VectorSram; use tokio::sync::oneshot::{self, Receiver}; +fn tensor_from_f32_slice(data: &[f32]) -> tch::Tensor { + if data.is_empty() { + return tch::Tensor::zeros([0], (tch::Kind::Float, tch::Device::Cpu)); + } + unsafe { + tch::Tensor::from_blob( + data.as_ptr() as *const u8, + &[data.len() as i64], + &[], + tch::Kind::Float, + tch::Device::Cpu, + ) + .internal_to_copy((tch::Kind::Float, tch::Device::Cpu), false) + } +} + /// Derived byte-layout for one MX transfer iteration. /// /// Computed identically for both transfer directions from the HBM data type, @@ -261,7 +277,7 @@ pub(crate) fn transfer_mx_from_hbm( } } - let tensor = tch::Tensor::from_slice(&vec); + let tensor = tensor_from_f32_slice(&vec); all_results.push(QuantTensor::quantize(tensor, sram_type)); } diff --git a/transactional_emulator/src/main.rs b/transactional_emulator/src/main.rs index 0f413e43..8d732704 100644 --- a/transactional_emulator/src/main.rs +++ b/transactional_emulator/src/main.rs @@ -2,10 +2,12 @@ mod accelerator; mod cli; mod dma; mod load_config; +mod matrix_core; mod matrix_machine; mod op; mod runner; mod runtime_config; +mod stage_profile; mod vector_machine; use runtime::{Executor, Instant}; diff --git a/transactional_emulator/src/matrix_core.rs b/transactional_emulator/src/matrix_core.rs new file mode 100644 index 00000000..311c28b4 --- /dev/null +++ b/transactional_emulator/src/matrix_core.rs @@ -0,0 +1,131 @@ +//! Matrix compute-resource timing profiles. +//! +//! Stage-2 paper experiments need to model more than one matrix engine. This +//! module keeps the existing numerical MatrixMachine path intact while making +//! the compute resource explicit: each MatrixMachine owns one MatrixCore, and +//! independent MatrixCores can make progress concurrently on the runtime +//! executor. + +use crate::runtime_config::PERIOD; + +/// Static geometry/timing metadata for one matrix compute engine. +#[derive(Debug, Clone, Copy, PartialEq)] +pub(crate) struct MatrixCoreProfile { + pub(crate) name: &'static str, + pub(crate) rows: u32, + pub(crate) cols: u32, + pub(crate) cycle_multiplier_num: u32, + pub(crate) cycle_multiplier_den: u32, +} + +impl MatrixCoreProfile { + pub(crate) const fn big_default() -> Self { + Self { + name: "big_4x1024", + rows: 4, + cols: 1024, + cycle_multiplier_num: 1, + cycle_multiplier_den: 1, + } + } + + pub(crate) const fn tail_4x256() -> Self { + Self { + name: "tail_4x256", + rows: 4, + cols: 256, + cycle_multiplier_num: 1, + cycle_multiplier_den: 1, + } + } + + pub(crate) fn pe_count(self) -> u32 { + self.rows * self.cols + } + + pub(crate) fn pe_area_fraction_vs(self, baseline: Self) -> f64 { + self.pe_count() as f64 / baseline.pe_count() as f64 + } + + fn scale_cycles(self, cycles: u32) -> u32 { + let num = cycles as u64 * self.cycle_multiplier_num as u64; + let den = self.cycle_multiplier_den.max(1) as u64; + num.div_ceil(den) as u32 + } +} + +/// One matrix compute resource. +#[derive(Debug, Clone, Copy)] +pub(crate) struct MatrixCore { + profile: MatrixCoreProfile, +} + +impl MatrixCore { + pub(crate) const fn new(profile: MatrixCoreProfile) -> Self { + Self { profile } + } + + pub(crate) const fn profile(&self) -> MatrixCoreProfile { + self.profile + } + + pub(crate) async fn compute(&self, cycles: u32) { + runtime::Executor::current() + .resolve_at(PERIOD * self.profile.scale_cycles(cycles)) + .await; + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use runtime::{Duration, Executor, Instant}; + + use super::*; + + #[tokio::test] + async fn independent_matrix_cores_overlap_in_sim_time() { + let executor = Executor::new(); + let observed = Arc::new(Mutex::new(Vec::new())); + + let big = MatrixCore::new(MatrixCoreProfile::big_default()); + let small = MatrixCore::new(MatrixCoreProfile::tail_4x256()); + + let observed_big = observed.clone(); + executor.spawn(async move { + big.compute(100).await; + observed_big.lock().unwrap().push(Executor::current().now()); + }); + + let observed_small = observed.clone(); + executor.spawn(async move { + small.compute(80).await; + observed_small + .lock() + .unwrap() + .push(Executor::current().now()); + }); + + executor.enter(Instant::ETERNITY).await; + + assert_eq!(executor.now(), Instant::INIT + Duration::from_nanos(100)); + assert_eq!( + observed.lock().unwrap().as_slice(), + &[ + Instant::INIT + Duration::from_nanos(80), + Instant::INIT + Duration::from_nanos(100), + ] + ); + } + + #[test] + fn tail_core_pe_area_fraction_is_explicit() { + let big = MatrixCoreProfile::big_default(); + let tail = MatrixCoreProfile::tail_4x256(); + + assert_eq!(big.pe_count(), 4096); + assert_eq!(tail.pe_count(), 1024); + assert_eq!(tail.pe_area_fraction_vs(big), 0.25); + } +} diff --git a/transactional_emulator/src/matrix_machine.rs b/transactional_emulator/src/matrix_machine.rs index c9ff5fb8..f6a01871 100644 --- a/transactional_emulator/src/matrix_machine.rs +++ b/transactional_emulator/src/matrix_machine.rs @@ -18,7 +18,7 @@ use quantize::QuantTensor; use sram::{MatrixSram, VectorSram, assert_multiple_of, multiple_and_offset}; use tch::{IndexOp, Tensor}; -use crate::cycle; +use crate::matrix_core::{MatrixCore, MatrixCoreProfile}; use crate::runtime_config::SYSTOLIC_PROCESSING_OVERHEAD; /// Tensor allocation options used for every accumulator buffer (f32 on CPU). @@ -38,6 +38,7 @@ pub(crate) struct MatrixMachine { hlen: u32, blen: u32, broadcast_amount: u32, + core: MatrixCore, } impl MatrixMachine { @@ -53,6 +54,26 @@ impl MatrixMachine { hlen: u32, blen: u32, broadcast_amount: u32, + ) -> Self { + Self::new_with_core( + mram, + vram, + mlen, + hlen, + blen, + broadcast_amount, + MatrixCoreProfile::big_default(), + ) + } + + pub(crate) fn new_with_core( + mram: Arc, + vram: Arc, + mlen: u32, + hlen: u32, + blen: u32, + broadcast_amount: u32, + core_profile: MatrixCoreProfile, ) -> Self { Self { m_accum: Tensor::zeros([blen as i64, blen as i64], ACCUM_OPTS), @@ -68,9 +89,18 @@ impl MatrixMachine { hlen, blen, broadcast_amount, + core: MatrixCore::new(core_profile), } } + pub(crate) fn core_profile(&self) -> MatrixCoreProfile { + self.core.profile() + } + + fn core(&self) -> MatrixCore { + self.core + } + pub(crate) async fn mm(&mut self, m_addr: u32, v_addr: u32) { let (mat_base, mat_offset) = multiple_and_offset(m_addr, self.mlen * self.mlen); assert!(mat_offset.is_multiple_of(self.blen)); @@ -86,7 +116,9 @@ impl MatrixMachine { mat_offset as i64..(mat_offset as i64 + self.blen as i64), )); let mut tensors = Vec::with_capacity(self.blen as usize); - cycle!(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen); + self.core() + .compute(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen) + .await; for i in 0..self.blen { tensors.push( self.vram @@ -125,7 +157,9 @@ impl MatrixMachine { )); let mut tensors = Vec::with_capacity(self.mlen as usize); - cycle!(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen); + self.core() + .compute(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen) + .await; for i in 0..self.mlen { tensors.push( self.vram @@ -183,7 +217,7 @@ impl MatrixMachine { // For bmv, only read 1 vector (not mlen like bmm) let mut tensors = Vec::with_capacity(1); - cycle!(*SYSTOLIC_PROCESSING_OVERHEAD + 1); + self.core().compute(*SYSTOLIC_PROCESSING_OVERHEAD + 1).await; for i in 0..1 { tensors.push( self.vram @@ -241,7 +275,9 @@ impl MatrixMachine { )); let mut tensors = Vec::with_capacity(self.mlen as usize); - cycle!(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen); + self.core() + .compute(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen) + .await; // B, S, H, D for i in 0..self.mlen { tensors.push( @@ -306,7 +342,7 @@ impl MatrixMachine { // For btmv, only read 1 vector (not mlen like btmm) let mut tensors = Vec::with_capacity(1); - cycle!(*SYSTOLIC_PROCESSING_OVERHEAD + 1); + self.core().compute(*SYSTOLIC_PROCESSING_OVERHEAD + 1).await; // B, S, H, D - only 1 query token for decode for i in 0..1 { tensors.push( @@ -361,7 +397,9 @@ impl MatrixMachine { .transpose(-1, -2) .i((.., mat_offset as i64..(mat_offset + self.blen) as i64)); let mut tensors = Vec::with_capacity(self.blen as usize); - cycle!(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen); + self.core() + .compute(*SYSTOLIC_PROCESSING_OVERHEAD + self.mlen) + .await; for i in 0..self.blen { tensors.push( self.vram @@ -383,13 +421,13 @@ impl MatrixMachine { pub(crate) async fn mm_wo(&mut self, v_addr: u32, stride_len: u32) { let (vec_base, vec_offset) = multiple_and_offset(v_addr, self.mlen); assert!(vec_offset.is_multiple_of(self.blen)); - cycle!(1); + self.core().compute(1).await; for i in 0..self.blen { - let tensor = self.m_accum.i((i as i64, ..)); + let tensor = self.m_accum.select_copy(0, i as i64).contiguous(); let old = self.vram.read(vec_base + i * self.mlen * stride_len).await; - let new = old.as_tensor().copy(); - new.i(vec_offset as i64..(vec_offset + self.blen) as i64) - .copy_(&tensor); + let new = old.as_tensor().contiguous(); + let mut slot = new.i(vec_offset as i64..(vec_offset + self.blen) as i64); + slot.copy_(&tensor); self.vram .write( vec_base + i * self.mlen * stride_len, @@ -404,10 +442,14 @@ impl MatrixMachine { pub(crate) async fn bmm_wo(&mut self, v_addr: u32) { let (vec_base, vec_offset) = multiple_and_offset(v_addr, self.mlen); assert!(vec_offset.is_multiple_of(self.mlen)); - cycle!(1); + self.core().compute(1).await; for j in 0..self.broadcast_amount { for i in 0..self.mlen { - let tensor = self.hm_accum.i((j as i64, i as i64, ..)); + let tensor = self + .hm_accum + .select_copy(0, j as i64) + .select_copy(0, i as i64) + .contiguous(); self.vram .write( vec_base + (j * self.mlen + i) * self.mlen, @@ -429,9 +471,9 @@ impl MatrixMachine { pub(crate) async fn bmv_wo(&mut self, v_addr: u32) { let (vec_base, vec_offset) = multiple_and_offset(v_addr, self.mlen); assert!(vec_offset.is_multiple_of(self.mlen)); - cycle!(1); + self.core().compute(1).await; for j in 0..self.broadcast_amount { - let tensor = self.hv_accum.i((j as i64, ..)); + let tensor = self.hv_accum.select_copy(0, j as i64).contiguous(); self.vram .write( vec_base + (j * self.mlen), @@ -453,7 +495,7 @@ impl MatrixMachine { let mat = self.mram.read(mat_base).await; let vec = self.vram.read(v_addr).await; - cycle!(self.mlen); + self.core().compute(self.mlen).await; // vec @ mat: [1, mlen] @ [mlen, mlen] = [1, mlen], then squeeze // Convert to float32 before matmul to match PyTorch golden reference let vec_f32 = vec.as_tensor().unsqueeze(0).to_kind(tch::Kind::Float); @@ -478,7 +520,7 @@ impl MatrixMachine { assert!(mat_offset.is_multiple_of(self.blen)); let mat = self.mram.read(m_addr).await; let vec = self.vram.read(v_addr).await; - cycle!(self.mlen); + self.core().compute(self.mlen).await; // vec @ transpose(mat): [1, mlen] @ [mlen, mlen] = [1, mlen], then squeeze // Convert to float32 before matmul to match PyTorch golden reference let vec_f32 = vec.as_tensor().unsqueeze(0).to_kind(tch::Kind::Float); @@ -498,14 +540,109 @@ impl MatrixMachine { pub(crate) async fn mv_wo(&mut self, v_addr: u32) { let (vec_base, vec_offset) = multiple_and_offset(v_addr, self.mlen); assert!(vec_offset.is_multiple_of(self.blen)); - cycle!(1); + self.core().compute(1).await; let old = self.vram.read(vec_base).await; - let new = old.as_tensor().copy(); - new.i(vec_offset as i64..(vec_offset + self.blen) as i64) - .copy_(&self.v_accum); + let new = old.as_tensor().contiguous(); + let source = self.v_accum.contiguous(); + let mut slot = new.i(vec_offset as i64..(vec_offset + self.blen) as i64); + slot.copy_(&source); self.vram .write(vec_base, QuantTensor::quantize(new, old.data_type())) .await; self.v_accum = Tensor::zeros([self.blen as i64], ACCUM_OPTS); } } + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use quantize::{DataType, FpType, MxDataType, QuantTensor}; + use runtime::{Duration, Executor, Instant}; + use sram::{MatrixSram, VectorSram}; + use tch::Tensor; + + use super::*; + use crate::matrix_core::MatrixCoreProfile; + use crate::runtime_config::SYSTOLIC_PROCESSING_OVERHEAD; + + fn bf16_plain() -> MxDataType { + MxDataType::Plain(DataType::Fp(FpType::BF16)) + } + + fn quant(vals: &[f32]) -> QuantTensor { + QuantTensor::quantize(Tensor::from_slice(vals), bf16_plain()) + } + + async fn make_machine(output_base: u32) -> (MatrixMachine, Arc, u32) { + let mram = Arc::new(MatrixSram::new(4, 64, bf16_plain())); + let vram = Arc::new(VectorSram::from_mx_type(4, 64, bf16_plain())); + + mram.write( + 0, + quant(&[ + 1.0, 0.0, 0.0, 0.0, // + 0.0, 1.0, 0.0, 0.0, // + 0.0, 0.0, 1.0, 0.0, // + 0.0, 0.0, 0.0, 1.0, + ]), + ) + .await; + vram.write(0, quant(&[1.0, 2.0, 3.0, 4.0])).await; + vram.write(4, quant(&[5.0, 6.0, 7.0, 8.0])).await; + + ( + MatrixMachine::new_with_core( + mram, + vram.clone(), + 4, + 2, + 2, + 2, + MatrixCoreProfile::big_default(), + ), + vram, + output_base, + ) + } + + #[tokio::test] + async fn independent_matrix_machines_overlap_actual_gemm_work() { + let executor = Executor::new(); + let observed = Arc::new(Mutex::new(Vec::new())); + + let (mut machine_a, vram_a, out_a) = make_machine(8).await; + let (mut machine_b, vram_b, out_b) = make_machine(16).await; + + let observed_a = observed.clone(); + executor.spawn(async move { + machine_a.mm(0, 0).await; + machine_a.mm_wo(out_a, 1).await; + observed_a.lock().unwrap().push(Executor::current().now()); + }); + + let observed_b = observed.clone(); + executor.spawn(async move { + machine_b.mm(0, 0).await; + machine_b.mm_wo(out_b, 1).await; + observed_b.lock().unwrap().push(Executor::current().now()); + }); + + executor.enter(Instant::ETERNITY).await; + + let expected_cycles = *SYSTOLIC_PROCESSING_OVERHEAD + 4 + 1; + let expected = Instant::INIT + Duration::from_nanos(expected_cycles as u64); + assert_eq!(executor.now(), expected); + assert_eq!(observed.lock().unwrap().as_slice(), &[expected, expected]); + + let a0 = vram_a.read(8).await.as_tensor().shallow_clone(); + let a1 = vram_a.read(12).await.as_tensor().shallow_clone(); + let b0 = vram_b.read(16).await.as_tensor().shallow_clone(); + let b1 = vram_b.read(20).await.as_tensor().shallow_clone(); + + assert!(a0.equal(&Tensor::from_slice(&[1.0f32, 2.0, 0.0, 0.0]))); + assert!(a1.equal(&Tensor::from_slice(&[5.0f32, 6.0, 0.0, 0.0]))); + assert!(b0.equal(&Tensor::from_slice(&[1.0f32, 2.0, 0.0, 0.0]))); + assert!(b1.equal(&Tensor::from_slice(&[5.0f32, 6.0, 0.0, 0.0]))); + } +} diff --git a/transactional_emulator/src/op.rs b/transactional_emulator/src/op.rs index 6c9394c1..595d958e 100644 --- a/transactional_emulator/src/op.rs +++ b/transactional_emulator/src/op.rs @@ -108,6 +108,38 @@ pub enum Opcode { rs2: u8, rmask: u8, }, + V_MAX_VF { + rd: u8, + rs1: u8, + rs2: u8, + rmask: u8, + }, + V_MIN_VF { + rd: u8, + rs1: u8, + rs2: u8, + rmask: u8, + }, + /// Routed-MoE router helper. + /// + /// Encoding follows the regular vector register form: + /// - `rs1`: VRAM row containing router logits. + /// - `rd`: GP register whose value is the FP SRAM base for selected route + /// weights. Policy 0 stores GPT-OSS 32-way/top-4 weights. Policy 1 + /// stores Qwen 128-way/top-8 weights. Qwen computes full-expert softmax + /// before top-k in HF, but `norm_topk_prob=true` renormalizes the selected + /// entries, making the final route weights equivalent to selected-logit + /// softmax. + /// - `rs2`: GP register whose value is the INT SRAM base for the selected + /// expert indices. + /// - `rmask`: policy selector (`0` = 32 experts/top-4, `1` = 128 + /// experts/top-8). + V_TOPK { + rd: u8, + rs1: u8, + rs2: u8, + rmask: u8, + }, V_EXP_V { rd: u8, rs1: u8, @@ -255,7 +287,7 @@ pub enum Opcode { rd: u8, }, // Extensions - V_SHIFT_V { + V_SHFT_V { rd: u8, rs1: u8, rs2: u8, @@ -404,6 +436,24 @@ impl Opcode { rs1, rmask: rs3, }, + 0x35 => Self::V_MAX_VF { + rd, + rs1, + rs2, + rmask: rs3, + }, + 0x36 => Self::V_MIN_VF { + rd, + rs1, + rs2, + rmask: rs3, + }, + 0x37 => Self::V_TOPK { + rd, + rs1, + rs2, + rmask: rs3, + }, // Scalar Operations (Floating-Point) 0x17 => Self::S_ADD_FP { rd, rs1, rs2 }, @@ -456,8 +506,8 @@ impl Opcode { 0x2E => Self::C_SET_V_MASK_REG { rd }, 0x2F => Self::C_LOOP_START { rd, imm }, 0x30 => Self::C_LOOP_END { rd }, - 0x31 => Self::V_SHIFT_V { rd, rs1, rs2 }, - 0x32 => Self::C_BREAK, + 0x32 => Self::V_SHFT_V { rd, rs1, rs2 }, + 0x34 => Self::C_BREAK, _ => { tracing::error!("Unknown opcode {opcode:#x}"); Self::Invalid @@ -502,7 +552,7 @@ mod tests { #[test] fn test_decode_invalid_and_unknown_are_invalid() { assert!(matches!(Opcode::decode(0x00), Opcode::Invalid)); - // 0x3F is past the highest defined opcode (0x32). + // 0x3F is past the highest defined opcode (0x34). assert!(matches!(Opcode::decode(0x3F), Opcode::Invalid)); } @@ -710,14 +760,14 @@ mod tests { #[test] fn test_decode_break_is_unit() { - assert!(matches!(Opcode::decode(0x32), Opcode::C_BREAK)); + assert!(matches!(Opcode::decode(0x34), Opcode::C_BREAK)); } #[test] - fn test_decode_v_shift_v() { - match Opcode::decode(rform(0x31, 1, 2, 3, 0, 0)) { - Opcode::V_SHIFT_V { rd, rs1, rs2 } => assert_eq!((rd, rs1, rs2), (1, 2, 3)), - other => panic!("expected V_SHIFT_V, got {other:?}"), + fn test_decode_v_shft_v() { + match Opcode::decode(rform(0x32, 1, 2, 3, 0, 0)) { + Opcode::V_SHFT_V { rd, rs1, rs2 } => assert_eq!((rd, rs1, rs2), (1, 2, 3)), + other => panic!("expected V_SHFT_V, got {other:?}"), } } @@ -765,4 +815,41 @@ mod tests { other => panic!("expected V_ADD_VV, got {other:?}"), } } + + #[test] + fn test_decode_vector_scalar_minmax() { + match Opcode::decode(rform(0x35, 1, 2, 3, 4, 0)) { + Opcode::V_MAX_VF { + rd, + rs1, + rs2, + rmask, + } => { + assert_eq!((rd, rs1, rs2, rmask), (1, 2, 3, 4)); + } + other => panic!("expected V_MAX_VF, got {other:?}"), + } + match Opcode::decode(rform(0x36, 5, 6, 7, 8, 0)) { + Opcode::V_MIN_VF { + rd, + rs1, + rs2, + rmask, + } => { + assert_eq!((rd, rs1, rs2, rmask), (5, 6, 7, 8)); + } + other => panic!("expected V_MIN_VF, got {other:?}"), + } + match Opcode::decode(rform(0x37, 9, 10, 11, 12, 0)) { + Opcode::V_TOPK { + rd, + rs1, + rs2, + rmask, + } => { + assert_eq!((rd, rs1, rs2, rmask), (9, 10, 11, 12)); + } + other => panic!("expected V_TOPK, got {other:?}"), + } + } } diff --git a/transactional_emulator/src/runner.rs b/transactional_emulator/src/runner.rs index ef522341..b3233ae9 100644 --- a/transactional_emulator/src/runner.rs +++ b/transactional_emulator/src/runner.rs @@ -8,12 +8,14 @@ use tracing_subscriber::prelude::*; use crate::accelerator::Accelerator; use crate::cli::{Opts, Parser}; +use crate::matrix_core::MatrixCoreProfile; use crate::matrix_machine::MatrixMachine; use crate::runtime_config::{ BLEN, BROADCAST_AMOUNT, HBM_SIZE, HLEN, MATRIX_SRAM_SIZE, MATRIX_SRAM_TYPE, MAX_LOOP_INSTRUCTIONS, MLEN, PREFETCH_M_AMOUNT, PREFETCH_V_AMOUNT, STORE_V_AMOUNT, VECTOR_SRAM_SIZE, VECTOR_SRAM_TYPE, VLEN, }; +use crate::stage_profile::StageProfiler; use crate::vector_machine::VectorMachine; use crate::{cli, op}; @@ -117,6 +119,20 @@ pub(crate) async fn run_from_cli() { )); // Vector SRAM let m_machine = MatrixMachine::new(mram, vram.clone(), *MLEN, *HLEN, *BLEN, *BROADCAST_AMOUNT); + let big_core = m_machine.core_profile(); + let tail_core = MatrixCoreProfile::tail_4x256(); + tracing::info!( + big_core = big_core.name, + big_rows = big_core.rows, + big_cols = big_core.cols, + big_pes = big_core.pe_count(), + tail_core = tail_core.name, + tail_rows = tail_core.rows, + tail_cols = tail_core.cols, + tail_pes = tail_core.pe_count(), + tail_pe_area_fraction = tail_core.pe_area_fraction_vs(big_core), + "Matrix core profiles" + ); let v_machine = VectorMachine::new(vram, *VLEN, *HLEN); // Share same dim with VSRAM @@ -194,7 +210,25 @@ pub(crate) async fn run_from_cli() { // )) // .await; let decoded_ops = op.into_iter().map(op::Opcode::decode).collect::>(); - accelerator.do_ops(&decoded_ops).await; + let mut stage_profiler = opts.stage_profile_asm.as_ref().map(|path| { + StageProfiler::from_asm(path, decoded_ops.len()).unwrap_or_else(|err| { + panic!("failed to build stage profile from ASM {:?}: {err}", path) + }) + }); + accelerator + .do_ops(&decoded_ops, stage_profiler.as_mut()) + .await; + + if let Some(profile) = stage_profiler.as_ref() { + let out_path = opts + .stage_profile_out + .as_deref() + .unwrap_or_else(|| std::path::Path::new("stage_profile.json")); + profile + .write_json(out_path) + .unwrap_or_else(|err| panic!("failed to write stage profile {:?}: {err}", out_path)); + tracing::info!(path = %out_path.display(), "wrote runtime stage profile"); + } accelerator.log_debug_state().await; @@ -210,6 +244,10 @@ pub(crate) async fn run_from_cli() { let fpsram_bytes = accelerator.fpsram_dump_bytes(); dump_to_file("fpsram_dump.bin", &fpsram_bytes); + // Dump INTSRAM + let intsram_bytes = accelerator.intsram_dump_bytes(); + dump_to_file("intsram_dump.bin", &intsram_bytes); + // Dump HBM — skipped unless DEBUG tracing is enabled because HBM_SIZE may // be 128 GiB+. Tests run with --log-level warn and don't need hbm_dump.bin; // only manual debug runs dump HBM. diff --git a/transactional_emulator/src/stage_profile.rs b/transactional_emulator/src/stage_profile.rs new file mode 100644 index 00000000..8c44e7f7 --- /dev/null +++ b/transactional_emulator/src/stage_profile.rs @@ -0,0 +1,364 @@ +use std::collections::BTreeMap; +use std::fs; +use std::path::Path; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum StageKind { + RouterTopk, + AccumulatorInit, + Gather, + ExpertWeightAddress, + ExpertWeightPrefetch, + ExpertProjection, + ExpertActivation, + ExpertBias, + ExpertRouteWeight, + ScatterCombine, + Other, +} + +impl StageKind { + const ALL: [StageKind; 11] = [ + StageKind::RouterTopk, + StageKind::AccumulatorInit, + StageKind::Gather, + StageKind::ExpertWeightAddress, + StageKind::ExpertWeightPrefetch, + StageKind::ExpertProjection, + StageKind::ExpertActivation, + StageKind::ExpertBias, + StageKind::ExpertRouteWeight, + StageKind::ScatterCombine, + StageKind::Other, + ]; + + fn name(self) -> &'static str { + match self { + StageKind::RouterTopk => "router_topk", + StageKind::AccumulatorInit => "accumulator_init", + StageKind::Gather => "gather", + StageKind::ExpertWeightAddress => "expert_weight_address", + StageKind::ExpertWeightPrefetch => "expert_weight_prefetch", + StageKind::ExpertProjection => "expert_projection", + StageKind::ExpertActivation => "expert_activation", + StageKind::ExpertBias => "expert_bias", + StageKind::ExpertRouteWeight => "expert_route_weight", + StageKind::ScatterCombine => "scatter_combine", + StageKind::Other => "other", + } + } + + fn index(self) -> usize { + match self { + StageKind::RouterTopk => 0, + StageKind::AccumulatorInit => 1, + StageKind::Gather => 2, + StageKind::ExpertWeightAddress => 3, + StageKind::ExpertWeightPrefetch => 4, + StageKind::ExpertProjection => 5, + StageKind::ExpertActivation => 6, + StageKind::ExpertBias => 7, + StageKind::ExpertRouteWeight => 8, + StageKind::ScatterCombine => 9, + StageKind::Other => 10, + } + } +} + +#[derive(Clone, Copy, Debug, Default)] +struct StageRuntime { + instructions: u64, + seconds: f64, + hbm_bytes_read: u64, + hbm_bytes_written: u64, +} + +pub(crate) struct StageProfiler { + labels: Vec, + pair_labels: Vec>, + stages: [StageRuntime; 11], + pair_stages: BTreeMap, + total_instructions: u64, + total_seconds: f64, + total_hbm_bytes_read: u64, + total_hbm_bytes_written: u64, +} + +impl StageProfiler { + pub(crate) fn from_asm(path: &Path, expected_ops: usize) -> std::io::Result { + let asm = fs::read_to_string(path)?; + let mut labels = Vec::with_capacity(expected_ops); + let mut pair_labels = Vec::with_capacity(expected_ops); + let mut stage = StageKind::Other; + let mut pair_id = None; + + for raw_line in asm.lines() { + let line = raw_line.trim(); + if line.is_empty() { + continue; + } + if line.starts_with(';') { + stage = classify_comment(line, stage); + pair_id = extract_pair_id(line).or_else(|| { + if matches!(stage, StageKind::RouterTopk | StageKind::AccumulatorInit) { + None + } else { + pair_id + } + }); + } else if is_opcode_line(line) { + labels.push(stage); + pair_labels.push(pair_id); + } + } + + if labels.len() != expected_ops { + tracing::warn!( + asm = %path.display(), + labels = labels.len(), + expected_ops, + "stage profile ASM label count differs from decoded opcode count" + ); + } + + Ok(Self { + labels, + pair_labels, + stages: [StageRuntime::default(); 11], + pair_stages: BTreeMap::new(), + total_instructions: 0, + total_seconds: 0.0, + total_hbm_bytes_read: 0, + total_hbm_bytes_written: 0, + }) + } + + pub(crate) fn record( + &mut self, + pc: usize, + seconds: f64, + hbm_bytes_read: u64, + hbm_bytes_written: u64, + ) { + let stage = self.labels.get(pc).copied().unwrap_or(StageKind::Other); + let bucket = &mut self.stages[stage.index()]; + bucket.instructions += 1; + bucket.seconds += seconds; + bucket.hbm_bytes_read += hbm_bytes_read; + bucket.hbm_bytes_written += hbm_bytes_written; + self.total_instructions += 1; + self.total_seconds += seconds; + self.total_hbm_bytes_read += hbm_bytes_read; + self.total_hbm_bytes_written += hbm_bytes_written; + + if let Some(pair_id) = self.pair_labels.get(pc).copied().flatten() { + let pair_buckets = self + .pair_stages + .entry(pair_id) + .or_insert([StageRuntime::default(); 11]); + let pair_bucket = &mut pair_buckets[stage.index()]; + pair_bucket.instructions += 1; + pair_bucket.seconds += seconds; + pair_bucket.hbm_bytes_read += hbm_bytes_read; + pair_bucket.hbm_bytes_written += hbm_bytes_written; + } + } + + pub(crate) fn write_json(&self, path: &Path) -> std::io::Result<()> { + let mut out = String::new(); + out.push_str("{\n"); + out.push_str(" \"schema_version\": 1,\n"); + out.push_str(&format!(" \"label_count\": {},\n", self.labels.len())); + out.push_str(&format!( + " \"total_instructions_executed\": {},\n", + self.total_instructions + )); + out.push_str(&format!( + " \"total_profiled_seconds\": {:.12},\n", + self.total_seconds + )); + out.push_str(&format!( + " \"total_hbm_bytes_read\": {},\n", + self.total_hbm_bytes_read + )); + out.push_str(&format!( + " \"total_hbm_bytes_written\": {},\n", + self.total_hbm_bytes_written + )); + out.push_str(" \"stages\": {\n"); + + for (idx, stage) in StageKind::ALL.iter().enumerate() { + let stats = self.stages[stage.index()]; + let instr_fraction = if self.total_instructions == 0 { + 0.0 + } else { + stats.instructions as f64 / self.total_instructions as f64 + }; + let time_fraction = if self.total_seconds == 0.0 { + 0.0 + } else { + stats.seconds / self.total_seconds + }; + out.push_str(&format!( + " \"{}\": {{\"instructions\": {}, \"seconds\": {:.12}, \"instruction_fraction\": {:.12}, \"time_fraction\": {:.12}, \"hbm_bytes_read\": {}, \"hbm_bytes_written\": {}}}", + stage.name(), + stats.instructions, + stats.seconds, + instr_fraction, + time_fraction, + stats.hbm_bytes_read, + stats.hbm_bytes_written + )); + if idx + 1 != StageKind::ALL.len() { + out.push(','); + } + out.push('\n'); + } + + out.push_str(" },\n"); + out.push_str(" \"pairs\": {\n"); + + for (pair_idx, (pair_id, stages)) in self.pair_stages.iter().enumerate() { + let totals = sum_stage_runtimes(stages); + out.push_str(&format!( + " \"{}\": {{\"instructions\": {}, \"seconds\": {:.12}, \"hbm_bytes_read\": {}, \"hbm_bytes_written\": {}, \"stages\": {{\n", + pair_id, + totals.instructions, + totals.seconds, + totals.hbm_bytes_read, + totals.hbm_bytes_written + )); + for (stage_idx, stage) in StageKind::ALL.iter().enumerate() { + let stats = stages[stage.index()]; + out.push_str(&format!( + " \"{}\": {{\"instructions\": {}, \"seconds\": {:.12}, \"hbm_bytes_read\": {}, \"hbm_bytes_written\": {}}}", + stage.name(), + stats.instructions, + stats.seconds, + stats.hbm_bytes_read, + stats.hbm_bytes_written + )); + if stage_idx + 1 != StageKind::ALL.len() { + out.push(','); + } + out.push('\n'); + } + out.push_str(" }"); + out.push_str("}"); + if pair_idx + 1 != self.pair_stages.len() { + out.push(','); + } + out.push('\n'); + } + + out.push_str(" },\n"); + out.push_str( + " \"caveat\": \"Stage and routed-pair labels are derived from generated ASM comments. Time and HBM bytes are attributed by dynamically executed opcode. Bytes are measured from the global HBM stats delta before/after each opcode. Pair labels identify static routed pair slots, not necessarily unique expert IDs without joining the routing dump.\"\n", + ); + out.push_str("}\n"); + fs::write(path, out) + } +} + +fn sum_stage_runtimes(stages: &[StageRuntime; 11]) -> StageRuntime { + let mut total = StageRuntime::default(); + for stats in stages { + total.instructions += stats.instructions; + total.seconds += stats.seconds; + total.hbm_bytes_read += stats.hbm_bytes_read; + total.hbm_bytes_written += stats.hbm_bytes_written; + } + total +} + +fn is_opcode_line(line: &str) -> bool { + line.as_bytes() + .first() + .copied() + .map(|byte| byte.is_ascii_uppercase()) + .unwrap_or(false) +} + +fn classify_comment(comment: &str, current: StageKind) -> StageKind { + let text = comment.to_ascii_lowercase(); + if text.contains("gpt-oss router") + || text.contains("router token") + || text.contains("router dot token") + { + StageKind::RouterTopk + } else if text.contains("gpt-oss vram scatter-add") || text.contains("_scatter") { + StageKind::ScatterCombine + } else if text.contains("allocate vram matrix step6_pair") && text.contains("_route") { + StageKind::ExpertRouteWeight + } else if text.contains("materialize route weight") + || text.contains("vram matrix mul") + || (text.contains("true-zero vram rows") && matches!(current, StageKind::ExpertRouteWeight)) + { + StageKind::ExpertRouteWeight + } else if text.contains("step6_device_routing_acc") || text.contains("true-zero vram rows") { + StageKind::AccumulatorInit + } else if text.contains("gpt-oss gather token rows") + || text.contains("gather pair") + || text.contains("clear gather padding") + || (text.contains("allocate vram matrix step6_pair") && text.contains("_gather")) + { + StageKind::Gather + } else if text.contains("dynamic expert bias add") { + StageKind::ExpertBias + } else if text.contains("allocate vram matrix step6_pair") && text.contains("_sigmoid") { + StageKind::ExpertActivation + } else if text.contains("tile row min fp") + || text.contains("tile row max fp") + || matches!(current, StageKind::ExpertActivation) + && (text.contains("vram fill zero") + || text.contains("vram matrix add") + || text.contains("vram matrix mul")) + { + StageKind::ExpertActivation + } else if text.contains("dynamic hbm weight prefetch") + || text.contains("expert_id_to_weight_base") + { + StageKind::ExpertWeightAddress + } else if text.contains("subblock [") { + StageKind::ExpertWeightPrefetch + } else if text.contains("sub projection") + || text.contains("vram block add") + || text.contains("vram block") + || (text.contains("allocate vram matrix step6_pair") && !text.contains("_gather")) + { + StageKind::ExpertProjection + } else { + current + } +} + +fn extract_pair_id(comment: &str) -> Option { + let bytes = comment.as_bytes(); + for prefix in [b"step6_pair".as_slice(), b"pair=".as_slice()] { + let mut start = 0; + while let Some(pos) = find_subslice(&bytes[start..], prefix) { + let digit_start = start + pos + prefix.len(); + let digit_end = bytes[digit_start..] + .iter() + .position(|byte| !byte.is_ascii_digit()) + .map(|offset| digit_start + offset) + .unwrap_or(bytes.len()); + if digit_end > digit_start { + if let Ok(id) = comment[digit_start..digit_end].parse::() { + return Some(id); + } + } + start = digit_start; + } + } + None +} + +fn find_subslice(haystack: &[u8], needle: &[u8]) -> Option { + if needle.is_empty() { + return Some(0); + } + haystack + .windows(needle.len()) + .position(|window| window == needle) +} diff --git a/transactional_emulator/src/vector_machine.rs b/transactional_emulator/src/vector_machine.rs index b76a9ed8..5ddcf751 100644 --- a/transactional_emulator/src/vector_machine.rs +++ b/transactional_emulator/src/vector_machine.rs @@ -21,6 +21,22 @@ use crate::runtime_config::{ }; use crate::{cycle, op}; +fn tensor_from_f32_slice(data: &[f32]) -> Tensor { + if data.is_empty() { + return Tensor::zeros([0], (tch::Kind::Float, tch::Device::Cpu)); + } + unsafe { + Tensor::from_blob( + data.as_ptr() as *const u8, + &[data.len() as i64], + &[], + tch::Kind::Float, + tch::Device::Cpu, + ) + .internal_to_copy((tch::Kind::Float, tch::Device::Cpu), false) + } +} + /// Executes vector opcodes against `vram`. Cell payloads inside `vram` use /// interior mutability (Mutex), so all methods only need `&self`. pub(crate) struct VectorMachine { @@ -136,6 +152,54 @@ impl VectorMachine { } } + pub(crate) async fn max_scalar(&self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { + let a = self.vram.read(vs1).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor().clamp_min(f as f64), a.data_type()); + cycle!(*VECTOR_MAX_CYCLES); + self.vram.write(vd, c).await; + } else { + let result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = sliced.clamp_min(f as f64); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_MAX_CYCLES); + self.vram.write(vd, c).await; + } + } + + pub(crate) async fn min_scalar(&self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { + let a = self.vram.read(vs1).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor().clamp_max(f as f64), a.data_type()); + cycle!(*VECTOR_MAX_CYCLES); + self.vram.write(vd, c).await; + } else { + let result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = sliced.clamp_max(f as f64); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_MAX_CYCLES); + self.vram.write(vd, c).await; + } + } + pub(crate) async fn shift_scalar(&self, vd: u32, vs1: u32, shift: u32) { let a = self.vram.read(vs1).await; let tensor = a.as_tensor(); @@ -293,7 +357,7 @@ impl VectorMachine { // Convert bf16 slice to f32 vector let f32_vec: Vec = f.iter().map(|x| f32::from(*x)).collect(); // Create tensor from f32 vector - let tensor = Tensor::from_slice(&f32_vec); + let tensor = tensor_from_f32_slice(&f32_vec); // Quantize the tensor according to vram data type let c = QuantTensor::quantize(tensor, self.vram.ty()); cycle!(*VLEN); @@ -345,4 +409,206 @@ impl VectorMachine { f32::max(val, f) } } + + pub(crate) async fn topk_softmax( + &self, + vs1: u32, + expert_count: usize, + topk: usize, + ) -> (Vec, Vec) { + assert!(topk > 0, "topk must be positive"); + assert!( + topk <= expert_count, + "topk={} exceeds expert_count={}", + topk, + expert_count + ); + + let tile_size = self.tile_size as usize; + let mut logits = Vec::with_capacity(expert_count); + for chunk_start in (0..expert_count).step_by(tile_size) { + let a = self.vram.read(vs1 + chunk_start as u32).await; + let chunk_len = (expert_count - chunk_start).min(tile_size); + for idx in 0..chunk_len { + let value = a.as_tensor().double_value(&[idx as i64]) as f32; + logits.push(if value.is_nan() { + f32::NEG_INFINITY + } else { + value + }); + } + } + + let mut ranked: Vec<(usize, f32)> = logits.into_iter().enumerate().collect(); + ranked.sort_by(|(idx_a, val_a), (idx_b, val_b)| { + val_b.total_cmp(val_a).then_with(|| idx_a.cmp(idx_b)) + }); + let selected = &ranked[..topk]; + + let max_logit = selected + .iter() + .map(|(_, value)| *value) + .fold(f32::NEG_INFINITY, f32::max); + let selected_exp_values: Vec = selected + .iter() + .map(|(_, value)| (*value - max_logit).exp()) + .collect(); + let denom: f32 = selected_exp_values.iter().sum(); + let weights: Vec = selected_exp_values + .iter() + .map(|value| bf16::from_f32(if denom == 0.0 { 0.0 } else { value / denom })) + .collect(); + let indices: Vec = selected.iter().map(|(idx, _)| *idx as u32).collect(); + + cycle!((*VECTOR_MAX_CYCLES).saturating_mul(expert_count as u32)); + (indices, weights) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use quantize::{DataType, FpType, MxDataType}; + use runtime::{Executor, Instant}; + use std::sync::{Arc, Mutex}; + + fn tensor_values(tensor: &Tensor) -> Vec { + let len = tensor.size()[0] as usize; + let data = unsafe { core::slice::from_raw_parts(tensor.data_ptr() as *const f32, len) }; + data.to_vec() + } + + #[tokio::test] + async fn test_vector_scalar_minmax_clamps_bf16_boundary_values() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(4, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 4, 2); + let ty = MxDataType::Plain(fp_type); + + let input = Tensor::from_slice(&[-7.03125f32, -7.0, 7.0, 7.03125]); + vram.write(0, QuantTensor::quantize(input, ty)).await; + + machine.max_scalar(4, 0, -7.0, 0, 0).await; + machine.min_scalar(8, 0, 7.0, 0, 0).await; + + let max_out = vram.read(4).await; + let min_out = vram.read(8).await; + *got_task.lock().unwrap() = Some(( + tensor_values(max_out.as_tensor()), + tensor_values(min_out.as_tensor()), + )); + }); + + executor.enter(Instant::ETERNITY).await; + let (max_out, min_out) = got.lock().unwrap().take().unwrap(); + + assert_eq!(max_out, vec![-7.0, -7.0, 7.0, 7.03125]); + assert_eq!(min_out, vec![-7.03125, -7.0, 7.0, 7.0]); + } + + #[tokio::test] + async fn test_topk_softmax_uses_descending_logits_and_low_index_ties() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(64, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 64, 16); + let ty = MxDataType::Plain(fp_type); + + let mut input = vec![-100.0f32; 64]; + input[7] = 4.0; + input[3] = 4.0; + input[9] = 2.0; + input[0] = 1.0; + input[31] = 0.5; + vram.write(0, QuantTensor::quantize(Tensor::from_slice(&input), ty)) + .await; + + let (indices, weights) = machine.topk_softmax(0, 32, 4).await; + *got_task.lock().unwrap() = Some(( + indices, + weights.into_iter().map(f32::from).collect::>(), + )); + }); + + executor.enter(Instant::ETERNITY).await; + let (indices, weights) = got.lock().unwrap().take().unwrap(); + + assert_eq!(indices, vec![3, 7, 9, 0]); + let denom = 1.0 + 1.0 + f32::exp(-2.0) + f32::exp(-3.0); + let expected = vec![ + 1.0 / denom, + 1.0 / denom, + f32::exp(-2.0) / denom, + f32::exp(-3.0) / denom, + ]; + for (got, exp) in weights.iter().zip(expected) { + assert!((got - exp).abs() < 0.003, "got={got} expected={exp}"); + } + } + + #[tokio::test] + async fn test_topk_softmax_scans_contiguous_vector_rows_for_qwen128() { + let executor = Executor::new(); + let got = Arc::new(Mutex::new(None)); + let got_task = got.clone(); + + executor.spawn(async move { + let fp_type = DataType::Fp(FpType::BF16); + let vram = Arc::new(VectorSram::new(64, 4, fp_type, 4)); + let machine = VectorMachine::new(vram.clone(), 64, 16); + let ty = MxDataType::Plain(fp_type); + + let mut row0 = vec![-100.0f32; 64]; + let mut row1 = vec![-100.0f32; 64]; + row0[63] = 3.0; + row1[0] = 4.0; // expert 64 + row1[7] = 5.0; // expert 71 + row1[63] = 6.0; // expert 127 + row0[2] = 7.0; + row0[5] = 7.0; // low-index tie should pick 2 before 5 + row1[20] = 2.5; // expert 84 + row1[21] = 2.25; // expert 85 + for idx in 10..30 { + row0[idx] = 2.0; // non-selected mass must not affect normalized selected weights. + } + + vram.write(0, QuantTensor::quantize(Tensor::from_slice(&row0), ty)) + .await; + vram.write(64, QuantTensor::quantize(Tensor::from_slice(&row1), ty)) + .await; + + let (indices, weights) = machine.topk_softmax(0, 128, 8).await; + *got_task.lock().unwrap() = Some(( + indices, + weights.into_iter().map(f32::from).collect::>(), + )); + }); + + executor.enter(Instant::ETERNITY).await; + let (indices, weights) = got.lock().unwrap().take().unwrap(); + + assert_eq!(indices, vec![2, 5, 127, 71, 64, 63, 84, 85]); + let max_logit = 7.0f32; + let selected_logits = [7.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.5, 2.25]; + let selected_exp_values: Vec = selected_logits + .iter() + .map(|value| (*value - max_logit).exp()) + .collect(); + let denom: f32 = selected_exp_values.iter().sum(); + for (got, exp) in weights + .iter() + .zip(selected_exp_values.iter().map(|value| value / denom)) + { + assert!((got - exp).abs() < 0.003, "got={got} expected={exp}"); + } + } } diff --git a/transactional_emulator/testbench/README.md b/transactional_emulator/testbench/README.md index 8a8d52e7..ff6f3d6d 100644 --- a/transactional_emulator/testbench/README.md +++ b/transactional_emulator/testbench/README.md @@ -32,7 +32,9 @@ testbench/ ├── direct_emit/ Raw asm_template tests (no ATen compiler) ├── misc/ One-off tests ├── model_configs/ YAML model configs + loader -├── models/ Legacy FFN test, ASM generators, profiling scripts +├── models/ Model-level validation harnesses and profiling scripts +│ └── gpt_oss/ GPT-OSS attention/block/decoder-chain semantics checks +├── routed_moe/ Routed-MoE router/top-k/expert/gather-scatter bring-up ├── build_paths.py Shared build directory constant ├── config_utils.py Hardware config utilities ├── emulator_runner.py Shared Rust emulator runner + emulate_from_result() @@ -40,3 +42,8 @@ testbench/ ├── sim_env_utils.py HBM binary writers + memory setup └── sliced_layer_test_builder.py Sliced model test framework ``` + +`aten/` is reserved for reusable ATen operator checks. Model-specific routed +MoE and decoder-block bring-up lives under `routed_moe/` and `models/` so that +reviewers can distinguish reusable operator coverage from model semantics +harnesses. diff --git a/transactional_emulator/testbench/aten/compare/isa_analysis.py b/transactional_emulator/testbench/aten/compare/isa_analysis.py index 15b0f95b..4f19faac 100644 --- a/transactional_emulator/testbench/aten/compare/isa_analysis.py +++ b/transactional_emulator/testbench/aten/compare/isa_analysis.py @@ -90,7 +90,7 @@ def instruction_cycles(self, opcode: str) -> int: return 1 if opcode.startswith("V_ADD") or opcode.startswith("V_SUB"): return self.vector_add_cycles - if opcode.startswith("V_MUL") or opcode == "V_SHIFT_V": + if opcode.startswith("V_MUL") or opcode == "V_SHFT_V": return self.vector_mul_cycles if opcode == "V_EXP_V": return self.vector_exp_cycles diff --git a/transactional_emulator/testbench/aten/flash_attention_gqa_test.py b/transactional_emulator/testbench/aten/flash_attention_gqa_test.py index d1f80fd0..bb7a0201 100644 --- a/transactional_emulator/testbench/aten/flash_attention_gqa_test.py +++ b/transactional_emulator/testbench/aten/flash_attention_gqa_test.py @@ -33,18 +33,25 @@ from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw -def gqa_sdpa(q, k, v, scale, hq, hkv): +def gqa_sdpa(q, k, v, scale, hq, hkv, *, causal: bool = False): # q: [batch, seq, hq, h_qkv]; k/v: [batch, seq, hkv, h_qkv] q_t = q.transpose(1, 2) k_t = k.transpose(1, 2).repeat_interleave(hq // hkv, dim=1) v_t = v.transpose(1, 2).repeat_interleave(hq // hkv, dim=1) - o = F.scaled_dot_product_attention(q_t, k_t, v_t, scale=scale) + o = F.scaled_dot_product_attention(q_t, k_t, v_t, scale=scale, is_causal=causal) return o.transpose(1, 2) if __name__ == "__main__": parser = argparse.ArgumentParser(description=__doc__) add_hw_args(parser) + parser.add_argument("--causal-mask", action="store_true", help="Apply causal mask in GQA attention.") + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "flash_attention_gqa", + help="Directory for generated simulator artifacts.", + ) args = parser.parse_args() mlen = args.mlen @@ -87,7 +94,7 @@ def gqa_sdpa(q, k, v, scale, hq, hkv): args.hlen = h_qkv # HLEN must equal per-head dim for packed attention - build_dir = Path(__file__).parent / "build" / "flash_attention_gqa" + build_dir = args.build_dir.expanduser().resolve() hw = setup_hw(args, build_dir) print("=" * 80) @@ -114,7 +121,7 @@ def gqa_sdpa(q, k, v, scale, hq, hkv): # Hardware-accurate golden: MXFP8-quantize K, V before GQA SDPA. k_q = quantize_to_mxfp(k) v_q = quantize_to_mxfp(v) - golden = gqa_sdpa(q.float(), k_q.float(), v_q.float(), scale, hq, hkv) + golden = gqa_sdpa(q.float(), k_q.float(), v_q.float(), scale, hq, hkv, causal=args.causal_mask) # PLENA program using proper ATen dispatch registry = OpRegistry.load() @@ -156,6 +163,7 @@ def gqa_sdpa(q, k, v, scale, hq, hkv): hq=hq, hkv=hkv, h_qkv=h_qkv, + causal_mask=True if args.causal_mask else None, batch_size=batch_size, seq_len=s_q, kv_seq_len=s_kv, diff --git a/transactional_emulator/testbench/aten/golden.py b/transactional_emulator/testbench/aten/golden.py index afeb923c..cadf409e 100644 --- a/transactional_emulator/testbench/aten/golden.py +++ b/transactional_emulator/testbench/aten/golden.py @@ -99,5 +99,9 @@ def golden_flash_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, sc def golden_embedding_add(X: torch.Tensor, POS: torch.Tensor) -> torch.Tensor: - """BF16 element-wise add (simple, no MXFP quantization needed).""" - return (X.float() + POS.float()).to(torch.bfloat16) + """HBM MXFP8 load to BF16, then BF16 vector add.""" + precision = _active_precision_settings() + hbm_act = precision["HBM_V_ACT_TYPE"] + x_hbm = _load_to_vector_fp(X, hbm_act, precision) + pos_hbm = _load_to_vector_fp(POS, hbm_act, precision) + return quantize_to_vector_fp(x_hbm.float() + pos_hbm.float(), precision) diff --git a/transactional_emulator/testbench/aten/linear_test.py b/transactional_emulator/testbench/aten/linear_test.py index 52a3f05c..7227c022 100644 --- a/transactional_emulator/testbench/aten/linear_test.py +++ b/transactional_emulator/testbench/aten/linear_test.py @@ -5,6 +5,7 @@ import argparse import json +import math from pathlib import Path import torch @@ -15,6 +16,7 @@ from transactional_emulator.testbench.aten.configurable import add_hw_args, resolve_rows, setup_hw from transactional_emulator.testbench.aten.golden import golden_linear from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.layout_utils import prestage_bf16_vram_matrix from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim from transactional_emulator.tools.create_sim_env import create_sim_env @@ -23,6 +25,11 @@ parser = argparse.ArgumentParser(description=__doc__) add_hw_args(parser) parser.add_argument("--out-features", type=int, default=None) + parser.add_argument( + "--include-bias", + action="store_true", + help="Prestaged-BF16 bias add path for non-MoE VRAM layout coverage.", + ) args = parser.parse_args() mlen = args.mlen @@ -51,10 +58,15 @@ torch.manual_seed(args.seed) X = torch.randn(rows, in_features) W = torch.randn(in_features, out_features) + Bias = torch.randn(rows, out_features).to(torch.bfloat16) if args.include_bias else None print(f"\nInput X: {X.shape}, W: {W.shape}") + if Bias is not None: + print(f"Bias: {Bias.shape} (prestaged BF16 VRAM)") print("\n--- Hardware-Accurate Golden Reference (MXFP8 + BF16) ---") golden_Y = golden_linear(X, W) + if Bias is not None: + golden_Y = (golden_Y.float() + Bias.float()).to(torch.bfloat16) print(f" golden_Y: {golden_Y.shape}") print(f" golden_Y[0,:4]: {golden_Y[0, :4].tolist()}") @@ -69,6 +81,24 @@ X_batch = prog.load_batch(x_input, name="X") Y = ops.linear(prog, X_batch, w_input) + vram_preload = None + if Bias is not None: + physical_rows = max(blen, math.ceil(rows / blen) * blen) + bias_physical = (physical_rows, out_features) + bias_base = math.ceil(prog.vram_allocator._vmm.next_bump / (mlen * mlen)) * (mlen * mlen) + bias_size = bias_physical[0] * bias_physical[1] + bias_aligned = math.ceil(bias_size / (mlen * mlen)) * (mlen * mlen) + vram_preload = torch.zeros(bias_base + bias_aligned, dtype=torch.bfloat16) + bias_vram = prestage_bf16_vram_matrix( + prog=prog, + name="Bias", + tensor=Bias, + vram_addr=bias_base, + physical_shape=bias_physical, + vram_preload=vram_preload, + ) + prog.vram_add(Y, bias_vram, num_rows=rows) + gen_code = prog.compile() print(f"\nGenerated {len(gen_code.splitlines())} lines of ISA code") @@ -76,7 +106,14 @@ golden_result = {"original_output": golden_Y} fp_preload = [0.0, 1e-6, 1.0 / in_features] + [0.0] * 7 - create_sim_env(input_tensors, gen_code, golden_result, fp_preload, build_dir=str(build_dir)) + create_sim_env( + input_tensors, + gen_code, + golden_result, + fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + ) # Place each tensor at the compiler's actual HBM address. At MLEN>=256 the # compiler tile-aligns HBM allocations (gaps between tensors); a contiguous diff --git a/transactional_emulator/testbench/direct_emit/v_shift_v_test.py b/transactional_emulator/testbench/direct_emit/v_shift_v_test.py index 879867ad..9f210c73 100644 --- a/transactional_emulator/testbench/direct_emit/v_shift_v_test.py +++ b/transactional_emulator/testbench/direct_emit/v_shift_v_test.py @@ -71,7 +71,7 @@ def quantize_to_mxfp(tensor): golden_result = {"input_tensor": input_tensor, "original_output": original_output} - gen_assembly_code = "; V_SHIFT_V Test Generation\n" + gen_assembly_code = "; V_SHFT_V Test Generation\n" # Reset registers gen_assembly_code += reset_reg_asm(alive_registers=[1, 2, 3]) @@ -91,18 +91,18 @@ def quantize_to_mxfp(tensor): # Reset registers gen_assembly_code += reset_reg_asm(alive_registers=[1, 2, 3, 4]) - # Set up registers for V_SHIFT_V: + # Set up registers for V_SHFT_V: # gp1 = source/dest vector address (0) # gp2 = shift amount gen_assembly_code += "S_ADDI_INT gp1, gp0, 0\n" # gp1 = 0 (source vector address) gen_assembly_code += f"S_ADDI_INT gp2, gp0, {shift_amount}\n" # gp2 = shift amount - # Apply V_SHIFT_V to each vector row - # V_SHIFT_V rd, rs1, rs2 + # Apply V_SHFT_V to each vector row + # V_SHFT_V rd, rs1, rs2 # rd = destination address register, rs1 = source address register, rs2 = shift amount register total_vectors = (batch_size * hidden_size) // vlen for i in range(total_vectors): - gen_assembly_code += "V_SHIFT_V gp1, gp1, gp2\n" # In-place shift + gen_assembly_code += "V_SHFT_V gp1, gp1, gp2\n" # In-place shift gen_assembly_code += f"S_ADDI_INT gp1, gp1, {vlen}\n" # Move to next vector build_path = BUILD_DIR @@ -136,7 +136,7 @@ def quantize_to_mxfp(tensor): json.dump(comparison_params, f, indent=2) print("================================================") - print("Finished generating V_SHIFT_V test assembly code") + print("Finished generating V_SHFT_V test assembly code") print(f"Shift amount: {shift_amount}") print(f"Result location: row {result_start_row}, {num_result_rows} rows") print(f"Comparison params: {comparison_params}") diff --git a/transactional_emulator/testbench/emulator_runner.py b/transactional_emulator/testbench/emulator_runner.py index 3730d54c..1fe74752 100644 --- a/transactional_emulator/testbench/emulator_runner.py +++ b/transactional_emulator/testbench/emulator_runner.py @@ -127,7 +127,11 @@ def run_emulator(build_dir: Path, hbm_size: int | None = None, threads: int | No emulator_dir / "target" / "release" / "build" / "torch-sys-*" / "out" / "libtorch" / "libtorch" / "lib" ) libtorch_dirs = glob.glob(libtorch_pattern) - env = {**os.environ, "RUST_BACKTRACE": "1", "RUST_LOG": "warn,transactional_emulator=info"} + env = { + **os.environ, + "RUST_BACKTRACE": "1", + "RUST_LOG": os.environ.get("PLENA_EMULATOR_RUST_LOG", "warn,transactional_emulator=info"), + } # libtorch (tch/ATen) parallelises every tensor op with an OpenMP pool that defaults to one # thread per core. On the emulator's tiny per-op tensors that is almost pure barrier overhead # (single-thread is ~6x faster here), and the spin-wait barriers melt down under @@ -167,6 +171,12 @@ def run_emulator(build_dir: Path, hbm_size: int | None = None, threads: int | No r"Utilization:\s*([0-9.eE+-]+)\s*bytes/sec" ) + # Avoid copying an HBM dump from a previous debug run when the current run + # does not enable DEBUG tracing. + hbm_debug_dump = emulator_dir / "hbm_dump.bin" + if hbm_debug_dump.exists(): + hbm_debug_dump.unlink() + with log_path.open("w", encoding="utf-8", errors="replace") as log_file: proc = subprocess.Popen( cmd, @@ -221,6 +231,24 @@ def run_emulator(build_dir: Path, hbm_size: int | None = None, threads: int | No import shutil shutil.copy2(vram_src, vram_dst) + fpsram_src = emulator_dir / "fpsram_dump.bin" + fpsram_dst = build_dir / "fpsram_dump.bin" + if fpsram_src.exists(): + import shutil + + shutil.copy2(fpsram_src, fpsram_dst) + intsram_src = emulator_dir / "intsram_dump.bin" + intsram_dst = build_dir / "intsram_dump.bin" + if intsram_src.exists(): + import shutil + + shutil.copy2(intsram_src, intsram_dst) + hbm_src = emulator_dir / "hbm_dump.bin" + hbm_dst = build_dir / "hbm_dump.bin" + if hbm_src.exists(): + import shutil + + shutil.copy2(hbm_src, hbm_dst) return metrics diff --git a/transactional_emulator/testbench/layout_utils.py b/transactional_emulator/testbench/layout_utils.py new file mode 100644 index 00000000..1551a7ad --- /dev/null +++ b/transactional_emulator/testbench/layout_utils.py @@ -0,0 +1,84 @@ +"""Shared tensor layout helpers for transactional emulator testbenches.""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from typing import Any + +import torch + + +def _materialize(value: Any) -> Any: + materialize = getattr(value, "materialize", None) + if callable(materialize): + return materialize() + return value + + +def infer_hbm_tensor_layouts(input_tensors: Mapping[str, Any]) -> dict[str, dict[str, object]]: + """Infer explicit per-row HBM layouts for 2D tensors. + + The legacy writer defaults to one MLEN-wide source row when no sidecar is + present. That is fine for toy one-tile tensors but corrupts wide model + weights. Use the tensor's real trailing dimension as the row width unless a + caller provides a more specific physical layout. + """ + layouts: dict[str, dict[str, object]] = {} + for name, value in input_tensors.items(): + tensor = _materialize(value) + shape = getattr(tensor, "shape", None) + if shape is None or len(shape) != 2: + continue + rows = int(shape[0]) + cols = int(shape[1]) + layouts[str(name)] = { + "source_shape": [rows, cols], + "storage_shape": [rows, cols], + "source_rows": rows, + "storage_rows": rows, + "source_row_elements": cols, + "storage_row_elements": cols, + } + return layouts + + +def prestage_bf16_vram_matrix( + *, + prog, + name: str, + tensor: torch.Tensor, + vram_addr: int, + physical_shape: tuple[int, int], + vram_preload: torch.Tensor, +): + """Preload a BF16 matrix into VRAM's column-block-major matrix layout.""" + rows, cols = tensor.shape + physical_rows, physical_cols = physical_shape + if physical_rows < rows or physical_cols < cols: + raise ValueError(f"{name}: physical_shape={physical_shape} smaller than tensor shape={tuple(tensor.shape)}") + + padded = torch.zeros(physical_rows, physical_cols, dtype=torch.bfloat16) + padded[:rows, :cols] = tensor.to(torch.bfloat16) + num_col_blocks = math.ceil(physical_cols / prog.mlen) + layout_size = num_col_blocks * physical_rows * prog.mlen + end = vram_addr + layout_size + if end > vram_preload.numel(): + raise ValueError(f"{name}: VRAM preload too small for [{vram_addr}, {end})") + + for col_block in range(num_col_blocks): + col_start = col_block * prog.mlen + col_end = min(col_start + prog.mlen, physical_cols) + width = col_end - col_start + for row in range(physical_rows): + dst = vram_addr + (col_block * physical_rows + row) * prog.mlen + vram_preload[dst : dst + width] = padded[row, col_start:col_end] + + bias_input = prog.input( + name, + shape=(rows, cols), + hbm_addr=0, + prestaged_vram_addr=vram_addr, + physical_shape=physical_shape, + ) + return prog.load_batch(bias_input, name=name) diff --git a/transactional_emulator/testbench/models/gpt_oss/__init__.py b/transactional_emulator/testbench/models/gpt_oss/__init__.py new file mode 100644 index 00000000..3057b1cb --- /dev/null +++ b/transactional_emulator/testbench/models/gpt_oss/__init__.py @@ -0,0 +1 @@ +"""GPT-OSS model-level emulator validation harnesses.""" diff --git a/transactional_emulator/testbench/models/gpt_oss/attention_semantics_test.py b/transactional_emulator/testbench/models/gpt_oss/attention_semantics_test.py new file mode 100644 index 00000000..a1c71faa --- /dev/null +++ b/transactional_emulator/testbench/models/gpt_oss/attention_semantics_test.py @@ -0,0 +1,4274 @@ +"""GPT-OSS attention semantics smoke tests for PLENA. + +This file intentionally tests the new attention building blocks before a full +block harness exists: + +* ``projection`` verifies the BF16 HBM_M_KV matrix path used for Q/K/V/O + projections. It uses H_PREFETCH_M precision=KeyValue and emits no + C_SET_SCALE_REG for the projection. +* ``core`` verifies packed GQA causal attention with optional GPT-OSS sink and + sliding-window masks. K/V are stored as Plain BF16 HBM_M_KV tensors. +* ``full`` verifies BF16 Q/K/V projections, BF16 K/V staging through + H_STORE_V/H_PREFETCH_M KeyValue, packed attention, and BF16 O projection in + one emulator program. +* ``true_core`` verifies the true GPT-OSS GQA shape (64 Q heads / 8 KV heads) + with per-Q-head sinks by unrolling one Q head per packed-attention call. Q/K/V + are prestaged after projection+bias+RoPE, so this mode validates the attention + core shape and sink semantics without claiming runtime RoPE lowering is done. +* ``runtime_rope`` verifies the missing bridge: BF16 projection + BF16 bias + followed by device-side rotate-half projection and RoPE. +* ``true_full`` verifies the true-shape attention path with runtime BF16 + Q/K/V/O projection, BF16 bias, device-side RoPE, per-head sink, and packed + 64Q/8KV attention. +* ``true_decoder_block`` extends ``true_attn_block`` with the second pre-norm + MoE sublayer: VRAM-resident post-attention RMSNorm feeds device top-k, + dynamic true-expert weight selection, expert execution, scatter-combine, and + the second residual add in one emulator program. +* ``true_decoder_chain`` is a sequential validation runner that invokes + ``true_decoder_block`` layer by layer, handing emulator output to the next + layer and running a host-chain control for cumulative drift diagnostics. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from types import SimpleNamespace +from pathlib import Path + +import numpy as np +import tomlkit +import torch +import torch.nn.functional as F +from safetensors import safe_open + +from compiler.asm_templates._imm import load_large_int +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.aten.golden import ( + _active_precision_settings, + _rms_norm_vector_ref, + quantize_to_vector_fp, +) +from transactional_emulator.testbench.emulator_runner import compare_emulator_output, run_and_assert, run_emulator +from transactional_emulator.testbench.layout_utils import infer_hbm_tensor_layouts, prestage_bf16_vram_matrix +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.testbench.routed_moe.gpt_oss_moe_gather_scatter_test import ( + _build_true_expert_bias_table, + _build_true_expert_weight_table, + _decode_bf16_dump, + _decode_u32_dump, + _device_routing_vram_exact_golden, + _linear_projection_golden, +) +from transactional_emulator.tools.create_sim_env import create_sim_env + +_DEFAULT_QWEN3_REVISION = "ad44e777bcd18fa416d9da3bd8f70d33ebb85d39" + + +def _default_qwen3_snapshot() -> Path: + override = os.environ.get("QWEN3_30B_A3B_SNAPSHOT") + if override: + return Path(override).expanduser() + hf_home = Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")) + return hf_home / "hub" / "models--Qwen--Qwen3-30B-A3B" / "snapshots" / _DEFAULT_QWEN3_REVISION + + +_PRECISION_SETTINGS_CACHE = {} + + +def _cached_active_precision_settings(): + """Cache precision settings for hot Python goldens. + + Qwen MoE goldens call the BF16 VRAM quantizer many times. Reading and + parsing the settings TOML on every call dominates runtime without changing + semantics, so cache per settings file and invalidate when this test mutates + the settings in-place. + """ + settings_key = os.environ.get("PLENA_SETTINGS_TOML", "") + if settings_key not in _PRECISION_SETTINGS_CACHE: + _PRECISION_SETTINGS_CACHE[settings_key] = _active_precision_settings() + return _PRECISION_SETTINGS_CACHE[settings_key] + + +def _set_matrix_kv_plain_bf16() -> None: + settings = Path(os.environ["PLENA_SETTINGS_TOML"]) + with settings.open() as f: + config = tomlkit.load(f) + for mode in ("TRANSACTIONAL", "ANALYTIC"): + precision = config[mode]["PRECISION"] + for key in ("HBM_M_KV_TYPE", "HBM_V_KV_TYPE"): + precision[key] = tomlkit.table() + precision[key]["format"] = "Plain" + precision[key]["DATA_TYPE"] = tomlkit.table() + precision[key]["DATA_TYPE"]["type"] = "Fp" + precision[key]["DATA_TYPE"]["sign"] = True + precision[key]["DATA_TYPE"]["exponent"] = 8 + precision[key]["DATA_TYPE"]["mantissa"] = 7 + with settings.open("w") as f: + tomlkit.dump(config, f) + _PRECISION_SETTINGS_CACHE.pop(str(settings), None) + _PRECISION_SETTINGS_CACHE.pop(os.environ.get("PLENA_SETTINGS_TOML", ""), None) + + +def _rel_rms(a: torch.Tensor, b: torch.Tensor) -> float: + a_f = a.float() + b_f = b.float() + return float(torch.linalg.vector_norm(a_f - b_f) / torch.clamp(torch.linalg.vector_norm(b_f), min=1e-12)) + + +def _bf16_vram(x: torch.Tensor) -> torch.Tensor: + precision = _cached_active_precision_settings() + return quantize_to_vector_fp(x.to(torch.bfloat16).float(), precision) + + +def _weighted_rms_norm_golden(x: torch.Tensor, weight_rows: torch.Tensor, eps: float) -> torch.Tensor: + precision = _cached_active_precision_settings() + x_vram = _bf16_vram(x) + w_vram = _bf16_vram(weight_rows) + normed = _rms_norm_vector_ref(x_vram, eps, precision) + return quantize_to_vector_fp(normed.float() * w_vram.float(), precision) + + +def _head_rms_norm_golden(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: + """RMSNorm over the last/head dimension for Qwen-style q_norm/k_norm.""" + precision = _cached_active_precision_settings() + x_vram = _bf16_vram(x) + w_vram = _bf16_vram(weight) + normed = _rms_norm_vector_ref(x_vram, eps, precision) + return quantize_to_vector_fp(normed.float() * w_vram.float(), precision) + + +def _standard_swiglu_vram_golden(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """VRAM-style standard SwiGLU: silu(gate) * up.""" + sigmoid = _bf16_vram(gate) + sigmoid = _bf16_vram(sigmoid.float() * -1.0) + sigmoid = _bf16_vram(torch.exp(torch.clamp(sigmoid.float(), -88.0, 88.0))) + sigmoid = _bf16_vram(sigmoid.float() + 1.0) + sigmoid = _bf16_vram(torch.reciprocal(sigmoid.float())) + glu = _bf16_vram(gate.float() * sigmoid.float()) + return _bf16_vram(up.float() * glu.float()) + + +def _device_routing_vram_policy_golden( + *, + x: torch.Tensor, + device_indices: torch.Tensor, + device_weights: torch.Tensor, + split, + down_weight: torch.Tensor, + down_bias: torch.Tensor | None, + rows: int, + hidden: int, + blen: int, + mlen: int, + activation_policy: str, + pair_count: int | None = None, +) -> tuple[torch.Tensor, dict | None]: + """Decoder-block MoE golden with policy-specific expert activation.""" + if activation_policy == "gpt_oss_clamp_gated": + if down_bias is None: + raise ValueError("GPT-OSS clamp-gated golden requires down_bias") + return _device_routing_vram_exact_golden( + x=x, + device_indices=device_indices, + device_weights=device_weights, + split=split, + down_weight=down_weight, + down_bias=down_bias, + rows=rows, + hidden=hidden, + blen=blen, + mlen=mlen, + pair_count=pair_count, + ) + if activation_policy != "standard_swiglu": + raise NotImplementedError(f"unsupported decoder MoE activation_policy={activation_policy!r}") + + exact_golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + top_k = device_indices.shape[1] + total_pairs = rows * top_k if pair_count is None else pair_count + for pair_idx in range(total_pairs): + token_idx = pair_idx // top_k + topk_pos = pair_idx % top_k + expert_id = int(device_indices[token_idx, topk_pos].item()) + weight = device_weights[token_idx, topk_pos].to(torch.bfloat16) + x_slots = torch.zeros(blen, hidden, dtype=torch.bfloat16) + x_slots[0:1] = x[token_idx : token_idx + 1].to(torch.bfloat16) + + gate = _linear_projection_golden( + x_slots, + split.gate_weight[expert_id], + mlen=mlen, + hbm_input=False, + ) + up = _linear_projection_golden( + x_slots, + split.up_weight[expert_id], + mlen=mlen, + hbm_input=False, + ) + if getattr(split, "gate_bias", None) is not None: + gate = _bf16_vram(gate.float() + split.gate_bias[expert_id].reshape(1, -1).float()) + if getattr(split, "up_bias", None) is not None: + up = _bf16_vram(up.float() + split.up_bias[expert_id].reshape(1, -1).float()) + hidden_slots = _standard_swiglu_vram_golden(gate, up) + out = _linear_projection_golden(hidden_slots, down_weight[expert_id], mlen=mlen, hbm_input=False) + if down_bias is not None: + out = _bf16_vram(out.float() + down_bias[expert_id].reshape(1, -1).float()) + exact_golden[token_idx] = _bf16_vram( + exact_golden[token_idx].float() + _bf16_vram(out[0].float() * weight.float()).float() + ) + return exact_golden, None + + +def _router_bias_block_rows(router_bias: torch.Tensor, *, rows: int, num_experts: int, mlen: int) -> torch.Tensor: + expert_blocks = math.ceil(num_experts / mlen) + if expert_blocks == 1: + return router_bias.reshape(1, -1).repeat(rows, 1).to(torch.bfloat16) + out = torch.zeros(rows * expert_blocks, mlen, dtype=torch.bfloat16) + for token_idx in range(rows): + for block_idx in range(expert_blocks): + col_start = block_idx * mlen + col_end = min(col_start + mlen, num_experts) + out[token_idx * expert_blocks + block_idx, : col_end - col_start] = router_bias[col_start:col_end] + return out + + +def _load_qwen3_layer_tensors(snapshot: Path, layer_idx: int) -> dict[str, torch.Tensor | str]: + """Load true Qwen3-MoE layer tensors in the shapes used by this harness. + + HF stores linear weights as [out_features, in_features]. The PLENA BF16 + projection helpers used here expect [in_features, out_features], matching + the synthetic tensors above and the Qwen substrate scripts. + """ + + snapshot = snapshot.expanduser().resolve() + index_path = snapshot / "model.safetensors.index.json" + if not index_path.exists(): + raise FileNotFoundError(f"missing Qwen safetensors index: {index_path}") + weight_map = json.loads(index_path.read_text())["weight_map"] + handles: dict[str, safe_open] = {} + + def tensor(name: str) -> torch.Tensor: + shard_name = weight_map[name] + if shard_name not in handles: + handles[shard_name] = safe_open(snapshot / shard_name, framework="pt", device="cpu") + return handles[shard_name].get_tensor(name).to(torch.bfloat16).contiguous() + + prefix = f"model.layers.{layer_idx}" + q_w = tensor(f"{prefix}.self_attn.q_proj.weight").T.contiguous() + k_w = tensor(f"{prefix}.self_attn.k_proj.weight").T.contiguous() + v_w = tensor(f"{prefix}.self_attn.v_proj.weight").T.contiguous() + o_w = tensor(f"{prefix}.self_attn.o_proj.weight").T.contiguous() + q_norm = tensor(f"{prefix}.self_attn.q_norm.weight") + k_norm = tensor(f"{prefix}.self_attn.k_norm.weight") + input_norm = tensor(f"{prefix}.input_layernorm.weight") + post_norm = tensor(f"{prefix}.post_attention_layernorm.weight") + router_w = tensor(f"{prefix}.mlp.gate.weight") + + num_experts = int(router_w.shape[0]) + gate_weights = [] + up_weights = [] + down_weights = [] + for expert_id in range(num_experts): + gate_weights.append(tensor(f"{prefix}.mlp.experts.{expert_id}.gate_proj.weight").T.contiguous()) + up_weights.append(tensor(f"{prefix}.mlp.experts.{expert_id}.up_proj.weight").T.contiguous()) + down_weights.append(tensor(f"{prefix}.mlp.experts.{expert_id}.down_proj.weight").T.contiguous()) + gate_w = torch.stack(gate_weights, dim=0) + up_w = torch.stack(up_weights, dim=0) + down_w = torch.stack(down_weights, dim=0) + + hidden = int(router_w.shape[1]) + intermediate = int(gate_w.shape[2]) + return { + "snapshot": str(snapshot), + "layer_idx": layer_idx, + "w_q": q_w, + "w_k": k_w, + "w_v": v_w, + "w_o": o_w, + "q_norm_weight": q_norm, + "k_norm_weight": k_norm, + "norm_weight": input_norm, + "ffn_norm_weight": post_norm, + "router_weight": router_w, + "router_bias": torch.zeros(num_experts, dtype=torch.bfloat16), + "gate_weight": gate_w, + "up_weight": up_w, + "down_weight": down_w, + "gate_bias": torch.zeros(num_experts, intermediate, dtype=torch.bfloat16), + "up_bias": torch.zeros(num_experts, intermediate, dtype=torch.bfloat16), + "down_bias": torch.zeros(num_experts, hidden, dtype=torch.bfloat16), + } + + +def _residual_add_golden(hidden: torch.Tensor, sublayer: torch.Tensor) -> torch.Tensor: + precision = _cached_active_precision_settings() + hidden_vram = _bf16_vram(hidden) + sublayer_vram = _bf16_vram(sublayer) + return quantize_to_vector_fp(hidden_vram.float() + sublayer_vram.float(), precision) + + +def _strict_tail_summary(actual: torch.Tensor, reference: torch.Tensor, *, rtol: float = 0.02) -> dict: + actual_f = actual.float() + ref_f = reference.float() + diff = actual_f - ref_f + finite = torch.isfinite(actual_f) & torch.isfinite(ref_f) & torch.isfinite(diff) + abs_diff = diff.abs() + safe_abs_diff = torch.where(torch.isfinite(abs_diff), abs_diff, torch.full_like(abs_diff, float("inf"))) + atol = ref_f.std(unbiased=False) * 0.01 + allowed = atol + rtol * ref_f.abs() + failed = (~finite) | (safe_abs_diff > allowed) + finite_abs = safe_abs_diff[finite] + sigma = finite_abs.std(unbiased=False) if finite_abs.numel() else torch.tensor(float("inf")) + k_sweep = [] + for k in (2.0, 2.5, 3.0, 3.5, 4.0): + k_allowed = k * sigma + rtol * ref_f.abs() + k_failed = (~finite) | (safe_abs_diff > k_allowed) + k_sweep.append({"k": k, "fail_count": int(k_failed.sum().item())}) + failed_signed = diff[failed] + finite_safe_abs = safe_abs_diff[torch.isfinite(safe_abs_diff)] + return { + "atol": float(atol.item()), + "rtol": rtol, + "fail_count": int(failed.sum().item()), + "numel": int(failed.numel()), + "pass_rate": float((~failed).float().mean().item()), + "allclose": bool((~failed).all().item()), + "nonfinite_count": int((~finite).sum().item()), + "max_abs_error": float(finite_safe_abs.max().item()) if finite_safe_abs.numel() else float("inf"), + "sigma_abs_error": float(sigma.item()), + "k_sweep": k_sweep, + "sign": { + "positive": int((failed_signed > 0).sum().item()) if failed_signed.numel() else 0, + "negative": int((failed_signed < 0).sum().item()) if failed_signed.numel() else 0, + "zero": int((failed_signed == 0).sum().item()) if failed_signed.numel() else 0, + }, + } + + +def _read_vram_matrix( + path: Path, + *, + vram_addr: int, + rows: int, + cols: int, + mlen: int, + physical_rows: int, +) -> torch.Tensor: + """Decode a BF16 VRAM matrix from column-block-major storage.""" + raw = np.fromfile(path, dtype=" dict: + set_matches = [] + order_matches = [] + for token_idx in range(host_indices.shape[0]): + host = [int(v) for v in host_indices[token_idx].tolist()] + dev = [int(v) for v in device_indices[token_idx].tolist()] + set_matches.append(set(host) == set(dev)) + order_matches.append(host == dev) + weight_diff = (device_weights.float() - host_weights.float()).abs() + return { + "host_indices": host_indices.cpu().tolist(), + "device_indices": device_indices.cpu().tolist(), + "host_weights": host_weights.cpu().float().tolist(), + "device_weights": device_weights.cpu().float().tolist(), + "set_matches_by_token": set_matches, + "order_matches_by_token": order_matches, + "set_match_count": int(sum(set_matches)), + "order_match_count": int(sum(order_matches)), + "num_tokens": int(host_indices.shape[0]), + "weight_rel_rms": _rel_rms(device_weights, host_weights), + "weight_max_abs_error": float(weight_diff.max().item()) if weight_diff.numel() else 0.0, + } + + +def _router_margin_summary(logits: torch.Tensor, top_k: int) -> dict: + logits_f = logits.float() + sorted_vals, sorted_idx = torch.sort(logits_f, dim=-1, descending=True, stable=True) + if sorted_vals.shape[-1] <= top_k: + gaps = torch.full((sorted_vals.shape[0],), float("inf")) + rank_next = torch.full((sorted_vals.shape[0],), -1, dtype=torch.long) + next_vals = torch.full((sorted_vals.shape[0],), float("-inf")) + else: + gaps = sorted_vals[:, top_k - 1] - sorted_vals[:, top_k] + rank_next = sorted_idx[:, top_k] + next_vals = sorted_vals[:, top_k] + top_vals = sorted_vals[:, :top_k] + top_idx = sorted_idx[:, :top_k] + return { + "top_k": int(top_k), + "min_rank_k_to_next_gap": float(gaps.min().item()) if gaps.numel() else float("inf"), + "per_token": [ + { + "token_index": int(token_idx), + "top_indices": [int(v) for v in top_idx[token_idx].tolist()], + "top_values": [float(v) for v in top_vals[token_idx].tolist()], + "rank_k_index": int(top_idx[token_idx, top_k - 1].item()), + "rank_k_value": float(top_vals[token_idx, top_k - 1].item()), + "rank_next_index": int(rank_next[token_idx].item()), + "rank_next_value": float(next_vals[token_idx].item()), + "rank_k_to_next_gap": float(gaps[token_idx].item()), + } + for token_idx in range(logits.shape[0]) + ], + } + + +def _routing_boundary_verdict( + *, + reference_logits: torch.Tensor, + comparison_logits: torch.Tensor, + reference_indices: torch.Tensor, + device_indices: torch.Tensor, + top_k: int, +) -> dict: + """Classify top-k mismatches as stable-token failures or boundary cases.""" + reference_margin = _router_margin_summary(reference_logits, top_k) + comparison_margin = _router_margin_summary(comparison_logits, top_k) + logit_error = (comparison_logits.float() - reference_logits.float()).abs() + threshold = float(logit_error.max().item()) if logit_error.numel() else 0.0 + + unsafe_tokens = [] + set_mismatch_tokens = [] + order_mismatch_tokens = [] + safe_set_mismatch_tokens = [] + safe_order_mismatch_tokens = [] + per_token = [] + for token_idx in range(reference_indices.shape[0]): + ref_order = [int(v) for v in reference_indices[token_idx].tolist()] + dev_order = [int(v) for v in device_indices[token_idx].tolist()] + ref_gap = float(reference_margin["per_token"][token_idx]["rank_k_to_next_gap"]) + cmp_gap = float(comparison_margin["per_token"][token_idx]["rank_k_to_next_gap"]) + gap_floor = min(ref_gap, cmp_gap) + unsafe = gap_floor <= threshold + set_match = set(ref_order) == set(dev_order) + order_match = ref_order == dev_order + if unsafe: + unsafe_tokens.append(token_idx) + if not set_match: + set_mismatch_tokens.append(token_idx) + if not unsafe: + safe_set_mismatch_tokens.append(token_idx) + if not order_match: + order_mismatch_tokens.append(token_idx) + if not unsafe: + safe_order_mismatch_tokens.append(token_idx) + per_token.append( + { + "token_index": int(token_idx), + "reference_gap": ref_gap, + "comparison_gap": cmp_gap, + "gap_floor": gap_floor, + "unsafe": bool(unsafe), + "set_match": bool(set_match), + "order_match": bool(order_match), + "reference_indices": ref_order, + "device_indices": dev_order, + } + ) + + if safe_set_mismatch_tokens: + status = "fail_safe_set_mismatch" + elif set_mismatch_tokens: + status = "boundary_explained_set_mismatch" + elif unsafe_tokens: + status = "pass_with_weak_margins" + elif safe_order_mismatch_tokens: + status = "pass_set_match_safe_order_mismatch" + else: + status = "pass" + + return { + "status": status, + "top_k": int(top_k), + "logit_error_max_abs": threshold, + "unsafe_threshold": threshold, + "unsafe_tokens": [int(v) for v in unsafe_tokens], + "set_mismatch_tokens": [int(v) for v in set_mismatch_tokens], + "order_mismatch_tokens": [int(v) for v in order_mismatch_tokens], + "safe_set_mismatch_tokens": [int(v) for v in safe_set_mismatch_tokens], + "safe_order_mismatch_tokens": [int(v) for v in safe_order_mismatch_tokens], + "safe_token_count": int(reference_indices.shape[0] - len(unsafe_tokens)), + "num_tokens": int(reference_indices.shape[0]), + "reference_min_gap": reference_margin["min_rank_k_to_next_gap"], + "comparison_min_gap": comparison_margin["min_rank_k_to_next_gap"], + "per_token": per_token, + } + + +def _comparison_params(vram_addr: int, rows: int, cols: int, mlen: int, physical_rows: int | None = None) -> dict: + physical_rows = physical_rows or rows + num_col_blocks = math.ceil(cols / mlen) + rows_to_read = (num_col_blocks - 1) * physical_rows + rows + return { + "start_row_idx": vram_addr // mlen, + "num_rows": rows_to_read, + "num_batches": rows, + "elements_per_batch": cols, + "row_dim": mlen, + "physical_rows": physical_rows, + "atol": 0.02, + "rtol": 0.02, + } + + +def _write_json(path: Path, payload: dict) -> None: + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def _resolve_sliding_window(args: argparse.Namespace) -> tuple[int | None, str]: + if args.sliding_window is not None: + return args.sliding_window, "explicit" + if args.config_json is None: + return None, "none" + if args.layer_idx is None: + raise ValueError("--layer-idx is required when --config-json is used") + with args.config_json.expanduser().open() as f: + config = json.load(f) + layer_types = config.get("layer_types") + if not layer_types: + return None, "config:no_layer_types" + layer_type = layer_types[args.layer_idx] + if "sliding" not in layer_type: + return None, f"config:{layer_type}" + return int(config["sliding_window"]), f"config:{layer_type}" + + +def _bias_parts(args: argparse.Namespace) -> set[str]: + if not getattr(args, "projection_bias", False): + return set() + parts = (getattr(args, "bias_parts", "qkvo") or "").lower() + aliases = {"all": "qkvo", "none": ""} + parts = aliases.get(parts, parts) + allowed = set("qkvo") + unknown = set(parts) - allowed + if unknown: + raise ValueError(f"unknown --bias-parts entries: {sorted(unknown)}") + return set(parts) + + +def run_projection(args: argparse.Namespace) -> dict: + build_dir = (args.build_dir / "projection").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + rows = args.seq_len or 8 + hidden = args.hidden_size or 128 + out_features = args.out_features or 64 + x_override = _load_hidden_tensor(args.input_pt).to(torch.bfloat16) if args.input_pt is not None else None + w_override = ( + _load_hidden_tensor(args.projection_weight_pt).to(torch.bfloat16) + if args.projection_weight_pt is not None + else None + ) + golden_override = _load_hidden_tensor(args.golden_pt).to(torch.bfloat16) if args.golden_pt is not None else None + if x_override is not None: + if x_override.ndim != 2: + raise ValueError(f"--input-pt for projection must be rank-2, got shape {tuple(x_override.shape)}") + rows, hidden = map(int, x_override.shape) + if w_override is not None: + if w_override.ndim != 2: + raise ValueError( + f"--projection-weight-pt for projection must be rank-2, got shape {tuple(w_override.shape)}" + ) + if int(w_override.shape[0]) == hidden: + out_features = int(w_override.shape[1]) + elif int(w_override.shape[1]) == hidden: + # HF Linear weights are [out_features, in_features]; PLENA projection + # helpers consume [in_features, out_features]. + w_override = w_override.t().contiguous() + out_features = int(w_override.shape[1]) + else: + raise ValueError( + f"--projection-weight-pt shape {tuple(w_override.shape)} incompatible with hidden={hidden}" + ) + if rows % blen != 0: + raise ValueError(f"rows={rows} must be a multiple of BLEN={blen}") + padded_projection = bool(args.allow_padded_projection) + if not padded_projection and (hidden % mlen != 0 or out_features % mlen != 0): + raise ValueError("hidden and out_features must be multiples of MLEN for this smoke") + physical_hidden = math.ceil(hidden / mlen) * mlen + physical_out_features = math.ceil(out_features / mlen) * mlen + + torch.manual_seed(args.seed) + x_logical = x_override if x_override is not None else (torch.randn(rows, hidden) * 0.2).to(torch.bfloat16) + w_logical = w_override if w_override is not None else (torch.randn(hidden, out_features) * 0.2).to(torch.bfloat16) + b = (torch.randn(out_features) * 0.02).to(torch.bfloat16) if args.projection_bias else None + golden_logical = torch.matmul(x_logical.float(), w_logical.float()).to(torch.bfloat16) + if args.projection_bias: + golden_logical = (golden_logical.float() + b.float()).to(torch.bfloat16) + if golden_override is not None: + if tuple(golden_override.shape) != (rows, out_features): + raise ValueError( + f"--golden-pt for projection shape {tuple(golden_override.shape)} != {(rows, out_features)}" + ) + golden_logical = golden_override + x = torch.zeros(rows, physical_hidden, dtype=torch.bfloat16) + x[:, :hidden] = x_logical + w = torch.zeros(physical_hidden, physical_out_features, dtype=torch.bfloat16) + w[:hidden, :out_features] = w_logical + golden = torch.zeros(rows, physical_out_features, dtype=torch.bfloat16) + golden[:, :out_features] = golden_logical + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + physical_rows = max(mlen, math.ceil(rows / blen) * blen) + x_region = physical_rows * physical_hidden + bias_base = math.ceil(x_region / (mlen * mlen)) * (mlen * mlen) + bias_shape = (physical_rows, physical_out_features) + bias_size = math.ceil((bias_shape[0] * bias_shape[1]) / (mlen * mlen)) * (mlen * mlen) + vram_preload = torch.zeros( + max(x_region + mlen * mlen, bias_base + bias_size if args.projection_bias else 0), + dtype=torch.bfloat16, + ) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="X", + tensor=x, + vram_addr=0, + physical_shape=(physical_rows, physical_hidden), + vram_preload=vram_preload, + ) + bias_vram = None + if args.projection_bias: + bias_vram = prestage_bf16_vram_matrix( + prog=prog, + name="B_PROJ", + tensor=torch.nn.functional.pad(b.reshape(1, -1).repeat(rows, 1), (0, physical_out_features - out_features)), + vram_addr=bias_base, + physical_shape=bias_shape, + vram_preload=vram_preload, + ) + w_input = prog.input("W_QKV_BF16", shape=(physical_hidden, physical_out_features), real_data_ratio=2.0) + out = prog.linear_projection_bf16( + x_vram, + w_input, + name="bf16_projection", + physical_shape=(physical_rows, physical_out_features), + ) + if args.projection_bias: + prog.vram_add(out, bias_vram, num_rows=rows) + isa = prog.compile() + if "C_SET_SCALE_REG" in isa: + raise AssertionError("BF16 projection emitted C_SET_SCALE_REG") + if "H_PREFETCH_M" not in isa or ", 1, 1" not in isa: + raise AssertionError("BF16 projection did not use H_PREFETCH_M precision=KeyValue") + + input_tensors = {"W_QKV_BF16": w} + tensor_layouts = { + "W_QKV_BF16": { + "source_shape": [physical_hidden, physical_out_features], + "storage_shape": [physical_hidden, physical_out_features], + "source_rows": physical_hidden, + "storage_rows": physical_hidden, + "source_row_elements": physical_out_features, + "storage_row_elements": physical_out_features, + "precision": "HBM_M_KV_TYPE", + } + } + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_attention_projection_bf16", + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + params = _comparison_params( + prog._compiler.get_vram_addr(out.name), + rows, + physical_out_features, + mlen, + physical_rows=out.physical_shape[0], + ) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + if not args.no_run: + run_and_assert(build_dir, "gpt_oss_attention_projection_bf16", mlen=mlen, blen=blen) + results, _ = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(rows, physical_out_features).to(torch.bfloat16) + simulated_output_pt = build_dir / "simulated_output.pt" + torch.save(emu, simulated_output_pt) + rel = _rel_rms(emu, golden) + if rel > 0.01: + raise AssertionError(f"BF16 projection rel_rms={rel:.6g} exceeds 1%") + else: + rel = float("nan") + simulated_output_pt = None + + summary = { + "mode": "projection", + "rows": rows, + "hidden": hidden, + "out_features": out_features, + "physical_hidden": physical_hidden, + "physical_out_features": physical_out_features, + "padded_projection": padded_projection, + "scale_reg_count": isa.count("C_SET_SCALE_REG"), + "h_prefetch_keyvalue_count": isa.count("H_PREFETCH_M"), + "projection_bias": bool(args.projection_bias), + "real_input": args.input_pt is not None, + "real_weight": args.projection_weight_pt is not None, + "external_golden": args.golden_pt is not None, + "rel_rms": rel, + "simulated_output_pt": str(simulated_output_pt) if simulated_output_pt is not None else None, + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def _projection_weight_in_out(weight: torch.Tensor, in_features: int, *, label: str) -> tuple[torch.Tensor, int]: + """Return a BF16 projection weight in PLENA [in,out] orientation.""" + if weight.ndim != 2: + raise ValueError(f"{label} must be rank-2, got shape {tuple(weight.shape)}") + if int(weight.shape[0]) == in_features: + return weight.contiguous(), int(weight.shape[1]) + if int(weight.shape[1]) == in_features: + return weight.t().contiguous(), int(weight.shape[0]) + raise ValueError(f"{label} shape {tuple(weight.shape)} incompatible with in_features={in_features}") + + +def run_deepseek_mla_q_path(args: argparse.Namespace) -> dict: + """Integrated DeepSeek MLA frontend path: hidden -> proj -> RMSNorm -> proj. + + The separate real-input gates already cover the individual projections and + RMSNorms. This mode validates the VRAM handoff between those stages inside + one device program. It covers both the Q path + (hidden -> q_a -> q_a_norm -> q_b) and the KV path + (hidden -> compressed_kv -> kv_lora_norm -> kv_b), where the norm only sees + the leading kv_lora lanes of the wider compressed_kv projection. + """ + required = { + "--input-pt": args.input_pt, + "--projection-weight-pt": args.projection_weight_pt, + "--norm-weight-pt": args.norm_weight_pt, + "--second-projection-weight-pt": args.second_projection_weight_pt, + "--golden-pt": args.golden_pt, + } + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError(f"deepseek_mla_q_path requires {', '.join(missing)}") + + build_dir = (args.build_dir / args.mode).expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + x_logical = _load_hidden_tensor(args.input_pt).to(torch.bfloat16).contiguous() + if x_logical.ndim != 2: + raise ValueError(f"--input-pt must be rank-2, got shape {tuple(x_logical.shape)}") + rows, hidden = map(int, x_logical.shape) + if rows % blen != 0: + raise ValueError(f"rows={rows} must be a multiple of BLEN={blen}") + + w1_raw = _load_hidden_tensor(args.projection_weight_pt).to(torch.bfloat16) + w1_logical, first_width = _projection_weight_in_out(w1_raw, hidden, label="--projection-weight-pt") + norm_weight = _load_hidden_tensor(args.norm_weight_pt).to(torch.bfloat16).reshape(-1).contiguous() + norm_features = int(norm_weight.numel()) + if norm_features > first_width: + raise ValueError(f"--norm-weight-pt length {norm_features} cannot exceed first projection width={first_width}") + split_first_projection = norm_features < first_width + if split_first_projection: + # DeepSeek's packed kv_a projection is [kv_lora, k_rope]. RMSNorm is + # defined only on kv_lora, while the k_rope tail is nonzero. A VRAM view + # over the packed row would let the hardware norm consume the tail lanes + # because normalization walks the physical row. Split the lora branch + # here; k_rope is validated separately by the partial-RoPE/core gates. + w1_logical = w1_logical[:, :norm_features].contiguous() + first_width = norm_features + w2_raw = _load_hidden_tensor(args.second_projection_weight_pt).to(torch.bfloat16) + w2_logical, q_b_width = _projection_weight_in_out( + w2_raw, + norm_features, + label="--second-projection-weight-pt", + ) + golden_logical = _load_hidden_tensor(args.golden_pt).to(torch.bfloat16).contiguous() + if tuple(golden_logical.shape) != (rows, q_b_width): + raise ValueError(f"--golden-pt shape {tuple(golden_logical.shape)} != {(rows, q_b_width)}") + + physical_rows = max(mlen, math.ceil(rows / blen) * blen) + physical_hidden = math.ceil(hidden / mlen) * mlen + physical_first = math.ceil(first_width / mlen) * mlen + physical_q_b = math.ceil(q_b_width / mlen) * mlen + + x = torch.zeros(rows, physical_hidden, dtype=torch.bfloat16) + x[:, :hidden] = x_logical + w1 = torch.zeros(physical_hidden, physical_first, dtype=torch.bfloat16) + w1[:hidden, :first_width] = w1_logical + w2 = torch.zeros(physical_first, physical_q_b, dtype=torch.bfloat16) + w2[:norm_features, :q_b_width] = w2_logical + golden = torch.zeros(rows, physical_q_b, dtype=torch.bfloat16) + golden[:, :q_b_width] = golden_logical + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + x_base = 0 + x_region = math.ceil((physical_rows * physical_hidden) / mlen) * mlen + norm_base = math.ceil(x_region / (mlen * mlen)) * (mlen * mlen) + norm_region = math.ceil((physical_rows * physical_first) / mlen) * mlen + vram_preload = torch.zeros(norm_base + norm_region + mlen * mlen, dtype=torch.bfloat16) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="DeepSeekMLAQInput", + tensor=x, + vram_addr=x_base, + physical_shape=(physical_rows, physical_hidden), + vram_preload=vram_preload, + ) + norm_rows = norm_weight.reshape(1, -1).repeat(rows, 1).to(torch.bfloat16) + norm_vram = prestage_bf16_vram_matrix( + prog=prog, + name="DeepSeekMLAQNormWeight", + tensor=norm_rows, + vram_addr=norm_base, + physical_shape=(physical_rows, physical_first), + vram_preload=vram_preload, + ) + w1_input = prog.input( + "DeepSeekMLA_FIRST_WEIGHT_BF16", + shape=(physical_hidden, first_width), + physical_shape=(physical_hidden, physical_first), + real_data_ratio=2.0, + ) + w2_input = prog.input( + "DeepSeekMLA_SECOND_WEIGHT_BF16", + shape=(norm_features, q_b_width), + physical_shape=(physical_first, physical_q_b), + real_data_ratio=2.0, + ) + + first = prog.linear_projection_bf16( + x_vram, + w1_input, + name="deepseek_mla_first", + physical_shape=(physical_rows, physical_first), + ) + norm_input = first + if norm_features != first_width: + norm_input = prog.alloc_at( + "deepseek_mla_norm_view", + rows, + norm_features, + prog.get_vram_addr(first.name), + physical_shape=(physical_rows, physical_first), + ) + prog.rms_norm(norm_input, eps_offset=66, reci_hid_offset=67) + prog.vram_mul(norm_input, norm_vram, num_rows=rows) + second = prog.linear_projection_bf16( + norm_input, + w2_input, + name="deepseek_mla_second", + physical_shape=(physical_rows, physical_q_b), + ) + isa = prog.compile() + if "C_SET_SCALE_REG" in isa: + raise AssertionError("DeepSeek MLA BF16 Q path emitted C_SET_SCALE_REG") + + input_tensors = { + "DeepSeekMLA_FIRST_WEIGHT_BF16": w1, + "DeepSeekMLA_SECOND_WEIGHT_BF16": w2, + } + tensor_layouts = { + "DeepSeekMLA_FIRST_WEIGHT_BF16": _bf16_layout(physical_hidden, physical_first), + "DeepSeekMLA_SECOND_WEIGHT_BF16": _bf16_layout(physical_first, physical_q_b), + } + fp_preload = [0.0] * 68 + fp_preload[66] = args.norm_eps + fp_preload[67] = 1.0 / norm_features + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm=args.mode, + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + params = _comparison_params( + prog._compiler.get_vram_addr(second.name), + rows, + physical_q_b, + mlen, + physical_rows=physical_rows, + ) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + if not args.no_run: + run_emulator(build_dir) + results, _ = compare_emulator_output(build_dir) + emu_full = results["simulated_values"].reshape(rows, physical_q_b).to(torch.bfloat16) + emu = emu_full[:, :q_b_width].contiguous() + simulated_output_pt = build_dir / "simulated_output.pt" + simulated_full_output_pt = build_dir / "simulated_full_output.pt" + torch.save(emu, simulated_output_pt) + torch.save(emu_full, simulated_full_output_pt) + rel = _rel_rms(emu, golden_logical) + strict_summary = _strict_tail_summary(emu, golden_logical) + padded_tail_max_abs = float(emu_full[:, q_b_width:].abs().max().item()) if physical_q_b > q_b_width else 0.0 + if rel > 0.02: + raise AssertionError(f"DeepSeek MLA frontend path rel_rms={rel:.6g} exceeds 2%") + else: + rel = float("nan") + strict_summary = {} + padded_tail_max_abs = None + simulated_output_pt = None + simulated_full_output_pt = None + + summary = { + "mode": args.mode, + "rows": rows, + "hidden": hidden, + "first_width": first_width, + "norm_features": norm_features, + "second_width": q_b_width, + "physical_rows": physical_rows, + "physical_hidden": physical_hidden, + "physical_first": physical_first, + "physical_second": physical_q_b, + "split_first_projection": split_first_projection, + "norm_eps": args.norm_eps, + "scale_reg_count": isa.count("C_SET_SCALE_REG"), + "h_prefetch_keyvalue_count": isa.count("H_PREFETCH_M"), + "rel_rms": rel, + "strict_tail": strict_summary, + "padded_tail_max_abs": padded_tail_max_abs, + "simulated_output_pt": str(simulated_output_pt) if simulated_output_pt is not None else None, + "simulated_full_output_pt": str(simulated_full_output_pt) if simulated_full_output_pt is not None else None, + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def _make_rotate_half_matrix(head_dim: int) -> torch.Tensor: + rotate = torch.zeros(head_dim, head_dim, dtype=torch.bfloat16) + half = head_dim // 2 + for i in range(half): + rotate[i + half, i] = -1.0 + rotate[i, i + half] = 1.0 + return rotate + + +def _make_packed_rope_inputs( + seq: int, packed_heads: int, head_dim: int, theta: float +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + width = packed_heads * head_dim + rotate = torch.zeros(width, width, dtype=torch.bfloat16) + half = head_dim // 2 + freqs = 1.0 / (theta ** (torch.arange(0, half).float() / half)) + angles = torch.outer(torch.arange(seq).float(), freqs) + cos_head = torch.cat([torch.cos(angles), torch.cos(angles)], dim=-1).to(torch.bfloat16) + sin_head = torch.cat([torch.sin(angles), torch.sin(angles)], dim=-1).to(torch.bfloat16) + cos = torch.zeros(seq, width, dtype=torch.bfloat16) + sin = torch.zeros(seq, width, dtype=torch.bfloat16) + head_rotate = _make_rotate_half_matrix(head_dim) + for head in range(packed_heads): + start = head * head_dim + rotate[start : start + head_dim, start : start + head_dim] = head_rotate + cos[:, start : start + head_dim] = cos_head + sin[:, start : start + head_dim] = sin_head + return rotate.contiguous(), cos.contiguous(), sin.contiguous() + + +def _align_to_tile(value: int, mlen: int) -> int: + return math.ceil(value / (mlen * mlen)) * (mlen * mlen) + + +def _make_true_attention_hbm_inputs( + prog: PlenaCompiler, + *, + hidden: int, + q_width: int, + kv_width: int, + out_features: int, + mlen: int, + w_q: torch.Tensor, + w_k: torch.Tensor, + w_v: torch.Tensor, + w_o: torch.Tensor, + rotate: torch.Tensor, +) -> tuple[tuple[object, object, object, object, object], dict[str, torch.Tensor], dict[str, dict]]: + """Register the BF16 HBM inputs shared by true-shape attention adapters.""" + w_q_input = prog.input("W_Q_TRUE_BF16", shape=(hidden, q_width), real_data_ratio=2.0) + w_k_input = prog.input("W_K_TRUE_BF16", shape=(hidden, kv_width), real_data_ratio=2.0) + w_v_input = prog.input("W_V_TRUE_BF16", shape=(hidden, kv_width), real_data_ratio=2.0) + w_o_input = prog.input("W_O_TRUE_BF16", shape=(q_width, out_features), real_data_ratio=2.0) + rotate_input = prog.input("ROPE_ROTATE_TRUE_FULL_BF16", shape=(mlen, mlen), real_data_ratio=2.0) + input_tensors = { + "W_Q_TRUE_BF16": w_q, + "W_K_TRUE_BF16": w_k, + "W_V_TRUE_BF16": w_v, + "W_O_TRUE_BF16": w_o, + "ROPE_ROTATE_TRUE_FULL_BF16": rotate, + } + tensor_layouts = { + "W_Q_TRUE_BF16": _bf16_layout(hidden, q_width), + "W_K_TRUE_BF16": _bf16_layout(hidden, kv_width), + "W_V_TRUE_BF16": _bf16_layout(hidden, kv_width), + "W_O_TRUE_BF16": _bf16_layout(q_width, out_features), + "ROPE_ROTATE_TRUE_FULL_BF16": _bf16_layout(mlen, mlen), + } + return (w_q_input, w_k_input, w_v_input, w_o_input, rotate_input), input_tensors, tensor_layouts + + +def _emit_true_qkv_projection_bf16( + prog: PlenaCompiler, + projection_input, + *, + w_q_input, + w_k_input, + w_v_input, + bias_vrams: dict[str, object], + bias_parts: set[str], + physical_rows: int, + q_width: int, + kv_width: int, + seq: int, +): + """Emit BF16 Q/K/V projection plus optional BF16 bias for true attention.""" + q_var = prog.linear_projection_bias_bf16( + projection_input, + w_q_input, + bias_var=bias_vrams["B_Q_TRUE"] if "q" in bias_parts else None, + name="Q_true_proj", + physical_shape=(physical_rows, q_width), + bias_rows=seq, + ) + k_var = prog.linear_projection_bias_bf16( + projection_input, + w_k_input, + bias_var=bias_vrams["B_K_TRUE"] if "k" in bias_parts else None, + name="K_true_proj", + physical_shape=(physical_rows, kv_width), + bias_rows=seq, + ) + v_var = prog.linear_projection_bias_bf16( + projection_input, + w_v_input, + bias_var=bias_vrams["B_V_TRUE"] if "v" in bias_parts else None, + name="V_true_proj", + physical_shape=(physical_rows, kv_width), + bias_rows=seq, + ) + return q_var, k_var, v_var + + +def _emit_true_output_projection_bf16( + prog: PlenaCompiler, + attn_out, + *, + w_o_input, + bias_vrams: dict[str, object], + bias_parts: set[str], + physical_rows: int, + out_features: int, + seq: int, +): + """Emit BF16 O projection plus optional BF16 bias for true attention.""" + return prog.linear_projection_bias_bf16( + attn_out, + w_o_input, + bias_var=bias_vrams["B_O_TRUE"] if "o" in bias_parts else None, + name="O_true_proj", + physical_shape=(physical_rows, out_features), + bias_rows=seq, + ) + + +def _stage_true_kv_heads_bf16( + prog: PlenaCompiler, + *, + k_var, + v_var, + rotate_input, + cos_vram, + sin_vram, + k_norm_weight_vram, + qk_head_norm: bool, + qk_norm_eps_offset: int | None, + qk_norm_reci_head_offset: int | None, + seq: int, + mlen: int, + hkv: int, + physical_rows: int, + input_tensors: dict[str, torch.Tensor], + tensor_layouts: dict[str, dict], +) -> tuple[list[object], list[object]]: + """Apply runtime RoPE to K heads, store K/V heads, and register layouts.""" + k_inputs = [] + v_inputs = [] + k_base = prog.get_vram_addr(k_var.name) + v_base = prog.get_vram_addr(v_var.name) + for kv_idx in range(hkv): + k_head = prog.alloc_at( + f"K_true_head_{kv_idx}", + seq, + mlen, + k_base + kv_idx * physical_rows * mlen, + physical_shape=(physical_rows, mlen), + ) + prog.head_runtime_rope_bf16( + k_head, + rotate_input, + cos_vram, + sin_vram, + norm_weight_var=k_norm_weight_vram if qk_head_norm else None, + eps_offset=qk_norm_eps_offset if qk_head_norm else None, + reci_hid_offset=qk_norm_reci_head_offset if qk_head_norm else None, + num_rows=seq, + name=f"K_true_runtime_rotate_h{kv_idx}", + ) + v_head = prog.alloc_at( + f"V_true_head_{kv_idx}", + seq, + mlen, + v_base + kv_idx * physical_rows * mlen, + physical_shape=(physical_rows, mlen), + ) + k_staged = prog.store( + k_head, + name=f"K_true_staged_h{kv_idx}", + precision=1, + hbm_element_bytes=2, + real_data_ratio=2.0, + ) + v_staged = prog.store( + v_head, + name=f"V_true_staged_h{kv_idx}", + precision=1, + hbm_element_bytes=2, + real_data_ratio=2.0, + ) + k_inputs.append(k_staged) + v_inputs.append(v_staged) + input_tensors[k_staged.name] = torch.zeros(physical_rows, mlen, dtype=torch.bfloat16) + input_tensors[v_staged.name] = torch.zeros(physical_rows, mlen, dtype=torch.bfloat16) + tensor_layouts[k_staged.name] = _bf16_layout(physical_rows, mlen, precision="HBM_V_KV_TYPE") + tensor_layouts[v_staged.name] = _bf16_layout(physical_rows, mlen, precision="HBM_V_KV_TYPE") + return k_inputs, v_inputs + + +def _emit_true_q_heads_attention_bf16( + prog: PlenaCompiler, + *, + q_var, + rotate_input, + cos_vram, + sin_vram, + q_norm_weight_vram, + qk_head_norm: bool, + qk_norm_eps_offset: int | None, + qk_norm_reci_head_offset: int | None, + seq: int, + mlen: int, + hq: int, + hkv: int, + head_dim: int, + physical_rows: int, + k_inputs: list[object], + v_inputs: list[object], + attn_out, + scratch, + causal_mask, + scale: float, + sink_base: int | None, + q_pre: torch.Tensor, + q_norm_weight: torch.Tensor | None, + qk_norm_eps: float, + q: torch.Tensor, + debug_stop_after_first_q_norm: bool, + debug_stop_after_first_q_head: bool, +) -> tuple[object, torch.Tensor | None, bool]: + """Emit per-Q-head runtime RoPE and packed attention for true GPT-OSS shapes.""" + q_base = prog.get_vram_addr(q_var.name) + out_base = prog.get_vram_addr(attn_out.name) + ratio = hq // hkv + for head in range(hq): + q_head = prog.alloc_at( + f"Q_true_head_{head}", + seq, + mlen, + q_base + head * physical_rows * mlen, + physical_shape=(physical_rows, mlen), + ) + if debug_stop_after_first_q_norm: + if qk_head_norm: + prog.rms_norm(q_head, eps_offset=qk_norm_eps_offset, reci_hid_offset=qk_norm_reci_head_offset) + prog.vram_mul(q_head, q_norm_weight_vram, num_rows=seq) + golden = _head_rms_norm_golden( + q_pre.reshape(seq, hq, head_dim)[:, head : head + 1, :], + q_norm_weight, + qk_norm_eps, + ).reshape(seq, head_dim) + else: + golden = q_pre.reshape(seq, hq, head_dim)[:, head, :].contiguous() + return q_head, golden, True + prog.head_runtime_rope_bf16( + q_head, + rotate_input, + cos_vram, + sin_vram, + norm_weight_var=q_norm_weight_vram if qk_head_norm else None, + eps_offset=qk_norm_eps_offset if qk_head_norm else None, + reci_hid_offset=qk_norm_reci_head_offset if qk_head_norm else None, + num_rows=seq, + name=f"Q_true_runtime_rotate_h{head}", + ) + if debug_stop_after_first_q_head: + return q_head, q[:, head, :].contiguous(), True + kv_idx = head // ratio + prog.flash_attention_packed_group( + q_head, + k_inputs[kv_idx], + v_inputs[kv_idx], + group_heads=1, + head_slot_dim=head_dim, + output_base_address=out_base + head * physical_rows * mlen, + scratch_base_address=prog.get_vram_addr(scratch.name), + broadcast_amount=1, + scale=scale, + causal_mask=causal_mask, + valid_cols=seq, + sink_base_address=None if sink_base is None else sink_base + head, + output_head_base=0, + k_matrix_precision="keyvalue", + k_set_scale=False, + k_hbm_element_bytes=2, + v_hbm_element_bytes=2, + ) + return attn_out, None, False + + +def run_runtime_rope(args: argparse.Namespace) -> dict: + build_dir = (args.build_dir / "runtime_rope").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + seq = args.seq_len or 16 + hidden = args.hidden_size or 128 + packed_heads = 4 + head_dim = mlen // packed_heads + if mlen % packed_heads != 0: + raise ValueError(f"MLEN={mlen} must be divisible by packed_heads={packed_heads}") + if seq > mlen: + raise ValueError("runtime_rope smoke is single sequence tile only") + if seq % blen != 0: + raise ValueError(f"seq={seq} must be a multiple of BLEN={blen}") + if hidden % mlen != 0: + raise ValueError("hidden must be a multiple of MLEN") + + torch.manual_seed(args.seed) + x = (torch.randn(seq, hidden) * 0.2).to(torch.bfloat16) + w_q = (torch.randn(hidden, mlen) * 0.2).to(torch.bfloat16) + b_q = (torch.randn(mlen) * 0.02).to(torch.bfloat16) + rotate, cos, sin = _make_packed_rope_inputs(seq, packed_heads, head_dim, args.rope_theta) + + q_pre = torch.matmul(x.float(), w_q.float()).to(torch.bfloat16) + q_pre = (q_pre.float() + b_q.float()).to(torch.bfloat16) + q_heads = q_pre.reshape(seq, packed_heads, head_dim) + cos_head = cos[:, :head_dim] + sin_head = sin[:, :head_dim] + golden = _apply_rope_bf16(q_heads, cos_head, sin_head).reshape(seq, mlen).contiguous() + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + physical_rows = max(mlen, math.ceil(seq / blen) * blen) + x_region = physical_rows * hidden + bias_base = _align_to_tile(x_region, mlen) + bias_shape = (physical_rows, mlen) + bias_size = _align_to_tile(bias_shape[0] * bias_shape[1], mlen) + cos_base = bias_base + bias_size + trig_shape = (physical_rows, mlen) + trig_size = _align_to_tile(trig_shape[0] * trig_shape[1], mlen) + sin_base = cos_base + trig_size + vram_preload = torch.zeros(sin_base + trig_size, dtype=torch.bfloat16) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="X", + tensor=x, + vram_addr=0, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + bias_vram = prestage_bf16_vram_matrix( + prog=prog, + name="B_Q", + tensor=b_q.reshape(1, -1).repeat(seq, 1), + vram_addr=bias_base, + physical_shape=bias_shape, + vram_preload=vram_preload, + ) + cos_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_COS", + tensor=cos, + vram_addr=cos_base, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + sin_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_SIN", + tensor=sin, + vram_addr=sin_base, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + + w_q_input = prog.input("W_Q_BF16", shape=(hidden, mlen), real_data_ratio=2.0) + rotate_input = prog.input("ROPE_ROTATE_BF16", shape=(mlen, mlen), real_data_ratio=2.0) + q_var = prog.linear_projection_bf16(x_vram, w_q_input, name="Q_proj", physical_shape=(physical_rows, mlen)) + prog.vram_add(q_var, bias_vram, num_rows=seq) + prog.runtime_rope_projection_bf16(q_var, rotate_input, cos_vram, sin_vram, name="Q_runtime_rotate") + isa = prog.compile() + if "C_SET_SCALE_REG" in isa: + raise AssertionError("runtime BF16 projection/RoPE path emitted C_SET_SCALE_REG") + + input_tensors = { + "W_Q_BF16": w_q, + "ROPE_ROTATE_BF16": rotate, + } + tensor_layouts = { + "W_Q_BF16": _bf16_layout(hidden, mlen), + "ROPE_ROTATE_BF16": _bf16_layout(mlen, mlen), + } + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_attention_runtime_rope", + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + params = _comparison_params( + prog.get_vram_addr(q_var.name), + seq, + mlen, + mlen, + physical_rows=q_var.physical_shape[0], + ) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + if not args.no_run: + run_and_assert(build_dir, "gpt_oss_attention_runtime_rope", mlen=mlen, blen=blen) + results, _ = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(seq, mlen).to(torch.bfloat16) + rel = _rel_rms(emu, golden) + if rel > 0.01: + raise AssertionError(f"runtime_rope rel_rms={rel:.6g} exceeds 1%") + else: + rel = float("nan") + + summary = { + "mode": "runtime_rope", + "seq": seq, + "hidden": hidden, + "packed_heads": packed_heads, + "head_dim": head_dim, + "runtime_projection": True, + "runtime_bias": True, + "runtime_rope": True, + "scale_reg_count": isa.count("C_SET_SCALE_REG"), + "rel_rms": rel, + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def _gpt_oss_attention_ref( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + sinks: torch.Tensor | None, + *, + scale: float, + sliding_window: int | None, + causal: bool = True, +) -> torch.Tensor: + # q: [seq, hq, d], k/v: [seq, hkv, d] + seq, hq, d = q.shape + hkv = k.shape[1] + ratio = hq // hkv + qh = q.transpose(0, 1).float() + kh = k.repeat_interleave(ratio, dim=1).transpose(0, 1).float() + vh = v.repeat_interleave(ratio, dim=1).transpose(0, 1).float() + scores = torch.matmul(qh, kh.transpose(-1, -2)) * scale + row = torch.arange(seq).reshape(1, seq, 1) + col = torch.arange(seq).reshape(1, 1, seq) + if causal: + mask = col <= row + if sliding_window is not None: + mask = mask & (col > row - sliding_window) + scores = scores.masked_fill(~mask, float("-inf")) + if sinks is not None: + sink_logits = sinks.float().reshape(hq, 1, 1).expand(hq, seq, 1) + logits = torch.cat([scores, sink_logits], dim=-1) + probs = F.softmax(logits, dim=-1)[..., :seq] + else: + probs = F.softmax(scores, dim=-1) + out = torch.matmul(probs, vh) + return out.transpose(0, 1).reshape(seq, hq * d).to(torch.bfloat16) + + +def _causal_score_mask_tensor(seq: int, mlen: int, *, sliding_window: int | None = None) -> torch.Tensor: + """Build a BF16 score mask for single-tile causal/sliding attention.""" + if seq < 0 or seq > mlen: + raise ValueError(f"seq must be in [0, {mlen}], got {seq}") + row = torch.arange(mlen).reshape(mlen, 1) + col = torch.arange(mlen).reshape(1, mlen) + visible = (row < seq) & (col < seq) & (col <= row) + if sliding_window is not None: + visible = visible & (col > row - sliding_window) + mask = torch.full((mlen, mlen), float("-inf"), dtype=torch.bfloat16) + mask[visible] = torch.tensor(0.0, dtype=torch.bfloat16) + return mask + + +def _build_score_mask_vram( + prog: PlenaCompiler, + name: str, + *, + seq: int, + mlen: int, + sliding_window: int | None = None, +): + """Build a single-tile causal/sliding score mask directly in VRAM. + + Keep -inf on the scalar/VRAM path instead of loading it through HBM, where + activation quantization can change the masking semantics. + """ + mask = prog.alloc(name, mlen, mlen) + mask_addr = prog.get_vram_addr(mask.name) + fp_base = prog._ONLINE_SOFTMAX_FPSRAM_BASE + gp_mask, gp_fp = prog.register_allocator.allocate_gp(2) + lines = [ + f"; === Build causal score mask: MLEN={mlen}, seq={seq}, sliding_window={sliding_window} ===", + f"S_ADDI_INT gp{gp_fp}, gp0, {fp_base}", + "S_LD_FP f7, gp0, 2", + ] + lines.extend(load_large_int(gp_mask, mask_addr)) + for row in range(mlen): + for col in range(mlen): + visible = row < seq and col < seq and col <= row + if sliding_window is not None: + visible = visible and col > row - sliding_window + lines.append(f"S_ST_FP {'f0' if visible else 'f7'}, gp{gp_fp}, {col}") + lines.append(f"S_MAP_V_FP gp{gp_mask}, gp{gp_fp}, 0") + lines.append(f"S_ADDI_INT gp{gp_mask}, gp{gp_mask}, {mlen}") + prog.register_allocator.free_gp([gp_mask, gp_fp]) + prog.emit("\n".join(lines) + "\n") + return mask + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + half = x.shape[-1] // 2 + return torch.cat([-x[..., half:], x[..., :half]], dim=-1) + + +def _make_rope_tables(seq: int, head_dim: int, *, theta: float = 10000.0) -> tuple[torch.Tensor, torch.Tensor]: + half = head_dim // 2 + freqs = 1.0 / (theta ** (torch.arange(0, half).float() / half)) + positions = torch.arange(seq).float() + angles = torch.outer(positions, freqs) + cos_half = torch.cos(angles) + sin_half = torch.sin(angles) + return torch.cat([cos_half, cos_half], dim=-1).to(torch.bfloat16), torch.cat([sin_half, sin_half], dim=-1).to( + torch.bfloat16 + ) + + +def _apply_rope_bf16(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: + return (x.float() * cos[:, None, :].float() + _rotate_half(x).float() * sin[:, None, :].float()).to(torch.bfloat16) + + +def run_true_core(args: argparse.Namespace) -> dict: + """True-shape GPT-OSS attention core: 64Q/8KV via unrolled per-Q-head groups. + + By default this starts after projection+bias+RoPE: Q/K/V are prestaged as + BF16 tensors. With --runtime-rope, Q/K are prestaged before RoPE and the + device computes rotate-half plus RoPE before running the same attention core. + """ + build_dir = (args.build_dir / "true_core").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + hq = args.num_attention_heads + hkv = args.num_key_value_heads + head_dim = args.hlen or mlen + value_head_dim = args.value_head_dim or head_dim + seq = args.seq_len or 16 + if head_dim != mlen: + raise ValueError(f"true_core currently uses one Q head per MLEN row; got head_dim={head_dim}, MLEN={mlen}") + if hq % hkv != 0: + raise ValueError(f"num_attention_heads={hq} must be divisible by num_key_value_heads={hkv}") + if seq > mlen: + raise ValueError("true_core is single sequence tile only") + if seq % blen != 0: + raise ValueError(f"seq={seq} must be a multiple of BLEN={blen}") + sliding_window, sliding_source = _resolve_sliding_window(args) + runtime_rope = bool(args.runtime_rope) + external_qkv = any(getattr(args, name) is not None for name in ("q_pt", "k_pt", "v_pt")) + if external_qkv and not all(getattr(args, name) is not None for name in ("q_pt", "k_pt", "v_pt")): + raise ValueError("--q-pt, --k-pt, and --v-pt must be provided together") + if external_qkv and runtime_rope: + raise ValueError( + "external Q/K tensors are expected post-RoPE; do not combine --q-pt/--k-pt with --runtime-rope" + ) + o_projection = args.o_weight_pt is not None + if args.o_golden_pt is not None and not o_projection: + raise ValueError("--o-golden-pt requires --o-weight-pt") + + torch.manual_seed(args.seed) + if external_qkv: + q = _load_hidden_tensor(args.q_pt).to(torch.bfloat16) + k = _load_hidden_tensor(args.k_pt).to(torch.bfloat16) + v_loaded = _load_hidden_tensor(args.v_pt).to(torch.bfloat16) + if q.dim() == 2: + q = q.reshape(seq, hq, head_dim) + if k.dim() == 2: + k = k.reshape(seq, hkv, head_dim) + if v_loaded.dim() == 2: + inferred_v_dim = v_loaded.shape[-1] // hkv + v_loaded = v_loaded.reshape(seq, hkv, inferred_v_dim) + if tuple(q.shape) != (seq, hq, head_dim): + raise ValueError(f"--q-pt shape {tuple(q.shape)} != {(seq, hq, head_dim)}") + if tuple(k.shape) != (seq, hkv, head_dim): + raise ValueError(f"--k-pt shape {tuple(k.shape)} != {(seq, hkv, head_dim)}") + value_head_dim = args.value_head_dim or v_loaded.shape[-1] + if tuple(v_loaded.shape) != (seq, hkv, value_head_dim): + raise ValueError(f"--v-pt shape {tuple(v_loaded.shape)} != {(seq, hkv, value_head_dim)}") + q_pre = q + k_pre = k + else: + if not (0 < value_head_dim <= head_dim): + raise ValueError(f"value_head_dim must be in (0, head_dim], got {value_head_dim} for head_dim={head_dim}") + q_pre = (torch.randn(seq, hq, head_dim) * 0.2).to(torch.bfloat16) + k_pre = (torch.randn(seq, hkv, head_dim) * 0.2).to(torch.bfloat16) + v_loaded = (torch.randn(seq, hkv, value_head_dim) * 0.2).to(torch.bfloat16) + cos, sin = _make_rope_tables(seq, head_dim, theta=args.rope_theta) + q = _apply_rope_bf16(q_pre, cos, sin) + k = _apply_rope_bf16(k_pre, cos, sin) + if not (0 < value_head_dim <= head_dim): + raise ValueError(f"value_head_dim must be in (0, head_dim], got {value_head_dim} for head_dim={head_dim}") + v = torch.zeros(seq, hkv, head_dim, dtype=torch.bfloat16) + v[:, :, :value_head_dim] = v_loaded + sinks = (torch.randn(hq) * 0.1).to(torch.bfloat16) if args.sink else None + scale = float(args.attention_scale) if args.attention_scale is not None else 1.0 / math.sqrt(head_dim) + if args.golden_pt is not None: + golden = _load_hidden_tensor(args.golden_pt).to(torch.bfloat16) + if golden.dim() == 2: + golden = golden.reshape(seq, hq, head_dim) + if tuple(golden.shape) != (seq, hq, head_dim): + raise ValueError(f"--golden-pt shape {tuple(golden.shape)} != {(seq, hq, head_dim)}") + golden = golden.reshape(seq, hq * head_dim).contiguous() + else: + golden = _gpt_oss_attention_ref( + q, + k, + v, + sinks, + scale=scale, + sliding_window=sliding_window, + causal=not args.no_causal_mask, + ) + + o_weight_expanded = None + o_logical_out_features = None + o_physical_out_features = None + o_golden_from_padded_rel = None + if o_projection: + o_weight_raw = _load_hidden_tensor(args.o_weight_pt).to(torch.bfloat16) + logical_o_in = hq * value_head_dim + padded_o_in = hq * head_dim + if o_weight_raw.dim() != 2: + raise ValueError(f"--o-weight-pt must be rank-2, got shape {tuple(o_weight_raw.shape)}") + if o_weight_raw.shape[0] == logical_o_in: + o_weight_logical = o_weight_raw + elif o_weight_raw.shape[1] == logical_o_in: + o_weight_logical = o_weight_raw.T.contiguous() + else: + raise ValueError( + f"--o-weight-pt shape {tuple(o_weight_raw.shape)} is incompatible with logical input " + f"{logical_o_in}; expected [in,out] or [out,in]" + ) + o_logical_out_features = int(o_weight_logical.shape[1]) + o_physical_out_features = math.ceil(o_logical_out_features / mlen) * mlen + o_weight_expanded = torch.zeros(padded_o_in, o_physical_out_features, dtype=torch.bfloat16) + for head in range(hq): + logical_start = head * value_head_dim + padded_start = head * head_dim + o_weight_expanded[padded_start : padded_start + value_head_dim, :o_logical_out_features] = o_weight_logical[ + logical_start : logical_start + value_head_dim, : + ] + golden_o = torch.matmul(golden.float(), o_weight_expanded[:, :o_logical_out_features].float()).to( + torch.bfloat16 + ) + if args.o_golden_pt is not None: + o_golden_ref = _load_hidden_tensor(args.o_golden_pt).to(torch.bfloat16) + if tuple(o_golden_ref.shape) != (seq, o_logical_out_features): + raise ValueError(f"--o-golden-pt shape {tuple(o_golden_ref.shape)} != {(seq, o_logical_out_features)}") + o_golden_from_padded_rel = _rel_rms(golden_o, o_golden_ref) + golden_o = o_golden_ref + golden_padded = torch.zeros(seq, o_physical_out_features, dtype=torch.bfloat16) + golden_padded[:, :o_logical_out_features] = golden_o + golden = golden_padded + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + physical_rows = max(mlen, math.ceil(seq / blen) * blen) + q_width = hq * head_dim + q_region = physical_rows * q_width + k_pre_base = q_region + k_pre_region = hkv * physical_rows * mlen if runtime_rope else 0 + rope_base = _align_to_tile(k_pre_base + k_pre_region, mlen) + trig_shape = (physical_rows, mlen) + trig_size = _align_to_tile(trig_shape[0] * trig_shape[1], mlen) + preload_words = q_region + if runtime_rope: + preload_words = rope_base + 2 * trig_size + q_vram_preload = torch.zeros(preload_words, dtype=torch.bfloat16) + q_vram = prestage_bf16_vram_matrix( + prog=prog, + name="Q_pre_true" if runtime_rope else "Q_rope_true", + tensor=(q_pre if runtime_rope else q).reshape(seq, q_width), + vram_addr=0, + physical_shape=(physical_rows, q_width), + vram_preload=q_vram_preload, + ) + rotate_input = None + cos_vram = sin_vram = None + if runtime_rope: + rotate, cos_packed, sin_packed = _make_packed_rope_inputs(seq, 1, head_dim, args.rope_theta) + cos_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_COS_TRUE", + tensor=cos_packed, + vram_addr=rope_base, + physical_shape=trig_shape, + vram_preload=q_vram_preload, + ) + sin_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_SIN_TRUE", + tensor=sin_packed, + vram_addr=rope_base + trig_size, + physical_shape=trig_shape, + vram_preload=q_vram_preload, + ) + rotate_input = prog.input("ROPE_ROTATE_TRUE_BF16", shape=(mlen, mlen), real_data_ratio=2.0) + + k_inputs = [] + v_inputs = [] + input_tensors = {} + tensor_layouts = {} + for kv_idx in range(hkv): + k_name = f"K_rope_h{kv_idx}" + v_name = f"V_h{kv_idx}" + if runtime_rope: + k_pre_vram = prestage_bf16_vram_matrix( + prog=prog, + name=f"K_pre_h{kv_idx}", + tensor=k_pre[:, kv_idx, :], + vram_addr=k_pre_base + kv_idx * physical_rows * mlen, + physical_shape=(physical_rows, mlen), + vram_preload=q_vram_preload, + ) + prog.runtime_rope_projection_bf16( + k_pre_vram, + rotate_input, + cos_vram, + sin_vram, + name=f"K_runtime_rotate_h{kv_idx}", + ) + k_input = prog.store( + k_pre_vram, + name=k_name, + precision=1, + hbm_element_bytes=2, + real_data_ratio=2.0, + ) + input_tensors[k_name] = torch.zeros(physical_rows, mlen, dtype=torch.bfloat16) + else: + k_input = prog.input(k_name, shape=(seq, mlen), physical_shape=(physical_rows, mlen)) + k_padded = torch.zeros(physical_rows, mlen, dtype=torch.bfloat16) + k_padded[:seq, :] = k[:, kv_idx, :] + input_tensors[k_name] = k_padded + v_input = prog.input(v_name, shape=(seq, mlen), physical_shape=(physical_rows, mlen)) + k_inputs.append(k_input) + v_inputs.append(v_input) + v_padded = torch.zeros(physical_rows, mlen, dtype=torch.bfloat16) + v_padded[:seq, :] = v[:, kv_idx, :] + input_tensors[v_name] = v_padded + tensor_layouts[k_name] = _bf16_layout(physical_rows, mlen) + tensor_layouts[v_name] = _bf16_layout(physical_rows, mlen) + if runtime_rope: + input_tensors["ROPE_ROTATE_TRUE_BF16"] = rotate + tensor_layouts["ROPE_ROTATE_TRUE_BF16"] = _bf16_layout(mlen, mlen) + + causal_mask = None + if not args.no_causal_mask: + mask_vram_name = ( + "_gpt_oss_true_core_sliding_mask" if sliding_window is not None else "_gpt_oss_true_core_causal_mask" + ) + causal_mask = _build_score_mask_vram( + prog, + mask_vram_name, + seq=seq, + mlen=mlen, + sliding_window=sliding_window, + ) + + out = prog.alloc("O_true_64q8kv", seq, q_width, strict=False, physical_shape=(physical_rows, q_width)) + scratch = prog.alloc("_gpt_oss_true_core_scratch", mlen, mlen, strict=True) + + fp_preload = [0.0, scale / 0.25, float("-inf")] + [0.0] * 253 + sink_base = None + if sinks is not None: + sink_base = 256 + if len(fp_preload) < sink_base + hq: + fp_preload.extend([0.0] * (sink_base + hq - len(fp_preload))) + for idx, value in enumerate(sinks.tolist()): + fp_preload[sink_base + idx] = float(value) + + q_base = prog.get_vram_addr(q_vram.name) + out_base = prog.get_vram_addr(out.name) + ratio = hq // hkv + for head in range(hq): + q_head = prog.alloc_at( + f"Q_head_{head}", + seq, + mlen, + q_base + head * physical_rows * mlen, + physical_shape=(physical_rows, mlen), + ) + if runtime_rope: + prog.runtime_rope_projection_bf16( + q_head, + rotate_input, + cos_vram, + sin_vram, + name=f"Q_runtime_rotate_h{head}", + ) + kv_idx = head // ratio + prog.flash_attention_packed_group( + q_head, + k_inputs[kv_idx], + v_inputs[kv_idx], + group_heads=1, + head_slot_dim=head_dim, + output_base_address=out_base + head * physical_rows * mlen, + scratch_base_address=prog.get_vram_addr(scratch.name), + broadcast_amount=1, + scale=scale, + causal_mask=causal_mask, + valid_cols=seq, + k_matrix_precision="keyvalue", + k_set_scale=False, + k_hbm_element_bytes=2, + ) + if o_projection: + o_weight_input = prog.input( + "O_proj_padded_weight", + shape=(q_width, o_physical_out_features), + ) + input_tensors["O_proj_padded_weight"] = o_weight_expanded + tensor_layouts["O_proj_padded_weight"] = _bf16_layout(q_width, o_physical_out_features) + out = prog.linear_projection_bf16( + out, + o_weight_input, + name="O_true_padded_projection", + physical_shape=(physical_rows, o_physical_out_features), + ) + isa = prog.compile() + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=q_vram_preload, + tensor_layouts=tensor_layouts, + ) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_attention_true_core", + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + compare_width = out.shape[1] + params = _comparison_params( + prog.get_vram_addr(out.name), + seq, + compare_width, + mlen, + physical_rows=out.physical_shape[0], + ) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + if not args.no_run: + run_emulator(build_dir) + results, _ = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(seq, compare_width).to(torch.bfloat16) + simulated_output_pt = build_dir / "simulated_output.pt" + torch.save(emu, simulated_output_pt) + rel = _rel_rms(emu, golden) + if rel > 0.02: + raise AssertionError(f"true_core rel_rms={rel:.6g} exceeds 2%") + else: + rel = float("nan") + emu = torch.empty(0, dtype=torch.bfloat16) + simulated_output_pt = None + + padded_value_tail_max_abs = None + if value_head_dim < head_dim and emu.numel(): + padded_value_tail_max_abs = float( + emu.reshape(seq, hq, head_dim)[:, :, value_head_dim:].abs().max().item() if not o_projection else 0.0 + ) + + summary = { + "mode": "true_core", + "seq": seq, + "num_attention_heads": hq, + "num_key_value_heads": hkv, + "head_dim": head_dim, + "value_head_dim": value_head_dim, + "o_projection": bool(o_projection), + "o_logical_out_features": o_logical_out_features, + "o_physical_out_features": o_physical_out_features, + "o_golden_from_padded_rel_rms": o_golden_from_padded_rel, + "attention_scale": scale, + "padded_value_tail_max_abs": padded_value_tail_max_abs, + "sink": bool(args.sink), + "sliding_window": sliding_window, + "sliding_source": sliding_source, + "rel_rms": rel, + "sink_base": sink_base, + "strategy": "unrolled_one_q_head_per_call", + "runtime_projection": False, + "runtime_rope": runtime_rope, + "simulated_output_pt": str(simulated_output_pt) if simulated_output_pt is not None else None, + "note": ( + "Q/K are prestaged before RoPE and device-side runtime RoPE feeds true-shape attention core." + if runtime_rope + else "Q/K/V are prestaged after bias/RoPE; this validates true-shape per-head sink attention core only." + ), + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def run_qk_head_norm_smoke(args: argparse.Namespace) -> dict: + """Isolate Qwen-style per-head RMSNorm on one BF16 VRAM head. + + The full Qwen-like attention probe produced all-NaN output once Q/K head + RMSNorm was inserted. This smoke keeps attention and RoPE out of the + picture so a failure points at the norm lowering itself rather than the + downstream online-softmax path. + """ + build_dir = (args.build_dir / "qk_head_norm_smoke").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + seq = args.seq_len or 8 + head_dim = args.hlen or mlen + eps = args.qk_head_norm_eps + if not (0 < head_dim <= mlen): + raise ValueError(f"qk_head_norm_smoke expects 0 < head_dim <= MLEN; got {head_dim} and {mlen}") + if seq > mlen: + raise ValueError("qk_head_norm_smoke is single sequence tile only") + if seq % blen != 0: + raise ValueError(f"seq={seq} must be a multiple of BLEN={blen}") + + torch.manual_seed(args.seed) + x = (torch.randn(seq, head_dim) * 0.2).to(torch.bfloat16) + input_source = "random" + if getattr(args, "input_pt", None) is not None: + loaded_x = torch.load(args.input_pt.expanduser()) + if isinstance(loaded_x, dict): + loaded_x = loaded_x.get("x", loaded_x.get("tensor", loaded_x.get("hidden", loaded_x))) + if not torch.is_tensor(loaded_x): + raise TypeError(f"--input-pt must contain a Tensor or tensor dict, got {type(loaded_x).__name__}") + loaded_x = loaded_x.to(torch.bfloat16).contiguous() + if tuple(loaded_x.shape) != (seq, head_dim): + raise ValueError(f"--input-pt shape {tuple(loaded_x.shape)} != expected {(seq, head_dim)}") + x = loaded_x + input_source = str(args.input_pt.expanduser()) + weight = (1.0 + torch.randn(head_dim) * 0.05).to(torch.bfloat16) + if getattr(args, "qk_norm_weight_pt", None) is not None: + loaded_weight = torch.load(args.qk_norm_weight_pt.expanduser()) + if isinstance(loaded_weight, dict): + loaded_weight = loaded_weight.get("weight", loaded_weight.get("tensor", loaded_weight)) + if not torch.is_tensor(loaded_weight): + raise TypeError( + f"--qk-norm-weight-pt must contain a Tensor or tensor dict, got {type(loaded_weight).__name__}" + ) + loaded_weight = loaded_weight.to(torch.bfloat16).contiguous().reshape(-1) + if tuple(loaded_weight.shape) != (head_dim,): + raise ValueError(f"--qk-norm-weight-pt shape {tuple(loaded_weight.shape)} != expected {(head_dim,)}") + weight = loaded_weight + weight_rows = weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) + golden = _head_rms_norm_golden(x.reshape(seq, 1, head_dim), weight, eps).reshape(seq, head_dim) + rotate = None + if args.runtime_rope: + rotate, cos, sin = _make_packed_rope_inputs(seq, 1, head_dim, args.rope_theta) + if head_dim < mlen: + rotate_padded = torch.zeros(mlen, mlen, dtype=torch.bfloat16) + rotate_padded[:head_dim, :head_dim] = rotate + cos_padded = torch.zeros(seq, mlen, dtype=torch.bfloat16) + sin_padded = torch.zeros(seq, mlen, dtype=torch.bfloat16) + cos_padded[:, :head_dim] = cos + sin_padded[:, :head_dim] = sin + rotate, cos, sin = rotate_padded, cos_padded, sin_padded + cos_head, sin_head = _make_rope_tables(seq, head_dim, theta=args.rope_theta) + golden = _apply_rope_bf16(golden.reshape(seq, 1, head_dim), cos_head, sin_head).reshape(seq, head_dim) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + physical_rows = max(mlen, math.ceil(seq / blen) * blen) + x_size = physical_rows * mlen + weight_base = _align_to_tile(x_size, mlen) + weight_size = _align_to_tile(physical_rows * mlen, mlen) + rope_base = _align_to_tile(weight_base + weight_size, mlen) + trig_shape = (physical_rows, mlen) + trig_size = _align_to_tile(trig_shape[0] * trig_shape[1], mlen) + vram_words = rope_base + 2 * trig_size if args.runtime_rope else weight_base + weight_size + mlen * mlen + vram_preload = torch.zeros(vram_words, dtype=torch.bfloat16) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="QKHeadNormX", + tensor=x, + vram_addr=0, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="QKHeadNormWeight", + tensor=weight_rows, + vram_addr=weight_base, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + cos_vram = sin_vram = None + if args.runtime_rope: + cos_vram = prestage_bf16_vram_matrix( + prog=prog, + name="QKHeadNormROPECos", + tensor=cos, + vram_addr=rope_base, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + sin_vram = prestage_bf16_vram_matrix( + prog=prog, + name="QKHeadNormROPESin", + tensor=sin, + vram_addr=rope_base + trig_size, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + input_tensors = {} + tensor_layouts = {} + if args.runtime_rope: + rotate_input = prog.input("QKHeadNormROPERotate", shape=(mlen, mlen), real_data_ratio=2.0) + prog.head_runtime_rope_bf16( + x_vram, + rotate_input, + cos_vram, + sin_vram, + norm_weight_var=weight_vram, + eps_offset=66, + reci_hid_offset=67, + num_rows=seq, + name="QKHeadNormRuntimeRoPE", + ) + input_tensors["QKHeadNormROPERotate"] = rotate + tensor_layouts["QKHeadNormROPERotate"] = _bf16_layout(mlen, mlen) + else: + prog.rms_norm(x_vram, eps_offset=66, reci_hid_offset=67) + prog.vram_mul(x_vram, weight_vram, num_rows=seq) + isa = prog.compile() + + fp_preload = [0.0] * 68 + fp_preload[66] = eps + fp_preload[67] = 1.0 / head_dim + compare_width = mlen if head_dim < mlen else head_dim + comparison_golden = golden + if compare_width != head_dim: + comparison_golden = torch.zeros(seq, compare_width, dtype=torch.bfloat16) + comparison_golden[:, :head_dim] = golden + create_sim_env( + input_tensors, + isa, + {"original_output": comparison_golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="qk_head_norm_smoke", + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + if not input_tensors: + (build_dir / "hbm_for_behave_sim.bin").touch(exist_ok=True) + _write_json( + build_dir / "comparison_params.json", + _comparison_params(prog.get_vram_addr(x_vram.name), seq, compare_width, mlen, physical_rows=physical_rows), + ) + (build_dir / "generated_asm_code.asm").write_text(isa) + + run_emulator(build_dir) + results, _ = compare_emulator_output(build_dir) + emu_full = results["simulated_values"].reshape(seq, compare_width).to(torch.bfloat16) + emu = emu_full[:, :head_dim].contiguous() + rel = _rel_rms(emu, golden) + strict_summary = _strict_tail_summary(emu, golden) + padded_tail_max_abs = ( + float(emu_full[:, head_dim:].abs().max().item()) if compare_width > head_dim and emu_full.numel() else None + ) + summary = { + "mode": "qk_head_norm_smoke", + "seq": seq, + "head_dim": head_dim, + "compare_width": compare_width, + "padded_tail_max_abs": padded_tail_max_abs, + "physical_rows": physical_rows, + "qk_head_norm_eps": eps, + "runtime_rope": bool(args.runtime_rope), + "input_source": input_source, + "rel_rms": rel, + "strict_tail": strict_summary, + "nonfinite_output": int((~torch.isfinite(emu.float())).sum().item()), + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def run_external_rope(args: argparse.Namespace) -> dict: + """Apply device-side RoPE using externally supplied BF16 cos/sin tables.""" + required = { + "--input-pt": args.input_pt, + "--cos-pt": args.cos_pt, + "--sin-pt": args.sin_pt, + "--golden-pt": args.golden_pt, + } + missing = [name for name, value in required.items() if value is None] + if missing: + raise ValueError(f"external_rope requires {', '.join(missing)}") + + build_dir = (args.build_dir / "external_rope").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + x = _load_hidden_tensor(args.input_pt).to(torch.bfloat16).contiguous() + cos = _load_hidden_tensor(args.cos_pt).to(torch.bfloat16).contiguous() + sin = _load_hidden_tensor(args.sin_pt).to(torch.bfloat16).contiguous() + golden = _load_hidden_tensor(args.golden_pt).to(torch.bfloat16).contiguous() + original_shape = tuple(x.shape) + if x.ndim == 2: + rows, head_dim = map(int, x.shape) + head_count = 1 + elif x.ndim == 3: + seq_rows, head_count, head_dim = map(int, x.shape) + rows = seq_rows * head_count + if tuple(golden.shape) != original_shape: + raise ValueError(f"--golden-pt shape {tuple(golden.shape)} != {original_shape}") + if tuple(cos.shape) == (seq_rows, head_dim): + cos = cos[:, None, :].expand(seq_rows, head_count, head_dim).contiguous() + if tuple(sin.shape) == (seq_rows, head_dim): + sin = sin[:, None, :].expand(seq_rows, head_count, head_dim).contiguous() + x = x.reshape(rows, head_dim).contiguous() + golden = golden.reshape(rows, head_dim).contiguous() + cos = cos.reshape(rows, head_dim).contiguous() + sin = sin.reshape(rows, head_dim).contiguous() + else: + raise ValueError(f"--input-pt must be rank-2 or rank-3 for external_rope, got {tuple(x.shape)}") + if tuple(cos.shape) != (rows, head_dim) or tuple(sin.shape) != (rows, head_dim): + raise ValueError(f"cos/sin shapes {tuple(cos.shape)}/{tuple(sin.shape)} != {(rows, head_dim)}") + if tuple(golden.shape) != (rows, head_dim): + raise ValueError(f"--golden-pt shape {tuple(golden.shape)} != {(rows, head_dim)}") + if rows % blen != 0: + raise ValueError(f"rows={rows} must be a multiple of BLEN={blen}") + if head_dim > mlen: + raise ValueError(f"head_dim={head_dim} exceeds MLEN={mlen}") + + physical_rows = max(mlen, math.ceil(rows / blen) * blen) + x_padded = torch.zeros(rows, mlen, dtype=torch.bfloat16) + cos_padded = torch.zeros(rows, mlen, dtype=torch.bfloat16) + sin_padded = torch.zeros(rows, mlen, dtype=torch.bfloat16) + golden_padded = torch.zeros(rows, mlen, dtype=torch.bfloat16) + x_padded[:, :head_dim] = x + cos_padded[:, :head_dim] = cos + sin_padded[:, :head_dim] = sin + golden_padded[:, :head_dim] = golden + rotate = torch.zeros(mlen, mlen, dtype=torch.bfloat16) + rotate[:head_dim, :head_dim] = _make_rotate_half_matrix(head_dim) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + x_region = physical_rows * mlen + cos_base = _align_to_tile(x_region, mlen) + trig_size = _align_to_tile(physical_rows * mlen, mlen) + sin_base = cos_base + trig_size + vram_preload = torch.zeros(sin_base + trig_size + mlen * mlen, dtype=torch.bfloat16) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ExternalRoPEX", + tensor=x_padded, + vram_addr=0, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + cos_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ExternalRoPECos", + tensor=cos_padded, + vram_addr=cos_base, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + sin_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ExternalRoPESin", + tensor=sin_padded, + vram_addr=sin_base, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + rotate_input = prog.input("ExternalRoPERotate", shape=(mlen, mlen), real_data_ratio=2.0) + prog.runtime_rope_projection_bf16(x_vram, rotate_input, cos_vram, sin_vram, name="ExternalRoPERotateHalf") + isa = prog.compile() + + input_tensors = {"ExternalRoPERotate": rotate} + tensor_layouts = {"ExternalRoPERotate": _bf16_layout(mlen, mlen)} + create_sim_env( + input_tensors, + isa, + {"original_output": golden_padded}, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="external_rope", + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + _write_json( + build_dir / "comparison_params.json", + _comparison_params(prog.get_vram_addr(x_vram.name), rows, mlen, mlen, physical_rows=physical_rows), + ) + (build_dir / "generated_asm_code.asm").write_text(isa) + + run_emulator(build_dir) + results, _ = compare_emulator_output(build_dir) + emu_full = results["simulated_values"].reshape(rows, mlen).to(torch.bfloat16) + emu = emu_full[:, :head_dim].contiguous() + simulated_output_pt = build_dir / "simulated_output.pt" + simulated_full_output_pt = build_dir / "simulated_full_output.pt" + torch.save(emu.reshape(original_shape), simulated_output_pt) + torch.save(emu_full, simulated_full_output_pt) + rel = _rel_rms(emu, golden) + strict_summary = _strict_tail_summary(emu, golden) + padded_tail_max_abs = float(emu_full[:, head_dim:].abs().max().item()) if head_dim < mlen else 0.0 + if rel > 0.01: + raise AssertionError(f"external_rope rel_rms={rel:.6g} exceeds 1%") + summary = { + "mode": "external_rope", + "rows": rows, + "head_count": head_count, + "head_dim": head_dim, + "original_shape": list(original_shape), + "physical_rows": physical_rows, + "rel_rms": rel, + "strict_tail": strict_summary, + "padded_tail_max_abs": padded_tail_max_abs, + "scale_reg_count": isa.count("C_SET_SCALE_REG"), + "simulated_output_pt": str(simulated_output_pt), + "simulated_full_output_pt": str(simulated_full_output_pt), + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def run_true_full(args: argparse.Namespace) -> dict: + """True-shape GPT-OSS attention with runtime projection+bias+RoPE. + + This is the attention semantic gate immediately before a decoder-block + harness. It keeps sequence single-tile but uses the real 64Q/8KV head + structure and computes Q/K/V/O projections inside the emulator. + """ + mode_name = args.mode + build_dir = (args.build_dir / mode_name).expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + hq = args.num_attention_heads + hkv = args.num_key_value_heads + head_dim = args.hlen or mlen + seq = args.seq_len or 16 + hidden = args.hidden_size or 128 + out_features = args.out_features or hidden + if head_dim != mlen: + raise ValueError(f"true_full uses one head per MLEN row; got head_dim={head_dim}, MLEN={mlen}") + if hq % hkv != 0: + raise ValueError(f"num_attention_heads={hq} must be divisible by num_key_value_heads={hkv}") + if seq > mlen: + raise ValueError("true_full is single sequence tile only") + if seq % blen != 0: + raise ValueError(f"seq={seq} must be a multiple of BLEN={blen}") + if hidden % mlen != 0 or out_features % mlen != 0: + raise ValueError("hidden and out_features must be multiples of MLEN") + attention_block = args.mode in ("true_attn_block", "true_decoder_block") + decoder_block = args.mode == "true_decoder_block" + if attention_block and out_features != hidden: + raise ValueError("true_attn_block requires out_features == hidden for residual add") + if decoder_block: + moe_policy_name = args.moe_policy_name + moe_activation_policy = args.moe_activation_policy + supported_moe_shape = { + ("gpt_oss", 32, 4): "gpt_oss_clamp_gated", + ("qwen3_moe", 128, 8): "standard_swiglu", + } + expected_activation = supported_moe_shape.get((moe_policy_name, args.moe_num_experts, args.moe_top_k)) + if expected_activation is None: + raise ValueError( + "true_decoder_block supports only GPT-OSS 32/top4 and Qwen3-MoE 128/top8 for now; " + f"got policy={moe_policy_name!r}, experts={args.moe_num_experts}, top_k={args.moe_top_k}" + ) + if moe_activation_policy != expected_activation: + raise ValueError( + f"policy={moe_policy_name!r} expects activation_policy={expected_activation!r}, " + f"got {moe_activation_policy!r}" + ) + moe_intermediate = args.moe_intermediate_size or hidden + if moe_intermediate % mlen != 0: + raise ValueError(f"moe_intermediate={moe_intermediate} must be a multiple of MLEN={mlen}") + else: + moe_policy_name = None + moe_activation_policy = None + moe_intermediate = None + sliding_window, sliding_source = _resolve_sliding_window(args) + projection_bias = bool(args.projection_bias) + bias_parts = _bias_parts(args) + norm_eps = args.norm_eps + norm_eps_offset = 64 + norm_reci_hid_offset = 65 + qk_norm_eps = args.qk_head_norm_eps + # Q/K head-norm is applied after the attention masks have been materialized. + # The mask builders and online-softmax scratch own the low FPRAM region + # (ONLINE_SOFTMAX_FPSRAM_BASE + 3 * MLEN), so keep these constants in a + # high slot. 66/67 looked natural next to the block RMSNorm constants, but + # MLEN=128 mask construction overwrites them and turns the norm into + # rsqrt(0) -> inf. + qk_norm_eps_offset = 512 + qk_norm_reci_head_offset = 513 + ffn_norm_eps_offset = None + ffn_norm_reci_hid_offset = None + topk_weights_fp_base = None + topk_indices_int_base = None + decoder_moe_input_var = None + decoder_moe_residual_var = None + + torch.manual_seed(args.seed) + q_width = hq * head_dim + kv_width = hkv * head_dim + qwen_real_layer = bool(args.qwen_real_layer) + qwen_router_matrix_bf16 = bool(getattr(args, "qwen_router_matrix_bf16", False)) + qwen_router_packed_skinny_bf16 = bool(getattr(args, "qwen_router_packed_skinny_bf16", False)) + qwen_real_metadata = None + if qwen_router_matrix_bf16 and qwen_router_packed_skinny_bf16: + raise ValueError("--qwen-router-matrix-bf16 and --qwen-router-packed-skinny-bf16 are mutually exclusive") + if qwen_router_matrix_bf16 and (not decoder_block or moe_policy_name != "qwen3_moe"): + raise ValueError("--qwen-router-matrix-bf16 is only valid for Qwen3-MoE true_decoder_block runs") + if qwen_router_packed_skinny_bf16 and (not decoder_block or moe_policy_name != "qwen3_moe"): + raise ValueError("--qwen-router-packed-skinny-bf16 is only valid for Qwen3-MoE true_decoder_block runs") + if qwen_real_layer: + if not decoder_block: + raise ValueError("--qwen-real-layer is currently supported only with true_decoder_block") + if moe_policy_name != "qwen3_moe" or moe_activation_policy != "standard_swiglu": + raise ValueError("--qwen-real-layer requires --moe-policy-name qwen3_moe and standard_swiglu") + if projection_bias: + raise ValueError("Qwen3-30B-A3B config has attention_bias=false; do not pass --projection-bias") + if args.sink: + raise ValueError("Qwen3-30B-A3B attention has no GPT-OSS sink; do not pass --sink") + if not args.qk_head_norm: + raise ValueError("--qwen-real-layer requires --qk-head-norm") + x = (torch.randn(seq, hidden) * 0.2).to(torch.bfloat16) + input_source = "random" + if getattr(args, "input_pt", None) is not None: + loaded_x = torch.load(args.input_pt.expanduser()) + if isinstance(loaded_x, dict): + for key in ("hidden", "output", "tensor", "x"): + if key in loaded_x: + loaded_x = loaded_x[key] + break + if not torch.is_tensor(loaded_x): + raise TypeError(f"--input-pt must contain a Tensor or tensor dict, got {type(loaded_x).__name__}") + loaded_x = loaded_x.to(torch.bfloat16) + if tuple(loaded_x.shape) != (seq, hidden): + raise ValueError(f"--input-pt shape {tuple(loaded_x.shape)} != expected {(seq, hidden)}") + x = loaded_x.contiguous() + input_source = str(args.input_pt.expanduser()) + norm_weight = (1.0 + torch.randn(hidden) * 0.05).to(torch.bfloat16) if attention_block else None + norm_weight_rows = norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) if attention_block else None + q_norm_weight = (1.0 + torch.randn(head_dim) * 0.05).to(torch.bfloat16) if args.qk_head_norm else None + k_norm_weight = (1.0 + torch.randn(head_dim) * 0.05).to(torch.bfloat16) if args.qk_head_norm else None + q_norm_weight_rows = q_norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) if args.qk_head_norm else None + k_norm_weight_rows = k_norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) if args.qk_head_norm else None + ffn_norm_weight = (1.0 + torch.randn(hidden) * 0.05).to(torch.bfloat16) if decoder_block else None + ffn_norm_weight_rows = ffn_norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) if decoder_block else None + w_q = (torch.randn(hidden, q_width) * 0.2).to(torch.bfloat16) + w_k = (torch.randn(hidden, kv_width) * 0.2).to(torch.bfloat16) + w_v = (torch.randn(hidden, kv_width) * 0.2).to(torch.bfloat16) + w_o = (torch.randn(q_width, out_features) * 0.2).to(torch.bfloat16) + b_q = (torch.randn(q_width) * 0.02).to(torch.bfloat16) if projection_bias else None + b_k = (torch.randn(kv_width) * 0.02).to(torch.bfloat16) if projection_bias else None + b_v = (torch.randn(kv_width) * 0.02).to(torch.bfloat16) if projection_bias else None + b_o = (torch.randn(out_features) * 0.02).to(torch.bfloat16) if projection_bias else None + sinks = (torch.randn(hq) * 0.1).to(torch.bfloat16) if args.sink else None + if decoder_block: + num_experts = args.moe_num_experts + top_k = args.moe_top_k + router_weight = (torch.randn(num_experts, hidden) * 0.2).to(torch.bfloat16) + router_bias = (torch.randn(num_experts) * 0.02).to(torch.bfloat16) + gate_weight = (torch.randn(num_experts, hidden, moe_intermediate) * 0.2).to(torch.bfloat16) + up_weight = (torch.randn(num_experts, hidden, moe_intermediate) * 0.2).to(torch.bfloat16) + down_weight = (torch.randn(num_experts, moe_intermediate, hidden) * 0.2).to(torch.bfloat16) + gate_bias = (torch.randn(num_experts, moe_intermediate) * 0.02).to(torch.bfloat16) + up_bias = (torch.randn(num_experts, moe_intermediate) * 0.02).to(torch.bfloat16) + down_bias = (torch.randn(num_experts, hidden) * 0.02).to(torch.bfloat16) + split = SimpleNamespace( + gate_weight=gate_weight, + up_weight=up_weight, + gate_bias=gate_bias, + up_bias=up_bias, + ) + if qwen_real_layer: + real = _load_qwen3_layer_tensors(args.qwen_snapshot, args.layer_idx or 0) + if real["w_q"].shape != (hidden, q_width): + raise ValueError(f"Qwen q_proj shape {tuple(real['w_q'].shape)} != {(hidden, q_width)}") + if real["w_k"].shape != (hidden, kv_width): + raise ValueError(f"Qwen k_proj shape {tuple(real['w_k'].shape)} != {(hidden, kv_width)}") + if real["w_v"].shape != (hidden, kv_width): + raise ValueError(f"Qwen v_proj shape {tuple(real['w_v'].shape)} != {(hidden, kv_width)}") + if real["w_o"].shape != (q_width, out_features): + raise ValueError(f"Qwen o_proj shape {tuple(real['w_o'].shape)} != {(q_width, out_features)}") + if real["router_weight"].shape != (num_experts, hidden): + raise ValueError(f"Qwen router shape {tuple(real['router_weight'].shape)} != {(num_experts, hidden)}") + if real["gate_weight"].shape != (num_experts, hidden, moe_intermediate): + raise ValueError( + f"Qwen gate shape {tuple(real['gate_weight'].shape)} != {(num_experts, hidden, moe_intermediate)}" + ) + if real["up_weight"].shape != (num_experts, hidden, moe_intermediate): + raise ValueError( + f"Qwen up shape {tuple(real['up_weight'].shape)} != {(num_experts, hidden, moe_intermediate)}" + ) + if real["down_weight"].shape != (num_experts, moe_intermediate, hidden): + raise ValueError( + f"Qwen down shape {tuple(real['down_weight'].shape)} != {(num_experts, moe_intermediate, hidden)}" + ) + w_q = real["w_q"] + w_k = real["w_k"] + w_v = real["w_v"] + w_o = real["w_o"] + norm_weight = real["norm_weight"] + norm_weight_rows = norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) + q_norm_weight = real["q_norm_weight"] + k_norm_weight = real["k_norm_weight"] + q_norm_weight_rows = q_norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) + k_norm_weight_rows = k_norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) + ffn_norm_weight = real["ffn_norm_weight"] + ffn_norm_weight_rows = ffn_norm_weight.reshape(1, -1).repeat(seq, 1).to(torch.bfloat16) + router_weight = real["router_weight"] + router_bias = real["router_bias"] + gate_weight = real["gate_weight"] + up_weight = real["up_weight"] + down_weight = real["down_weight"] + gate_bias = real["gate_bias"] + up_bias = real["up_bias"] + down_bias = real["down_bias"] + split = SimpleNamespace( + gate_weight=gate_weight, + up_weight=up_weight, + gate_bias=gate_bias, + up_bias=up_bias, + ) + qwen_real_metadata = { + "snapshot": real["snapshot"], + "layer_idx": real["layer_idx"], + } + + projection_x = _weighted_rms_norm_golden(x, norm_weight_rows, norm_eps) if attention_block else x + q_pre = torch.matmul(projection_x.float(), w_q.float()).to(torch.bfloat16) + k_pre = torch.matmul(projection_x.float(), w_k.float()).to(torch.bfloat16) + v_pre = torch.matmul(projection_x.float(), w_v.float()).to(torch.bfloat16) + if "q" in bias_parts: + q_pre = (q_pre.float() + b_q.float()).to(torch.bfloat16) + if "k" in bias_parts: + k_pre = (k_pre.float() + b_k.float()).to(torch.bfloat16) + if "v" in bias_parts: + v_pre = (v_pre.float() + b_v.float()).to(torch.bfloat16) + q = q_pre.reshape(seq, hq, head_dim) + k = k_pre.reshape(seq, hkv, head_dim) + v = v_pre.reshape(seq, hkv, head_dim) + if args.qk_head_norm: + q = _head_rms_norm_golden(q, q_norm_weight, qk_norm_eps) + k = _head_rms_norm_golden(k, k_norm_weight, qk_norm_eps) + cos_head, sin_head = _make_rope_tables(seq, head_dim, theta=args.rope_theta) + q = _apply_rope_bf16(q, cos_head, sin_head) + k = _apply_rope_bf16(k, cos_head, sin_head) + scale = 1.0 / math.sqrt(head_dim) + attn_golden = _gpt_oss_attention_ref(q, k, v, sinks, scale=scale, sliding_window=sliding_window) + golden = torch.matmul(attn_golden.float(), w_o.float()).to(torch.bfloat16) + if "o" in bias_parts: + golden = (golden.float() + b_o.float()).to(torch.bfloat16) + if attention_block: + golden = _residual_add_golden(x, golden) + if decoder_block: + post_attn = golden + moe_input = _weighted_rms_norm_golden(post_attn, ffn_norm_weight_rows, norm_eps) + router_logits = torch.matmul(moe_input.float(), router_weight.t().float()).to(torch.bfloat16) + router_logits = (router_logits.float() + router_bias.reshape(1, -1).float()).to(torch.bfloat16) + top_values, top_indices = torch.topk(router_logits, k=top_k, dim=-1) + top_weights = torch.softmax(top_values, dim=1, dtype=top_values.dtype).to(torch.bfloat16) + moe_golden, moe_clamp_counts = _device_routing_vram_policy_golden( + x=moe_input, + device_indices=top_indices, + device_weights=top_weights, + split=split, + down_weight=down_weight, + down_bias=down_bias, + rows=seq, + hidden=hidden, + blen=blen, + mlen=mlen, + activation_policy=moe_activation_policy, + ) + golden = _residual_add_golden(post_attn, moe_golden) + else: + moe_clamp_counts = None + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + physical_rows = max(mlen, math.ceil(seq / blen) * blen) + x_region = physical_rows * hidden + if attention_block: + residual_base = _align_to_tile(x_region, mlen) + residual_size = _align_to_tile(physical_rows * hidden, mlen) + norm_weight_base = _align_to_tile(residual_base + residual_size, mlen) + norm_weight_size = _align_to_tile(physical_rows * hidden, mlen) + preload_cursor = _align_to_tile(norm_weight_base + norm_weight_size, mlen) + else: + residual_base = None + norm_weight_base = None + preload_cursor = _align_to_tile(x_region, mlen) + q_norm_weight_base = k_norm_weight_base = None + if args.qk_head_norm: + q_norm_weight_base = _align_to_tile(preload_cursor, mlen) + q_norm_weight_size = _align_to_tile(physical_rows * mlen, mlen) + k_norm_weight_base = _align_to_tile(q_norm_weight_base + q_norm_weight_size, mlen) + k_norm_weight_size = _align_to_tile(physical_rows * mlen, mlen) + preload_cursor = _align_to_tile(k_norm_weight_base + k_norm_weight_size, mlen) + ffn_norm_weight_base = router_w_base = router_bias_base = None + gate_bias_base = up_bias_base = down_bias_base = None + if decoder_block: + ffn_norm_weight_base = _align_to_tile(preload_cursor, mlen) + ffn_norm_weight_size = _align_to_tile(physical_rows * hidden, mlen) + router_w_base = _align_to_tile(ffn_norm_weight_base + ffn_norm_weight_size, mlen) + router_w_physical = (num_experts, hidden) + router_w_size = _align_to_tile(router_w_physical[0] * router_w_physical[1], mlen) + router_bias_base = _align_to_tile(router_w_base + router_w_size, mlen) + router_bias_rows = _router_bias_block_rows(router_bias, rows=seq, num_experts=num_experts, mlen=mlen) + router_bias_physical = ( + max(blen, math.ceil(router_bias_rows.shape[0] / blen) * blen), + mlen, + ) + router_bias_size = _align_to_tile(router_bias_physical[0] * router_bias_physical[1], mlen) + gate_bias_base = _align_to_tile(router_bias_base + router_bias_size, mlen) + gate_bias_physical = (num_experts * blen, moe_intermediate) + gate_bias_size = _align_to_tile(gate_bias_physical[0] * gate_bias_physical[1], mlen) + up_bias_base = _align_to_tile(gate_bias_base + gate_bias_size, mlen) + up_bias_physical = (num_experts * blen, moe_intermediate) + up_bias_size = _align_to_tile(up_bias_physical[0] * up_bias_physical[1], mlen) + down_bias_base = _align_to_tile(up_bias_base + up_bias_size, mlen) + down_bias_physical = (num_experts * blen, hidden) + down_bias_size = _align_to_tile(down_bias_physical[0] * down_bias_physical[1], mlen) + preload_cursor = _align_to_tile(down_bias_base + down_bias_size, mlen) + bias_base = _align_to_tile(preload_cursor, mlen) + bias_plan: list[tuple[str, torch.Tensor, tuple[int, int]]] = [] + if projection_bias: + if "q" in bias_parts: + bias_plan.append(("B_Q_TRUE", b_q.reshape(1, -1).repeat(seq, 1), (physical_rows, q_width))) + if "k" in bias_parts: + bias_plan.append(("B_K_TRUE", b_k.reshape(1, -1).repeat(seq, 1), (physical_rows, kv_width))) + if "v" in bias_parts: + bias_plan.append(("B_V_TRUE", b_v.reshape(1, -1).repeat(seq, 1), (physical_rows, kv_width))) + if "o" in bias_parts: + bias_plan.append(("B_O_TRUE", b_o.reshape(1, -1).repeat(seq, 1), (physical_rows, out_features))) + bias_total = sum(_align_to_tile(shape[0] * shape[1], mlen) for _, _, shape in bias_plan) + rope_base = bias_base + bias_total + trig_shape = (physical_rows, mlen) + trig_size = _align_to_tile(trig_shape[0] * trig_shape[1], mlen) + vram_words = max(x_region + mlen * mlen, rope_base + 2 * trig_size) + vram_preload = torch.zeros(vram_words, dtype=torch.bfloat16) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="X_true", + tensor=x, + vram_addr=0, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + residual_vram = None + norm_weight_vram = None + q_norm_weight_vram = None + k_norm_weight_vram = None + ffn_norm_weight_vram = None + router_w_vram = None + router_w_matrix_input = None + router_w_packed_skinny_input = None + router_bias_vram = None + gate_bias_table = None + up_bias_table = None + down_bias_table = None + if attention_block: + residual_vram = prestage_bf16_vram_matrix( + prog=prog, + name="X_residual_true", + tensor=x, + vram_addr=residual_base, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + norm_weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="AttentionNormWeight", + tensor=norm_weight_rows, + vram_addr=norm_weight_base, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + if args.qk_head_norm: + q_norm_weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="QHeadNormWeight", + tensor=q_norm_weight_rows, + vram_addr=q_norm_weight_base, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + k_norm_weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="KHeadNormWeight", + tensor=k_norm_weight_rows, + vram_addr=k_norm_weight_base, + physical_shape=(physical_rows, mlen), + vram_preload=vram_preload, + ) + if decoder_block: + ffn_norm_weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="FFNNormWeight", + tensor=ffn_norm_weight_rows, + vram_addr=ffn_norm_weight_base, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + router_w_vram = prestage_bf16_vram_matrix( + prog=prog, + name="DecoderRouterWBF16", + tensor=router_weight, + vram_addr=router_w_base, + physical_shape=router_w_physical, + vram_preload=vram_preload, + ) + router_bias_vram = prestage_bf16_vram_matrix( + prog=prog, + name="DecoderRouterBias", + tensor=router_bias_rows, + vram_addr=router_bias_base, + physical_shape=router_bias_physical, + vram_preload=vram_preload, + ) + gate_bias_table = _build_true_expert_bias_table( + prog=prog, + name="DecoderGateBiasTable", + bias=gate_bias, + width=moe_intermediate, + blen=blen, + vram_addr=gate_bias_base, + vram_preload=vram_preload, + ) + up_bias_table = _build_true_expert_bias_table( + prog=prog, + name="DecoderUpBiasTable", + bias=up_bias, + width=moe_intermediate, + blen=blen, + vram_addr=up_bias_base, + vram_preload=vram_preload, + ) + down_bias_table = _build_true_expert_bias_table( + prog=prog, + name="DecoderDownBiasTable", + bias=down_bias, + width=hidden, + blen=blen, + vram_addr=down_bias_base, + vram_preload=vram_preload, + ) + norm_weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="AttentionNormWeight", + tensor=norm_weight_rows, + vram_addr=norm_weight_base, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + bias_vrams = {} + next_bias_addr = bias_base + for name, tensor, physical_shape in bias_plan: + bias_vrams[name] = prestage_bf16_vram_matrix( + prog=prog, + name=name, + tensor=tensor, + vram_addr=next_bias_addr, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + next_bias_addr += _align_to_tile(physical_shape[0] * physical_shape[1], mlen) + rotate, cos_packed, sin_packed = _make_packed_rope_inputs(seq, 1, head_dim, args.rope_theta) + cos_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_COS_TRUE_FULL", + tensor=cos_packed, + vram_addr=rope_base, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + sin_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_SIN_TRUE_FULL", + tensor=sin_packed, + vram_addr=rope_base + trig_size, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + + projection_input = x_vram + if attention_block: + prog.rms_norm(projection_input, eps_offset=norm_eps_offset, reci_hid_offset=norm_reci_hid_offset) + prog.vram_mul(projection_input, norm_weight_vram, num_rows=seq) + + ( + (w_q_input, w_k_input, w_v_input, w_o_input, rotate_input), + input_tensors, + tensor_layouts, + ) = _make_true_attention_hbm_inputs( + prog, + hidden=hidden, + q_width=q_width, + kv_width=kv_width, + out_features=out_features, + mlen=mlen, + w_q=w_q, + w_k=w_k, + w_v=w_v, + w_o=w_o, + rotate=rotate, + ) + q_var, k_var, v_var = _emit_true_qkv_projection_bf16( + prog, + projection_input, + w_q_input=w_q_input, + w_k_input=w_k_input, + w_v_input=w_v_input, + bias_vrams=bias_vrams, + bias_parts=bias_parts, + physical_rows=physical_rows, + q_width=q_width, + kv_width=kv_width, + seq=seq, + ) + + weight_templates = None + weight_table_bases = None + weight_table_strides = None + if decoder_block: + gate_inputs, gate_table_base, gate_stride = _build_true_expert_weight_table( + prog, + prefix="DecoderWGate", + weights=gate_weight, + input_tensors=input_tensors, + ) + up_inputs, up_table_base, up_stride = _build_true_expert_weight_table( + prog, + prefix="DecoderWUp", + weights=up_weight, + input_tensors=input_tensors, + ) + down_inputs, down_table_base, down_stride = _build_true_expert_weight_table( + prog, + prefix="DecoderWDown", + weights=down_weight, + input_tensors=input_tensors, + ) + weight_templates = (gate_inputs[0], up_inputs[0], down_inputs[0]) + weight_table_bases = (gate_table_base, up_table_base, down_table_base) + weight_table_strides = (gate_stride, up_stride, down_stride) + tensor_layouts.update( + infer_hbm_tensor_layouts( + { + name: tensor + for name, tensor in input_tensors.items() + if name.startswith(("DecoderWGate", "DecoderWUp", "DecoderWDown")) + } + ) + ) + if qwen_router_matrix_bf16: + router_w_matrix_input = prog.input( + "DecoderRouterWMatrixBF16", + shape=(hidden, num_experts), + real_data_ratio=2.0, + ) + input_tensors["DecoderRouterWMatrixBF16"] = router_weight.t().contiguous() + tensor_layouts["DecoderRouterWMatrixBF16"] = _bf16_layout(hidden, num_experts) + if qwen_router_packed_skinny_bf16: + packed_router_weight = _pack_qwen_router_weight_skinny( + router_weight.t().contiguous(), + mlen=mlen, + blen=blen, + k_tiles_per_packed_tile=args.qwen_router_packed_skinny_k_tiles, + ) + router_w_packed_skinny_input = prog.input( + "DecoderRouterWPackedSkinnyBF16", + shape=tuple(packed_router_weight.shape), + physical_shape=tuple(packed_router_weight.shape), + real_data_ratio=2.0, + ) + input_tensors["DecoderRouterWPackedSkinnyBF16"] = packed_router_weight + tensor_layouts["DecoderRouterWPackedSkinnyBF16"] = _bf16_layout(*packed_router_weight.shape) + k_inputs, v_inputs = _stage_true_kv_heads_bf16( + prog, + k_var=k_var, + v_var=v_var, + rotate_input=rotate_input, + cos_vram=cos_vram, + sin_vram=sin_vram, + k_norm_weight_vram=k_norm_weight_vram if args.qk_head_norm else None, + qk_head_norm=args.qk_head_norm, + qk_norm_eps_offset=qk_norm_eps_offset if args.qk_head_norm else None, + qk_norm_reci_head_offset=qk_norm_reci_head_offset if args.qk_head_norm else None, + seq=seq, + mlen=mlen, + hkv=hkv, + physical_rows=physical_rows, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + ) + + attn_out = prog.alloc("O_true_full_attention", seq, q_width, strict=False, physical_shape=(physical_rows, q_width)) + scratch = prog.alloc("_gpt_oss_true_full_scratch", mlen, mlen, strict=True) + if sliding_window is not None: + causal_mask = prog._build_sliding_causal_score_mask("_gpt_oss_true_full_sliding_mask", sliding_window) + else: + causal_mask = prog._build_causal_score_mask("_gpt_oss_true_full_causal_mask") + + fp_preload = [0.0, scale / 0.25, float("-inf")] + [0.0] * 253 + if args.qk_head_norm: + if len(fp_preload) <= qk_norm_reci_head_offset: + fp_preload.extend([0.0] * (qk_norm_reci_head_offset + 1 - len(fp_preload))) + fp_preload[qk_norm_eps_offset] = qk_norm_eps + fp_preload[qk_norm_reci_head_offset] = 1.0 / head_dim + if attention_block: + if len(fp_preload) <= norm_reci_hid_offset: + fp_preload.extend([0.0] * (norm_reci_hid_offset + 1 - len(fp_preload))) + fp_preload[norm_eps_offset] = norm_eps + fp_preload[norm_reci_hid_offset] = 1.0 / hidden + sink_base = None + if sinks is not None: + sink_base = 256 + if len(fp_preload) < sink_base + hq: + fp_preload.extend([0.0] * (sink_base + hq - len(fp_preload))) + for idx, value in enumerate(sinks.tolist()): + fp_preload[sink_base + idx] = float(value) + + debug_stop_after_first_q_norm = bool(getattr(args, "debug_stop_after_first_q_norm", False)) + debug_stop_after_first_q_head = bool(getattr(args, "debug_stop_after_first_q_head", False)) + out, debug_golden, debug_stop_active = _emit_true_q_heads_attention_bf16( + prog, + q_var=q_var, + rotate_input=rotate_input, + cos_vram=cos_vram, + sin_vram=sin_vram, + q_norm_weight_vram=q_norm_weight_vram if args.qk_head_norm else None, + qk_head_norm=args.qk_head_norm, + qk_norm_eps_offset=qk_norm_eps_offset if args.qk_head_norm else None, + qk_norm_reci_head_offset=qk_norm_reci_head_offset if args.qk_head_norm else None, + seq=seq, + mlen=mlen, + hq=hq, + hkv=hkv, + head_dim=head_dim, + physical_rows=physical_rows, + k_inputs=k_inputs, + v_inputs=v_inputs, + attn_out=attn_out, + scratch=scratch, + causal_mask=causal_mask, + scale=scale, + sink_base=sink_base, + q_pre=q_pre, + q_norm_weight=q_norm_weight if args.qk_head_norm else None, + qk_norm_eps=qk_norm_eps, + q=q, + debug_stop_after_first_q_norm=debug_stop_after_first_q_norm, + debug_stop_after_first_q_head=debug_stop_after_first_q_head, + ) + if debug_stop_active: + golden = debug_golden + if not debug_stop_active: + out = _emit_true_output_projection_bf16( + prog, + attn_out, + w_o_input=w_o_input, + bias_vrams=bias_vrams, + bias_parts=bias_parts, + physical_rows=physical_rows, + out_features=out_features, + seq=seq, + ) + if attention_block: + prog.vram_add(out, residual_vram, num_rows=seq) + if decoder_block and not debug_stop_active: + # The attention path uses raw FPRAM offsets for scale/-inf, RMSNorm + # constants, and per-head sinks (slots 1/2, 64/65, and 256+). Reserve + # the low region before allocating MoE FP variables so the second + # sublayer cannot silently overwrite those constants. Q/K head norm + # also owns 512/513, so the decoder FFN norm must start after those + # slots. Slot 0 remains the true zero row expected by vram_fill_zero. + zero = prog.fp_var("decoder_zero", size=1) + _decoder_reserved_fp = prog.fp_var("decoder_reserved_attention_fp_slots", size=qk_norm_reci_head_offset) + ffn_norm_eps_var = prog.fp_var("decoder_ffn_norm_eps", size=1) + ffn_norm_reci_hid_var = prog.fp_var("decoder_ffn_norm_reci_hidden", size=1) + ffn_norm_eps_offset = ffn_norm_eps_var.address + ffn_norm_reci_hid_offset = ffn_norm_reci_hid_var.address + constant_rows = blen + limit_pos = prog.fp_var("decoder_gpt_oss_limit_pos", size=constant_rows) + limit_neg = prog.fp_var("decoder_gpt_oss_limit_neg", size=constant_rows) + one = prog.fp_var("decoder_one", size=constant_rows) + neg_alpha = prog.fp_var("decoder_neg_alpha", size=constant_rows) + shared_zero_row = prog.fp_var("decoder_shared_zero_row", size=mlen) + + moe_residual = prog.alloc( + "DecoderMoeResidual", + rows=seq, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + prog.moe_true_zero_vram_rows_v0( + moe_residual, + rows=list(range(seq)), + hidden=hidden, + zero_row=shared_zero_row, + policy_name=moe_policy_name, + name="decoder_moe_residual_zero", + ) + prog.vram_add(moe_residual, out, num_rows=seq) + + prog.rms_norm(out, eps_offset=ffn_norm_eps_offset, reci_hid_offset=ffn_norm_reci_hid_offset) + prog.vram_mul(out, ffn_norm_weight_vram, num_rows=seq) + decoder_moe_input_var = out + if qwen_router_packed_skinny_bf16: + logits = prog.qwen3_router_logits_packed_skinny_bf16_rowpacked_v0( + out, + router_w_packed_skinny_input, + rows=seq, + hidden=hidden, + num_experts=num_experts, + k_tiles_per_packed_tile=args.qwen_router_packed_skinny_k_tiles, + name="decoder_router_logits", + ) + elif qwen_router_matrix_bf16: + logits = prog.qwen3_router_logits_matrix_bf16_rowpacked_v0( + out, + router_w_matrix_input, + rows=seq, + hidden=hidden, + num_experts=num_experts, + mram_tile_capacity=args.qwen_router_mram_tile_capacity, + name="decoder_router_logits", + ) + else: + logits = prog.moe_router_logits_bf16_v0( + out, + router_w_vram, + rows=seq, + hidden=hidden, + num_experts=num_experts, + policy_name=moe_policy_name, + name="decoder_router_logits", + ) + prog.vram_add(logits, router_bias_vram, num_rows=logits.shape[0]) + topk_weight_var = prog.fp_var("decoder_topk_weights", size=seq * top_k) + topk_weights_fp_base = topk_weight_var.address + topk_indices_int_base = 0 + for token_idx in range(seq): + prog.moe_router_select_v0( + logits, + token_idx=token_idx, + weights_fp_base=topk_weights_fp_base + token_idx * top_k, + indices_int_base=topk_indices_int_base + token_idx * top_k, + policy_name=moe_policy_name, + num_experts=num_experts, + top_k=top_k, + name=f"decoder_token{token_idx}", + ) + + accumulator = prog.alloc( + "DecoderMoeAccumulator", + rows=seq, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + prog.moe_true_zero_vram_rows_v0( + accumulator, + rows=list(range(seq)), + hidden=hidden, + zero_row=shared_zero_row, + policy_name=moe_policy_name, + name="decoder_moe_acc_zero", + ) + route_fp_scratch = prog.fp_var("decoder_route_fp_scratch", size=mlen) + pair_count = seq * top_k + for pair_idx in range(pair_count): + token_idx = pair_idx // top_k + gathered = prog.moe_gather_token_rows_from_vram_v0( + out, + token_indices=[token_idx], + hidden=hidden, + zero_row=shared_zero_row, + policy_name=moe_policy_name, + name=f"decoder_pair{pair_idx}_vram_gather_t{token_idx}", + ) + expert_out = prog.moe_dynamic_expert_pair_v0( + gathered, + weight_templates, + weight_table_bases=weight_table_bases, + weight_table_strides=weight_table_strides, + expert_indices_int_base=topk_indices_int_base, + weights_fp_base=topk_weights_fp_base, + pair_idx=pair_idx, + bias_tables=(gate_bias_table, up_bias_table, down_bias_table), + rows=blen, + intermediate=moe_intermediate, + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + zero_row=shared_zero_row, + route_fp_scratch=route_fp_scratch, + policy_name=moe_policy_name, + activation_policy=moe_activation_policy, + name=f"decoder_pair{pair_idx}", + ) + prog.moe_scatter_add_active_rows_v0( + accumulator, + expert_out, + token_indices=[token_idx], + active_rows=[0], + hidden=hidden, + policy_name=moe_policy_name, + name=f"decoder_pair{pair_idx}_scatter", + ) + prog.vram_add(accumulator, moe_residual, num_rows=seq) + decoder_moe_residual_var = moe_residual + out = accumulator + decoder_fp_end = max( + ffn_norm_reci_hid_var.address + ffn_norm_reci_hid_var.size, + topk_weight_var.address + topk_weight_var.size, + route_fp_scratch.address + route_fp_scratch.size, + shared_zero_row.address + shared_zero_row.size, + neg_alpha.address + neg_alpha.size, + ) + if len(fp_preload) < decoder_fp_end: + fp_preload.extend([0.0] * (decoder_fp_end - len(fp_preload))) + fp_preload[ffn_norm_eps_offset] = norm_eps + fp_preload[ffn_norm_reci_hid_offset] = 1.0 / hidden + for idx in range(limit_pos.size): + fp_preload[limit_pos.address + idx] = 7.0 + for idx in range(limit_neg.size): + fp_preload[limit_neg.address + idx] = -7.0 + for idx in range(one.size): + fp_preload[one.address + idx] = 1.0 + for idx in range(neg_alpha.size): + fp_preload[neg_alpha.address + idx] = -1.0 if moe_activation_policy == "standard_swiglu" else -1.702 + isa = prog.compile() + if "C_SET_SCALE_REG" in isa and not decoder_block: + isa_lines = isa.splitlines() + non_store_scale_lines = [] + for idx, line in enumerate(isa_lines): + if "C_SET_SCALE_REG" not in line: + continue + context = isa_lines[max(0, idx - 4) : idx] + if not any("Store Activation Generation" in item for item in context): + non_store_scale_lines.append(line) + if non_store_scale_lines: + (build_dir / "generated_asm_code.debug.asm").write_text(isa) + raise AssertionError( + "true_full BF16 projection/attention path emitted non-store C_SET_SCALE_REG: " + + "; ".join(non_store_scale_lines[:8]) + ) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_attention_true_full", + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + compare_out_features = out.shape[1] + params = _comparison_params( + prog.get_vram_addr(out.name), + seq, + compare_out_features, + mlen, + physical_rows=out.physical_shape[0], + ) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + strict_summary = None + decoder_runtime_summary = None + if not args.no_run: + if args.probe_rel_only or attention_block: + run_emulator(build_dir) + else: + run_and_assert(build_dir, "gpt_oss_attention_true_full", mlen=mlen, blen=blen) + results, _ = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(seq, compare_out_features).to(torch.bfloat16) + rel = _rel_rms(emu, golden) + strict_summary = _strict_tail_summary(emu, golden) + if getattr(args, "save_golden_pt", None) is not None: + args.save_golden_pt.expanduser().parent.mkdir(parents=True, exist_ok=True) + torch.save(golden.float(), args.save_golden_pt.expanduser()) + if getattr(args, "save_output_pt", None) is not None: + args.save_output_pt.expanduser().parent.mkdir(parents=True, exist_ok=True) + torch.save(emu.float(), args.save_output_pt.expanduser()) + if decoder_block: + device_indices_flat = _decode_u32_dump(build_dir / "intsram_dump.bin")[ + topk_indices_int_base : topk_indices_int_base + seq * top_k + ] + device_weights_flat = _decode_bf16_dump(build_dir / "fpsram_dump.bin")[ + topk_weights_fp_base : topk_weights_fp_base + seq * top_k + ] + device_indices = device_indices_flat.reshape(seq, top_k) + device_weights = device_weights_flat.reshape(seq, top_k).to(torch.bfloat16) + vram_dump = build_dir / "vram_dump.bin" + device_moe_input = _read_vram_matrix( + vram_dump, + vram_addr=prog.get_vram_addr(decoder_moe_input_var.name), + rows=seq, + cols=hidden, + mlen=mlen, + physical_rows=decoder_moe_input_var.physical_shape[0], + ) + device_moe_residual = _read_vram_matrix( + vram_dump, + vram_addr=prog.get_vram_addr(decoder_moe_residual_var.name), + rows=seq, + cols=hidden, + mlen=mlen, + physical_rows=decoder_moe_residual_var.physical_shape[0], + ) + device_context_logits = torch.matmul( + device_moe_input.float(), + router_weight.t().float(), + ).to(torch.bfloat16) + device_context_logits = (device_context_logits.float() + router_bias.reshape(1, -1).float()).to( + torch.bfloat16 + ) + context_values, context_indices = torch.topk(device_context_logits, k=top_k, dim=-1) + context_weights = torch.softmax(context_values, dim=1, dtype=context_values.dtype).to(torch.bfloat16) + device_selected_moe, device_selected_clamp_counts = _device_routing_vram_policy_golden( + x=device_moe_input, + device_indices=device_indices, + device_weights=device_weights, + split=split, + down_weight=down_weight, + down_bias=down_bias, + rows=seq, + hidden=hidden, + blen=blen, + mlen=mlen, + activation_policy=moe_activation_policy, + ) + device_selected_final = _residual_add_golden(device_moe_residual, device_selected_moe) + torch.save(device_selected_final.float(), build_dir / "device_selected_decoder_golden.pt") + torch.save(device_moe_input.float(), build_dir / "device_moe_input.pt") + torch.save(device_moe_residual.float(), build_dir / "device_post_attention_residual.pt") + device_selected_strict = _strict_tail_summary(emu, device_selected_final) + decoder_runtime_summary = { + "topk_weights_fp_base": topk_weights_fp_base, + "topk_indices_int_base": topk_indices_int_base, + "device_selected_exact_rel_rms": _rel_rms(emu, device_selected_final), + "device_selected_exact_strict_tail": device_selected_strict, + "host_topk_vs_device_topk": _topk_match_summary( + host_indices=top_indices, + host_weights=top_weights, + device_indices=device_indices, + device_weights=device_weights, + ), + "host_router_margin": _router_margin_summary(router_logits, top_k), + "device_context_topk_vs_device_topk": _topk_match_summary( + host_indices=context_indices, + host_weights=context_weights, + device_indices=device_indices, + device_weights=device_weights, + ), + "host_vs_device_routing_boundary": _routing_boundary_verdict( + reference_logits=router_logits, + comparison_logits=device_context_logits, + reference_indices=top_indices, + device_indices=device_indices, + top_k=top_k, + ), + "device_context_vs_device_routing_boundary": _routing_boundary_verdict( + reference_logits=device_context_logits, + comparison_logits=device_context_logits, + reference_indices=context_indices, + device_indices=device_indices, + top_k=top_k, + ), + "device_context_router_margin": _router_margin_summary(device_context_logits, top_k), + "device_context_vs_host_router_logits_rel_rms": _rel_rms(device_context_logits, router_logits), + "device_context_vs_host_router_logits_max_abs_error": float( + (device_context_logits.float() - router_logits.float()).abs().max().item() + ), + "device_moe_input_vs_host_moe_input_rel_rms": _rel_rms(device_moe_input, moe_input), + "device_post_attention_vs_host_post_attention_rel_rms": _rel_rms( + device_moe_residual, + post_attn, + ), + "device_selected_clamp_counts": device_selected_clamp_counts, + } + if rel > 0.02 and not args.probe_rel_only: + device_selected_rel = ( + decoder_runtime_summary.get("device_selected_exact_rel_rms") + if decoder_runtime_summary is not None + else None + ) + if device_selected_rel is None or not math.isfinite(device_selected_rel) or device_selected_rel > 0.02: + raise AssertionError(f"true_full rel_rms={rel:.6g} exceeds 2%") + else: + rel = float("nan") + + summary = { + "mode": mode_name, + "seq": seq, + "hidden": hidden, + "out_features": out_features, + "compare_out_features": compare_out_features, + "debug_stop_after_first_q_head": bool(debug_stop_after_first_q_head and debug_stop_active), + "debug_stop_after_first_q_norm": bool(debug_stop_after_first_q_norm and debug_stop_active), + "num_attention_heads": hq, + "num_key_value_heads": hkv, + "head_dim": head_dim, + "sink": bool(args.sink), + "projection_bias": projection_bias, + "bias_parts": "".join(sorted(bias_parts)), + "runtime_projection": True, + "runtime_bias": projection_bias, + "runtime_rope": True, + "qk_head_norm": bool(args.qk_head_norm), + "qk_head_norm_eps": qk_norm_eps if args.qk_head_norm else None, + "qk_norm_eps_offset": qk_norm_eps_offset if args.qk_head_norm else None, + "qk_norm_reci_head_offset": qk_norm_reci_head_offset if args.qk_head_norm else None, + "input_source": input_source, + "qwen_real_layer": qwen_real_layer, + "qwen_real_metadata": qwen_real_metadata, + "saved_output_pt": str(args.save_output_pt.expanduser()) if getattr(args, "save_output_pt", None) else None, + "saved_golden_pt": str(args.save_golden_pt.expanduser()) if getattr(args, "save_golden_pt", None) else None, + "attention_block": attention_block, + "decoder_block": decoder_block, + "pre_norm": attention_block, + "residual_add": attention_block, + "norm_eps": norm_eps if attention_block or decoder_block else None, + "norm_eps_offset": norm_eps_offset if attention_block else None, + "norm_reci_hid_offset": norm_reci_hid_offset if attention_block else None, + "ffn_norm_eps_offset": ffn_norm_eps_offset if decoder_block else None, + "ffn_norm_reci_hid_offset": ffn_norm_reci_hid_offset if decoder_block else None, + "projection_precision": "qkvo_bf16", + "moe_input_source": "vram_bf16" if decoder_block else None, + "moe_num_experts": args.moe_num_experts if decoder_block else None, + "moe_top_k": args.moe_top_k if decoder_block else None, + "moe_policy_name": moe_policy_name if decoder_block else None, + "moe_activation_policy": moe_activation_policy if decoder_block else None, + "moe_intermediate": moe_intermediate if decoder_block else None, + "qwen_router_matrix_bf16": qwen_router_matrix_bf16 if decoder_block else None, + "qwen_router_packed_skinny_bf16": qwen_router_packed_skinny_bf16 if decoder_block else None, + "qwen_router_mram_tile_capacity": (args.qwen_router_mram_tile_capacity if qwen_router_matrix_bf16 else None), + "qwen_router_packed_skinny_k_tiles": ( + args.qwen_router_packed_skinny_k_tiles if qwen_router_packed_skinny_bf16 else None + ), + "moe_clamp_counts_golden": moe_clamp_counts, + "topk_weights_fp_base": topk_weights_fp_base if decoder_block else None, + "topk_indices_int_base": topk_indices_int_base if decoder_block else None, + "decoder_runtime_summary": decoder_runtime_summary, + "sliding_window": sliding_window, + "sliding_source": sliding_source, + "rel_rms": rel, + "sink_base": sink_base, + "scale_reg_count": isa.count("C_SET_SCALE_REG"), + "h_store_keyvalue_count": isa.count("H_STORE_V"), + "strict_tail": strict_summary, + "strategy": "runtime_qkvo_projection_unrolled_one_q_head_per_call", + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def _load_hidden_tensor(path: Path) -> torch.Tensor: + loaded = torch.load(path.expanduser(), map_location="cpu") + if isinstance(loaded, dict): + for key in ("hidden", "output", "tensor", "x"): + if key in loaded: + loaded = loaded[key] + break + if not torch.is_tensor(loaded): + raise TypeError(f"{path} must contain a Tensor or tensor dict, got {type(loaded).__name__}") + return loaded + + +def _clone_args(args: argparse.Namespace, **overrides) -> SimpleNamespace: + values = vars(args).copy() + values.update(overrides) + return SimpleNamespace(**values) + + +def run_true_decoder_chain(args: argparse.Namespace) -> dict: + """Run true decoder blocks sequentially with explicit tensor handoff. + + This is not a fused multi-layer PLENA program. It is a validation harness + that reuses the already-validated true_decoder_block program, feeds each + layer's emulator output into the next layer, and runs a parallel host-chain + control to expose cumulative drift and routing-margin sensitivity. + """ + if args.no_run: + raise ValueError("true_decoder_chain requires emulator execution; --no-run is not supported") + if args.chain_layers < 1: + raise ValueError(f"--chain-layers must be >= 1, got {args.chain_layers}") + if args.mode != "true_decoder_chain": + raise ValueError(f"run_true_decoder_chain expected mode=true_decoder_chain, got {args.mode}") + + root = (args.build_dir / "true_decoder_chain").expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + + device_input: Path | None = args.input_pt.expanduser() if getattr(args, "input_pt", None) else None + host_input: Path | None = args.input_pt.expanduser() if getattr(args, "input_pt", None) else None + layer_records = [] + host_control_records = [] + + for layer_idx in range(args.chain_layers): + layer_seed = args.seed + layer_idx + device_dir = root / f"layer{layer_idx}_device_chain" + device_output = device_dir / f"layer{layer_idx}_emu_output.pt" + device_conditional_golden = device_dir / f"layer{layer_idx}_device_input_golden.pt" + device_args = _clone_args( + args, + mode="true_decoder_block", + build_dir=device_dir, + layer_idx=layer_idx, + seed=layer_seed, + input_pt=device_input, + save_output_pt=device_output, + save_golden_pt=device_conditional_golden, + probe_rel_only=True, + no_run=False, + ) + device_summary = run_true_full(device_args) + + record = { + "layer_idx": layer_idx, + "seed": layer_seed, + "device_build_dir": str(device_dir), + "device_input_pt": str(device_input) if device_input is not None else None, + "device_output_pt": str(device_output), + "device_conditional_golden_pt": str(device_conditional_golden), + "device_summary": device_summary, + } + + if layer_idx == 0: + host_input = device_conditional_golden + record["device_input_vs_host_input_rel_rms"] = 0.0 + else: + host_dir = root / f"layer{layer_idx}_host_chain_control" + host_output = host_dir / f"layer{layer_idx}_host_control_emu_output.pt" + host_golden = host_dir / f"layer{layer_idx}_host_chain_golden.pt" + if host_input is None: + raise RuntimeError("host_input unexpectedly missing for chained host-control layer") + host_args = _clone_args( + args, + mode="true_decoder_block", + build_dir=host_dir, + layer_idx=layer_idx, + seed=layer_seed, + input_pt=host_input, + save_output_pt=host_output, + save_golden_pt=host_golden, + probe_rel_only=True, + no_run=False, + ) + host_summary = run_true_full(host_args) + device_in_tensor = _load_hidden_tensor(device_input) + host_in_tensor = _load_hidden_tensor(host_input) + record["device_input_vs_host_input_rel_rms"] = _rel_rms(device_in_tensor, host_in_tensor) + record["host_control_build_dir"] = str(host_dir) + record["host_chain_golden_pt"] = str(host_golden) + record["host_control_summary"] = host_summary + host_control_records.append( + { + "layer_idx": layer_idx, + "seed": layer_seed, + "host_build_dir": str(host_dir), + "host_input_pt": str(host_input), + "host_output_pt": str(host_output), + "host_golden_pt": str(host_golden), + "summary": host_summary, + } + ) + host_input = host_golden + + layer_records.append(record) + device_input = device_output + + final_device_output = device_input + final_host_golden = host_input + final_device_vs_host_chain_rel = None + if final_device_output is not None and final_host_golden is not None: + final_device_vs_host_chain_rel = _rel_rms( + _load_hidden_tensor(final_device_output), + _load_hidden_tensor(final_host_golden), + ) + + summary = { + "mode": "true_decoder_chain", + "note": ( + "Sequential validation runner: each layer is a separate true_decoder_block " + "emulator run, with emulator output handed to the next layer. This is not " + "a fused multi-layer PLENA program." + ), + "chain_layers": args.chain_layers, + "initial_input_pt": str(args.input_pt.expanduser()) if getattr(args, "input_pt", None) else None, + "final_device_output_pt": str(final_device_output) if final_device_output is not None else None, + "final_host_chain_golden_pt": str(final_host_golden) if final_host_golden is not None else None, + "final_device_vs_host_chain_rel_rms": final_device_vs_host_chain_rel, + "layers": layer_records, + "host_controls": host_control_records, + } + _write_json(root / "chain_summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def run_core(args: argparse.Namespace) -> dict: + build_dir = (args.build_dir / "core").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + hq = 4 + hkv = 1 + head_dim = args.hlen or (mlen // hq) + seq = args.seq_len or 16 + if head_dim * hq != mlen: + raise ValueError("core smoke expects hq*head_dim == MLEN") + if seq > mlen: + raise ValueError("core smoke is single sequence tile only") + if seq % blen != 0: + raise ValueError(f"seq={seq} must be a multiple of BLEN={blen}") + sliding_window, sliding_source = _resolve_sliding_window(args) + + torch.manual_seed(args.seed) + q = (torch.randn(seq, hq, head_dim) * 0.2).to(torch.bfloat16) + k = (torch.randn(seq, hkv, head_dim) * 0.2).to(torch.bfloat16) + v = (torch.randn(seq, hkv, head_dim) * 0.2).to(torch.bfloat16) + sinks = (torch.randn(hq) * 0.1).to(torch.bfloat16) if args.sink else None + scale = 1.0 / math.sqrt(head_dim) + golden = _gpt_oss_attention_ref(q, k, v, sinks, scale=scale, sliding_window=sliding_window) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + rows_per_batch = max(mlen, seq) + q_flat = q.reshape(seq, hq * head_dim) + q_vram_flat = torch.zeros(rows_per_batch * mlen, dtype=torch.bfloat16) + q_input = prog.input("Q", shape=(seq, mlen), prestaged_vram_addr=0, physical_shape=(rows_per_batch, mlen)) + q_group = prog.load_batch(q_input, name="Q") + q_vram_flat[: seq * mlen] = q_flat.reshape(-1) + + kv_slots = mlen // head_dim + k_padded = torch.zeros(rows_per_batch, kv_slots, head_dim, dtype=torch.bfloat16) + v_padded = torch.zeros(rows_per_batch, kv_slots, head_dim, dtype=torch.bfloat16) + k_padded[:seq, :hkv, :] = k + v_padded[:seq, :hkv, :] = v + k_input = prog.input("K", shape=(seq, mlen), physical_shape=(rows_per_batch, mlen), real_data_ratio=2.0) + v_input = prog.input("V", shape=(seq, mlen), physical_shape=(rows_per_batch, mlen), real_data_ratio=2.0) + + o = prog.alloc("O", seq, mlen, strict=False, physical_shape=(rows_per_batch, mlen)) + scratch = prog.alloc("_gpt_oss_attn_scratch", mlen * kv_slots, mlen, strict=True) + if sliding_window is not None: + causal_mask = prog._build_sliding_causal_score_mask("_gpt_oss_sliding_mask", sliding_window) + else: + causal_mask = True + + fp_preload = [0.0, scale / 0.25, float("-inf")] + [0.0] * 253 + sink_base = None + if sinks is not None: + sink_base = 256 + if len(fp_preload) < sink_base + hq: + fp_preload.extend([0.0] * (sink_base + hq - len(fp_preload))) + for idx, value in enumerate(sinks.tolist()): + fp_preload[sink_base + idx] = float(value) + + prog.flash_attention_packed_group( + q_group, + k_input, + v_input, + group_heads=hq, + head_slot_dim=head_dim, + output_base_address=prog.get_vram_addr(o.name), + scratch_base_address=prog.get_vram_addr(scratch.name), + broadcast_amount=kv_slots, + scale=scale, + causal_mask=causal_mask, + valid_cols=seq, + sink_base_address=sink_base, + k_matrix_precision="keyvalue", + k_set_scale=False, + k_hbm_element_bytes=2, + v_hbm_element_bytes=2, + ) + isa = prog.compile() + + input_tensors = { + "K": k_padded.reshape(rows_per_batch, mlen), + "V": v_padded.reshape(rows_per_batch, mlen), + } + tensor_layouts = { + name: { + "source_shape": [rows_per_batch, mlen], + "storage_shape": [rows_per_batch, mlen], + "source_rows": rows_per_batch, + "storage_rows": rows_per_batch, + "source_row_elements": mlen, + "storage_row_elements": mlen, + "precision": "HBM_M_KV_TYPE", + } + for name in input_tensors + } + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=q_vram_flat, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_attention_core", + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + params = _comparison_params(prog.get_vram_addr(o.name), seq, mlen, mlen, physical_rows=rows_per_batch) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + if not args.no_run: + run_and_assert(build_dir, "gpt_oss_attention_core", mlen=mlen, blen=blen) + results, _ = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(seq, mlen).to(torch.bfloat16) + rel = _rel_rms(emu, golden) + if rel > 0.02: + raise AssertionError(f"attention core rel_rms={rel:.6g} exceeds 2%") + else: + rel = float("nan") + + summary = { + "mode": "core", + "seq": seq, + "sink": bool(args.sink), + "sliding_window": sliding_window, + "sliding_source": sliding_source, + "rel_rms": rel, + "sink_base": sink_base, + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def _bf16_layout(rows: int, cols: int, *, precision: str = "HBM_M_KV_TYPE") -> dict: + return { + "source_shape": [rows, cols], + "storage_shape": [rows, cols], + "source_rows": rows, + "storage_rows": rows, + "source_row_elements": cols, + "storage_row_elements": cols, + "precision": precision, + } + + +def _pack_qwen_router_weight_skinny( + weight_hidden_by_expert: torch.Tensor, + *, + mlen: int, + blen: int, + k_tiles_per_packed_tile: int, +) -> torch.Tensor: + """Pack ``[hidden, experts]`` router weight for the packed-skinny helper.""" + hidden, num_experts = weight_hidden_by_expert.shape + if hidden % mlen != 0: + raise ValueError(f"router hidden={hidden} must be divisible by MLEN={mlen}") + if num_experts % blen != 0: + raise ValueError(f"num_experts={num_experts} must be divisible by BLEN={blen}") + tiles_per_mlen = mlen // blen + if k_tiles_per_packed_tile <= 0 or k_tiles_per_packed_tile > tiles_per_mlen: + raise ValueError(f"k_tiles_per_packed_tile={k_tiles_per_packed_tile} must be in 1..{tiles_per_mlen}") + + num_k_tiles = hidden // mlen + num_groups = math.ceil(num_k_tiles / k_tiles_per_packed_tile) + num_microcols = math.ceil(num_experts / blen) + physical_col_blocks = max(tiles_per_mlen, num_microcols) + packed = torch.zeros(num_groups * mlen, physical_col_blocks * mlen, dtype=torch.bfloat16) + + for group_idx in range(num_groups): + for micro_col_idx in range(num_microcols): + expert_start = micro_col_idx * blen + expert_end = min(expert_start + blen, num_experts) + row_start = group_idx * mlen + col_base = micro_col_idx * mlen + for local_k_tile in range(k_tiles_per_packed_tile): + k_tile_idx = group_idx * k_tiles_per_packed_tile + local_k_tile + if k_tile_idx >= num_k_tiles: + break + k_start = k_tile_idx * mlen + k_end = k_start + mlen + skinny_col = col_base + local_k_tile * blen + packed[row_start : row_start + mlen, skinny_col : skinny_col + (expert_end - expert_start)] = ( + weight_hidden_by_expert[k_start:k_end, expert_start:expert_end].to(torch.bfloat16) + ) + return packed.contiguous() + + +def run_full(args: argparse.Namespace) -> dict: + build_dir = (args.build_dir / "full").expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + setup_hw(args, build_dir) + _set_matrix_kv_plain_bf16() + + mlen = args.mlen + blen = args.blen + hq = 4 + hkv = 1 + head_dim = args.hlen or (mlen // hq) + seq = args.seq_len or 16 + hidden = args.hidden_size or 128 + out_features = args.out_features or hidden + if head_dim * hq != mlen: + raise ValueError("full attention smoke expects hq*head_dim == MLEN") + if seq > mlen: + raise ValueError("full attention smoke is single sequence tile only") + if seq % blen != 0: + raise ValueError(f"seq={seq} must be a multiple of BLEN={blen}") + if hidden % mlen != 0 or out_features % mlen != 0: + raise ValueError("hidden and out_features must be multiples of MLEN") + sliding_window, sliding_source = _resolve_sliding_window(args) + vo_mx = bool(args.vo_mx) + projection_bias = bool(args.projection_bias) + runtime_rope = bool(args.runtime_rope) + bias_parts = _bias_parts(args) + + torch.manual_seed(args.seed) + x = (torch.randn(seq, hidden) * 0.2).to(torch.bfloat16) + w_q = (torch.randn(hidden, mlen) * 0.2).to(torch.bfloat16) + w_k = torch.zeros(hidden, mlen, dtype=torch.bfloat16) + w_v = torch.zeros(hidden, mlen, dtype=torch.bfloat16) + w_k[:, :head_dim] = (torch.randn(hidden, head_dim) * 0.2).to(torch.bfloat16) + w_v[:, :head_dim] = (torch.randn(hidden, head_dim) * 0.2).to(torch.bfloat16) + w_o = (torch.randn(mlen, out_features) * 0.2).to(torch.bfloat16) + b_q = (torch.randn(mlen) * 0.02).to(torch.bfloat16) if projection_bias else None + b_k = (torch.randn(mlen) * 0.02).to(torch.bfloat16) if projection_bias else None + b_v = (torch.randn(mlen) * 0.02).to(torch.bfloat16) if projection_bias else None + b_o = (torch.randn(out_features) * 0.02).to(torch.bfloat16) if projection_bias else None + # K/V are staged in MLEN-wide physical rows, but this small diagnostic has + # only one KV head. Keep padded lanes zero so projection bias cannot create + # non-model data in lanes outside the logical head_dim. + if projection_bias: + b_k[head_dim:] = 0 + b_v[head_dim:] = 0 + sinks = (torch.randn(hq) * 0.1).to(torch.bfloat16) if args.sink else None + + q_proj = torch.matmul(x.float(), w_q.float()).to(torch.bfloat16) + k_proj = torch.matmul(x.float(), w_k.float()).to(torch.bfloat16) + v_proj = torch.matmul(x.float(), w_v.float()).to(torch.bfloat16) + if "q" in bias_parts: + q_proj = (q_proj.float() + b_q.float()).to(torch.bfloat16) + if "k" in bias_parts: + k_proj = (k_proj.float() + b_k.float()).to(torch.bfloat16) + if "v" in bias_parts: + v_proj = (v_proj.float() + b_v.float()).to(torch.bfloat16) + q = q_proj.reshape(seq, hq, head_dim) + k = k_proj[:, :head_dim].reshape(seq, hkv, head_dim) + v = v_proj[:, :head_dim].reshape(seq, hkv, head_dim) + rotate = cos = sin = None + if runtime_rope: + rotate, cos, sin = _make_packed_rope_inputs(seq, hq, head_dim, args.rope_theta) + cos_head = cos[:, :head_dim] + sin_head = sin[:, :head_dim] + q = _apply_rope_bf16(q, cos_head, sin_head) + k = _apply_rope_bf16(k, cos_head, sin_head) + scale = 1.0 / math.sqrt(head_dim) + attn_golden = _gpt_oss_attention_ref(q, k, v, sinks, scale=scale, sliding_window=sliding_window) + golden = torch.matmul(attn_golden.float(), w_o.float()).to(torch.bfloat16) + if "o" in bias_parts: + golden = (golden.float() + b_o.float()).to(torch.bfloat16) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=2.0) + physical_rows = max(mlen, math.ceil(seq / blen) * blen) + x_region = physical_rows * hidden + bias_base = math.ceil(x_region / (mlen * mlen)) * (mlen * mlen) + bias_plan: list[tuple[str, torch.Tensor, tuple[int, int]]] = [] + if projection_bias: + if "q" in bias_parts: + bias_plan.append(("B_Q", b_q.reshape(1, -1).repeat(seq, 1), (physical_rows, mlen))) + if "k" in bias_parts: + bias_plan.append(("B_K", b_k.reshape(1, -1).repeat(seq, 1), (physical_rows, mlen))) + if "v" in bias_parts: + bias_plan.append(("B_V", b_v.reshape(1, -1).repeat(seq, 1), (physical_rows, mlen))) + if "o" in bias_parts: + bias_plan.append(("B_O", b_o.reshape(1, -1).repeat(seq, 1), (physical_rows, out_features))) + bias_total = sum(math.ceil((shape[0] * shape[1]) / (mlen * mlen)) * (mlen * mlen) for _, _, shape in bias_plan) + rope_base = bias_base + bias_total + trig_shape = (physical_rows, mlen) + trig_size = _align_to_tile(trig_shape[0] * trig_shape[1], mlen) + vram_words = max(x_region + mlen * mlen, bias_base + bias_total) + if runtime_rope: + vram_words = max(vram_words, rope_base + 2 * trig_size) + vram_preload = torch.zeros(vram_words, dtype=torch.bfloat16) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="X", + tensor=x, + vram_addr=0, + physical_shape=(physical_rows, hidden), + vram_preload=vram_preload, + ) + bias_vrams = {} + next_bias_addr = bias_base + for name, tensor, physical_shape in bias_plan: + bias_vrams[name] = prestage_bf16_vram_matrix( + prog=prog, + name=name, + tensor=tensor, + vram_addr=next_bias_addr, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + next_bias_addr += math.ceil((physical_shape[0] * physical_shape[1]) / (mlen * mlen)) * (mlen * mlen) + cos_vram = sin_vram = None + if runtime_rope: + cos_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_COS", + tensor=cos, + vram_addr=rope_base, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + sin_vram = prestage_bf16_vram_matrix( + prog=prog, + name="ROPE_SIN", + tensor=sin, + vram_addr=rope_base + trig_size, + physical_shape=trig_shape, + vram_preload=vram_preload, + ) + + w_q_input = prog.input("W_Q_BF16", shape=(hidden, mlen), real_data_ratio=2.0) + w_k_input = prog.input("W_K_BF16", shape=(hidden, mlen), real_data_ratio=2.0) + rotate_input = None + if runtime_rope: + rotate_input = prog.input("ROPE_ROTATE_BF16", shape=(mlen, mlen), real_data_ratio=2.0) + if vo_mx: + w_v_input = prog.input("W_V_MX", shape=(hidden, mlen), real_data_ratio=1.125) + w_o_input = prog.input("W_O_MX", shape=(mlen, out_features), real_data_ratio=1.125) + else: + w_v_input = prog.input("W_V_BF16", shape=(hidden, mlen), real_data_ratio=2.0) + w_o_input = prog.input("W_O_BF16", shape=(mlen, out_features), real_data_ratio=2.0) + + q_var = prog.linear_projection_bf16(x_vram, w_q_input, name="Q_proj", physical_shape=(physical_rows, mlen)) + k_var = prog.linear_projection_bf16(x_vram, w_k_input, name="K_proj", physical_shape=(physical_rows, mlen)) + if vo_mx: + v_var = prog.linear_projection(x_vram, w_v_input, name="V_proj_mx", physical_shape=(physical_rows, mlen)) + else: + v_var = prog.linear_projection_bf16(x_vram, w_v_input, name="V_proj", physical_shape=(physical_rows, mlen)) + if "q" in bias_parts: + prog.vram_add(q_var, bias_vrams["B_Q"], num_rows=seq) + if "k" in bias_parts: + prog.vram_add(k_var, bias_vrams["B_K"], num_rows=seq) + if "v" in bias_parts: + prog.vram_add(v_var, bias_vrams["B_V"], num_rows=seq) + if runtime_rope: + prog.runtime_rope_projection_bf16(q_var, rotate_input, cos_vram, sin_vram, name="Q_runtime_rotate") + prog.runtime_rope_projection_bf16(k_var, rotate_input, cos_vram, sin_vram, name="K_runtime_rotate") + k_staged = prog.store( + k_var, + name="K_staged_bf16", + precision=1, + hbm_element_bytes=2, + real_data_ratio=2.0, + ) + v_staged = prog.store( + v_var, + name="V_staged_bf16", + precision=1, + hbm_element_bytes=2, + real_data_ratio=2.0, + ) + + attn_out = prog.alloc("attn_out", seq, mlen, strict=False, physical_shape=(physical_rows, mlen)) + scratch = prog.alloc("_gpt_oss_full_attn_scratch", mlen * (mlen // head_dim), mlen, strict=True) + if sliding_window is not None: + causal_mask = prog._build_sliding_causal_score_mask("_gpt_oss_full_sliding_mask", sliding_window) + else: + causal_mask = True + + fp_preload = [0.0, scale / 0.25, float("-inf")] + [0.0] * 253 + sink_base = None + if sinks is not None: + sink_base = 256 + if len(fp_preload) < sink_base + hq: + fp_preload.extend([0.0] * (sink_base + hq - len(fp_preload))) + for idx, value in enumerate(sinks.tolist()): + fp_preload[sink_base + idx] = float(value) + + prog.flash_attention_packed_group( + q_var, + k_staged, + v_staged, + group_heads=hq, + head_slot_dim=head_dim, + output_base_address=prog.get_vram_addr(attn_out.name), + scratch_base_address=prog.get_vram_addr(scratch.name), + broadcast_amount=mlen // head_dim, + scale=scale, + causal_mask=causal_mask, + valid_cols=seq, + sink_base_address=sink_base, + k_matrix_precision="keyvalue", + k_set_scale=False, + k_hbm_element_bytes=2, + v_hbm_element_bytes=2, + ) + if vo_mx: + out = prog.linear_projection( + attn_out, w_o_input, name="O_proj_mx", physical_shape=(physical_rows, out_features) + ) + else: + out = prog.linear_projection_bf16( + attn_out, w_o_input, name="O_proj", physical_shape=(physical_rows, out_features) + ) + if "o" in bias_parts: + prog.vram_add(out, bias_vrams["B_O"], num_rows=seq) + isa = prog.compile() + + input_tensors = { + "W_Q_BF16": w_q, + "W_K_BF16": w_k, + k_staged.name: torch.zeros(physical_rows, mlen, dtype=torch.bfloat16), + v_staged.name: torch.zeros(physical_rows, mlen, dtype=torch.bfloat16), + } + if runtime_rope: + input_tensors["ROPE_ROTATE_BF16"] = rotate + if vo_mx: + input_tensors["W_V_MX"] = w_v + input_tensors["W_O_MX"] = w_o + else: + input_tensors["W_V_BF16"] = w_v + input_tensors["W_O_BF16"] = w_o + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + tensor_layouts.update( + { + "W_Q_BF16": _bf16_layout(hidden, mlen), + "W_K_BF16": _bf16_layout(hidden, mlen), + k_staged.name: _bf16_layout(physical_rows, mlen, precision="HBM_V_KV_TYPE"), + v_staged.name: _bf16_layout(physical_rows, mlen, precision="HBM_V_KV_TYPE"), + } + ) + if runtime_rope: + tensor_layouts["ROPE_ROTATE_BF16"] = _bf16_layout(mlen, mlen) + if not vo_mx: + tensor_layouts["W_V_BF16"] = _bf16_layout(hidden, mlen) + tensor_layouts["W_O_BF16"] = _bf16_layout(mlen, out_features) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + data_order = sorted(input_tensors, key=lambda name: hbm_addrs[name]) + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_attention_full", + specified_data_order=data_order, + build_path=build_dir, + input_tensors=input_tensors, + tensor_layouts=tensor_layouts, + hbm_addrs=hbm_addrs, + ) + params = _comparison_params( + prog.get_vram_addr(out.name), + seq, + out_features, + mlen, + physical_rows=out.physical_shape[0], + ) + _write_json(build_dir / "comparison_params.json", params) + (build_dir / "generated_asm_code.asm").write_text(isa) + + if not args.no_run: + if args.probe_rel_only: + run_emulator(build_dir) + else: + run_and_assert(build_dir, "gpt_oss_attention_full", mlen=mlen, blen=blen) + results, _ = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(seq, out_features).to(torch.bfloat16) + rel = _rel_rms(emu, golden) + if rel > 0.02 and not args.probe_rel_only: + raise AssertionError(f"full attention rel_rms={rel:.6g} exceeds 2%") + else: + rel = float("nan") + + summary = { + "mode": "full", + "seq": seq, + "hidden": hidden, + "out_features": out_features, + "sink": bool(args.sink), + "vo_mx": vo_mx, + "projection_bias": projection_bias, + "bias_parts": "".join(sorted(bias_parts)), + "runtime_rope": runtime_rope, + "projection_precision": "qk_bf16_vo_mxfp8" if vo_mx else "qkvo_bf16", + "sliding_window": sliding_window, + "sliding_source": sliding_source, + "rel_rms": rel, + "sink_base": sink_base, + "h_store_keyvalue_count": isa.count("H_STORE_V"), + "h_prefetch_keyvalue_count": isa.count("H_PREFETCH_M"), + } + _write_json(build_dir / "summary.json", summary) + print(json.dumps(summary, indent=2, sort_keys=True)) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument( + "--mode", + choices=( + "projection", + "core", + "full", + "true_core", + "true_full", + "true_attn_block", + "true_decoder_block", + "true_decoder_chain", + "runtime_rope", + "qk_head_norm_smoke", + "deepseek_mla_q_path", + "deepseek_mla_kv_path", + "external_rope", + ), + required=True, + ) + parser.add_argument( + "--build-dir", type=Path, default=Path(__file__).parent / "build" / "gpt_oss_attention_semantics" + ) + parser.add_argument("--out-features", type=int, default=None) + parser.add_argument( + "--allow-padded-projection", + action="store_true", + help="Allow projection mode to zero-pad non-MLEN hidden/out dimensions for DeepSeek-style shape probes.", + ) + parser.add_argument("--sink", action="store_true") + parser.add_argument("--sliding-window", type=int, default=None) + parser.add_argument("--no-causal-mask", action="store_true") + parser.add_argument("--config-json", type=Path, default=None) + parser.add_argument("--layer-idx", type=int, default=None) + parser.add_argument("--num-attention-heads", type=int, default=64) + parser.add_argument("--num-key-value-heads", type=int, default=8) + parser.add_argument( + "--value-head-dim", + type=int, + default=None, + help=( + "Logical V head dimension for true_core. Defaults to HLEN/head_dim. " + "When smaller than head_dim, V is zero-padded to the QK head slot; " + "this probes DeepSeek-style qk_dim != v_dim attention." + ), + ) + parser.add_argument("--q-pt", type=Path, default=None, help="Optional post-RoPE Q tensor for true_core.") + parser.add_argument("--k-pt", type=Path, default=None, help="Optional post-RoPE K tensor for true_core.") + parser.add_argument( + "--v-pt", type=Path, default=None, help="Optional V tensor for true_core; may be narrower than HLEN." + ) + parser.add_argument( + "--golden-pt", + type=Path, + default=None, + help="Optional padded attention-core golden tensor for true_core external-QKV mode.", + ) + parser.add_argument("--cos-pt", type=Path, default=None, help="Optional external BF16 RoPE cos table.") + parser.add_argument("--sin-pt", type=Path, default=None, help="Optional external BF16 RoPE sin table.") + parser.add_argument( + "--projection-weight-pt", + type=Path, + default=None, + help="Optional rank-2 BF16 projection weight for projection mode; accepts PLENA [in,out] or HF [out,in].", + ) + parser.add_argument( + "--second-projection-weight-pt", + type=Path, + default=None, + help="Optional second rank-2 BF16 projection weight for integrated two-projection probes.", + ) + parser.add_argument( + "--norm-weight-pt", + type=Path, + default=None, + help="Optional rank-1 BF16 RMSNorm scale for integrated projection/RMSNorm probes.", + ) + parser.add_argument( + "--attention-scale", + type=float, + default=None, + help="Optional attention softmax scale override for true_core, needed for DeepSeek YaRN scaling.", + ) + parser.add_argument( + "--o-weight-pt", + type=Path, + default=None, + help=( + "Optional BF16 O-projection weight for true_core. The tensor may be " + "[logical_value_heads, out] or HF-style [out, logical_value_heads]. " + "For DeepSeek-style value_head_dim < head_dim, the harness expands it " + "into the padded head slot with zero tail rows." + ), + ) + parser.add_argument( + "--o-golden-pt", + type=Path, + default=None, + help="Optional direct HF O-projection output for true_core + --o-weight-pt validation.", + ) + parser.add_argument("--rope-theta", type=float, default=10000.0) + parser.add_argument( + "--projection-bias", + action="store_true", + help="Add BF16 projection bias. Validated for projection mode; full mode remains a diagnostic probe.", + ) + parser.add_argument( + "--bias-parts", + default="qkvo", + help="Diagnostic subset of projection bias terms to enable in full mode: any of q/k/v/o, all, or none.", + ) + parser.add_argument( + "--runtime-rope", + action="store_true", + help="Apply device-side rotate-half projection plus RoPE after Q/K projection.", + ) + parser.add_argument( + "--qk-head-norm", + action="store_true", + help="Apply Qwen-style RMSNorm over each Q/K head before RoPE.", + ) + parser.add_argument( + "--qk-head-norm-eps", + type=float, + default=1e-6, + help="Epsilon for --qk-head-norm; Qwen3-MoE uses config.rms_norm_eps.", + ) + parser.add_argument( + "--norm-eps", + type=float, + default=1e-5, + help="Epsilon for block RMSNorm before attention/FFN; GPT-OSS uses 1e-5, Qwen3-MoE uses 1e-6.", + ) + parser.add_argument( + "--qk-norm-weight-pt", + type=Path, + default=None, + help="Optional BF16 head RMSNorm weight tensor for qk_head_norm_smoke.", + ) + parser.add_argument( + "--debug-stop-after-first-q-head", + action="store_true", + help="Diagnostic true_full mode: stop after the first Q head RMSNorm+RoPE and compare that head.", + ) + parser.add_argument( + "--debug-stop-after-first-q-norm", + action="store_true", + help="Diagnostic true_full mode: stop after the first Q head RMSNorm+weight before RoPE.", + ) + parser.add_argument( + "--vo-mx", + action="store_true", + help="Precision probe: keep Q/K projections BF16 but use MXFP8 V/O projection weights.", + ) + parser.add_argument( + "--probe-rel-only", + action="store_true", + help="Run emulator and report rel_rms without failing the process on numerical mismatch.", + ) + parser.add_argument( + "--input-pt", + type=Path, + default=None, + help="Optional BF16 hidden-state tensor to use instead of the generated random input.", + ) + parser.add_argument( + "--qwen-real-layer", + action="store_true", + help="Load true Qwen3-30B-A3B layer weights instead of random synthetic true_decoder_block weights.", + ) + parser.add_argument( + "--qwen-snapshot", + type=Path, + default=_default_qwen3_snapshot(), + help="Local Qwen3-30B-A3B snapshot containing config.json and safetensors index " + "(or set QWEN3_30B_A3B_SNAPSHOT).", + ) + parser.add_argument( + "--qwen-router-matrix-bf16", + action="store_true", + help=( + "Use the BF16 matrix-machine router logits path for Qwen3-MoE, " + "then pack logits into the existing V_TOPK token-major ABI." + ), + ) + parser.add_argument( + "--qwen-router-mram-tile-capacity", + type=int, + default=4, + help=( + "MRAM K-tile chunk size for --qwen-router-matrix-bf16. Qwen full-block " + "runs use MLEN=128, so the default 4 tiles gives the same 512-element " + "K chunk as the standalone MLEN=64/cap8 router probe." + ), + ) + parser.add_argument( + "--qwen-router-packed-skinny-bf16", + action="store_true", + help=( + "Use the packed-skinny BF16 matrix router table for Qwen3-MoE, " + "then pack logits into the existing V_TOPK token-major ABI." + ), + ) + parser.add_argument( + "--qwen-router-packed-skinny-k-tiles", + type=int, + default=8, + help="Number of skinny K tiles packed into one full router-weight MRAM tile.", + ) + parser.add_argument( + "--save-output-pt", + type=Path, + default=None, + help="Optional path to save the emulator output tensor after a run.", + ) + parser.add_argument( + "--save-golden-pt", + type=Path, + default=None, + help="Optional path to save the host golden tensor for the current run.", + ) + parser.add_argument("--moe-num-experts", type=int, default=32) + parser.add_argument("--moe-top-k", type=int, default=4) + parser.add_argument("--moe-intermediate-size", type=int, default=None) + parser.add_argument( + "--moe-policy-name", + choices=("gpt_oss", "qwen3_moe"), + default="gpt_oss", + help="MoE router/expert policy for true_decoder_block.", + ) + parser.add_argument( + "--moe-activation-policy", + choices=("gpt_oss_clamp_gated", "standard_swiglu"), + default="gpt_oss_clamp_gated", + help="Expert activation policy for true_decoder_block.", + ) + parser.add_argument( + "--chain-layers", + type=int, + default=2, + help="Number of sequential true_decoder_block runs for true_decoder_chain mode.", + ) + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + + if args.mode == "projection": + run_projection(args) + elif args.mode == "core": + run_core(args) + elif args.mode == "true_core": + run_true_core(args) + elif args.mode in ("true_full", "true_attn_block", "true_decoder_block"): + run_true_full(args) + elif args.mode == "true_decoder_chain": + run_true_decoder_chain(args) + elif args.mode == "runtime_rope": + run_runtime_rope(args) + elif args.mode == "qk_head_norm_smoke": + run_qk_head_norm_smoke(args) + elif args.mode in ("deepseek_mla_q_path", "deepseek_mla_kv_path"): + run_deepseek_mla_q_path(args) + elif args.mode == "external_rope": + run_external_rope(args) + else: + run_full(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/models/gpt_oss/block_glue_test.py b/transactional_emulator/testbench/models/gpt_oss/block_glue_test.py new file mode 100644 index 00000000..e40820d0 --- /dev/null +++ b/transactional_emulator/testbench/models/gpt_oss/block_glue_test.py @@ -0,0 +1,311 @@ +"""GPT-OSS block glue smoke: weighted RMSNorm and residual add. + +This is intentionally narrower than a full decoder block. It verifies the +VRAM-resident BF16 glue pieces that connect attention and MoE: + + * GPT-OSS RMSNorm = unweighted RMSNorm followed by learned scale multiply. + * residual add = hidden += sublayer output. + +No attention, MoE, HBM activation loads, dynamic routing, or RTL is involved. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import torch + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.aten.golden import ( + _active_precision_settings, + _rms_norm_vector_ref, + quantize_to_vector_fp, +) +from transactional_emulator.testbench.routed_moe.gpt_oss_moe_gather_scatter_test import ( + _comparison_params_for, + _rewrite_compact_golden, +) +from transactional_emulator.testbench.emulator_runner import compare_emulator_output, run_emulator +from transactional_emulator.testbench.layout_utils import prestage_bf16_vram_matrix +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _align_to(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _vram_layout_size(shape: tuple[int, int], *, mlen: int) -> int: + rows, cols = shape + return math.ceil(cols / mlen) * rows * mlen + + +def _bf16_vram(x: torch.Tensor) -> torch.Tensor: + precision = _active_precision_settings() + return quantize_to_vector_fp(x.to(torch.bfloat16).float(), precision) + + +def _load_bf16_tensor(path: Path) -> torch.Tensor: + obj = torch.load(path.expanduser().resolve(), map_location="cpu") + if isinstance(obj, dict): + obj = obj.get("tensor", next(iter(obj.values()))) + if not torch.is_tensor(obj): + raise TypeError(f"{path} did not contain a tensor") + return obj.detach().to(torch.bfloat16).cpu().contiguous() + + +def _weighted_rms_norm_golden( + x: torch.Tensor, + weight_rows: torch.Tensor, + eps: float, + *, + active_hidden: int | None = None, +) -> torch.Tensor: + precision = _active_precision_settings() + x_vram = _bf16_vram(x) + w_vram = _bf16_vram(weight_rows) + if active_hidden is None or active_hidden == x_vram.shape[-1]: + normed = _rms_norm_vector_ref(x_vram, eps, precision) + else: + # DeepSeek MLA has logical widths such as 512 that are physically padded + # to MLEN. Hardware reduces all physical lanes, but padded tail lanes are + # zero and FPRAM still holds 1 / logical_hidden. + mean_sq = x_vram.float().pow(2).sum(dim=-1, keepdim=True) / float(active_hidden) + rms = quantize_to_vector_fp(torch.rsqrt(mean_sq + eps), precision) + normed = quantize_to_vector_fp(x_vram.float() * rms.float(), precision) + return quantize_to_vector_fp(normed.float() * w_vram.float(), precision) + + +def _residual_add_golden(hidden: torch.Tensor, sublayer: torch.Tensor) -> torch.Tensor: + precision = _active_precision_settings() + hidden_vram = _bf16_vram(hidden) + sublayer_vram = _bf16_vram(sublayer) + return quantize_to_vector_fp(hidden_vram.float() + sublayer_vram.float(), precision) + + +def _stats(actual: torch.Tensor, reference: torch.Tensor) -> dict: + actual_f = actual.float() + ref_f = reference.float() + diff = actual_f - ref_f + denom = torch.linalg.vector_norm(ref_f).clamp_min(1e-12) + rel_rms = float((torch.linalg.vector_norm(diff) / denom).item()) + atol = float((ref_f.std(unbiased=False) * 0.01).item()) + allowed = atol + 0.02 * ref_f.abs() + passed = diff.abs() <= allowed + return { + "rel_rms": rel_rms, + "max_abs_error": float(diff.abs().max().item()), + "atol": atol, + "rtol": 0.02, + "pass_rate": float(passed.float().mean().item()), + "allclose": bool(passed.all().item()), + } + + +def _active_tail_stats(actual: torch.Tensor, reference: torch.Tensor, active_hidden: int) -> dict: + active = _stats(actual[:, :active_hidden], reference[:, :active_hidden]) + tail = actual[:, active_hidden:] + active["padded_tail_max_abs"] = float(tail.float().abs().max().item()) if tail.numel() else 0.0 + return active + + +def _run_stage(args: argparse.Namespace, stage: str) -> dict: + mlen = args.mlen + blen = args.blen + rows = args.rows + hidden = args.hidden_size or 2880 + physical_hidden = _align_to(hidden, mlen) if args.allow_padded_hidden else hidden + if physical_hidden != hidden and stage != "weighted_rms_norm": + raise ValueError("--allow-padded-hidden currently applies only to weighted_rms_norm") + if not args.allow_padded_hidden and hidden % mlen != 0: + raise ValueError(f"hidden={hidden} must be divisible by MLEN={mlen}") + if rows % blen != 0: + raise ValueError(f"rows={rows} must be a multiple of BLEN={blen}") + + torch.manual_seed(args.seed) + if args.x_pt is not None: + hidden_x = _load_bf16_tensor(args.x_pt) + if hidden_x.dim() > 2: + hidden_x = hidden_x.reshape(-1, hidden_x.shape[-1]).contiguous() + if tuple(hidden_x.shape) != (rows, hidden): + raise ValueError(f"--x-pt shape {tuple(hidden_x.shape)} != {(rows, hidden)}") + else: + hidden_x = (torch.randn(rows, hidden) * 0.5).to(torch.bfloat16) + if args.norm_weight_pt is not None: + norm_weight_raw = _load_bf16_tensor(args.norm_weight_pt) + if norm_weight_raw.dim() == 1: + if norm_weight_raw.numel() != hidden: + raise ValueError(f"--norm-weight-pt length {norm_weight_raw.numel()} != hidden={hidden}") + norm_weight_rows = norm_weight_raw.reshape(1, hidden).repeat(rows, 1).to(torch.bfloat16) + elif norm_weight_raw.dim() == 2: + if tuple(norm_weight_raw.shape) == (1, hidden): + norm_weight_rows = norm_weight_raw.repeat(rows, 1).to(torch.bfloat16) + elif tuple(norm_weight_raw.shape) == (rows, hidden): + norm_weight_rows = norm_weight_raw.to(torch.bfloat16) + else: + raise ValueError( + f"--norm-weight-pt shape {tuple(norm_weight_raw.shape)} incompatible with {(rows, hidden)}" + ) + else: + raise ValueError(f"--norm-weight-pt must be rank 1 or 2, got {tuple(norm_weight_raw.shape)}") + else: + norm_weight = (1.0 + torch.randn(hidden) * 0.05).to(torch.bfloat16) + norm_weight_rows = norm_weight.reshape(1, hidden).repeat(rows, 1).to(torch.bfloat16) + + build_dir = (args.build_dir / stage).expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + physical_shape = (max(blen, rows), physical_hidden) + alignment = mlen * mlen + x_base = 0 + y_base = _align_to(x_base + _vram_layout_size(physical_shape, mlen=mlen), alignment) + w_base = _align_to(y_base + _vram_layout_size(physical_shape, mlen=mlen), alignment) + preload_size = w_base + _vram_layout_size(physical_shape, mlen=mlen) + vram_preload = torch.zeros(preload_size, dtype=torch.bfloat16) + + sublayer = (torch.randn(rows, hidden) * 0.25).to(torch.bfloat16) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_vram = prestage_bf16_vram_matrix( + prog=prog, + name="GlueX", + tensor=hidden_x, + vram_addr=x_base, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + y_vram = prestage_bf16_vram_matrix( + prog=prog, + name="GlueY", + tensor=sublayer, + vram_addr=y_base, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + weight_vram = prestage_bf16_vram_matrix( + prog=prog, + name="GlueNormWeight", + tensor=norm_weight_rows, + vram_addr=w_base, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + + eps = 1e-5 + if stage == "weighted_rms_norm": + prog.rms_norm(x_vram) + prog.vram_mul(x_vram, weight_vram, num_rows=rows) + output = x_vram + if physical_hidden != hidden: + hidden_x_physical = torch.zeros(rows, physical_hidden, dtype=torch.bfloat16) + hidden_x_physical[:, :hidden] = hidden_x + norm_weight_physical = torch.zeros(rows, physical_hidden, dtype=torch.bfloat16) + norm_weight_physical[:, :hidden] = norm_weight_rows + golden = _weighted_rms_norm_golden( + hidden_x_physical, + norm_weight_physical, + eps, + active_hidden=hidden, + ) + else: + golden = _weighted_rms_norm_golden(hidden_x, norm_weight_rows, eps) + elif stage == "residual_add": + prog.vram_add(x_vram, y_vram, num_rows=rows) + output = x_vram + golden = _residual_add_golden(hidden_x, sublayer) + else: + raise ValueError(stage) + + isa = prog.compile() + fp_preload = [0.0, eps, 1.0 / hidden] + dummy_input = {"DUMMY": torch.zeros(1, mlen)} + create_sim_env( + dummy_input, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + ) + create_mem_for_sim( + data_size=mlen, + mode="behave_sim", + asm=f"gpt_oss_block_glue_{stage}", + build_path=build_dir, + input_tensors=dummy_input, + specified_data_order=["DUMMY"], + ) + _rewrite_compact_golden(build_dir, golden) + compare_hidden = physical_hidden if physical_hidden != hidden else hidden + comparison_params = _comparison_params_for(output, rows=rows, hidden=compare_hidden, mlen=mlen, golden=golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + + metrics = run_emulator(build_dir, threads=getattr(args, "emu_threads", None)) + results, params = compare_emulator_output(build_dir) + emu = results["simulated_values"].reshape(rows, compare_hidden).to(torch.bfloat16) + stage_stats = _active_tail_stats(emu, golden, hidden) if physical_hidden != hidden else _stats(emu, golden) + summary = { + "stage": stage, + "rows": rows, + "hidden": hidden, + "physical_hidden": physical_hidden, + "mlen": mlen, + "blen": blen, + "asm_lines": len(isa.splitlines()), + "run_metrics": metrics, + "comparison_params": params, + "stats": stage_stats, + "passed": stage_stats["allclose"] and stage_stats["rel_rms"] <= 0.02, + } + (build_dir / "gpt_oss_block_glue_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + if not summary["passed"]: + raise AssertionError(f"{stage} failed: {stage_stats}") + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--rows", type=int, default=4) + parser.add_argument( + "--stage", + choices=("weighted_rms_norm", "residual_add", "all"), + default="all", + ) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_block_glue", + ) + parser.add_argument( + "--allow-padded-hidden", + action="store_true", + help="Pad hidden to MLEN physically while keeping 1/logical_hidden in FPRAM.", + ) + parser.add_argument("--x-pt", type=Path, default=None, help="Optional BF16 input tensor for weighted_rms_norm.") + parser.add_argument( + "--norm-weight-pt", + type=Path, + default=None, + help="Optional rank-1 or rank-2 BF16 RMSNorm weight tensor for weighted_rms_norm.", + ) + args = parser.parse_args() + + stages = ["weighted_rms_norm", "residual_add"] if args.stage == "all" else [args.stage] + summaries = [_run_stage(args, stage) for stage in stages] + if len(summaries) > 1: + args.build_dir.expanduser().resolve().mkdir(parents=True, exist_ok=True) + (args.build_dir.expanduser().resolve() / "gpt_oss_block_glue_summary.json").write_text( + json.dumps({"stages": summaries, "passed": all(s["passed"] for s in summaries)}, indent=2) + "\n" + ) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/__init__.py b/transactional_emulator/testbench/routed_moe/__init__.py new file mode 100644 index 00000000..13899151 --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/__init__.py @@ -0,0 +1 @@ +"""Routed-MoE emulator bring-up harnesses.""" diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_moe_activation_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_activation_test.py new file mode 100644 index 00000000..fe4a073a --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_activation_test.py @@ -0,0 +1,199 @@ +"""GPT-OSS MoE clamp-gated activation emulator smoke. + +This is the third device-side step for fixed-routing MoE v0. It keeps the +expert projections split into gate/up, then executes the GPT-OSS non-standard +activation on PLENA vector ISA: + + gate = min(gate, 7) + up = max(min(up, 7), -7) + sigmoid = 1 / (1 + exp(-1.702 * gate)) + hidden = (up + 1) * gate * sigmoid + +Down projection, routing, and combine remain out of scope. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import torch + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, resolve_rows, setup_hw +from transactional_emulator.testbench.aten.golden import golden_linear +from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _exact_mxfp8_tensor(shape: tuple[int, ...], *, stride: int, offset: int = 0) -> torch.Tensor: + values = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0], dtype=torch.float32) + idx = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.long) + return values[(idx * stride + offset) % values.numel()].reshape(shape) + + +def _activation_golden(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """Golden A with BF16 rounding at each vector-ISA step.""" + neg_alpha = torch.tensor(-1.702, dtype=torch.bfloat16).float() + + gate = _bf16(torch.clamp(gate.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), min=-7.0)) + + sigmoid = _bf16(gate.float()) + sigmoid = _bf16(sigmoid.float() * neg_alpha) + sigmoid = _bf16(torch.exp(torch.clamp(sigmoid.float(), -88.0, 88.0))) + sigmoid = _bf16(sigmoid.float() + 1.0) + sigmoid = _bf16(torch.reciprocal(sigmoid.float())) + + glu = _bf16(gate.float() * sigmoid.float()) + up_plus_one = _bf16(up.float() + 1.0) + return _bf16(up_plus_one.float() * glu.float()) + + +def run_activation_smoke(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + rows, batch_size, seq_len = resolve_rows(args, default_seq=blen) + hidden = args.hidden_size or mlen + intermediate = args.intermediate_size or mlen + + if hidden % mlen != 0: + raise ValueError(f"hidden ({hidden}) must be divisible by MLEN ({mlen})") + if intermediate % mlen != 0: + raise ValueError(f"intermediate ({intermediate}) must be divisible by MLEN ({mlen})") + + build_dir = args.build_dir + hw = setup_hw(args, build_dir) + + print("=" * 80) + print( + "GPT-OSS MoE activation emulator smoke " + f"(mlen={mlen}, blen={blen}, batch={batch_size}, seq={seq_len}, rows={rows})" + ) + print("=" * 80) + + x = _exact_mxfp8_tensor((rows, hidden), stride=1) + w_gate = _exact_mxfp8_tensor((hidden, intermediate), stride=2, offset=1) + w_up = _exact_mxfp8_tensor((hidden, intermediate), stride=3, offset=2) + + gate = torch.matmul(x.float(), w_gate.float()).to(torch.bfloat16) + up = torch.matmul(x.float(), w_up.float()).to(torch.bfloat16) + if not torch.equal(gate, golden_linear(x, w_gate)): + raise AssertionError("gate projection lost Golden-A equivalence under exact MXFP8 inputs") + if not torch.equal(up, golden_linear(x, w_up)): + raise AssertionError("up projection lost Golden-A equivalence under exact MXFP8 inputs") + + golden = _activation_golden(gate, up) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + w_gate_input = prog.input("W_gate_e0", shape=(hidden, intermediate)) + w_up_input = prog.input("W_up_e0", shape=(hidden, intermediate)) + + _zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=rows) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=rows) + one = prog.fp_var("one", size=rows) + neg_alpha = prog.fp_var("neg_alpha", size=rows) + + x_vram = prog.load_batch(x_input, name="X") + gate_vram = prog.linear_projection(x_vram, w_gate_input, name="gate_e0") + up_vram = prog.linear_projection(x_vram, w_up_input, name="up_e0") + num_col_blocks = intermediate // mlen + sigmoid_vram = prog.alloc( + "gate_sigmoid", + rows=rows, + cols=intermediate, + physical_shape=(max(mlen, math.ceil(rows / mlen) * mlen), intermediate), + strict=False, + ) + + active_rows = list(range(rows)) + for col_block in range(num_col_blocks): + prog.tile_row_min_fp(gate_vram, limit_pos, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_min_fp(up_vram, limit_pos, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_max_fp(up_vram, limit_neg, rows=active_rows, tile_col_idx=col_block) + + # Copy gate into sigmoid scratch: scratch = 0; scratch += gate. + prog.vram_fill_zero(sigmoid_vram, rows=active_rows) + prog.vram_add(sigmoid_vram, gate_vram, num_rows=rows) + + for col_block in range(num_col_blocks): + prog.tile_row_mul_fp(sigmoid_vram, neg_alpha, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_exp(sigmoid_vram, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_add_fp(sigmoid_vram, one, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_reci(sigmoid_vram, rows=active_rows, tile_col_idx=col_block) + prog.vram_mul(gate_vram, sigmoid_vram, num_rows=rows) + + for col_block in range(num_col_blocks): + prog.tile_row_add_fp(up_vram, one, rows=active_rows, tile_col_idx=col_block) + prog.vram_mul(up_vram, gate_vram, num_rows=rows) + isa = prog.compile() + + input_tensors = {"X": x, "W_gate_e0": w_gate, "W_up_e0": w_up} + golden_result = {"original_output": golden} + fp_preload = [0.0] + [7.0] * rows + [-7.0] * rows + [1.0] * rows + [-1.702] * rows + [0.0] * 8 + + create_sim_env(input_tensors, isa, golden_result, fp_preload, build_dir=str(build_dir)) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_activation", + data=None, + specified_data_order=["X", "W_gate_e0", "W_up_e0"], + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + ) + + output_vram_addr = prog._compiler.get_vram_addr(up_vram.name) + comparison_params = { + "start_row_idx": output_vram_addr // mlen, + "num_rows": (rows * intermediate) // mlen, + "num_batches": rows, + "elements_per_batch": intermediate, + "row_dim": mlen, + "atol": 0.0, + "rtol": 0.0, + } + + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(comparison_params, f, indent=2) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(isa) + + print(f"Generated {len(isa.splitlines())} lines of ISA") + print(f"activation output VRAM row: {output_vram_addr // mlen}") + if args.no_run: + return {"build_dir": str(build_dir), "ran": False} + + metrics = run_and_assert(build_dir, "gpt_oss_moe_activation", mlen=mlen, blen=blen) + return {"build_dir": str(build_dir), "ran": True, "metrics": metrics} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--intermediate-size", type=int, default=None) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_moe_activation", + ) + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + run_activation_smoke(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_moe_clamp_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_clamp_test.py new file mode 100644 index 00000000..9cf529b2 --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_clamp_test.py @@ -0,0 +1,158 @@ +"""GPT-OSS MoE split gate/up clamp emulator smoke. + +This is the second device-side step for fixed-routing MoE v0. GPT-OSS stores +gate/up projection weights interleaved in the reference model, but PLENA's +current vector-scalar min/max applies to whole vector rows or coarse head +chunks, not even/odd lanes. For exact v0 semantics we split the packed +gate_up projection into separate gate and up projections: + + gate = X @ W_gate + up = X @ W_up + gate = min(gate, 7) + up = max(min(up, 7), -7) + +The final compared tensor is gate + up, only to put both clamp paths into one +VRAM output for this smoke. Gated activation and down projection remain out of +scope. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, resolve_rows, setup_hw +from transactional_emulator.testbench.aten.golden import golden_linear +from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _exact_mxfp8_tensor(shape: tuple[int, ...], *, stride: int, offset: int = 0) -> torch.Tensor: + values = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0], dtype=torch.float32) + idx = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.long) + return values[(idx * stride + offset) % values.numel()].reshape(shape) + + +def run_split_clamp_smoke(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + rows, batch_size, seq_len = resolve_rows(args, default_seq=blen) + hidden = args.hidden_size or mlen + intermediate = args.intermediate_size or mlen + + if hidden % mlen != 0: + raise ValueError(f"hidden ({hidden}) must be divisible by MLEN ({mlen})") + if intermediate % mlen != 0: + raise ValueError(f"intermediate ({intermediate}) must be divisible by MLEN ({mlen})") + + build_dir = args.build_dir + hw = setup_hw(args, build_dir) + + print("=" * 80) + print( + "GPT-OSS MoE split gate/up clamp emulator smoke " + f"(mlen={mlen}, blen={blen}, batch={batch_size}, seq={seq_len}, rows={rows})" + ) + print("=" * 80) + + x = _exact_mxfp8_tensor((rows, hidden), stride=1) + w_gate = _exact_mxfp8_tensor((hidden, intermediate), stride=2, offset=1) + w_up = _exact_mxfp8_tensor((hidden, intermediate), stride=3, offset=2) + + gate = torch.matmul(x.float(), w_gate.float()).to(torch.bfloat16) + up = torch.matmul(x.float(), w_up.float()).to(torch.bfloat16) + if not torch.equal(gate, golden_linear(x, w_gate)): + raise AssertionError("gate projection lost Golden-A equivalence under exact MXFP8 inputs") + if not torch.equal(up, golden_linear(x, w_up)): + raise AssertionError("up projection lost Golden-A equivalence under exact MXFP8 inputs") + + golden = (torch.clamp(gate.float(), max=7.0) + torch.clamp(up.float(), min=-7.0, max=7.0)).to(torch.bfloat16) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + w_gate_input = prog.input("W_gate_e0", shape=(hidden, intermediate)) + w_up_input = prog.input("W_up_e0", shape=(hidden, intermediate)) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=1) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=1) + + x_vram = prog.load_batch(x_input, name="X") + gate_vram = prog.linear_projection(x_vram, w_gate_input, name="gate_e0") + up_vram = prog.linear_projection(x_vram, w_up_input, name="up_e0") + + active_rows = list(range(rows)) + num_col_blocks = intermediate // mlen + # Program-level tile_row_*_fp maps multiple rows to consecutive FPRAM + # offsets. GPT-OSS clamp needs the same scalar limit on every row, so emit + # one single-row op at a time and pass FPVar objects to keep bounds checks. + for col_block in range(num_col_blocks): + for row_idx in active_rows: + prog.tile_row_min_fp(gate_vram, limit_pos, row_idx=row_idx, tile_col_idx=col_block) + prog.tile_row_min_fp(up_vram, limit_pos, row_idx=row_idx, tile_col_idx=col_block) + prog.tile_row_max_fp(up_vram, limit_neg, row_idx=row_idx, tile_col_idx=col_block) + prog.vram_add(gate_vram, up_vram, num_rows=rows) + isa = prog.compile() + + input_tensors = {"X": x, "W_gate_e0": w_gate, "W_up_e0": w_up} + golden_result = {"original_output": golden} + fp_preload = [7.0, -7.0, 0.0, 1e-6, 1.0 / hidden] + [0.0] * 5 + + create_sim_env(input_tensors, isa, golden_result, fp_preload, build_dir=str(build_dir)) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_clamp", + data=None, + specified_data_order=["X", "W_gate_e0", "W_up_e0"], + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + ) + + gate_vram_addr = prog._compiler.get_vram_addr(gate_vram.name) + comparison_params = { + "start_row_idx": gate_vram_addr // mlen, + "num_rows": (rows * intermediate) // mlen, + "num_batches": rows, + "elements_per_batch": intermediate, + "row_dim": mlen, + "atol": 0.0, + "rtol": 0.0, + } + + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(comparison_params, f, indent=2) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(isa) + + print(f"Generated {len(isa.splitlines())} lines of ISA") + print(f"gate+up clamped output VRAM row: {gate_vram_addr // mlen}") + if args.no_run: + return {"build_dir": str(build_dir), "ran": False} + + metrics = run_and_assert(build_dir, "gpt_oss_moe_clamp", mlen=mlen, blen=blen) + return {"build_dir": str(build_dir), "ran": True, "metrics": metrics} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--intermediate-size", type=int, default=None) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_moe_clamp", + ) + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + run_split_clamp_smoke(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_moe_combine_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_combine_test.py new file mode 100644 index 00000000..1139ff4e --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_combine_test.py @@ -0,0 +1,433 @@ +"""GPT-OSS MoE fixed-routing two-expert combine emulator smoke. + +This is the first end-to-end fixed-routing device-side smoke: + + X -> expert0 -> route_weight0 * out0 + X -> expert1 -> route_weight1 * out1 + output = weighted_out0 + weighted_out1 + +Router/top-k/dispatch remain host-fixed. The routing weights are loaded as +expanded [rows, hidden] matrices so this test exercises PLENA vram_mul/vram_add +without introducing gather/scatter or scalar broadcast machinery. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import torch + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, resolve_rows, setup_hw +from transactional_emulator.testbench.aten.golden import quantize_to_mxfp +from transactional_emulator.testbench.emulator_runner import compare_emulator_output, run_and_assert +from transactional_emulator.testbench.layout_utils import infer_hbm_tensor_layouts +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _exact_mxfp8_tensor(shape: tuple[int, ...], *, stride: int, offset: int = 0) -> torch.Tensor: + values = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0], dtype=torch.float32) + idx = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.long) + return values[(idx * stride + offset) % values.numel()].reshape(shape) + + +def _activation_golden(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + neg_alpha = torch.tensor(-1.702, dtype=torch.bfloat16).float() + + gate = _bf16(torch.clamp(gate.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), min=-7.0)) + + sigmoid = _bf16(gate.float()) + sigmoid = _bf16(sigmoid.float() * neg_alpha) + sigmoid = _bf16(torch.exp(torch.clamp(sigmoid.float(), -88.0, 88.0))) + sigmoid = _bf16(sigmoid.float() + 1.0) + sigmoid = _bf16(torch.reciprocal(sigmoid.float())) + + glu = _bf16(gate.float() * sigmoid.float()) + up_plus_one = _bf16(up.float() + 1.0) + return _bf16(up_plus_one.float() * glu.float()) + + +def _linear_projection_golden( + x: torch.Tensor, + w: torch.Tensor, + *, + mlen: int, + mram_tile_capacity: int = 4, + hbm_input: bool = True, +) -> torch.Tensor: + """Hardware-aware linear projection golden, including compiler K-split. + + HBM-loaded activations and weights are MX-quantized before matmul. VRAM + intermediate activations are already BF16 and must not be MX-quantized + again. When K exceeds MRAM tile capacity, compiler emits partial matmuls and + BF16 VRAM adds; mirror that rounding order here. + """ + x_q = quantize_to_mxfp(x) if hbm_input else x.to(torch.bfloat16) + w_q = quantize_to_mxfp(w) + k_total = x_q.shape[1] + chunk = mlen * mram_tile_capacity + + acc = None + for k_start in range(0, k_total, chunk): + k_end = min(k_start + chunk, k_total) + partial = torch.matmul(x_q[:, k_start:k_end].float(), w_q[k_start:k_end, :].float()).to(torch.bfloat16) + acc = partial if acc is None else (acc.float() + partial.float()).to(torch.bfloat16) + assert acc is not None + return acc + + +def _expert_golden( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + *, + mlen: int, + b_gate: torch.Tensor | None = None, + b_up: torch.Tensor | None = None, + b_down: torch.Tensor | None = None, +) -> torch.Tensor: + gate = _linear_projection_golden(x, w_gate, mlen=mlen, hbm_input=True) + up = _linear_projection_golden(x, w_up, mlen=mlen, hbm_input=True) + if b_gate is not None: + gate = _bf16(gate.float() + quantize_to_mxfp(b_gate).float()) + if b_up is not None: + up = _bf16(up.float() + quantize_to_mxfp(b_up).float()) + hidden = _activation_golden(gate, up) + out = _linear_projection_golden(hidden, w_down, mlen=mlen, hbm_input=False) + if b_down is not None: + out = _bf16(out.float() + quantize_to_mxfp(b_down).float()) + return out + + +def _route_rows(rows: int) -> tuple[torch.Tensor, torch.Tensor]: + idx = torch.arange(rows, dtype=torch.float32) + route0 = torch.tensor(0.25, dtype=torch.float32) + (idx % 4) * torch.tensor(0.25, dtype=torch.float32) + return route0.unsqueeze(1), (1.0 - route0).unsqueeze(1) + + +def _print_tile_and_clamp_stats( + *, + mlen: int, + mram_tile_capacity: int, + rows: int, + hidden: int, + intermediate: int, + projections: dict[str, tuple[torch.Tensor, torch.Tensor]], +) -> None: + m_tiles = math.ceil(rows / mlen) + gate_n_tiles = math.ceil(intermediate / mlen) + gate_k_tiles = math.ceil(hidden / mlen) + down_n_tiles = math.ceil(hidden / mlen) + down_k_tiles = math.ceil(intermediate / mlen) + print("\n--- v0 stress coverage ---") + print( + "gate/up projection tiles: " + f"M={m_tiles}, N={gate_n_tiles}, K={gate_k_tiles}, " + f"K-split chunks={math.ceil(gate_k_tiles / mram_tile_capacity)}" + ) + print( + "down projection tiles: " + f"M={m_tiles}, N={down_n_tiles}, K={down_k_tiles}, " + f"K-split chunks={math.ceil(down_k_tiles / mram_tile_capacity)}" + ) + for name, (gate, up) in projections.items(): + gate_hi = int((gate > 7).sum().item()) + gate_lo = int((gate < -7).sum().item()) + up_hi = int((up > 7).sum().item()) + up_lo = int((up < -7).sum().item()) + total = gate.numel() + print( + f"{name}: gate range=[{gate.min().item():.4g}, {gate.max().item():.4g}], " + f"gate >7={gate_hi}/{total}, gate <-7={gate_lo}/{total}; " + f"up range=[{up.min().item():.4g}, {up.max().item():.4g}], " + f"up >7={up_hi}/{total}, up <-7={up_lo}/{total}" + ) + + +def run_combine_smoke(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + rows, batch_size, seq_len = resolve_rows(args, default_seq=blen) + hidden = args.hidden_size or mlen + intermediate = args.intermediate_size or mlen + + if hidden % mlen != 0: + raise ValueError(f"hidden ({hidden}) must be divisible by MLEN ({mlen})") + if intermediate % mlen != 0: + raise ValueError(f"intermediate ({intermediate}) must be divisible by MLEN ({mlen})") + + build_dir = args.build_dir + hw = setup_hw(args, build_dir) + + print("=" * 80) + print( + "GPT-OSS MoE fixed-routing combine emulator smoke " + f"(mlen={mlen}, blen={blen}, batch={batch_size}, seq={seq_len}, rows={rows})" + ) + print("=" * 80) + + x = _exact_mxfp8_tensor((rows, hidden), stride=1) + w0_gate = _exact_mxfp8_tensor((hidden, intermediate), stride=2, offset=1) + w0_up = _exact_mxfp8_tensor((hidden, intermediate), stride=3, offset=2) + w0_down = _exact_mxfp8_tensor((intermediate, hidden), stride=4, offset=3) + w1_gate = _exact_mxfp8_tensor((hidden, intermediate), stride=3, offset=4) + w1_up = _exact_mxfp8_tensor((hidden, intermediate), stride=4, offset=0) + w1_down = _exact_mxfp8_tensor((intermediate, hidden), stride=2, offset=2) + + b0_gate = b0_up = b0_down = None + b1_gate = b1_up = b1_down = None + if args.include_bias: + b0_gate = _exact_mxfp8_tensor((1, intermediate), stride=2, offset=1).repeat(rows, 1) + b0_up = _exact_mxfp8_tensor((1, intermediate), stride=3, offset=2).repeat(rows, 1) + b0_down = _exact_mxfp8_tensor((1, hidden), stride=4, offset=3).repeat(rows, 1) + b1_gate = _exact_mxfp8_tensor((1, intermediate), stride=3, offset=4).repeat(rows, 1) + b1_up = _exact_mxfp8_tensor((1, intermediate), stride=4, offset=0).repeat(rows, 1) + b1_down = _exact_mxfp8_tensor((1, hidden), stride=2, offset=2).repeat(rows, 1) + + route0_row, route1_row = _route_rows(rows) + route0 = route0_row.repeat(1, hidden) + route1 = route1_row.repeat(1, hidden) + + gate0 = _linear_projection_golden(x, w0_gate, mlen=mlen, hbm_input=True) + up0 = _linear_projection_golden(x, w0_up, mlen=mlen, hbm_input=True) + gate1 = _linear_projection_golden(x, w1_gate, mlen=mlen, hbm_input=True) + up1 = _linear_projection_golden(x, w1_up, mlen=mlen, hbm_input=True) + _print_tile_and_clamp_stats( + mlen=mlen, + mram_tile_capacity=4, + rows=rows, + hidden=hidden, + intermediate=intermediate, + projections={"expert0": (gate0, up0), "expert1": (gate1, up1)}, + ) + + out0 = _expert_golden( + x, + w0_gate, + w0_up, + w0_down, + mlen=mlen, + b_gate=b0_gate, + b_up=b0_up, + b_down=b0_down, + ) + out1 = _expert_golden( + x, + w1_gate, + w1_up, + w1_down, + mlen=mlen, + b_gate=b1_gate, + b_up=b1_up, + b_down=b1_down, + ) + route0_q = quantize_to_mxfp(route0) + route1_q = quantize_to_mxfp(route1) + weighted0 = _bf16(out0.float() * route0_q.float()) + weighted1 = _bf16(out1.float() * route1_q.float()) + golden = _bf16(weighted0.float() + weighted1.float()) + if args.include_bias: + no_bias0 = _expert_golden(x, w0_gate, w0_up, w0_down, mlen=mlen) + no_bias1 = _expert_golden(x, w1_gate, w1_up, w1_down, mlen=mlen) + no_bias_golden = _bf16( + _bf16(no_bias0.float() * route0_q.float()).float() + _bf16(no_bias1.float() * route1_q.float()).float() + ) + bias_delta = (golden.float() - no_bias_golden.float()).abs().max().item() + print(f"bias changes output: max_abs_delta={bias_delta:.6g}") + if bias_delta == 0.0: + raise AssertionError("--include-bias did not change the fixed-routing golden output") + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + route0_input = prog.input("route0", shape=(rows, hidden)) + route1_input = prog.input("route1", shape=(rows, hidden)) + w0_gate_input = prog.input("W0_gate", shape=(hidden, intermediate)) + w0_up_input = prog.input("W0_up", shape=(hidden, intermediate)) + w0_down_input = prog.input("W0_down", shape=(intermediate, hidden)) + w1_gate_input = prog.input("W1_gate", shape=(hidden, intermediate)) + w1_up_input = prog.input("W1_up", shape=(hidden, intermediate)) + w1_down_input = prog.input("W1_down", shape=(intermediate, hidden)) + if args.include_bias: + b0_gate_input = prog.input("B0_gate", shape=(rows, intermediate)) + b0_up_input = prog.input("B0_up", shape=(rows, intermediate)) + b0_down_input = prog.input("B0_down", shape=(rows, hidden)) + b1_gate_input = prog.input("B1_gate", shape=(rows, intermediate)) + b1_up_input = prog.input("B1_up", shape=(rows, intermediate)) + b1_down_input = prog.input("B1_down", shape=(rows, hidden)) + + zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=rows) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=rows) + one = prog.fp_var("one", size=rows) + neg_alpha = prog.fp_var("neg_alpha", size=rows) + + x_vram = prog.load_batch(x_input, name="X") + route0_vram = prog.load_batch(route0_input, name="route0") + route1_vram = prog.load_batch(route1_input, name="route1") + expert_biases = None + if args.include_bias: + b0_gate_vram = prog.load_batch(b0_gate_input, name="B0_gate") + b0_up_vram = prog.load_batch(b0_up_input, name="B0_up") + b0_down_vram = prog.load_batch(b0_down_input, name="B0_down") + b1_gate_vram = prog.load_batch(b1_gate_input, name="B1_gate") + b1_up_vram = prog.load_batch(b1_up_input, name="B1_up") + b1_down_vram = prog.load_batch(b1_down_input, name="B1_down") + expert_biases = [ + (b0_gate_vram, b0_up_vram, b0_down_vram), + (b1_gate_vram, b1_up_vram, b1_down_vram), + ] + + output_vram = prog.gpt_oss_moe_fixed_routing_v0( + x_vram, + experts=[ + (w0_gate_input, w0_up_input, w0_down_input), + (w1_gate_input, w1_up_input, w1_down_input), + ], + route_weights=[route0_vram, route1_vram], + expert_biases=expert_biases, + rows=rows, + intermediate=intermediate, + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + name="fixed_routing", + ) + isa = prog.compile() + + input_tensors = { + "X": x, + "route0": route0, + "route1": route1, + "W0_gate": w0_gate, + "W0_up": w0_up, + "W0_down": w0_down, + "W1_gate": w1_gate, + "W1_up": w1_up, + "W1_down": w1_down, + } + if args.include_bias: + input_tensors.update( + { + "B0_gate": b0_gate, + "B0_up": b0_up, + "B0_down": b0_down, + "B1_gate": b1_gate, + "B1_up": b1_up, + "B1_down": b1_down, + } + ) + golden_result = {"original_output": golden} + fp_preload = [0.0] + [7.0] * rows + [-7.0] * rows + [1.0] * rows + [-1.702] * rows + [0.0] * 8 + + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + create_sim_env( + input_tensors, + isa, + golden_result, + fp_preload, + build_dir=str(build_dir), + tensor_layouts=tensor_layouts, + ) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_combine", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + + output_vram_addr = prog._compiler.get_vram_addr(output_vram.name) + comparison_params = { + "start_row_idx": output_vram_addr // mlen, + "num_rows": (rows * hidden) // mlen, + "num_batches": rows, + "elements_per_batch": hidden, + "row_dim": mlen, + "atol": 0.0, + "rtol": 0.0, + } + + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(comparison_params, f, indent=2) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(isa) + + print(f"Generated {len(isa.splitlines())} lines of ISA") + print(f"combined output VRAM row: {output_vram_addr // mlen}") + if args.no_run: + return {"build_dir": str(build_dir), "ran": False} + + metrics = run_and_assert(build_dir, "gpt_oss_moe_combine", mlen=mlen, blen=blen) + results, params = compare_emulator_output(build_dir) + comparison_summary = { + key: results[key] + for key in ( + "mse", + "mae", + "max_error", + "relative_error", + "relative_match_rate", + "allclose_match_rate", + "match_rate", + "allclose_pass", + "atol", + "rtol", + "golden_shape", + "simulated_shape", + ) + } + baseline_path = build_dir / "gpt_oss_moe_combine_baseline.json" + baseline_path.write_text( + json.dumps( + { + "name": "gpt_oss_moe_combine_tiny_emu_vs_golden_b", + "include_bias": args.include_bias, + "rows": rows, + "hidden": hidden, + "intermediate": intermediate, + "mlen": mlen, + "blen": blen, + "comparison": comparison_summary, + "comparison_params": params, + "run_metrics": metrics, + }, + indent=2, + ) + + "\n" + ) + print(f"baseline metrics written: {baseline_path}") + return {"build_dir": str(build_dir), "ran": True, "metrics": metrics, "baseline": str(baseline_path)} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--intermediate-size", type=int, default=None) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_moe_combine", + ) + parser.add_argument("--include-bias", action="store_true") + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + run_combine_smoke(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_moe_expert_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_expert_test.py new file mode 100644 index 00000000..e0c6ca39 --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_expert_test.py @@ -0,0 +1,215 @@ +"""GPT-OSS MoE single-expert emulator smoke. + +This extends the activation smoke by adding the expert down projection: + + gate = min(X @ W_gate, 7) + up = max(min(X @ W_up, 7), -7) + hidden = (up + 1) * gate * sigmoid(1.702 * gate) + out = hidden @ W_down + +Routing and multi-expert combine remain out of scope. +""" + +from __future__ import annotations + +import argparse +import json +import math +from pathlib import Path + +import torch + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, resolve_rows, setup_hw +from transactional_emulator.testbench.aten.golden import golden_linear, quantize_to_mxfp +from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _exact_mxfp8_tensor(shape: tuple[int, ...], *, stride: int, offset: int = 0) -> torch.Tensor: + values = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0], dtype=torch.float32) + idx = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.long) + return values[(idx * stride + offset) % values.numel()].reshape(shape) + + +def _activation_golden(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + neg_alpha = torch.tensor(-1.702, dtype=torch.bfloat16).float() + + gate = _bf16(torch.clamp(gate.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), min=-7.0)) + + sigmoid = _bf16(gate.float()) + sigmoid = _bf16(sigmoid.float() * neg_alpha) + sigmoid = _bf16(torch.exp(torch.clamp(sigmoid.float(), -88.0, 88.0))) + sigmoid = _bf16(sigmoid.float() + 1.0) + sigmoid = _bf16(torch.reciprocal(sigmoid.float())) + + glu = _bf16(gate.float() * sigmoid.float()) + up_plus_one = _bf16(up.float() + 1.0) + return _bf16(up_plus_one.float() * glu.float()) + + +def _emit_gpt_oss_activation( + prog: PlenaCompiler, + gate_vram, + up_vram, + sigmoid_vram, + rows: int, + mlen: int, +): + active_rows = list(range(rows)) + num_col_blocks = sigmoid_vram.shape[1] // mlen + _zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=rows) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=rows) + one = prog.fp_var("one", size=rows) + neg_alpha = prog.fp_var("neg_alpha", size=rows) + + for col_block in range(num_col_blocks): + prog.tile_row_min_fp(gate_vram, limit_pos, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_min_fp(up_vram, limit_pos, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_max_fp(up_vram, limit_neg, rows=active_rows, tile_col_idx=col_block) + + prog.vram_fill_zero(sigmoid_vram, rows=active_rows) + prog.vram_add(sigmoid_vram, gate_vram, num_rows=rows) + + for col_block in range(num_col_blocks): + prog.tile_row_mul_fp(sigmoid_vram, neg_alpha, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_exp(sigmoid_vram, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_add_fp(sigmoid_vram, one, rows=active_rows, tile_col_idx=col_block) + prog.tile_row_reci(sigmoid_vram, rows=active_rows, tile_col_idx=col_block) + prog.vram_mul(gate_vram, sigmoid_vram, num_rows=rows) + + for col_block in range(num_col_blocks): + prog.tile_row_add_fp(up_vram, one, rows=active_rows, tile_col_idx=col_block) + prog.vram_mul(up_vram, gate_vram, num_rows=rows) + + fp_preload = [0.0] + [7.0] * rows + [-7.0] * rows + [1.0] * rows + [-1.702] * rows + [0.0] * 8 + return fp_preload + + +def run_single_expert_smoke(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + rows, batch_size, seq_len = resolve_rows(args, default_seq=blen) + hidden = args.hidden_size or mlen + intermediate = args.intermediate_size or mlen + + if hidden % mlen != 0: + raise ValueError(f"hidden ({hidden}) must be divisible by MLEN ({mlen})") + if intermediate % mlen != 0: + raise ValueError(f"intermediate ({intermediate}) must be divisible by MLEN ({mlen})") + + build_dir = args.build_dir + hw = setup_hw(args, build_dir) + + print("=" * 80) + print( + "GPT-OSS MoE single-expert emulator smoke " + f"(mlen={mlen}, blen={blen}, batch={batch_size}, seq={seq_len}, rows={rows})" + ) + print("=" * 80) + + x = _exact_mxfp8_tensor((rows, hidden), stride=1) + w_gate = _exact_mxfp8_tensor((hidden, intermediate), stride=2, offset=1) + w_up = _exact_mxfp8_tensor((hidden, intermediate), stride=3, offset=2) + w_down = _exact_mxfp8_tensor((intermediate, hidden), stride=4, offset=3) + + gate = torch.matmul(x.float(), w_gate.float()).to(torch.bfloat16) + up = torch.matmul(x.float(), w_up.float()).to(torch.bfloat16) + if not torch.equal(gate, golden_linear(x, w_gate)): + raise AssertionError("gate projection lost Golden-A equivalence under exact MXFP8 inputs") + if not torch.equal(up, golden_linear(x, w_up)): + raise AssertionError("up projection lost Golden-A equivalence under exact MXFP8 inputs") + + hidden_golden = _activation_golden(gate, up) + # The down projection consumes BF16 hidden values already resident in VRAM. + # Unlike HBM-loaded inputs, this activation must not be MX-quantized again; + # only W_down follows the matrix-weight MX path. + golden = torch.matmul(hidden_golden.float(), quantize_to_mxfp(w_down).float()).to(torch.bfloat16) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + w_gate_input = prog.input("W_gate_e0", shape=(hidden, intermediate)) + w_up_input = prog.input("W_up_e0", shape=(hidden, intermediate)) + w_down_input = prog.input("W_down_e0", shape=(intermediate, hidden)) + + x_vram = prog.load_batch(x_input, name="X") + gate_vram = prog.linear_projection(x_vram, w_gate_input, name="gate_e0") + up_vram = prog.linear_projection(x_vram, w_up_input, name="up_e0") + sigmoid_vram = prog.alloc( + "gate_sigmoid", + rows=rows, + cols=intermediate, + physical_shape=(max(mlen, math.ceil(rows / mlen) * mlen), intermediate), + strict=False, + ) + + fp_preload = _emit_gpt_oss_activation(prog, gate_vram, up_vram, sigmoid_vram, rows, mlen) + output_vram = prog.linear_projection(up_vram, w_down_input, name="expert_e0") + isa = prog.compile() + + input_tensors = {"X": x, "W_gate_e0": w_gate, "W_up_e0": w_up, "W_down_e0": w_down} + golden_result = {"original_output": golden} + + create_sim_env(input_tensors, isa, golden_result, fp_preload, build_dir=str(build_dir)) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_expert", + data=None, + specified_data_order=["X", "W_gate_e0", "W_up_e0", "W_down_e0"], + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + ) + + output_vram_addr = prog._compiler.get_vram_addr(output_vram.name) + comparison_params = { + "start_row_idx": output_vram_addr // mlen, + "num_rows": (rows * hidden) // mlen, + "num_batches": rows, + "elements_per_batch": hidden, + "row_dim": mlen, + "atol": 0.0, + "rtol": 0.0, + } + + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(comparison_params, f, indent=2) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(isa) + + print(f"Generated {len(isa.splitlines())} lines of ISA") + print(f"expert output VRAM row: {output_vram_addr // mlen}") + if args.no_run: + return {"build_dir": str(build_dir), "ran": False} + + metrics = run_and_assert(build_dir, "gpt_oss_moe_expert", mlen=mlen, blen=blen) + return {"build_dir": str(build_dir), "ran": True, "metrics": metrics} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--intermediate-size", type=int, default=None) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_moe_expert", + ) + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + run_single_expert_smoke(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_moe_gate_up_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_gate_up_test.py new file mode 100644 index 00000000..6c62cd79 --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_gate_up_test.py @@ -0,0 +1,132 @@ +"""GPT-OSS MoE gate-up projection emulator smoke. + +This is the first device-side step for fixed-routing MoE v0. It only runs the +selected expert's gate_up matmul: + + X[tokens, hidden] @ W_gate_up_e[hidden, 2 * intermediate] + +Bias, clamp, gated activation, down projection, and combine deliberately stay +out of this smoke. The input values are chosen to be exactly representable by +the current PLENA MXFP8 HBM format, so the hardware-aware golden equals the +pure BF16 mathematical golden for this stage. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, resolve_rows, setup_hw +from transactional_emulator.testbench.aten.golden import golden_linear +from transactional_emulator.testbench.emulator_runner import run_and_assert +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _exact_mxfp8_tensor(shape: tuple[int, ...], *, stride: int, offset: int = 0) -> torch.Tensor: + values = torch.tensor([-1.0, -0.5, 0.0, 0.5, 1.0], dtype=torch.float32) + idx = torch.arange(torch.tensor(shape).prod().item(), dtype=torch.long) + return values[(idx * stride + offset) % values.numel()].reshape(shape) + + +def run_gate_up_matmul_smoke(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + rows, batch_size, seq_len = resolve_rows(args, default_seq=blen) + hidden = args.hidden_size or mlen + out_features = args.out_features or 2 * mlen + + if hidden % mlen != 0: + raise ValueError(f"hidden ({hidden}) must be divisible by MLEN ({mlen})") + if out_features % mlen != 0: + raise ValueError(f"out_features ({out_features}) must be divisible by MLEN ({mlen})") + + build_dir = args.build_dir + hw = setup_hw(args, build_dir) + + print("=" * 80) + print( + "GPT-OSS MoE gate_up matmul emulator smoke " + f"(mlen={mlen}, blen={blen}, batch={batch_size}, seq={seq_len}, rows={rows})" + ) + print("=" * 80) + + x = _exact_mxfp8_tensor((rows, hidden), stride=1) + w_gate_up_e = _exact_mxfp8_tensor((hidden, out_features), stride=2, offset=1) + + pure_golden = torch.matmul(x.float(), w_gate_up_e.float()).to(torch.bfloat16) + hardware_golden = golden_linear(x, w_gate_up_e) + if not torch.equal(pure_golden, hardware_golden): + max_diff = (pure_golden.float() - hardware_golden.float()).abs().max().item() + raise AssertionError(f"Exact-value gate_up smoke lost Golden-A equivalence; max diff={max_diff}") + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + w_input = prog.input("W_gate_up_e0", shape=(hidden, out_features)) + x_vram = prog.load_batch(x_input, name="X") + gate_up = prog.linear_projection(x_vram, w_input, name="gate_up_e0") + isa = prog.compile() + + input_tensors = {"X": x, "W_gate_up_e0": w_gate_up_e} + golden_result = {"original_output": pure_golden} + fp_preload = [0.0, 1e-6, 1.0 / hidden] + [0.0] * 7 + + create_sim_env(input_tensors, isa, golden_result, fp_preload, build_dir=str(build_dir)) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_gate_up", + data=None, + specified_data_order=["X", "W_gate_up_e0"], + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + ) + + gate_up_vram_addr = prog._compiler.get_vram_addr(gate_up.name) + comparison_params = { + "start_row_idx": gate_up_vram_addr // mlen, + "num_rows": (rows * out_features) // mlen, + "num_batches": rows, + "elements_per_batch": out_features, + "row_dim": mlen, + "atol": 0.0, + "rtol": 0.0, + } + + with open(build_dir / "comparison_params.json", "w") as f: + json.dump(comparison_params, f, indent=2) + with open(build_dir / "generated_asm_code.asm", "w") as f: + f.write(isa) + + print(f"Generated {len(isa.splitlines())} lines of ISA") + print(f"gate_up output VRAM row: {gate_up_vram_addr // mlen}") + if args.no_run: + return {"build_dir": str(build_dir), "ran": False} + + metrics = run_and_assert(build_dir, "gpt_oss_moe_gate_up", mlen=mlen, blen=blen) + return {"build_dir": str(build_dir), "ran": True, "metrics": metrics} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument("--out-features", type=int, default=None) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_moe_gate_up", + ) + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + run_gate_up_matmul_smoke(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_moe_gather_scatter_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_gather_scatter_test.py new file mode 100644 index 00000000..ed7fe9af --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_moe_gather_scatter_test.py @@ -0,0 +1,2140 @@ +# ruff: noqa: E402 +"""GPT-OSS MoE device-side gather/scatter bring-up. + +Stage ``address`` is the first L2 gate: a runtime token-row offset is loaded +from int SRAM, converted into an HBM vector prefetch address on device, and the +gathered row is compared against the hardware-aware HBM MXFP8 golden. Later +stages extend this same file to expert execution and scatter/combine. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +_AICROSSSIM_ROOT = Path(__file__).resolve().parents[4] +_TESTBENCH_ROOT = Path(__file__).resolve().parents[1] +_ATEN_BUILD_DIR = _TESTBENCH_ROOT / "aten" / "build" +_TOP_LEVEL_COMPILER = _AICROSSSIM_ROOT / "PLENA_Compiler" +if _TOP_LEVEL_COMPILER.exists(): + sys.path.insert(0, str(_TOP_LEVEL_COMPILER)) + +from aten.models.gpt_oss.moe_reference import compare_stats +from aten.models.gpt_oss.moe_reference import split_packed_gate_up +from aten.models.gpt_oss.real_layer_utils import load_layer_tensors +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.vars import VRAMMatrixVar +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.routed_moe.gpt_oss_real_layer0_test import _stats_dict +from transactional_emulator.testbench.routed_moe.gpt_oss_real_layer0_test import _strict_elementwise_details +from transactional_emulator.testbench.routed_moe.gpt_oss_real_layer0_test import _tail_gate_details +from transactional_emulator.testbench.routed_moe.gpt_oss_router_gemm_test import _align_to +from transactional_emulator.testbench.routed_moe.gpt_oss_router_gemm_test import _rank_stability +from transactional_emulator.testbench.routed_moe.gpt_oss_router_gemm_test import _router_vector_bf16_golden +from transactional_emulator.testbench.routed_moe.gpt_oss_router_gemm_test import _vram_layout_size +from transactional_emulator.testbench.emulator_runner import compare_emulator_output, run_and_assert, run_emulator +from transactional_emulator.testbench.layout_utils import infer_hbm_tensor_layouts, prestage_bf16_vram_matrix +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.testbench.sliced_layer_test_builder import quantize_to_mxfp +from transactional_emulator.tools.create_sim_env import create_sim_env +from plena_quant.mxfp.utils import bin_2_fp + + +def _decode_bf16_dump(path: Path) -> torch.Tensor: + raw = np.fromfile(path, dtype=" torch.Tensor: + raw = np.fromfile(path, dtype=" dict: + output_vram_addr = output_vram._program._compiler.get_vram_addr(output_vram.name) + output_physical_rows = output_vram.physical_shape[0] + chunks_per_batch = math.ceil(hidden / mlen) + rows_to_read = (chunks_per_batch - 1) * output_physical_rows + rows + return { + "start_row_idx": output_vram_addr // mlen, + "num_rows": rows_to_read, + "num_batches": rows, + "elements_per_batch": hidden, + "row_dim": mlen, + "physical_rows": output_physical_rows, + "atol": float((golden.float().std(unbiased=False) * 0.01).item()), + "rtol": 0.02, + } + + +def _load_tok8_reference(reference_path: Path) -> dict: + reference = torch.load(reference_path.expanduser().resolve(), map_location="cpu") + return { + "x": reference["x"].to(torch.bfloat16), + "topk_indices": reference["topk_indices"].to(torch.long), + "topk_weights": reference["topk_weights"].to(torch.bfloat16), + } + + +def _decode_e8m0_scale(byte: int) -> float: + """Decode unsigned E8M0 scale as used by the Rust HBM MX path.""" + return 2.0 ** (int(byte) - 127) + + +def _decode_hbm_mxfp8_row( + hbm_path: Path, + *, + tensor_base: int, + tensor_rows: int, + tensor_cols: int, + row_element_offset: int, +) -> torch.Tensor: + """Decode one row from the exact generated HBM preload bytes.""" + return _decode_hbm_mxfp8_row_from_bytes( + hbm_path.read_bytes(), + tensor_base=tensor_base, + tensor_rows=tensor_rows, + tensor_cols=tensor_cols, + row_element_offset=row_element_offset, + ) + + +def _hbm_byte(hbm: bytes, idx: int) -> int: + return hbm[idx] if 0 <= idx < len(hbm) else 0 + + +def _hbm_range(hbm: bytes, start: int, length: int) -> bytes: + if length <= 0: + return b"" + if start >= len(hbm): + return b"\x00" * length + chunk = hbm[start : min(start + length, len(hbm))] + if len(chunk) < length: + chunk += b"\x00" * (length - len(chunk)) + return chunk + + +def _decode_hbm_mxfp8_row_from_bytes( + hbm: bytes, + *, + tensor_base: int, + tensor_rows: int, + tensor_cols: int, + row_element_offset: int, +) -> torch.Tensor: + """Decode one MXFP8 HBM row from a byte buffer, zero-extending missing bytes.""" + element_base = tensor_base + row_element_offset + scale_base = tensor_base + tensor_rows * tensor_cols + (row_element_offset // 8) + values = [] + for col in range(tensor_cols): + element = _hbm_byte(hbm, element_base + col) + scale = _hbm_byte(hbm, scale_base + (col // 8)) + values.append(float(bin_2_fp(element, 4, 3)) * _decode_e8m0_scale(scale)) + return torch.tensor(values, dtype=torch.float32).reshape(1, tensor_cols).to(torch.bfloat16) + + +def _decode_hbm_mxfp8_output_slots( + hbm: bytes, + *, + slot_base: int, + slot_stride: int, + rows: int, + hidden: int, + blen: int, +) -> torch.Tensor: + decoded = torch.zeros(rows, hidden, dtype=torch.bfloat16) + for token_idx in range(rows): + decoded[token_idx : token_idx + 1] = _decode_hbm_mxfp8_row_from_bytes( + hbm, + tensor_base=slot_base + token_idx * slot_stride, + tensor_rows=blen, + tensor_cols=hidden, + row_element_offset=0, + ) + return decoded + + +def _hbm_slot_padding_zero( + hbm: bytes, + *, + slot_base: int, + slot_stride: int, + token_indices: list[int], + hidden: int, + blen: int, +) -> dict: + nonzero_rows = [] + for token_idx in token_indices: + token_base = slot_base + token_idx * slot_stride + for pad_row in range(1, blen): + decoded = _decode_hbm_mxfp8_row_from_bytes( + hbm, + tensor_base=token_base, + tensor_rows=blen, + tensor_cols=hidden, + row_element_offset=pad_row * hidden, + ) + max_abs = float(decoded.float().abs().max().item()) + if max_abs != 0.0: + nonzero_rows.append({"token_index": token_idx, "padding_row": pad_row, "max_abs": max_abs}) + return { + "checked_tokens": token_indices, + "padding_rows_per_token": blen - 1, + "nonzero_rows": nonzero_rows, + "passed": not nonzero_rows, + } + + +def _hbm_unwritten_slots_unchanged( + before: bytes, + after: bytes, + *, + slot_base: int, + slot_stride: int, + rows: int, + written_tokens: list[int], +) -> dict: + written = set(written_tokens) + changed = [] + for token_idx in range(rows): + if token_idx in written: + continue + start = slot_base + token_idx * slot_stride + before_slot = _hbm_range(before, start, slot_stride) + after_slot = _hbm_range(after, start, slot_stride) + if before_slot != after_slot: + first_diff = next((idx for idx, (a, b) in enumerate(zip(before_slot, after_slot)) if a != b), None) + changed.append( + { + "token_index": token_idx, + "first_diff_byte": first_diff, + "absolute_hbm_byte": None if first_diff is None else start + first_diff, + } + ) + return { + "written_tokens": written_tokens, + "unchanged_tokens_checked": [token for token in range(rows) if token not in written], + "changed_unwritten_slots": changed, + "passed": not changed, + } + + +def _rewrite_compact_golden(build_dir: Path, golden: torch.Tensor) -> None: + torch.save(golden.detach().cpu().float(), build_dir / "golden_output.pt") + summary = { + "format": "compact", + "reason": "exact HBM-byte decoded gather golden", + "golden_output_file": "golden_output.pt", + "original_output": { + "shape": list(golden.shape), + "dtype": str(golden.dtype), + "numel": int(golden.numel()), + }, + } + (build_dir / "golden_result.txt").write_text("Compact Golden Result\n" + json.dumps(summary, indent=2) + "\n") + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _activation_golden(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + neg_alpha = torch.tensor(-1.702, dtype=torch.bfloat16).float() + gate = _bf16(torch.clamp(gate.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), min=-7.0)) + + sigmoid = _bf16(gate.float()) + sigmoid = _bf16(sigmoid.float() * neg_alpha) + sigmoid = _bf16(torch.exp(torch.clamp(sigmoid.float(), -88.0, 88.0))) + sigmoid = _bf16(sigmoid.float() + 1.0) + sigmoid = _bf16(torch.reciprocal(sigmoid.float())) + + glu = _bf16(gate.float() * sigmoid.float()) + up_plus_one = _bf16(up.float() + 1.0) + return _bf16(up_plus_one.float() * glu.float()) + + +def _linear_projection_golden( + x: torch.Tensor, + w: torch.Tensor, + *, + mlen: int, + mram_tile_capacity: int = 4, + hbm_input: bool, +) -> torch.Tensor: + x_q = quantize_to_mxfp(x) if hbm_input else x.to(torch.bfloat16) + w_q = quantize_to_mxfp(w) + k_total = x_q.shape[1] + chunk = mlen * mram_tile_capacity + acc = None + for k_start in range(0, k_total, chunk): + k_end = min(k_start + chunk, k_total) + partial = torch.matmul(x_q[:, k_start:k_end].float(), w_q[k_start:k_end, :].float()).to(torch.bfloat16) + acc = partial if acc is None else (acc.float() + partial.float()).to(torch.bfloat16) + assert acc is not None + return acc + + +def _expert_golden_from_gathered_vram( + x_slots: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + *, + b_gate: torch.Tensor, + b_up: torch.Tensor, + b_down: torch.Tensor, + mlen: int, +) -> torch.Tensor: + gate = _linear_projection_golden(x_slots, w_gate, mlen=mlen, hbm_input=False) + up = _linear_projection_golden(x_slots, w_up, mlen=mlen, hbm_input=False) + gate = _bf16(gate.float() + b_gate.float()) + up = _bf16(up.float() + b_up.float()) + hidden = _activation_golden(gate, up) + out = _linear_projection_golden(hidden, w_down, mlen=mlen, hbm_input=False) + return _bf16(out.float() + b_down.float()) + + +def _expanded_bias(row_bias: torch.Tensor, rows: int) -> torch.Tensor: + return row_bias.to(torch.bfloat16).reshape(1, -1).repeat(rows, 1) + + +def _active_row_indices(pair_count: int, blen: int) -> list[int]: + return [pair_idx * blen for pair_idx in range(pair_count)] + + +def _slot_route_matrix( + *, + slot_rows: int, + hidden: int, + blen: int, + token_indices: list[int], + topk_positions: list[int], + topk_weights: torch.Tensor, +) -> torch.Tensor: + route = torch.zeros(slot_rows, hidden, dtype=torch.bfloat16) + for pair_idx, (token_idx, topk_pos) in enumerate(zip(token_indices, topk_positions, strict=True)): + route[pair_idx * blen] = topk_weights[token_idx, topk_pos].to(torch.bfloat16) + return route + + +def _expert_hbm_stride(prog: PlenaCompiler, shape: tuple[int, int]) -> int: + raw_size = int(shape[0] * shape[1] * prog.real_data_ratio) + return _align_to(raw_size, prog.mlen) + + +def _build_true_expert_weight_table( + prog: PlenaCompiler, + *, + prefix: str, + weights: torch.Tensor, + input_tensors: dict[str, torch.Tensor], +) -> tuple[list, int, int]: + """Register all 32 true expert weights as one stride-addressable HBM table.""" + num_experts, rows, cols = weights.shape + stride = _expert_hbm_stride(prog, (rows, cols)) + base = prog._allocate_hbm(stride * num_experts) + inputs = [] + for expert_id in range(num_experts): + name = f"{prefix}_true_e{expert_id}" + inputs.append(prog.input(name, shape=(rows, cols), hbm_addr=base + expert_id * stride)) + input_tensors[name] = weights[expert_id].to(torch.bfloat16).contiguous().clone() + return inputs, base, stride + + +def _build_true_expert_bias_table( + *, + prog: PlenaCompiler, + name: str, + bias: torch.Tensor, + width: int, + blen: int, + vram_addr: int, + vram_preload: torch.Tensor, +) -> VRAMMatrixVar: + num_experts = int(bias.shape[0]) + table = torch.zeros(num_experts * blen, width, dtype=torch.bfloat16) + for expert_id in range(num_experts): + table[expert_id * blen : (expert_id + 1) * blen] = bias[expert_id].to(torch.bfloat16).reshape(1, -1) + return prestage_bf16_vram_matrix( + prog=prog, + name=name, + tensor=table, + vram_addr=vram_addr, + physical_shape=(num_experts * blen, width), + vram_preload=vram_preload, + ) + + +def _device_routing_exact_golden( + *, + hbm_path: Path, + x_input, + x: torch.Tensor, + device_indices: torch.Tensor, + device_weights: torch.Tensor, + split, + down_weight: torch.Tensor, + down_bias: torch.Tensor, + rows: int, + hidden: int, + blen: int, + mlen: int, + pair_count: int | None = None, +) -> tuple[torch.Tensor, dict]: + exact_golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + clamp_counts = { + "gate_gt_7": 0, + "gate_lt_neg_7": 0, + "up_gt_7": 0, + "up_lt_neg_7": 0, + "per_token": [ + {"token_index": token_idx, "gate_gt_7": 0, "gate_lt_neg_7": 0, "up_gt_7": 0, "up_lt_neg_7": 0} + for token_idx in range(rows) + ], + } + hbm_bytes = hbm_path.read_bytes() + top_k = device_indices.shape[1] + total_pairs = rows * top_k if pair_count is None else pair_count + for pair_idx in range(total_pairs): + token_idx = pair_idx // top_k + topk_pos = pair_idx % top_k + expert_id = int(device_indices[token_idx, topk_pos].item()) + weight = device_weights[token_idx, topk_pos].to(torch.bfloat16) + x_slots = torch.zeros(blen, hidden, dtype=torch.bfloat16) + x_slots[0:1] = _decode_hbm_mxfp8_row_from_bytes( + hbm_bytes, + tensor_base=x_input.hbm_addr, + tensor_rows=rows, + tensor_cols=hidden, + row_element_offset=token_idx * hidden, + ) + + b_gate = _expanded_bias(split.gate_bias[expert_id], blen) + b_up = _expanded_bias(split.up_bias[expert_id], blen) + b_down = _expanded_bias(down_bias[expert_id], blen) + gate_pre = _linear_projection_golden( + x_slots, + split.gate_weight[expert_id], + mlen=mlen, + hbm_input=False, + ) + up_pre = _linear_projection_golden( + x_slots, + split.up_weight[expert_id], + mlen=mlen, + hbm_input=False, + ) + gate_pre = _bf16(gate_pre.float() + b_gate.float()) + up_pre = _bf16(up_pre.float() + b_up.float()) + local_clamp = _clamp_counts_from_slots( + token_indices=[token_idx], + gate=gate_pre, + up=up_pre, + blen=blen, + rows=rows, + ) + for key in ("gate_gt_7", "gate_lt_neg_7", "up_gt_7", "up_lt_neg_7"): + clamp_counts[key] += local_clamp[key] + for idx, item in enumerate(local_clamp["per_token"]): + for key in ("gate_gt_7", "gate_lt_neg_7", "up_gt_7", "up_lt_neg_7"): + clamp_counts["per_token"][idx][key] += item[key] + + expert_out = _expert_golden_from_gathered_vram( + x_slots, + split.gate_weight[expert_id], + split.up_weight[expert_id], + down_weight[expert_id], + b_gate=b_gate, + b_up=b_up, + b_down=b_down, + mlen=mlen, + ) + exact_golden[token_idx] = _bf16( + exact_golden[token_idx].float() + _bf16(expert_out[0].float() * weight.float()).float() + ) + return exact_golden, clamp_counts + + +def _device_routing_vram_exact_golden( + *, + x: torch.Tensor, + device_indices: torch.Tensor, + device_weights: torch.Tensor, + split, + down_weight: torch.Tensor, + down_bias: torch.Tensor, + rows: int, + hidden: int, + blen: int, + mlen: int, + pair_count: int | None = None, +) -> tuple[torch.Tensor, dict]: + """Exact golden for decoder-block MoE input already resident in BF16 VRAM.""" + exact_golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + clamp_counts = { + "gate_gt_7": 0, + "gate_lt_neg_7": 0, + "up_gt_7": 0, + "up_lt_neg_7": 0, + "per_token": [ + {"token_index": token_idx, "gate_gt_7": 0, "gate_lt_neg_7": 0, "up_gt_7": 0, "up_lt_neg_7": 0} + for token_idx in range(rows) + ], + } + top_k = device_indices.shape[1] + total_pairs = rows * top_k if pair_count is None else pair_count + for pair_idx in range(total_pairs): + token_idx = pair_idx // top_k + topk_pos = pair_idx % top_k + expert_id = int(device_indices[token_idx, topk_pos].item()) + weight = device_weights[token_idx, topk_pos].to(torch.bfloat16) + x_slots = torch.zeros(blen, hidden, dtype=torch.bfloat16) + x_slots[0:1] = x[token_idx : token_idx + 1].to(torch.bfloat16) + + b_gate = _expanded_bias(split.gate_bias[expert_id], blen) + b_up = _expanded_bias(split.up_bias[expert_id], blen) + b_down = _expanded_bias(down_bias[expert_id], blen) + gate_pre = _linear_projection_golden( + x_slots, + split.gate_weight[expert_id], + mlen=mlen, + hbm_input=False, + ) + up_pre = _linear_projection_golden( + x_slots, + split.up_weight[expert_id], + mlen=mlen, + hbm_input=False, + ) + gate_pre = _bf16(gate_pre.float() + b_gate.float()) + up_pre = _bf16(up_pre.float() + b_up.float()) + local_clamp = _clamp_counts_from_slots( + token_indices=[token_idx], + gate=gate_pre, + up=up_pre, + blen=blen, + rows=rows, + ) + for key in ("gate_gt_7", "gate_lt_neg_7", "up_gt_7", "up_lt_neg_7"): + clamp_counts[key] += local_clamp[key] + for idx, item in enumerate(local_clamp["per_token"]): + for key in ("gate_gt_7", "gate_lt_neg_7", "up_gt_7", "up_lt_neg_7"): + clamp_counts["per_token"][idx][key] += item[key] + + expert_out = _expert_golden_from_gathered_vram( + x_slots, + split.gate_weight[expert_id], + split.up_weight[expert_id], + down_weight[expert_id], + b_gate=b_gate, + b_up=b_up, + b_down=b_down, + mlen=mlen, + ) + exact_golden[token_idx] = _bf16( + exact_golden[token_idx].float() + _bf16(expert_out[0].float() * weight.float()).float() + ) + return exact_golden, clamp_counts + + +def _device_topk_contract( + *, + hf_indices: torch.Tensor, + hf_weights: torch.Tensor, + device_indices: torch.Tensor, + device_weights: torch.Tensor, + device_logits: torch.Tensor, + rank_gate: dict, +) -> dict: + boundary_safe = [bool(v) for v in rank_gate["gap_gt_measured_error_by_token"]] + internal_order_safe = [bool(v) for v in rank_gate["internal_order_gap_gt_measured_error_by_token"]] + boundary_ambiguous_tokens = {idx for idx, safe in enumerate(boundary_safe) if not safe} + internal_order_ambiguous_tokens = {idx for idx, safe in enumerate(internal_order_safe) if not safe} + set_matches = [] + order_matches = [] + weights_by_expert_close = [] + tie_break_matches = [] + for token_idx in range(hf_indices.shape[0]): + hf_row = [int(v) for v in hf_indices[token_idx].tolist()] + dev_row = [int(v) for v in device_indices[token_idx].tolist()] + # V_TOPK contract: descending logit value, exact ties by smaller expert id. + sorted_by_device_context = sorted( + range(device_logits.shape[1]), + key=lambda expert_id: (-float(device_logits[token_idx, expert_id].float().item()), expert_id), + )[: device_indices.shape[1]] + set_match = set(hf_row) == set(dev_row) + order_match = hf_row == dev_row + tie_break_matches.append(dev_row == [int(v) for v in sorted_by_device_context]) + set_matches.append(set_match) + order_matches.append(order_match) + hf_map = {int(expert): hf_weights[token_idx, pos].float() for pos, expert in enumerate(hf_row)} + dev_map = {int(expert): device_weights[token_idx, pos].float() for pos, expert in enumerate(dev_row)} + common = sorted(set(hf_map) & set(dev_map)) + if common: + hf_vals = torch.stack([hf_map[e] for e in common]) + dev_vals = torch.stack([dev_map[e] for e in common]) + weights_by_expert_close.append(bool(torch.allclose(dev_vals, hf_vals, atol=0.003, rtol=0.0))) + else: + weights_by_expert_close.append(False) + safe_set_pass = all( + match for token_idx, match in enumerate(set_matches) if token_idx not in boundary_ambiguous_tokens + ) + safe_order_pass = all( + match + for token_idx, match in enumerate(order_matches) + if token_idx not in boundary_ambiguous_tokens | internal_order_ambiguous_tokens + ) + weight_diff = (device_weights.float() - hf_weights.float()).abs() + weight_rel_rms = torch.linalg.vector_norm(device_weights.float() - hf_weights.float()) / torch.clamp( + torch.linalg.vector_norm(hf_weights.float()), + min=1e-12, + ) + return { + "boundary_ambiguous_tokens": sorted(boundary_ambiguous_tokens), + "internal_order_ambiguous_tokens": sorted(internal_order_ambiguous_tokens), + "hf_indices": hf_indices.cpu().tolist(), + "device_indices": device_indices.cpu().tolist(), + "hf_weights": hf_weights.cpu().float().tolist(), + "device_weights": device_weights.cpu().float().tolist(), + "weight_rel_rms": float(weight_rel_rms.item()), + "weight_max_abs_error": float(weight_diff.max().item()), + "weight_max_abs_error_by_token": weight_diff.max(dim=1).values.cpu().float().tolist(), + "set_matches_by_token": set_matches, + "order_matches_by_token": order_matches, + "tie_break_matches_device_logits_by_token": tie_break_matches, + "tie_break_matches_device_logits": all(tie_break_matches), + "weights_by_expert_close_by_token": weights_by_expert_close, + "rank4_rank5_gap": rank_gate["rank4_rank5_gap"], + "min_internal_topk_gap": rank_gate["min_internal_topk_gap"], + "device_max_abs_error_by_token": rank_gate["device_max_abs_error_by_token"], + "gap_safe_by_token": boundary_safe, + "internal_order_safe_by_token": internal_order_safe, + "ambiguity_policy": ( + "Tokens are classified per sample/layer. A token is gap-safe when " + "rank4-rank5 gap exceeds that token's measured BF16 router logit " + "error; safe tokens must match HF top-4 sets. Internal order is " + "required only when the minimum adjacent in-top-k gap exceeds the " + "same measured error. Unsafe tokens are recorded and checked for " + "deterministic device-logit tie-break behavior, not forced to match HF." + ), + "gap_safe_set_pass": safe_set_pass, + "gap_safe_order_pass": safe_order_pass, + "passed": safe_set_pass and safe_order_pass and all(tie_break_matches), + } + + +def _clamp_counts_from_slots( + *, + token_indices: list[int], + gate: torch.Tensor, + up: torch.Tensor, + blen: int, + rows: int, +) -> dict: + per_token = [ + { + "token_index": token_idx, + "gate_gt_7": 0, + "gate_lt_neg_7": 0, + "up_gt_7": 0, + "up_lt_neg_7": 0, + } + for token_idx in range(rows) + ] + for pair_idx, token_idx in enumerate(token_indices): + active = pair_idx * blen + per_token[token_idx]["gate_gt_7"] += int((gate[active] > 7).sum().item()) + per_token[token_idx]["gate_lt_neg_7"] += int((gate[active] < -7).sum().item()) + per_token[token_idx]["up_gt_7"] += int((up[active] > 7).sum().item()) + per_token[token_idx]["up_lt_neg_7"] += int((up[active] < -7).sum().item()) + return { + "gate_gt_7": sum(item["gate_gt_7"] for item in per_token), + "gate_lt_neg_7": sum(item["gate_lt_neg_7"] for item in per_token), + "up_gt_7": sum(item["up_gt_7"] for item in per_token), + "up_lt_neg_7": sum(item["up_lt_neg_7"] for item in per_token), + "per_token": per_token, + } + + +def run_address_stage(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + build_dir = args.build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + ref = _load_tok8_reference(args.reference_path) + x = ref["x"] + rows, hidden = x.shape + if hidden % mlen != 0: + raise ValueError(f"hidden={hidden} must be divisible by MLEN={mlen}") + if args.token_index < 0 or args.token_index >= rows: + raise ValueError(f"token_index={args.token_index} out of range [0, {rows})") + + slot_rows = blen + token_offsets = torch.tensor([args.token_index * hidden], dtype=torch.int32) + golden = quantize_to_mxfp(x[args.token_index : args.token_index + 1]).to(torch.bfloat16) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + _zero = prog.fp_var("zero", size=1) + _limit_pos = prog.fp_var("gpt_oss_limit_pos", size=slot_rows) + _limit_neg = prog.fp_var("gpt_oss_limit_neg", size=slot_rows) + _one = prog.fp_var("one", size=slot_rows) + _neg_alpha = prog.fp_var("neg_alpha", size=slot_rows) + + x_input = prog.input("X", shape=(rows, hidden)) + gathered = prog.moe_gather_token_rows_from_hbm_v0( + x_input, + token_offsets_int_base=0, + pair_count=1, + hidden=hidden, + name="gather_address_strict_x", + ) + isa = prog.compile() + scale_reg_count = isa.count("C_SET_SCALE_REG") + if scale_reg_count <= 0: + raise AssertionError("gather path must set C_SET_SCALE_REG for MXFP8 HBM activation loads") + + input_tensors = {"X": x} + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=[0.0], + int_preload=token_offsets, + build_dir=str(build_dir), + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_gather_address", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + + # The strict address gate should compare against the exact HBM bytes the + # DMA reads, not a parallel high-level quantizer. A few E4M3/E8M0 boundary + # values differ between those views; the generated preload is the source of + # truth for this address test. + golden = _decode_hbm_mxfp8_row( + build_dir / "hbm_for_behave_sim.bin", + tensor_base=x_input.hbm_addr, + tensor_rows=rows, + tensor_cols=hidden, + row_element_offset=int(token_offsets[0].item()), + ) + _rewrite_compact_golden(build_dir, golden) + + comparison_params = _comparison_params_for(gathered, rows=1, hidden=hidden, mlen=mlen, golden=golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + + manifest = { + "stage": "address", + "reference_path": str(args.reference_path.expanduser().resolve()), + "token_index": args.token_index, + "token_offsets_int_base": 0, + "token_offsets": token_offsets.tolist(), + "rows": rows, + "hidden": hidden, + "mlen": mlen, + "blen": blen, + "h_prefetch_v_granularity_elements": blen * mlen, + "scale_reg_count": scale_reg_count, + "x_hbm_base_addr": x_input.hbm_addr, + "expected_hbm_element_offset": int(token_offsets[0].item()), + "expected_hbm_absolute_element_addr": int(x_input.hbm_addr + token_offsets[0].item()), + "output_vram_row": prog._compiler.get_vram_addr(gathered.name) // mlen, + "comparison_params": comparison_params, + "asm_lines": len(isa.splitlines()), + } + (build_dir / "gather_scatter_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + if args.no_run: + return {"build_dir": str(build_dir), "ran": False, "manifest": manifest} + + metrics = run_and_assert( + build_dir, + "gpt_oss_moe_gather_address", + mlen=mlen, + blen=blen, + threads=args.emu_threads, + ) + results, params = compare_emulator_output(build_dir) + emu_output = results["simulated_values"].reshape(1, hidden).to(torch.bfloat16) + stats = compare_stats(emu_output, golden, rtol=0.0) + if stats.max_abs_error != 0.0 or not stats.allclose: + raise AssertionError( + "single-token gather address strict failed: " + f"max_abs={stats.max_abs_error:.6g}, rel_rms={stats.rel_rms:.6g}, allclose={stats.allclose}" + ) + + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params_runtime": params, + "emulator_compare_raw": { + key: results[key] + for key in ( + "mse", + "mae", + "max_error", + "relative_error", + "relative_match_rate", + "allclose_match_rate", + "match_rate", + "allclose_pass", + "atol", + "rtol", + "golden_shape", + "simulated_shape", + ) + if key in results + }, + "emulator_vs_gather_golden": { + "rel_rms": stats.rel_rms, + "allclose": stats.allclose, + "pass_rate": stats.pass_rate, + "max_abs_error": stats.max_abs_error, + "atol": stats.atol, + "rtol": stats.rtol, + }, + "address_gate": { + "device_address_matches_host_formula": True, + "gathered_row_strict": stats.max_abs_error == 0.0 and stats.allclose, + "scale_reg_count_nonzero": scale_reg_count > 0, + "passed": stats.max_abs_error == 0.0 and stats.allclose and scale_reg_count > 0, + }, + } + (build_dir / "gather_scatter_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + return {"build_dir": str(build_dir), "ran": True, "summary": summary} + + +def run_four_pair_stage(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + build_dir = args.build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + ref = _load_tok8_reference(args.reference_path) + x = ref["x"] + topk_indices = ref["topk_indices"] + rows, hidden = x.shape + tensors = load_layer_tensors(args.layer_index) + split = split_packed_gate_up(tensors["gate_up_weight"], tensors["gate_up_bias"]) + down_weight = tensors["down_weight"] + down_bias = tensors["down_bias"] + + token_idx, _ = torch.where(topk_indices == args.expert_id) + token_indices = [int(v) for v in token_idx.tolist()] + if not token_indices: + raise ValueError(f"expert {args.expert_id} has no routed rows in this reference") + + pair_count = len(token_indices) + slot_rows = pair_count * blen + token_offsets = torch.tensor([token * hidden for token in token_indices], dtype=torch.int32) + + w_gate = split.gate_weight[args.expert_id].to(torch.bfloat16).contiguous().clone() + w_up = split.up_weight[args.expert_id].to(torch.bfloat16).contiguous().clone() + w_down = down_weight[args.expert_id].to(torch.bfloat16).contiguous().clone() + b_gate = _expanded_bias(split.gate_bias[args.expert_id], slot_rows) + b_up = _expanded_bias(split.up_bias[args.expert_id], slot_rows) + b_down = _expanded_bias(down_bias[args.expert_id], slot_rows) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=slot_rows) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=slot_rows) + one = prog.fp_var("one", size=slot_rows) + neg_alpha = prog.fp_var("neg_alpha", size=slot_rows) + + x_input = prog.input("X", shape=(rows, hidden)) + gathered = prog.moe_gather_token_rows_from_hbm_v0( + x_input, + token_offsets_int_base=0, + pair_count=pair_count, + hidden=hidden, + name=f"gather_expert{args.expert_id}_x", + ) + + w_gate_input = prog.input("W_gate", shape=(hidden, w_gate.shape[1])) + w_up_input = prog.input("W_up", shape=(hidden, w_up.shape[1])) + w_down_input = prog.input("W_down", shape=(w_down.shape[0], hidden)) + + physical_rows = max(blen, math.ceil(slot_rows / blen) * blen) + bias_physical_inter = (physical_rows, w_gate.shape[1]) + bias_physical_hidden = (physical_rows, hidden) + bias_base = math.ceil(prog.vram_allocator._vmm.next_bump / (mlen * mlen)) * (mlen * mlen) + bias_sizes = [ + bias_physical_inter[0] * bias_physical_inter[1], + bias_physical_inter[0] * bias_physical_inter[1], + bias_physical_hidden[0] * bias_physical_hidden[1], + ] + bias_total = sum(math.ceil(size / (mlen * mlen)) * (mlen * mlen) for size in bias_sizes) + vram_preload = torch.zeros(bias_base + bias_total, dtype=torch.bfloat16) + next_bias_addr = bias_base + bias_vrams = [] + for suffix, tensor, physical_shape in ( + ("gate", b_gate, bias_physical_inter), + ("up", b_up, bias_physical_inter), + ("down", b_down, bias_physical_hidden), + ): + aligned_size = math.ceil((physical_shape[0] * physical_shape[1]) / (mlen * mlen)) * (mlen * mlen) + bias_vrams.append( + prestage_bf16_vram_matrix( + prog=prog, + name=f"B_{suffix}", + tensor=tensor, + vram_addr=next_bias_addr, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + ) + next_bias_addr += aligned_size + + output_vram = prog.moe_expert_v0( + gathered, + (w_gate_input, w_up_input, w_down_input), + biases=tuple(bias_vrams), + rows=slot_rows, + intermediate=w_gate.shape[1], + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + name=f"gather_expert{args.expert_id}", + ) + isa = prog.compile() + scale_reg_count = isa.count("C_SET_SCALE_REG") + if scale_reg_count <= 0: + raise AssertionError("four-pair expert path must emit C_SET_SCALE_REG for gathered activations/weights") + + input_tensors = { + "X": x, + "W_gate": w_gate, + "W_up": w_up, + "W_down": w_down, + } + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + placeholder_golden = torch.zeros(slot_rows, hidden, dtype=torch.bfloat16) + fp_preload = [0.0] + [7.0] * slot_rows + [-7.0] * slot_rows + [1.0] * slot_rows + [-1.702] * slot_rows + create_sim_env( + input_tensors, + isa, + {"original_output": placeholder_golden}, + fp_preload=fp_preload, + int_preload=token_offsets, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_gather_four", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + + x_slots = torch.zeros(slot_rows, hidden, dtype=torch.bfloat16) + for pair_idx, token_offset in enumerate(token_offsets.tolist()): + x_slots[pair_idx * blen : pair_idx * blen + 1] = _decode_hbm_mxfp8_row( + build_dir / "hbm_for_behave_sim.bin", + tensor_base=x_input.hbm_addr, + tensor_rows=rows, + tensor_cols=hidden, + row_element_offset=int(token_offset), + ) + golden = _expert_golden_from_gathered_vram( + x_slots, + w_gate, + w_up, + w_down, + b_gate=b_gate, + b_up=b_up, + b_down=b_down, + mlen=mlen, + ) + _rewrite_compact_golden(build_dir, golden) + + comparison_params = _comparison_params_for(output_vram, rows=slot_rows, hidden=hidden, mlen=mlen, golden=golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + manifest = { + "stage": "four", + "expert_id": args.expert_id, + "token_indices": token_indices, + "token_offsets": token_offsets.tolist(), + "slot_rows": slot_rows, + "active_rows": [i * blen for i in range(pair_count)], + "hidden": hidden, + "intermediate": int(w_gate.shape[1]), + "mlen": mlen, + "blen": blen, + "scale_reg_count": scale_reg_count, + "output_vram_row": prog._compiler.get_vram_addr(output_vram.name) // mlen, + "comparison_params": comparison_params, + "asm_lines": len(isa.splitlines()), + } + (build_dir / "gather_scatter_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + if args.no_run: + return {"build_dir": str(build_dir), "ran": False, "manifest": manifest} + + metrics = run_and_assert( + build_dir, + "gpt_oss_moe_gather_four", + mlen=mlen, + blen=blen, + threads=args.emu_threads, + ) + results, params = compare_emulator_output(build_dir) + emu_output = results["simulated_values"].reshape(slot_rows, hidden).to(torch.bfloat16) + stats = compare_stats(emu_output, golden, rtol=0.02) + if stats.rel_rms > 0.02 or not stats.allclose: + raise AssertionError( + "four-pair gather+expert failed: " + f"rel_rms={stats.rel_rms:.6g}, max_abs={stats.max_abs_error:.6g}, allclose={stats.allclose}" + ) + + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params_runtime": params, + "emulator_compare_raw": { + key: results[key] + for key in ( + "mse", + "mae", + "max_error", + "relative_error", + "relative_match_rate", + "allclose_match_rate", + "match_rate", + "allclose_pass", + "atol", + "rtol", + "golden_shape", + "simulated_shape", + ) + if key in results + }, + "emulator_vs_four_pair_golden": { + "rel_rms": stats.rel_rms, + "allclose": stats.allclose, + "pass_rate": stats.pass_rate, + "max_abs_error": stats.max_abs_error, + "atol": stats.atol, + "rtol": stats.rtol, + }, + "four_pair_gate": { + "scale_reg_count_nonzero": scale_reg_count > 0, + "passed": scale_reg_count > 0 and stats.rel_rms <= 0.02 and stats.allclose, + }, + } + (build_dir / "gather_scatter_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + return {"build_dir": str(build_dir), "ran": True, "summary": summary} + + +def run_full_vram_stage(args: argparse.Namespace) -> dict: + """Run Step5 segment A, optionally followed by HBM storeback validation.""" + mlen = args.mlen + blen = args.blen + build_dir = args.build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + store_mode = args.stage in ("hbm_store_single", "hbm_store_full") + store_single = args.stage == "hbm_store_single" + + ref = _load_tok8_reference(args.reference_path) + x = ref["x"] + topk_indices = ref["topk_indices"] + topk_weights = ref["topk_weights"] + rows, hidden = x.shape + if hidden % mlen != 0: + raise ValueError(f"hidden={hidden} must be divisible by MLEN={mlen}") + if topk_indices.shape != topk_weights.shape: + raise ValueError(f"topk indices/weights shape mismatch: {topk_indices.shape} vs {topk_weights.shape}") + + tensors = load_layer_tensors(args.layer_index) + split = split_packed_gate_up(tensors["gate_up_weight"], tensors["gate_up_bias"]) + down_weight = tensors["down_weight"] + down_bias = tensors["down_bias"] + intermediate = int(split.gate_weight.shape[-1]) + + selected_experts = [int(v) for v in torch.unique(topk_indices).sort().values.tolist()] + if args.only_experts: + requested = [int(part) for part in args.only_experts.split(",") if part.strip()] + missing = sorted(set(requested) - set(selected_experts)) + if missing: + raise ValueError(f"--only-experts requested unrouted experts: {missing}") + selected_experts = requested + expert_specs = [] + max_slot_rows = blen + int_preload_values: list[int] = [] + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=max(blen, rows * blen)) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=max(blen, rows * blen)) + one = prog.fp_var("one", size=max(blen, rows * blen)) + neg_alpha = prog.fp_var("neg_alpha", size=max(blen, rows * blen)) + shared_zero_row = prog.fp_var("gpt_oss_shared_zero_row", size=mlen) + + x_input = prog.input("X", shape=(rows, hidden)) + acc_physical_rows = max(blen, math.ceil(rows / blen) * blen) + accumulator = prog.alloc( + "gather_full_vram_acc", + rows=rows, + cols=hidden, + strict=False, + physical_shape=(acc_physical_rows, hidden), + ) + prog.moe_true_zero_vram_rows_v0( + accumulator, + rows=list(range(rows)), + hidden=hidden, + zero_row=shared_zero_row, + name="gather_full_acc", + ) + + input_tensors: dict[str, torch.Tensor] = {"X": x} + route_inputs = {} + weight_inputs = {} + gathered_by_expert = {} + + for expert_id in selected_experts: + token_idx, topk_pos = torch.where(topk_indices == expert_id) + token_indices = [int(v) for v in token_idx.tolist()] + topk_positions = [int(v) for v in topk_pos.tolist()] + if not token_indices: + continue + pair_count = len(token_indices) + slot_rows = pair_count * blen + max_slot_rows = max(max_slot_rows, slot_rows) + active_rows = _active_row_indices(pair_count, blen) + token_offsets_base = len(int_preload_values) + int_preload_values.extend([token * hidden for token in token_indices]) + + gathered_by_expert[expert_id] = prog.moe_gather_token_rows_from_hbm_v0( + x_input, + token_offsets_int_base=token_offsets_base, + pair_count=pair_count, + hidden=hidden, + zero_row=shared_zero_row, + name=f"full_gather_e{expert_id}_x", + ) + + w_gate = split.gate_weight[expert_id].to(torch.bfloat16).contiguous().clone() + w_up = split.up_weight[expert_id].to(torch.bfloat16).contiguous().clone() + w_down = down_weight[expert_id].to(torch.bfloat16).contiguous().clone() + w_gate_name = f"W_e{expert_id}_gate" + w_up_name = f"W_e{expert_id}_up" + w_down_name = f"W_e{expert_id}_down" + w_gate_input = prog.input(w_gate_name, shape=(hidden, intermediate)) + w_up_input = prog.input(w_up_name, shape=(hidden, intermediate)) + w_down_input = prog.input(w_down_name, shape=(intermediate, hidden)) + weight_inputs[expert_id] = (w_gate_input, w_up_input, w_down_input) + input_tensors[w_gate_name] = w_gate + input_tensors[w_up_name] = w_up + input_tensors[w_down_name] = w_down + + route = _slot_route_matrix( + slot_rows=slot_rows, + hidden=hidden, + blen=blen, + token_indices=token_indices, + topk_positions=topk_positions, + topk_weights=topk_weights, + ) + route_name = f"route_e{expert_id}" + route_input = prog.input(route_name, shape=(slot_rows, hidden)) + route_inputs[expert_id] = route_input + input_tensors[route_name] = route + + expert_specs.append( + { + "expert_id": expert_id, + "token_indices": token_indices, + "topk_positions": topk_positions, + "pair_count": pair_count, + "slot_rows": slot_rows, + "active_rows": active_rows, + "token_offsets_int_base": token_offsets_base, + "token_offsets": [token * hidden for token in token_indices], + } + ) + + route_vrams = { + expert_id: prog.load_batch(route_input, name=route_input.display_name) + for expert_id, route_input in route_inputs.items() + } + + bias_physical: dict[int, tuple[tuple[int, int], tuple[int, int]]] = {} + bias_sizes = [] + for spec in expert_specs: + slot_rows = int(spec["slot_rows"]) + physical_rows = max(blen, math.ceil(slot_rows / blen) * blen) + inter_shape = (physical_rows, intermediate) + hidden_shape = (physical_rows, hidden) + bias_physical[int(spec["expert_id"])] = (inter_shape, hidden_shape) + bias_sizes.extend( + [ + inter_shape[0] * inter_shape[1], + inter_shape[0] * inter_shape[1], + hidden_shape[0] * hidden_shape[1], + ] + ) + bias_base = math.ceil(prog.vram_allocator._vmm.next_bump / (mlen * mlen)) * (mlen * mlen) + bias_total = sum(math.ceil(size / (mlen * mlen)) * (mlen * mlen) for size in bias_sizes) + vram_preload = torch.zeros(bias_base + bias_total, dtype=torch.bfloat16) + next_bias_addr = bias_base + bias_vrams = {} + for spec in expert_specs: + expert_id = int(spec["expert_id"]) + slot_rows = int(spec["slot_rows"]) + inter_shape, hidden_shape = bias_physical[expert_id] + b_gate = _expanded_bias(split.gate_bias[expert_id], slot_rows) + b_up = _expanded_bias(split.up_bias[expert_id], slot_rows) + b_down = _expanded_bias(down_bias[expert_id], slot_rows) + local_bias_vrams = [] + for suffix, tensor, physical_shape in ( + ("gate", b_gate, inter_shape), + ("up", b_up, inter_shape), + ("down", b_down, hidden_shape), + ): + size = physical_shape[0] * physical_shape[1] + aligned_size = math.ceil(size / (mlen * mlen)) * (mlen * mlen) + local_bias_vrams.append( + prestage_bf16_vram_matrix( + prog=prog, + name=f"B_e{expert_id}_{suffix}", + tensor=tensor, + vram_addr=next_bias_addr, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + ) + next_bias_addr += aligned_size + bias_vrams[expert_id] = tuple(local_bias_vrams) + + for spec in expert_specs: + expert_id = int(spec["expert_id"]) + output_vram = prog.moe_expert_v0( + gathered_by_expert[expert_id], + weight_inputs[expert_id], + biases=bias_vrams[expert_id], + rows=int(spec["slot_rows"]), + intermediate=intermediate, + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + name=f"full_vram_e{expert_id}", + ) + prog.vram_mul(output_vram, route_vrams[expert_id], num_rows=int(spec["slot_rows"])) + prog.moe_scatter_add_active_rows_v0( + accumulator, + output_vram, + token_indices=list(spec["token_indices"]), + active_rows=list(spec["active_rows"]), + hidden=hidden, + name=f"expert{expert_id}", + ) + + store_token_indices: list[int] = [] + output_slot_hbm_base = None + output_slot_hbm_stride = None + expected_h_store_v_dynamic_count = 0 + if store_mode: + if args.only_experts: + raise ValueError("HBM store stages require the full routed expert set; do not pass --only-experts") + if args.store_token_index < 0 or args.store_token_index >= rows: + raise ValueError(f"store_token_index={args.store_token_index} out of range [0, {rows})") + store_token_indices = [args.store_token_index] if store_single else list(range(rows)) + output_slot_hbm_stride = int(blen * hidden * prog.real_data_ratio) + output_slot_hbm_base = prog._allocate_hbm(output_slot_hbm_stride * rows) + store_slot = prog.alloc( + "gather_full_hbm_store_slot", + rows=blen, + cols=hidden, + strict=False, + physical_shape=(blen, hidden), + ) + for token_idx in store_token_indices: + prog.moe_true_zero_vram_rows_v0( + store_slot, + rows=list(range(blen)), + hidden=hidden, + zero_row=shared_zero_row, + name=f"hbm_store_t{token_idx}_zero", + ) + prog.moe_scatter_add_active_rows_v0( + store_slot, + accumulator, + token_indices=[0], + active_rows=[token_idx], + hidden=hidden, + name=f"hbm_store_t{token_idx}_copy", + ) + prog.store( + store_slot, + name=f"gather_full_output_slot_t{token_idx}", + hbm_addr=output_slot_hbm_base + token_idx * output_slot_hbm_stride, + ) + expected_h_store_v_dynamic_count = len(store_token_indices) * math.ceil(hidden / mlen) + + isa = prog.compile() + scale_reg_count = isa.count("C_SET_SCALE_REG") + h_store_count = isa.count("H_STORE_V") + if scale_reg_count <= 0: + raise AssertionError("full VRAM gather/expert path must emit C_SET_SCALE_REG for MXFP8 tensors") + expected_h_store_v_source_count = len(store_token_indices) + if not store_mode and h_store_count != 0: + raise AssertionError(f"segment A must not emit H_STORE_V, got {h_store_count}") + if store_mode and h_store_count != expected_h_store_v_source_count: + raise AssertionError( + f"HBM store stage expected {expected_h_store_v_source_count} source H_STORE_V instructions, " + f"got {h_store_count}" + ) + + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + l1_golden_path = args.l1_golden_path.expanduser().resolve() + l1_golden = torch.load(l1_golden_path, map_location="cpu").to(torch.bfloat16) + if tuple(l1_golden.shape) != (rows, hidden): + raise ValueError(f"L1 golden shape {tuple(l1_golden.shape)} does not match {(rows, hidden)}") + fp_preload = [0.0] + [7.0] * max(blen, rows * blen) + [-7.0] * max(blen, rows * blen) + fp_preload += [1.0] * max(blen, rows * blen) + [-1.702] * max(blen, rows * blen) + int_preload = torch.tensor(int_preload_values, dtype=torch.int32) + create_sim_env( + input_tensors, + isa, + {"original_output": l1_golden}, + fp_preload=fp_preload, + int_preload=int_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_gather_full_vram", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + if store_mode: + assert output_slot_hbm_base is not None + assert output_slot_hbm_stride is not None + required_hbm_size = output_slot_hbm_base + output_slot_hbm_stride * rows + required_hbm_size = ((required_hbm_size + 63) // 64) * 64 + (build_dir / "hbm_size.txt").write_text(f"{required_hbm_size}\n") + + exact_golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + clamp_counts = { + "gate_gt_7": 0, + "gate_lt_neg_7": 0, + "up_gt_7": 0, + "up_lt_neg_7": 0, + "per_token": [ + {"token_index": token_idx, "gate_gt_7": 0, "gate_lt_neg_7": 0, "up_gt_7": 0, "up_lt_neg_7": 0} + for token_idx in range(rows) + ], + } + for spec in expert_specs: + expert_id = int(spec["expert_id"]) + slot_rows = int(spec["slot_rows"]) + token_indices = list(spec["token_indices"]) + active_rows = list(spec["active_rows"]) + w_gate = input_tensors[f"W_e{expert_id}_gate"] + w_up = input_tensors[f"W_e{expert_id}_up"] + w_down = input_tensors[f"W_e{expert_id}_down"] + b_gate = _expanded_bias(split.gate_bias[expert_id], slot_rows) + b_up = _expanded_bias(split.up_bias[expert_id], slot_rows) + b_down = _expanded_bias(down_bias[expert_id], slot_rows) + + x_slots = torch.zeros(slot_rows, hidden, dtype=torch.bfloat16) + for pair_idx, token_offset in enumerate(spec["token_offsets"]): + x_slots[pair_idx * blen : pair_idx * blen + 1] = _decode_hbm_mxfp8_row( + build_dir / "hbm_for_behave_sim.bin", + tensor_base=x_input.hbm_addr, + tensor_rows=rows, + tensor_cols=hidden, + row_element_offset=int(token_offset), + ) + + gate_pre = _linear_projection_golden(x_slots, w_gate, mlen=mlen, hbm_input=False) + up_pre = _linear_projection_golden(x_slots, w_up, mlen=mlen, hbm_input=False) + gate_pre = _bf16(gate_pre.float() + b_gate.float()) + up_pre = _bf16(up_pre.float() + b_up.float()) + local_clamp = _clamp_counts_from_slots( + token_indices=token_indices, + gate=gate_pre, + up=up_pre, + blen=blen, + rows=rows, + ) + for key in ("gate_gt_7", "gate_lt_neg_7", "up_gt_7", "up_lt_neg_7"): + clamp_counts[key] += local_clamp[key] + for token_idx, item in enumerate(local_clamp["per_token"]): + for key in ("gate_gt_7", "gate_lt_neg_7", "up_gt_7", "up_lt_neg_7"): + clamp_counts["per_token"][token_idx][key] += item[key] + + expert_out = _expert_golden_from_gathered_vram( + x_slots, + w_gate, + w_up, + w_down, + b_gate=b_gate, + b_up=b_up, + b_down=b_down, + mlen=mlen, + ) + + route_exact = torch.zeros(slot_rows, hidden, dtype=torch.bfloat16) + route_input = route_inputs[expert_id] + for active_row in active_rows: + route_exact[active_row : active_row + 1] = _decode_hbm_mxfp8_row( + build_dir / "hbm_for_behave_sim.bin", + tensor_base=route_input.hbm_addr, + tensor_rows=slot_rows, + tensor_cols=hidden, + row_element_offset=active_row * hidden, + ) + weighted = _bf16(expert_out.float() * route_exact.float()) + for pair_idx, token_idx in enumerate(token_indices): + active_row = pair_idx * blen + exact_golden[token_idx] = _bf16(exact_golden[token_idx].float() + weighted[active_row].float()) + + exact_vs_l1 = compare_stats(exact_golden, l1_golden, rtol=0.02) + torch.save(exact_golden.float(), build_dir / "exact_hbm_decode_full_vram_golden.pt") + # Segment A's mechanism anchor is the exact tensor produced by the same + # HBM-decode semantics as the device gather path. The archived L1 golden + # remains recorded below as an external continuity check; it carries the + # already adjudicated BF16 tail residual from the direct L1 proof. + target_golden = exact_golden + target_golden_name = "exact_hbm_decode_subset_golden" if args.only_experts else "exact_hbm_decode_full_golden" + _rewrite_compact_golden(build_dir, target_golden) + + comparison_params = _comparison_params_for(accumulator, rows=rows, hidden=hidden, mlen=mlen, golden=target_golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + manifest = { + "stage": args.stage, + "reference_path": str(args.reference_path.expanduser().resolve()), + "l1_golden_path": str(l1_golden_path), + "selected_experts": selected_experts, + "only_experts": args.only_experts, + "target_golden": target_golden_name, + "expert_specs": expert_specs, + "rows": rows, + "hidden": hidden, + "intermediate": intermediate, + "mlen": mlen, + "blen": blen, + "scale_reg_count": scale_reg_count, + "h_store_v_count": h_store_count, + "expected_h_store_v_source_count": expected_h_store_v_source_count, + "expected_h_store_v_dynamic_count": expected_h_store_v_dynamic_count, + "segment_a_no_h_store": (not store_mode) and h_store_count == 0, + "store_mode": args.stage if store_mode else None, + "store_token_indices": store_token_indices, + "output_slot_hbm_base": output_slot_hbm_base, + "output_slot_hbm_stride": output_slot_hbm_stride, + "output_slot_hbm_total_bytes": None if output_slot_hbm_stride is None else output_slot_hbm_stride * rows, + "accumulator_vram_row": prog._compiler.get_vram_addr(accumulator.name) // mlen, + "comparison_params": comparison_params, + "clamp_counts": clamp_counts, + "exact_hbm_decode_golden_vs_l1_golden": _stats_dict(exact_vs_l1), + "asm_lines": len(isa.splitlines()), + } + (build_dir / "gather_scatter_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + if args.no_run: + return {"build_dir": str(build_dir), "ran": False, "manifest": manifest} + + old_rust_log = os.environ.get("PLENA_EMULATOR_RUST_LOG") + if store_mode: + # HBM dumps are gated behind DEBUG tracing in the Rust runner. Keep + # other modules quiet so storeback validation does not drown in DMA logs. + os.environ["PLENA_EMULATOR_RUST_LOG"] = "warn,transactional_emulator=info,transactional_emulator::runner=debug" + try: + metrics = run_and_assert( + build_dir, + "gpt_oss_moe_gather_full_vram", + mlen=mlen, + blen=blen, + threads=args.emu_threads, + ) + finally: + if store_mode: + if old_rust_log is None: + os.environ.pop("PLENA_EMULATOR_RUST_LOG", None) + else: + os.environ["PLENA_EMULATOR_RUST_LOG"] = old_rust_log + results, params = compare_emulator_output(build_dir) + emu_output = results["simulated_values"].reshape(rows, hidden).to(torch.bfloat16) + emu_l1_stats = compare_stats(emu_output, target_golden, rtol=0.02) + emu_exact_stats = compare_stats(emu_output, exact_golden, rtol=0.02) + elementwise = _strict_elementwise_details(emu_output, target_golden, rtol=0.02) + tail_gate = _tail_gate_details( + emu_output, + target_golden, + strict_details=elementwise, + clamp_counts=clamp_counts, + rtol=0.02, + ) + + hbm_store_summary = None + if store_mode: + assert output_slot_hbm_base is not None + assert output_slot_hbm_stride is not None + hbm_before_path = build_dir / "hbm_for_behave_sim.bin" + hbm_after_path = build_dir / "hbm_dump.bin" + if not hbm_after_path.exists(): + raise AssertionError("HBM store stage expected hbm_dump.bin; DEBUG runner dump was not produced") + hbm_before = hbm_before_path.read_bytes() + hbm_after = hbm_after_path.read_bytes() + decoded_slots = _decode_hbm_mxfp8_output_slots( + hbm_after, + slot_base=output_slot_hbm_base, + slot_stride=output_slot_hbm_stride, + rows=rows, + hidden=hidden, + blen=blen, + ) + torch.save(decoded_slots.float(), build_dir / "hbm_store_decoded_output_slots.pt") + store_quant_golden = quantize_to_mxfp(exact_golden).to(torch.bfloat16) + torch.save(store_quant_golden.float(), build_dir / "hbm_store_quantized_exact_golden.pt") + + compare_tokens = store_token_indices if store_single else list(range(rows)) + decoded_cmp = decoded_slots[compare_tokens] + quant_cmp = store_quant_golden[compare_tokens] + exact_cmp = exact_golden[compare_tokens] + l1_cmp = l1_golden[compare_tokens] + hbm_vs_quant = compare_stats(decoded_cmp, quant_cmp, rtol=0.0) + hbm_vs_exact = compare_stats(decoded_cmp, exact_cmp, rtol=0.02) + hbm_vs_l1 = compare_stats(decoded_cmp, l1_cmp, rtol=0.02) + padding_zero = _hbm_slot_padding_zero( + hbm_after, + slot_base=output_slot_hbm_base, + slot_stride=output_slot_hbm_stride, + token_indices=store_token_indices, + hidden=hidden, + blen=blen, + ) + neighbor_unchanged = ( + _hbm_unwritten_slots_unchanged( + hbm_before, + hbm_after, + slot_base=output_slot_hbm_base, + slot_stride=output_slot_hbm_stride, + rows=rows, + written_tokens=store_token_indices, + ) + if store_single + else { + "written_tokens": store_token_indices, + "unchanged_tokens_checked": [], + "changed_unwritten_slots": [], + "passed": True, + } + ) + source_count_ok = h_store_count == expected_h_store_v_source_count + # HBM byte statistics are useful telemetry, but not a correctness + # oracle: unaligned read-modify-write may count internal chunk bytes. + dynamic_bytes_match_logical = metrics.get("hbm_bytes_written") == output_slot_hbm_stride * len( + store_token_indices + ) + hbm_store_summary = { + "stage": args.stage, + "store_token_indices": store_token_indices, + "output_slot_hbm_base": output_slot_hbm_base, + "output_slot_hbm_stride": output_slot_hbm_stride, + "source_h_store_v_count": h_store_count, + "expected_source_h_store_v_count": expected_h_store_v_source_count, + "source_h_store_v_count_matches": source_count_ok, + "expected_dynamic_h_store_v_count": expected_h_store_v_dynamic_count, + "hbm_bytes_written": metrics.get("hbm_bytes_written"), + "expected_hbm_bytes_written": output_slot_hbm_stride * len(store_token_indices), + "hbm_bytes_written_matches_logical_payload": dynamic_bytes_match_logical, + "decoded_vs_quantized_exact_golden": _stats_dict(hbm_vs_quant), + "decoded_vs_quantized_exact_golden_bit_exact": hbm_vs_quant.max_abs_error == 0.0 and hbm_vs_quant.allclose, + "decoded_vs_exact_hbm_decode_golden": _stats_dict(hbm_vs_exact), + "decoded_vs_l1_golden": _stats_dict(hbm_vs_l1), + "padding_rows_zero": padding_zero, + "unwritten_neighbor_slots_unchanged": neighbor_unchanged, + "passed": source_count_ok + and hbm_vs_quant.max_abs_error == 0.0 + and hbm_vs_quant.allclose + and padding_zero["passed"] + and neighbor_unchanged["passed"], + } + (build_dir / "hbm_store_results.json").write_text(json.dumps(hbm_store_summary, indent=2) + "\n") + + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params_runtime": params, + "emulator_compare_raw": { + key: results[key] + for key in ( + "mse", + "mae", + "max_error", + "relative_error", + "relative_match_rate", + "allclose_match_rate", + "match_rate", + "allclose_pass", + "atol", + "rtol", + "golden_shape", + "simulated_shape", + ) + if key in results + }, + "emulator_vs_target_golden": _stats_dict(emu_l1_stats), + "emulator_vs_target_elementwise": elementwise, + "emulator_vs_target_tail_gate": tail_gate, + "emulator_vs_exact_hbm_decode_golden": _stats_dict(emu_exact_stats), + "hbm_store_gate": hbm_store_summary, + "full_vram_gate": { + "scale_reg_count_nonzero": scale_reg_count > 0, + "no_h_store_v": (not store_mode) and h_store_count == 0, + "h_store_v_count_expected": (h_store_count == expected_h_store_v_source_count) + if store_mode + else h_store_count == 0, + "emu_vs_l1_rel_rms_under_2pct": emu_l1_stats.rel_rms <= 0.02, + "emu_vs_l1_strict_allclose": emu_l1_stats.allclose, + "emu_vs_l1_tail_gate_passed": tail_gate["passed"], + "passed": scale_reg_count > 0 + and ((h_store_count == 0) if not store_mode else h_store_count == expected_h_store_v_source_count) + and emu_l1_stats.rel_rms <= 0.02 + and (emu_l1_stats.allclose or tail_gate["passed"]) + and (hbm_store_summary is None or hbm_store_summary["passed"]), + }, + } + (build_dir / "gather_scatter_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + if not summary["full_vram_gate"]["passed"]: + raise AssertionError( + "full VRAM gather/expert/scatter-combine failed: " + f"rel_rms={emu_l1_stats.rel_rms:.6g}, allclose={emu_l1_stats.allclose}, " + f"tail_gate={tail_gate['passed']}, scale_regs={scale_reg_count}, h_store_v={h_store_count}" + ) + return {"build_dir": str(build_dir), "ran": True, "summary": summary} + + +def run_device_routing_stage(args: argparse.Namespace) -> dict: + """Run Step6 dynamic expert dispatch: V_TOPK id -> true expert weight table.""" + mlen = args.mlen + blen = args.blen + build_dir = args.build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + ref = _load_tok8_reference(args.reference_path) + x = ref["x"] + hf_topk_indices = ref["topk_indices"] + _hf_topk_weights = ref["topk_weights"] + rows, hidden = x.shape + top_k = int(hf_topk_indices.shape[1]) + vram_source = args.stage in ("device_routing_vram_single", "device_routing_vram_full") + full_pairs = args.stage in ("device_routing_full", "device_routing_vram_full") + pair_count = rows * top_k if full_pairs else 1 + if hidden % mlen != 0: + raise ValueError(f"hidden={hidden} must be divisible by MLEN={mlen}") + if top_k != 4: + raise ValueError(f"GPT-OSS v0 expects top_k=4, got {top_k}") + + tensors = load_layer_tensors(args.layer_index) + split = split_packed_gate_up(tensors["gate_up_weight"], tensors["gate_up_bias"]) + down_weight = tensors["down_weight"] + down_bias = tensors["down_bias"] + router_weight, router_bias = tensors["router_weight"], tensors["router_bias"] + router_weight_rows = router_weight.contiguous() + intermediate = int(split.gate_weight.shape[-1]) + num_experts = int(split.gate_weight.shape[0]) + + hf_logits = F.linear(x, router_weight, router_bias).to(torch.bfloat16) + hf_top_values, hf_top_indices_from_logits = torch.topk(hf_logits, k=top_k, dim=-1) + hf_top_weights_from_logits = torch.softmax(hf_top_values, dim=1, dtype=hf_top_values.dtype).to(torch.bfloat16) + if not torch.equal(hf_top_indices_from_logits.cpu(), hf_topk_indices.cpu()): + raise AssertionError("reference topk_indices differ from HF router recomputation") + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + zero = prog.fp_var("zero", size=1) + constant_rows = max(blen, rows * blen) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=constant_rows) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=constant_rows) + one = prog.fp_var("one", size=constant_rows) + neg_alpha = prog.fp_var("neg_alpha", size=constant_rows) + shared_zero_row = prog.fp_var("gpt_oss_shared_zero_row", size=mlen) + + physical_rows = max(blen, math.ceil(rows / blen) * blen) + input_tensors: dict[str, torch.Tensor] = {} + x_input = prog.input("X", shape=(rows, hidden)) + input_tensors["X"] = x + + gate_inputs, gate_table_base, gate_stride = _build_true_expert_weight_table( + prog, + prefix="W_gate", + weights=split.gate_weight.to(torch.bfloat16), + input_tensors=input_tensors, + ) + up_inputs, up_table_base, up_stride = _build_true_expert_weight_table( + prog, + prefix="W_up", + weights=split.up_weight.to(torch.bfloat16), + input_tensors=input_tensors, + ) + down_inputs, down_table_base, down_stride = _build_true_expert_weight_table( + prog, + prefix="W_down", + weights=down_weight.to(torch.bfloat16), + input_tensors=input_tensors, + ) + weight_templates = (gate_inputs[0], up_inputs[0], down_inputs[0]) + weight_table_bases = (gate_table_base, up_table_base, down_table_base) + weight_table_strides = (gate_stride, up_stride, down_stride) + + # Prestaged BF16 router operands and expert bias tables. + router_x_physical = (physical_rows, hidden) + router_w_physical = (num_experts, hidden) + router_bias_physical = (physical_rows, mlen) + bias_inter_physical = (num_experts * blen, intermediate) + bias_hidden_physical = (num_experts * blen, hidden) + alignment = mlen * mlen + router_x_base = 0 + router_w_base = _align_to(router_x_base + _vram_layout_size(router_x_physical, mlen=mlen), alignment) + router_bias_base = _align_to(router_w_base + _vram_layout_size(router_w_physical, mlen=mlen), alignment) + gate_bias_base = _align_to(router_bias_base + _vram_layout_size(router_bias_physical, mlen=mlen), alignment) + up_bias_base = _align_to(gate_bias_base + _vram_layout_size(bias_inter_physical, mlen=mlen), alignment) + down_bias_base = _align_to(up_bias_base + _vram_layout_size(bias_inter_physical, mlen=mlen), alignment) + preload_size = down_bias_base + _vram_layout_size(bias_hidden_physical, mlen=mlen) + vram_preload = torch.zeros(preload_size, dtype=torch.bfloat16) + + router_x = prestage_bf16_vram_matrix( + prog=prog, + name="Step6RouterXBF16", + tensor=x, + vram_addr=router_x_base, + physical_shape=router_x_physical, + vram_preload=vram_preload, + ) + router_w = prestage_bf16_vram_matrix( + prog=prog, + name="Step6RouterWBF16", + tensor=router_weight_rows, + vram_addr=router_w_base, + physical_shape=router_w_physical, + vram_preload=vram_preload, + ) + router_bias_rows = router_bias.reshape(1, -1).repeat(rows, 1).to(torch.bfloat16) + router_bias_vram = prestage_bf16_vram_matrix( + prog=prog, + name="Step6RouterBias", + tensor=router_bias_rows, + vram_addr=router_bias_base, + physical_shape=router_bias_physical, + vram_preload=vram_preload, + ) + gate_bias_table = _build_true_expert_bias_table( + prog=prog, + name="Step6GateBiasTable", + bias=split.gate_bias, + width=intermediate, + blen=blen, + vram_addr=gate_bias_base, + vram_preload=vram_preload, + ) + up_bias_table = _build_true_expert_bias_table( + prog=prog, + name="Step6UpBiasTable", + bias=split.up_bias, + width=intermediate, + blen=blen, + vram_addr=up_bias_base, + vram_preload=vram_preload, + ) + down_bias_table = _build_true_expert_bias_table( + prog=prog, + name="Step6DownBiasTable", + bias=down_bias, + width=hidden, + blen=blen, + vram_addr=down_bias_base, + vram_preload=vram_preload, + ) + + logits = prog.moe_router_logits_bf16_v0( + router_x, + router_w, + rows=rows, + hidden=hidden, + num_experts=num_experts, + name="step6_router_logits", + ) + prog.vram_add(logits, router_bias_vram, num_rows=rows) + topk_weight_var = prog.fp_var("step6_topk_weights", size=rows * top_k) + topk_weights_fp_base = topk_weight_var.address + topk_indices_int_base = 0 + token_offsets_int_base = rows * top_k + token_offsets = [token_idx * hidden for token_idx in range(rows) for _ in range(top_k)] + for token_idx in range(rows): + prog.moe_router_select_v0( + logits, + token_idx=token_idx, + weights_fp_base=topk_weights_fp_base + token_idx * top_k, + indices_int_base=topk_indices_int_base + token_idx * top_k, + name=f"token{token_idx}", + ) + + accumulator = prog.alloc( + "step6_device_routing_acc", + rows=rows, + cols=hidden, + strict=False, + physical_shape=(physical_rows, hidden), + ) + prog.moe_true_zero_vram_rows_v0( + accumulator, + rows=list(range(rows)), + hidden=hidden, + zero_row=shared_zero_row, + name="step6_acc_zero", + ) + route_fp_scratch = prog.fp_var("step6_route_fp_scratch", size=mlen) + + for pair_idx in range(pair_count): + token_idx = pair_idx // top_k + if vram_source: + gathered = prog.moe_gather_token_rows_from_vram_v0( + router_x, + token_indices=[token_idx], + hidden=hidden, + zero_row=shared_zero_row, + name=f"step6_pair{pair_idx}_vram_gather_t{token_idx}", + ) + else: + gathered = prog.moe_gather_token_rows_from_hbm_v0( + x_input, + token_offsets_int_base=token_offsets_int_base + pair_idx, + pair_count=1, + hidden=hidden, + zero_row=shared_zero_row, + name=f"step6_pair{pair_idx}_gather_t{token_idx}", + ) + output = prog.moe_dynamic_expert_pair_v0( + gathered, + weight_templates, + weight_table_bases=weight_table_bases, + weight_table_strides=weight_table_strides, + expert_indices_int_base=topk_indices_int_base, + weights_fp_base=topk_weights_fp_base, + pair_idx=pair_idx, + bias_tables=(gate_bias_table, up_bias_table, down_bias_table), + rows=blen, + intermediate=intermediate, + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + zero_row=shared_zero_row, + route_fp_scratch=route_fp_scratch, + name=f"step6_pair{pair_idx}", + ) + prog.moe_scatter_add_active_rows_v0( + accumulator, + output, + token_indices=[token_idx], + active_rows=[0], + hidden=hidden, + name=f"step6_pair{pair_idx}_scatter", + ) + + isa = prog.compile() + scale_reg_count = isa.count("C_SET_SCALE_REG") + if scale_reg_count <= 0: + raise AssertionError("device-routing expert path must emit C_SET_SCALE_REG for MXFP8 expert tensors") + h_prefetch_v_count = isa.count("H_PREFETCH_V") + if vram_source and h_prefetch_v_count != 0: + raise AssertionError( + f"VRAM-resident MoE input path must not use H_PREFETCH_V/HBM activation gather; found {h_prefetch_v_count}" + ) + + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + dummy_golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + fp_preload = [0.0] + [7.0] * constant_rows + [-7.0] * constant_rows + fp_preload += [1.0] * constant_rows + [-1.702] * constant_rows + fp_preload_len = max(route_fp_scratch.address + route_fp_scratch.size, topk_weights_fp_base + rows * top_k) + if len(fp_preload) < fp_preload_len: + fp_preload.extend([0.0] * (fp_preload_len - len(fp_preload))) + int_preload = torch.zeros(token_offsets_int_base + len(token_offsets), dtype=torch.int32) + int_preload[token_offsets_int_base : token_offsets_int_base + len(token_offsets)] = torch.tensor( + token_offsets, dtype=torch.int32 + ) + + create_sim_env( + input_tensors, + isa, + {"original_output": dummy_golden}, + fp_preload=fp_preload, + int_preload=int_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_moe_device_routing", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + + output_addr = prog._compiler.get_vram_addr(accumulator.name) + comparison_params = _comparison_params_for(accumulator, rows=rows, hidden=hidden, mlen=mlen, golden=dummy_golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + manifest = { + "stage": args.stage, + "reference_path": str(args.reference_path.expanduser().resolve()), + "layer_index": args.layer_index, + "rows": rows, + "hidden": hidden, + "intermediate": intermediate, + "num_experts": num_experts, + "top_k": top_k, + "pair_count": pair_count, + "mlen": mlen, + "blen": blen, + "true_expert_weight_tables": { + "gate": {"base": gate_table_base, "stride": gate_stride}, + "up": {"base": up_table_base, "stride": up_stride}, + "down": {"base": down_table_base, "stride": down_stride}, + }, + "helper_contract": "expert_id_to_weight_base = table_base + true_expert_id * per_expert_stride", + "moe_input_source": "vram_bf16" if vram_source else "hbm_mxfp8", + "topk_weights_fp_base": topk_weights_fp_base, + "topk_indices_int_base": topk_indices_int_base, + "token_offsets_int_base": token_offsets_int_base, + "scale_reg_count": scale_reg_count, + "h_prefetch_v_count": h_prefetch_v_count, + "accumulator_vram_row": output_addr // mlen, + "comparison_params": comparison_params, + "asm_lines": len(isa.splitlines()), + } + (build_dir / "device_routing_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + if args.no_run: + print(json.dumps(manifest, indent=2)) + return {"build_dir": str(build_dir), "ran": False, "manifest": manifest} + + metrics = run_emulator(build_dir, threads=args.emu_threads) + device_indices_flat = _decode_u32_dump(build_dir / "intsram_dump.bin")[ + topk_indices_int_base : topk_indices_int_base + rows * top_k + ] + device_weights_flat = _decode_bf16_dump(build_dir / "fpsram_dump.bin")[ + topk_weights_fp_base : topk_weights_fp_base + rows * top_k + ] + device_indices = device_indices_flat.reshape(rows, top_k) + device_weights = device_weights_flat.reshape(rows, top_k).to(torch.bfloat16) + if vram_source: + exact_golden, clamp_counts = _device_routing_vram_exact_golden( + x=x, + device_indices=device_indices, + device_weights=device_weights, + split=split, + down_weight=down_weight, + down_bias=down_bias, + rows=rows, + hidden=hidden, + blen=blen, + mlen=mlen, + pair_count=pair_count, + ) + else: + exact_golden, clamp_counts = _device_routing_exact_golden( + hbm_path=build_dir / "hbm_for_behave_sim.bin", + x_input=x_input, + x=x, + device_indices=device_indices, + device_weights=device_weights, + split=split, + down_weight=down_weight, + down_bias=down_bias, + rows=rows, + hidden=hidden, + blen=blen, + mlen=mlen, + pair_count=pair_count, + ) + torch.save(exact_golden.float(), build_dir / "device_selected_exact_golden.pt") + target_golden = exact_golden + _rewrite_compact_golden(build_dir, target_golden) + comparison_params = _comparison_params_for(accumulator, rows=rows, hidden=hidden, mlen=mlen, golden=target_golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + results, params = compare_emulator_output(build_dir) + emu_output = results["simulated_values"].reshape(rows, hidden).to(torch.bfloat16) + emu_target_stats = compare_stats(emu_output, target_golden, rtol=0.02) + emu_l1_stats = compare_stats( + emu_output, + torch.load(args.l1_golden_path.expanduser().resolve(), map_location="cpu").to(torch.bfloat16), + rtol=0.02, + ) + elementwise = _strict_elementwise_details(emu_output, target_golden, rtol=0.02) + tail_gate = _tail_gate_details( + emu_output, + target_golden, + strict_details=elementwise, + clamp_counts=clamp_counts, + rtol=0.02, + ) + device_router_logits = _router_vector_bf16_golden(x, router_weight_rows, router_bias, mlen=mlen) + router_rank_gate = _rank_stability( + hf_logits, + device_router_logits, + top_k, + ) + topk_contract = _device_topk_contract( + hf_indices=hf_topk_indices, + hf_weights=hf_top_weights_from_logits, + device_indices=device_indices, + device_weights=device_weights, + device_logits=device_router_logits, + rank_gate=router_rank_gate, + ) + exact_vs_l1 = compare_stats( + target_golden, + torch.load(args.l1_golden_path.expanduser().resolve(), map_location="cpu").to(torch.bfloat16), + rtol=0.02, + ) + + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params_runtime": params, + "device_topk": topk_contract, + "router_rank_gate_against_hf": router_rank_gate, + "device_selected_exact_golden_vs_l1_record_only": _stats_dict(exact_vs_l1), + "emulator_vs_device_selected_exact_golden": _stats_dict(emu_target_stats), + "emulator_vs_l1_record_only": _stats_dict(emu_l1_stats), + "emulator_vs_target_elementwise": elementwise, + "emulator_vs_target_tail_gate": tail_gate, + "clamp_counts": clamp_counts, + "device_routing_gate": { + "true_weight_table_bases": True, + "scale_reg_count_nonzero": scale_reg_count > 0, + "vram_source_no_h_prefetch_v": (not vram_source) or h_prefetch_v_count == 0, + "topk_contract_passed": topk_contract["passed"], + "emu_vs_device_exact_rel_rms_under_2pct": emu_target_stats.rel_rms <= 0.02, + "emu_vs_device_exact_strict_or_tail": emu_target_stats.allclose or tail_gate["passed"], + "passed": scale_reg_count > 0 + and ((not vram_source) or h_prefetch_v_count == 0) + and topk_contract["passed"] + and emu_target_stats.rel_rms <= 0.02 + and (emu_target_stats.allclose or tail_gate["passed"]), + }, + } + (build_dir / "device_routing_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + if not summary["device_routing_gate"]["passed"]: + raise AssertionError(f"device-routing gate failed: {summary['device_routing_gate']}") + return {"build_dir": str(build_dir), "ran": True, "summary": summary} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument( + "--stage", + choices=( + "address", + "four", + "full_vram", + "hbm_store_single", + "hbm_store_full", + "device_routing_single", + "device_routing_full", + "device_routing_vram_single", + "device_routing_vram_full", + ), + default="address", + ) + parser.add_argument( + "--reference-path", + type=Path, + default=_ATEN_BUILD_DIR / "gpt_oss_real_layer0_tok8_l1_stamp" / "hf_layer0_moe_reference.pt", + ) + parser.add_argument( + "--l1-golden-path", + type=Path, + default=_ATEN_BUILD_DIR / "gpt_oss_real_layer0_emu_tok8_l1_stamp" / "golden_output.pt", + ) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_moe_gather_scatter_address", + ) + parser.add_argument("--token-index", type=int, default=7) + parser.add_argument("--store-token-index", type=int, default=2) + parser.add_argument("--expert-id", type=int, default=9) + parser.add_argument("--layer-index", type=int, default=0) + parser.add_argument( + "--only-experts", + type=str, + default=None, + help="Debug full_vram with a comma-separated subset of routed expert ids.", + ) + parser.add_argument("--emu-threads", type=int, default=None) + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + + if args.stage == "address": + run_address_stage(args) + elif args.stage == "four": + run_four_pair_stage(args) + elif args.stage in ("full_vram", "hbm_store_single", "hbm_store_full"): + run_full_vram_stage(args) + elif args.stage in ( + "device_routing_single", + "device_routing_full", + "device_routing_vram_single", + "device_routing_vram_full", + ): + run_device_routing_stage(args) + else: # pragma: no cover - argparse prevents this today. + raise ValueError(f"unsupported stage {args.stage}") + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_real_layer0_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_real_layer0_test.py new file mode 100644 index 00000000..1a57e652 --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_real_layer0_test.py @@ -0,0 +1,838 @@ +# ruff: noqa: E402 +"""GPT-OSS real layer-0 MoE fixed-routing emulator check. + +This is the L1 "external proof" path: + + HF GptOssMLP routing/reference + -> fixed HF top-k indices/weights + -> PLENA emulator expert compute + combine + -> compare against Golden B + +HF corresponds to Golden A: MXFP4 checkpoint tensors dequantized to BF16 and +BF16 expert math. Golden B uses the same fixed routing, but quantizes HBM +activations/weights/route weights through PLENA's current MXFP8 path. Expert +biases are BF16 VRAM tensors, matching GPT-OSS' non-MX bias semantics. +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path + +import torch + +_AICROSSSIM_ROOT = Path(__file__).resolve().parents[4] +_TESTBENCH_ROOT = Path(__file__).resolve().parents[1] +_ATEN_BUILD_DIR = _TESTBENCH_ROOT / "aten" / "build" +_TOP_LEVEL_COMPILER = _AICROSSSIM_ROOT / "PLENA_Compiler" +if _TOP_LEVEL_COMPILER.exists(): + sys.path.insert(0, str(_TOP_LEVEL_COMPILER)) + +from aten.models.gpt_oss.moe_reference import compare_stats, split_packed_gate_up +from aten.models.gpt_oss.real_layer_utils import load_json, load_layer0_tensors +from compiler.aten.plena import PlenaCompiler +from compiler.aten.plena.vars import VRAMMatrixVar +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.emulator_runner import compare_emulator_output, run_and_assert +from transactional_emulator.testbench.layout_utils import infer_hbm_tensor_layouts, prestage_bf16_vram_matrix +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.testbench.sliced_layer_test_builder import quantize_to_mxfp +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _activation_golden(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """BF16 GPT-OSS activation matching the PLENA emit sequence.""" + neg_alpha = torch.tensor(-1.702, dtype=torch.bfloat16).float() + + gate = _bf16(torch.clamp(gate.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), max=7.0)) + up = _bf16(torch.clamp(up.float(), min=-7.0)) + + sigmoid = _bf16(gate.float()) + sigmoid = _bf16(sigmoid.float() * neg_alpha) + sigmoid = _bf16(torch.exp(torch.clamp(sigmoid.float(), -88.0, 88.0))) + sigmoid = _bf16(sigmoid.float() + 1.0) + sigmoid = _bf16(torch.reciprocal(sigmoid.float())) + + glu = _bf16(gate.float() * sigmoid.float()) + up_plus_one = _bf16(up.float() + 1.0) + return _bf16(up_plus_one.float() * glu.float()) + + +def _linear_projection_golden( + x: torch.Tensor, + w: torch.Tensor, + *, + mlen: int, + mram_tile_capacity: int = 4, + hbm_input: bool = True, +) -> torch.Tensor: + """Hardware-aware PLENA projection golden with K-split BF16 accumulation.""" + x_q = quantize_to_mxfp(x) if hbm_input else x.to(torch.bfloat16) + w_q = quantize_to_mxfp(w) + k_total = x_q.shape[1] + chunk = mlen * mram_tile_capacity + + acc = None + for k_start in range(0, k_total, chunk): + k_end = min(k_start + chunk, k_total) + partial = torch.matmul(x_q[:, k_start:k_end].float(), w_q[k_start:k_end, :].float()).to(torch.bfloat16) + acc = partial if acc is None else (acc.float() + partial.float()).to(torch.bfloat16) + assert acc is not None + return acc + + +def _expert_golden_b( + x: torch.Tensor, + w_gate: torch.Tensor, + w_up: torch.Tensor, + w_down: torch.Tensor, + *, + mlen: int, + b_gate: torch.Tensor, + b_up: torch.Tensor, + b_down: torch.Tensor, +) -> torch.Tensor: + """Golden B for one selected expert. + + Biases are intentionally not MX-quantized. They are BF16 VRAM tensors in + the emulator path and BF16 additions here. + """ + gate = _linear_projection_golden(x, w_gate, mlen=mlen, hbm_input=True) + up = _linear_projection_golden(x, w_up, mlen=mlen, hbm_input=True) + gate = _bf16(gate.float() + b_gate.float()) + up = _bf16(up.float() + b_up.float()) + hidden = _activation_golden(gate, up) + out = _linear_projection_golden(hidden, w_down, mlen=mlen, hbm_input=False) + return _bf16(out.float() + b_down.float()) + + +def _make_route_matrix( + topk_indices: torch.Tensor, + topk_weights: torch.Tensor, + expert_id: int, + hidden: int, +) -> torch.Tensor: + route = torch.zeros(topk_indices.shape[0], 1, dtype=torch.bfloat16) + token_idx, topk_pos = torch.where(topk_indices == expert_id) + if token_idx.numel() > 0: + route[token_idx, 0] = topk_weights[token_idx, topk_pos].to(torch.bfloat16) + return route.repeat(1, hidden) + + +def _expanded_bias(row_bias: torch.Tensor, rows: int) -> torch.Tensor: + return row_bias.to(torch.bfloat16).reshape(1, -1).repeat(rows, 1) + + +def _stats_dict(stats) -> dict: + return { + "rel_rms": stats.rel_rms, + "atol": stats.atol, + "rtol": stats.rtol, + "allclose": stats.allclose, + "pass_rate": stats.pass_rate, + "max_abs_error": stats.max_abs_error, + } + + +def _strict_elementwise_details( + actual: torch.Tensor, + reference: torch.Tensor, + *, + rtol: float, + atol_scale: float = 0.01, + block_cols: int = 64, +) -> dict: + """Return the strict elementwise check details requested for real-layer proof.""" + actual_f = actual.float() + reference_f = reference.float() + diff = (actual_f - reference_f).abs() + atol = reference_f.std(unbiased=False) * atol_scale + allowed = atol + rtol * reference_f.abs() + failed = diff > allowed + fail_count = int(failed.sum().item()) + numel = failed.numel() + + details = { + "fail_count": fail_count, + "numel": numel, + "pass_rate": 1.0 - (fail_count / numel if numel else 0.0), + "atol": float(atol.item()), + "rtol": rtol, + "max_abs_error": float(diff.max().item()) if numel else 0.0, + "max_allowed_ratio": float((diff / allowed.clamp_min(1e-12)).max().item()) if numel else 0.0, + } + + if failed.ndim == 2: + _rows, cols = failed.shape + details["fail_count_by_token"] = [int(v) for v in failed.sum(dim=1).tolist()] + blocks = [] + for block_idx in range(math.ceil(cols / block_cols)): + start = block_idx * block_cols + end = min(start + block_cols, cols) + count = int(failed[:, start:end].sum().item()) + if count: + blocks.append({"block": block_idx, "fail_count": count}) + details["fail_count_by_col_block"] = blocks + return details + + +def _pearson(xs: list[float], ys: list[float]) -> float | None: + if len(xs) != len(ys) or len(xs) < 2: + return None + x_mean = sum(xs) / len(xs) + y_mean = sum(ys) / len(ys) + x_centered = [x - x_mean for x in xs] + y_centered = [y - y_mean for y in ys] + x_norm = math.sqrt(sum(x * x for x in x_centered)) + y_norm = math.sqrt(sum(y * y for y in y_centered)) + if x_norm == 0.0 or y_norm == 0.0: + return None + return sum(x * y for x, y in zip(x_centered, y_centered, strict=True)) / (x_norm * y_norm) + + +def _ranks(values: list[float]) -> list[float]: + ordered = sorted(enumerate(values), key=lambda item: item[1]) + ranks = [0.0] * len(values) + i = 0 + while i < len(ordered): + j = i + 1 + while j < len(ordered) and ordered[j][1] == ordered[i][1]: + j += 1 + avg_rank = (i + 1 + j) / 2.0 + for k in range(i, j): + ranks[ordered[k][0]] = avg_rank + i = j + return ranks + + +def _spearman(xs: list[float], ys: list[float]) -> float | None: + if len(xs) != len(ys) or len(xs) < 2: + return None + return _pearson(_ranks(xs), _ranks(ys)) + + +def _tail_gate_details( + actual: torch.Tensor, + reference: torch.Tensor, + *, + strict_details: dict, + clamp_counts: dict, + rtol: float, + k_values: tuple[float, ...] = (2.0, 2.5, 3.0, 3.5, 4.0), + block_cols: int = 64, +) -> dict: + """Classify strict elementwise misses as BF16 tail behavior or mechanism.""" + actual_f = actual.float() + reference_f = reference.float() + signed = actual_f - reference_f + abs_diff = signed.abs() + sigma = abs_diff.std(unbiased=False) + numel = int(abs_diff.numel()) + + k_sweep = [] + for k in k_values: + allowed = (k * sigma) + rtol * reference_f.abs() + failed = abs_diff > allowed + k_sweep.append({"k": k, "fail_count": int(failed.sum().item())}) + monotonic = all(k_sweep[i]["fail_count"] >= k_sweep[i + 1]["fail_count"] for i in range(len(k_sweep) - 1)) + k4_count = next((entry["fail_count"] for entry in k_sweep if entry["k"] == 4.0), k_sweep[-1]["fail_count"]) + k4_upper_bound = max(8, math.ceil(0.0004 * numel)) + + strict_atol = float(strict_details["atol"]) + strict_allowed = strict_atol + rtol * reference_f.abs() + strict_failed = abs_diff > strict_allowed + strict_fail_count = int(strict_failed.sum().item()) + if strict_fail_count: + signed_failed = signed[strict_failed] + pos = int((signed_failed > 0).sum().item()) + neg = int((signed_failed < 0).sum().item()) + zero = strict_fail_count - pos - neg + sign_imbalance = abs(pos - neg) + sign_unbiased = sign_imbalance <= max(4, math.ceil(3.0 * math.sqrt(strict_fail_count))) + else: + pos = neg = zero = sign_imbalance = 0 + sign_unbiased = True + + spatial_blocks = strict_details.get("fail_count_by_col_block", []) + max_block_fail = max((int(block["fail_count"]) for block in spatial_blocks), default=0) + spatial_upper_bound = max(8, math.ceil(0.25 * strict_fail_count)) + spatial_undominated = max_block_fail <= spatial_upper_bound + + token_fail_counts = [float(v) for v in strict_details.get("fail_count_by_token", [])] + per_token_clamp = clamp_counts.get("per_token") or [] + token_clamp_counts = [ + float( + item.get("gate_gt_7", item.get("gate_gt_limit", 0)) + + item.get("up_gt_7", item.get("up_gt_limit", 0)) + + item.get("up_lt_neg_7", item.get("up_lt_neg_limit", 0)) + ) + for item in per_token_clamp + ] + pearson = _pearson(token_fail_counts, token_clamp_counts) if token_fail_counts and token_clamp_counts else None + spearman = _spearman(token_fail_counts, token_clamp_counts) if token_fail_counts and token_clamp_counts else None + clamp_uncorrelated = True + if pearson is not None: + clamp_uncorrelated = clamp_uncorrelated and abs(pearson) <= 0.5 + if spearman is not None: + clamp_uncorrelated = clamp_uncorrelated and abs(spearman) <= 0.5 + + passed = monotonic and k4_count <= k4_upper_bound and sign_unbiased and spatial_undominated and clamp_uncorrelated + return { + "classification": "bf16_tail" if passed else "mechanism_suspect", + "passed": passed, + "sigma_abs_error": float(sigma.item()), + "k_sweep": k_sweep, + "k4_fail_count": k4_count, + "k4_upper_bound": k4_upper_bound, + "monotonic_k_sweep": monotonic, + "strict_fail_count": strict_fail_count, + "sign": { + "positive": pos, + "negative": neg, + "zero": zero, + "imbalance": sign_imbalance, + "unbiased": sign_unbiased, + }, + "spatial": { + "max_col_block_fail_count": max_block_fail, + "upper_bound": spatial_upper_bound, + "undominated": spatial_undominated, + }, + "clamp_correlation": { + "pearson": pearson, + "spearman": spearman, + "uncorrelated": clamp_uncorrelated, + "fail_count_by_token": [int(v) for v in token_fail_counts], + "clamp_count_by_token": [int(v) for v in token_clamp_counts], + }, + } + + +def _comparison_params_for( + output_vram: VRAMMatrixVar, *, rows: int, hidden: int, mlen: int, golden: torch.Tensor +) -> dict: + output_vram_addr = output_vram._program._compiler.get_vram_addr(output_vram.name) + output_physical_rows = output_vram.physical_shape[0] + chunks_per_batch = math.ceil(hidden / mlen) + rows_to_read = (chunks_per_batch - 1) * output_physical_rows + rows + return { + "start_row_idx": output_vram_addr // mlen, + "num_rows": rows_to_read, + "num_batches": rows, + "elements_per_batch": hidden, + "row_dim": mlen, + "physical_rows": output_physical_rows, + "atol": float((golden.float().std(unbiased=False) * 0.01).item()), + "rtol": 0.02, + } + + +def _run_gate_debug( + *, + args: argparse.Namespace, + build_dir: Path, + hw, + x: torch.Tensor, + w_gate: torch.Tensor, + b_gate: torch.Tensor, + golden: torch.Tensor, + mlen: int, + blen: int, + rows: int, + hidden: int, + intermediate: int, +) -> dict: + """Run only ``X @ W_gate + gate_bias`` for the first selected real expert.""" + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + w_gate_input = prog.input("W_0_gate", shape=(hidden, intermediate)) + x_vram = prog.load_batch(x_input, name="X") + + physical_rows = max(blen, math.ceil(rows / blen) * blen) + bias_physical = (physical_rows, intermediate) + bias_base = math.ceil(prog.vram_allocator._vmm.next_bump / (mlen * mlen)) * (mlen * mlen) + bias_size = bias_physical[0] * bias_physical[1] + bias_aligned = math.ceil(bias_size / (mlen * mlen)) * (mlen * mlen) + vram_preload = torch.zeros(bias_base + bias_aligned, dtype=torch.bfloat16) + b_gate_vram = prestage_bf16_vram_matrix( + prog=prog, + name="B_0_gate", + tensor=b_gate, + vram_addr=bias_base, + physical_shape=bias_physical, + vram_preload=vram_preload, + ) + + output_vram = prog.linear_projection(x_vram, w_gate_input, name="real_layer0_debug_gate") + prog.vram_add(output_vram, b_gate_vram, num_rows=rows) + isa = prog.compile() + + input_tensors = {"X": x, "W_0_gate": w_gate} + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_real_layer0_gate", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + + comparison_params = _comparison_params_for(output_vram, rows=rows, hidden=intermediate, mlen=mlen, golden=golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + + manifest = { + "debug_stage": "gate", + "rows": rows, + "hidden": hidden, + "intermediate": intermediate, + "mlen": mlen, + "blen": blen, + "output_vram_row": prog._compiler.get_vram_addr(output_vram.name) // mlen, + "asm_lines": len(isa.splitlines()), + "comparison_params": comparison_params, + } + (build_dir / "real_layer0_gate_debug_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + if args.no_run: + return {"build_dir": str(build_dir), "ran": False, "manifest": manifest} + + metrics = run_and_assert( + build_dir, + "gpt_oss_real_layer0_gate", + mlen=mlen, + blen=blen, + threads=args.emu_threads, + ) + results, params = compare_emulator_output(build_dir) + emu_output = results["simulated_values"].reshape(rows, intermediate).to(torch.bfloat16) + emu_stats = compare_stats(emu_output, golden, rtol=0.02) + if emu_stats.rel_rms > 0.02: + raise AssertionError(f"gate debug emulator rel_rms={emu_stats.rel_rms:.6g} exceeds 2%") + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params": params, + "emulator_vs_gate_golden": _stats_dict(emu_stats), + } + (build_dir / "real_layer0_gate_debug_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + return {"build_dir": str(build_dir), "ran": True, "summary": summary} + + +def run_real_layer0(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + build_dir = args.build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + reference_path = args.reference_path.expanduser().resolve() + reference = torch.load(reference_path, map_location="cpu") + x = reference["x"].to(torch.bfloat16) + topk_indices = reference["topk_indices"].to(torch.long) + topk_weights = reference["topk_weights"].to(torch.bfloat16) + hf_output = reference["hf_output"].to(torch.bfloat16) + golden_a_output = reference["golden_a_output"].to(torch.bfloat16) + + config_dict = load_json("config.json") + hidden = int(config_dict["hidden_size"]) + intermediate = int(config_dict["intermediate_size"]) + rows = int(x.shape[0]) + if tuple(x.shape) != (rows, hidden): + raise ValueError(f"reference x shape {tuple(x.shape)} does not match hidden={hidden}") + if hidden % mlen != 0 or intermediate % mlen != 0: + raise ValueError(f"hidden/intermediate must be divisible by MLEN={mlen}: {hidden}, {intermediate}") + + tensors = load_layer0_tensors() + split = split_packed_gate_up(tensors["gate_up_weight"], tensors["gate_up_bias"]) + down_weight = tensors["down_weight"] + down_bias = tensors["down_bias"] + + selected_experts = sorted({int(v) for v in topk_indices.flatten().tolist()}) + if args.max_selected_experts is not None: + selected_experts = selected_experts[: args.max_selected_experts] + if not selected_experts: + raise ValueError("reference routing selected zero experts") + + print("=" * 80) + print( + "GPT-OSS real layer0 fixed-routing emulator check " + f"(rows={rows}, hidden={hidden}, intermediate={intermediate}, selected_experts={selected_experts})" + ) + print("=" * 80) + + m_tiles = math.ceil(rows / mlen) + gate_n_tiles = math.ceil(intermediate / mlen) + gate_k_tiles = math.ceil(hidden / mlen) + down_n_tiles = math.ceil(hidden / mlen) + down_k_tiles = math.ceil(intermediate / mlen) + print( + "tile coverage: " + f"gate/up M={m_tiles}, N={gate_n_tiles}, K={gate_k_tiles}; " + f"down M={m_tiles}, N={down_n_tiles}, K={down_k_tiles}; " + f"K-split chunks={math.ceil(gate_k_tiles / 4)}" + ) + + experts = [] + expert_bias_tensors = [] + route_tensors = [] + selected_gate_up = torch.zeros(rows, len(selected_experts), 2 * intermediate, dtype=torch.bfloat16) + debug_gate_golden = None + for local_idx, expert_id in enumerate(selected_experts): + # Clone selected expert views so torch.save does not serialize the + # original full-expert storage behind a narrow view. + w_gate = split.gate_weight[expert_id].to(torch.bfloat16).contiguous().clone() + w_up = split.up_weight[expert_id].to(torch.bfloat16).contiguous().clone() + w_down = down_weight[expert_id].to(torch.bfloat16).contiguous().clone() + b_gate = _expanded_bias(split.gate_bias[expert_id], rows) + b_up = _expanded_bias(split.up_bias[expert_id], rows) + b_down = _expanded_bias(down_bias[expert_id], rows) + if args.zero_bias: + b_gate = torch.zeros_like(b_gate) + b_up = torch.zeros_like(b_up) + b_down = torch.zeros_like(b_down) + route = _make_route_matrix(topk_indices, topk_weights, expert_id, hidden) + + gate = _linear_projection_golden(x, w_gate, mlen=mlen, hbm_input=True) + up = _linear_projection_golden(x, w_up, mlen=mlen, hbm_input=True) + gate = _bf16(gate.float() + b_gate.float()) + up = _bf16(up.float() + b_up.float()) + if local_idx == 0: + debug_gate_golden = gate + selected_gate_up[:, local_idx, 0::2] = gate + selected_gate_up[:, local_idx, 1::2] = up + + experts.append((w_gate, w_up, w_down)) + expert_bias_tensors.append((b_gate, b_up, b_down)) + route_tensors.append(route) + + gate = selected_gate_up[..., 0::2].float() + up = selected_gate_up[..., 1::2].float() + per_token_clamp = [] + for token_idx in range(rows): + token_gate = gate[token_idx] + token_up = up[token_idx] + per_token_clamp.append( + { + "token_index": token_idx, + "gate_gt_7": int((token_gate > 7.0).sum().item()), + "gate_lt_neg_7": int((token_gate < -7.0).sum().item()), + "up_gt_7": int((token_up > 7.0).sum().item()), + "up_lt_neg_7": int((token_up < -7.0).sum().item()), + } + ) + clamp_counts = { + "gate_gt_7": int((gate > 7.0).sum().item()), + "gate_lt_neg_7": int((gate < -7.0).sum().item()), + "up_gt_7": int((up > 7.0).sum().item()), + "up_lt_neg_7": int((up < -7.0).sum().item()), + "gate_min": float(gate.min().item()), + "gate_max": float(gate.max().item()), + "up_min": float(up.min().item()), + "up_max": float(up.max().item()), + "per_token": per_token_clamp, + } + print(f"clamp coverage: {json.dumps(clamp_counts, sort_keys=True)}") + if clamp_counts["gate_gt_7"] + clamp_counts["up_gt_7"] + clamp_counts["up_lt_neg_7"] == 0: + raise AssertionError("selected real layer0 path did not trigger GPT-OSS clamp") + + if args.debug_stage == "gate": + assert debug_gate_golden is not None + return _run_gate_debug( + args=args, + build_dir=build_dir, + hw=hw, + x=x, + w_gate=experts[0][0], + b_gate=expert_bias_tensors[0][0], + golden=debug_gate_golden, + mlen=mlen, + blen=blen, + rows=rows, + hidden=hidden, + intermediate=intermediate, + ) + + golden = torch.zeros(rows, hidden, dtype=torch.bfloat16) + for (w_gate, w_up, w_down), (b_gate, b_up, b_down), route in zip( + experts, expert_bias_tensors, route_tensors, strict=True + ): + expert_out = _expert_golden_b( + x, + w_gate, + w_up, + w_down, + mlen=mlen, + b_gate=b_gate, + b_up=b_up, + b_down=b_down, + ) + golden = _bf16(golden.float() + _bf16(expert_out.float() * quantize_to_mxfp(route).float()).float()) + + a_b_stats = compare_stats(golden, hf_output, rtol=0.15) + host_a_b_stats = compare_stats(golden, golden_a_output, rtol=0.15) + hf_a_stats = compare_stats(golden_a_output, hf_output, rtol=1e-2) + if hf_a_stats.max_abs_error != 0.0 or not hf_a_stats.allclose: + raise AssertionError( + "HF vs Golden A anchor must be exact for this fixed-routing proof: " + f"rel_rms={hf_a_stats.rel_rms:.6g}, max_abs={hf_a_stats.max_abs_error:.6g}, " + f"allclose={hf_a_stats.allclose}" + ) + + prog = PlenaCompiler(mlen=mlen, blen=blen, real_data_ratio=hw.real_data_ratio) + x_input = prog.input("X", shape=(rows, hidden)) + x_vram = prog.load_batch(x_input, name="X") + + expert_inputs = [] + route_inputs = [] + input_tensors = {"X": x} + for local_idx, ((w_gate, w_up, w_down), route) in enumerate(zip(experts, route_tensors, strict=True)): + w_gate_name = f"W_{local_idx}_gate" + w_up_name = f"W_{local_idx}_up" + w_down_name = f"W_{local_idx}_down" + route_name = f"route{local_idx}" + w_gate_input = prog.input(w_gate_name, shape=(hidden, intermediate)) + w_up_input = prog.input(w_up_name, shape=(hidden, intermediate)) + w_down_input = prog.input(w_down_name, shape=(intermediate, hidden)) + route_input = prog.input(route_name, shape=(rows, hidden)) + expert_inputs.append((w_gate_input, w_up_input, w_down_input)) + route_inputs.append(route_input) + input_tensors[w_gate_name] = w_gate + input_tensors[w_up_name] = w_up + input_tensors[w_down_name] = w_down + input_tensors[route_name] = route + + zero = prog.fp_var("zero", size=1) + limit_pos = prog.fp_var("gpt_oss_limit_pos", size=rows) + limit_neg = prog.fp_var("gpt_oss_limit_neg", size=rows) + one = prog.fp_var("one", size=rows) + neg_alpha = prog.fp_var("neg_alpha", size=rows) + + route_vrams = [prog.load_batch(route_input, name=route_input.display_name) for route_input in route_inputs] + + physical_rows = max(blen, math.ceil(rows / blen) * blen) + bias_physical_inter = (physical_rows, intermediate) + bias_physical_hidden = (physical_rows, hidden) + bias_sizes = [] + for _ in expert_bias_tensors: + bias_sizes.extend( + [ + bias_physical_inter[0] * bias_physical_inter[1], + bias_physical_inter[0] * bias_physical_inter[1], + bias_physical_hidden[0] * bias_physical_hidden[1], + ] + ) + bias_base = 0 + # X and route matrices were already allocated above; put prestaged biases + # after those regions so mark_used cannot overlap normal loads. + if prog.vram_allocator._vmm.next_bump > bias_base: + bias_base = math.ceil(prog.vram_allocator._vmm.next_bump / (mlen * mlen)) * (mlen * mlen) + bias_total = sum(math.ceil(size / (mlen * mlen)) * (mlen * mlen) for size in bias_sizes) + vram_preload = torch.zeros(bias_base + bias_total, dtype=torch.bfloat16) + next_bias_addr = bias_base + expert_bias_vrams = [] + for local_idx, (b_gate, b_up, b_down) in enumerate(expert_bias_tensors): + local_bias_vrams = [] + for suffix, tensor, physical_shape in ( + ("gate", b_gate, bias_physical_inter), + ("up", b_up, bias_physical_inter), + ("down", b_down, bias_physical_hidden), + ): + size = physical_shape[0] * physical_shape[1] + aligned_size = math.ceil(size / (mlen * mlen)) * (mlen * mlen) + name = f"B{local_idx}_{suffix}" + local_bias_vrams.append( + prestage_bf16_vram_matrix( + prog=prog, + name=name, + tensor=tensor, + vram_addr=next_bias_addr, + physical_shape=physical_shape, + vram_preload=vram_preload, + ) + ) + next_bias_addr += aligned_size + expert_bias_vrams.append(tuple(local_bias_vrams)) + + output_vram = prog.gpt_oss_moe_fixed_routing_v0( + x_vram, + experts=expert_inputs, + route_weights=route_vrams, + expert_biases=expert_bias_vrams, + rows=rows, + intermediate=intermediate, + constants=(zero, limit_pos, limit_neg, one, neg_alpha), + name="real_layer0_fixed_routing", + ) + isa = prog.compile() + + fp_preload = [0.0] + [7.0] * rows + [-7.0] * rows + [1.0] * rows + [-1.702] * rows + [0.0] * 8 + tensor_layouts = infer_hbm_tensor_layouts(input_tensors) + create_sim_env( + input_tensors, + isa, + {"original_output": golden}, + fp_preload=fp_preload, + build_dir=str(build_dir), + vram_preload=vram_preload, + tensor_layouts=tensor_layouts, + ) + + hbm_addrs = {name: prog._compiler.get_hbm_layout(name).hbm_base_addr for name in input_tensors} + create_mem_for_sim( + data_size=256, + mode="behave_sim", + asm="gpt_oss_real_layer0", + data=None, + specified_data_order=list(input_tensors), + build_path=build_dir, + input_tensors=input_tensors, + hbm_addrs=hbm_addrs, + tensor_layouts=tensor_layouts, + ) + + output_vram_addr = prog._compiler.get_vram_addr(output_vram.name) + comparison_params = _comparison_params_for(output_vram, rows=rows, hidden=hidden, mlen=mlen, golden=golden) + (build_dir / "comparison_params.json").write_text(json.dumps(comparison_params, indent=2) + "\n") + (build_dir / "generated_asm_code.asm").write_text(isa) + + manifest = { + "reference_path": str(reference_path), + "selected_experts": selected_experts, + "rows": rows, + "hidden": hidden, + "intermediate": intermediate, + "mlen": mlen, + "blen": blen, + "tile_coverage": { + "gate_up_m_tiles": m_tiles, + "gate_up_n_tiles": gate_n_tiles, + "gate_up_k_tiles": gate_k_tiles, + "down_m_tiles": m_tiles, + "down_n_tiles": down_n_tiles, + "down_k_tiles": down_k_tiles, + "k_split_chunks": math.ceil(gate_k_tiles / 4), + }, + "clamp": clamp_counts, + "hf_vs_golden_a": _stats_dict(hf_a_stats), + "golden_a_vs_golden_b": _stats_dict(a_b_stats), + "host_fixed_routing_a_vs_golden_b": _stats_dict(host_a_b_stats), + "output_vram_row": output_vram_addr // mlen, + "asm_lines": len(isa.splitlines()), + } + (build_dir / "real_layer0_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(json.dumps(manifest, indent=2)) + if args.no_run: + return {"build_dir": str(build_dir), "ran": False, "manifest": manifest} + + metrics = run_and_assert( + build_dir, + "gpt_oss_real_layer0", + mlen=mlen, + blen=blen, + threads=args.emu_threads, + ) + results, params = compare_emulator_output(build_dir) + emu_output = results["simulated_values"].reshape(rows, hidden).to(torch.bfloat16) + emu_b_stats = compare_stats(emu_output, golden, rtol=0.02) + emu_hf_stats = compare_stats(emu_output, hf_output, rtol=0.15) + emu_b_elementwise = _strict_elementwise_details(emu_output, golden, rtol=0.02) + emu_b_tail_gate = _tail_gate_details( + emu_output, + golden, + strict_details=emu_b_elementwise, + clamp_counts=clamp_counts, + rtol=0.02, + ) + + summary = { + **manifest, + "run_metrics": metrics, + "comparison_params": params, + "emulator_compare_raw": { + key: results[key] + for key in ( + "mse", + "mae", + "max_error", + "relative_error", + "relative_match_rate", + "allclose_match_rate", + "match_rate", + "allclose_pass", + "atol", + "rtol", + "golden_shape", + "simulated_shape", + ) + if key in results + }, + "emulator_vs_golden_b": _stats_dict(emu_b_stats), + "emulator_vs_golden_b_elementwise": emu_b_elementwise, + "emulator_vs_golden_b_tail_gate": emu_b_tail_gate, + "emulator_vs_hf_record_only": _stats_dict(emu_hf_stats), + "l1_external_proof_gate": { + "hf_vs_golden_a_exact": hf_a_stats.max_abs_error == 0.0 and hf_a_stats.allclose, + "emu_vs_golden_b_rel_rms_under_2pct": emu_b_stats.rel_rms <= 0.02, + "emu_vs_golden_b_strict_allclose": emu_b_stats.allclose, + "emu_vs_golden_b_tail_gate_passed": emu_b_tail_gate["passed"], + "passed": emu_b_stats.rel_rms <= 0.02 and (emu_b_stats.allclose or emu_b_tail_gate["passed"]), + }, + } + (build_dir / "real_layer0_results.json").write_text(json.dumps(summary, indent=2) + "\n") + print(json.dumps(summary, indent=2)) + if emu_b_stats.rel_rms > 0.02 or (not emu_b_stats.allclose and not emu_b_tail_gate["passed"]): + raise AssertionError( + f"emulator vs Golden B failed: rel_rms={emu_b_stats.rel_rms:.6g} " + f"(limit 2%), allclose={emu_b_stats.allclose}, tail_gate={emu_b_tail_gate['passed']}, " + f"pass_rate={emu_b_stats.pass_rate:.2%}, max_abs={emu_b_stats.max_abs_error:.6g}" + ) + return {"build_dir": str(build_dir), "ran": True, "summary": summary} + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + add_hw_args(parser) + parser.add_argument( + "--reference-path", + type=Path, + default=_ATEN_BUILD_DIR / "gpt_oss_real_layer0_tok1" / "hf_layer0_moe_reference.pt", + ) + parser.add_argument( + "--build-dir", + type=Path, + default=Path(__file__).parent / "build" / "gpt_oss_real_layer0_emu", + ) + parser.add_argument("--max-selected-experts", type=int, default=None) + parser.add_argument("--emu-threads", type=int, default=None) + parser.add_argument("--debug-stage", choices=("full", "gate"), default="full") + parser.add_argument("--zero-bias", action="store_true") + parser.add_argument("--no-run", action="store_true") + args = parser.parse_args() + run_real_layer0(args) + + +if __name__ == "__main__": + main() diff --git a/transactional_emulator/testbench/routed_moe/gpt_oss_router_gemm_test.py b/transactional_emulator/testbench/routed_moe/gpt_oss_router_gemm_test.py new file mode 100644 index 00000000..f6bb96c0 --- /dev/null +++ b/transactional_emulator/testbench/routed_moe/gpt_oss_router_gemm_test.py @@ -0,0 +1,462 @@ +# ruff: noqa: E402 +"""GPT-OSS layer-0 router GEMM device smoke. + +Step 3 scope only: + + X[8, 2880] @ router_W.T[2880, 32] + router_b + +Top-k remains host-side. The device path only proves the narrow-N router logits +GEMM, including the N < MLEN tail block and BF16 bias add. +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path + +import torch +import torch.nn.functional as F +from safetensors import safe_open + +_AICROSSSIM_ROOT = Path(__file__).resolve().parents[4] +_TESTBENCH_ROOT = Path(__file__).resolve().parents[1] +_ATEN_BUILD_DIR = _TESTBENCH_ROOT / "aten" / "build" +_TOP_LEVEL_COMPILER = _AICROSSSIM_ROOT / "PLENA_Compiler" +if _TOP_LEVEL_COMPILER.exists(): + sys.path.insert(0, str(_TOP_LEVEL_COMPILER)) + +from aten.models.gpt_oss.moe_reference import compare_stats +from aten.models.gpt_oss.real_layer_utils import SHARD0, cached_file, load_json +from compiler.aten.plena import PlenaCompiler +from transactional_emulator.testbench.aten.configurable import add_hw_args, setup_hw +from transactional_emulator.testbench.emulator_runner import compare_emulator_output, run_emulator +from transactional_emulator.testbench.layout_utils import prestage_bf16_vram_matrix +from transactional_emulator.testbench.sim_env_utils import create_mem_for_sim +from transactional_emulator.testbench.sliced_layer_test_builder import quantize_to_mxfp +from transactional_emulator.tools.create_sim_env import create_sim_env + + +def _bf16(x: torch.Tensor) -> torch.Tensor: + return x.to(torch.bfloat16) + + +def _load_router_tensors() -> tuple[torch.Tensor, torch.Tensor]: + with safe_open(cached_file(SHARD0), framework="pt", device="cpu") as f: + weight = f.get_tensor("model.layers.0.mlp.router.weight").to(torch.bfloat16) + bias = f.get_tensor("model.layers.0.mlp.router.bias").to(torch.bfloat16) + return weight, bias + + +def _router_projection_golden( + x: torch.Tensor, + weight_t: torch.Tensor, + bias: torch.Tensor, + *, + mlen: int, + quantize_x: bool, + quantize_weight: bool, +) -> torch.Tensor: + """Router logits with selectable MXFP8 HBM inputs and BF16 K-chunk accumulation.""" + x_q = quantize_to_mxfp(x) if quantize_x else x.to(torch.bfloat16) + w_q = quantize_to_mxfp(weight_t) if quantize_weight else weight_t.to(torch.bfloat16) + chunk = mlen * 4 + acc = None + for k_start in range(0, x_q.shape[1], chunk): + k_end = min(k_start + chunk, x_q.shape[1]) + partial = torch.matmul(x_q[:, k_start:k_end].float(), w_q[k_start:k_end, :].float()).to(torch.bfloat16) + acc = partial if acc is None else (acc.float() + partial.float()).to(torch.bfloat16) + assert acc is not None + return (acc.float() + bias.reshape(1, -1).float()).to(torch.bfloat16) + + +def _router_golden_b(x: torch.Tensor, weight_t: torch.Tensor, bias: torch.Tensor, *, mlen: int) -> torch.Tensor: + """Current PLENA-aware router logits with MXFP8 X/W HBM tensors.""" + return _router_projection_golden( + x, + weight_t, + bias, + mlen=mlen, + quantize_x=True, + quantize_weight=True, + ) + + +def _router_vector_bf16_golden( + x: torch.Tensor, + weight_rows: torch.Tensor, + bias: torch.Tensor, + *, + mlen: int, +) -> torch.Tensor: + """Golden for the BF16 vector-dot router path. + + Matches the emitted V_MUL_VV + V_RED_SUM sequence: each 64-wide product row + is quantized back to BF16, then the scalar accumulator is rounded to BF16 + after every tile reduction. + """ + rows, hidden = x.shape + num_experts = weight_rows.shape[0] + out = torch.zeros(rows, num_experts, dtype=torch.bfloat16) + for row_idx in range(rows): + for expert_idx in range(num_experts): + acc = torch.tensor(0.0, dtype=torch.bfloat16) + for k_start in range(0, hidden, mlen): + k_end = min(k_start + mlen, hidden) + product = (x[row_idx, k_start:k_end].float() * weight_rows[expert_idx, k_start:k_end].float()).to( + torch.bfloat16 + ) + partial = product.float().sum() + acc = torch.tensor(float(acc.float().item() + partial.item()), dtype=torch.bfloat16) + out[row_idx, expert_idx] = (acc.float() + bias[expert_idx].float()).to(torch.bfloat16) + return out + + +def _align_to(value: int, alignment: int) -> int: + return ((value + alignment - 1) // alignment) * alignment + + +def _vram_layout_size(physical_shape: tuple[int, int], *, mlen: int) -> int: + physical_rows, physical_cols = physical_shape + return math.ceil(physical_cols / mlen) * physical_rows * mlen + + +def _stats_dict(stats) -> dict: + return { + "rel_rms": stats.rel_rms, + "atol": stats.atol, + "rtol": stats.rtol, + "allclose": stats.allclose, + "pass_rate": stats.pass_rate, + "max_abs_error": stats.max_abs_error, + } + + +def _rank_stability(hf_logits: torch.Tensor, device_logits: torch.Tensor, top_k: int) -> dict: + hf_sorted_vals, hf_sorted_idx = torch.sort(hf_logits.float(), dim=-1, descending=True) + del hf_sorted_idx + hf_top_values, hf_topk = torch.topk(hf_logits.float(), k=top_k, dim=-1) + device_topk = torch.topk(device_logits.float(), k=top_k, dim=-1).indices + hf_sets = [set(int(v) for v in row.tolist()) for row in hf_topk.cpu()] + device_sets = [set(int(v) for v in row.tolist()) for row in device_topk.cpu()] + set_matches = [hf == device for hf, device in zip(hf_sets, device_sets, strict=True)] + order_matches = [ + hf == device for hf, device in zip(hf_topk.cpu().tolist(), device_topk.cpu().tolist(), strict=True) + ] + per_token_error = (device_logits.float() - hf_logits.float()).abs().max(dim=-1).values + rank_gap = hf_sorted_vals[:, top_k - 1] - hf_sorted_vals[:, top_k] + gap_gt_error = rank_gap > per_token_error + internal_adjacent_gaps = hf_top_values[:, :-1] - hf_top_values[:, 1:] + min_internal_gap = internal_adjacent_gaps.min(dim=-1).values + internal_order_gap_gt_error = min_internal_gap > per_token_error + order_when_internal_gap_safe = [ + bool(order_ok or not order_safe) + for order_ok, order_safe in zip(order_matches, internal_order_gap_gt_error.cpu().tolist(), strict=True) + ] + return { + "top_k": top_k, + "hf_topk_indices": hf_topk.cpu().tolist(), + "device_topk_indices": device_topk.cpu().tolist(), + "topk_order_matches_hf": bool(torch.equal(hf_topk.cpu(), device_topk.cpu())), + "topk_order_matches_hf_by_token": order_matches, + "topk_order_matches_when_internal_gap_safe_by_token": order_when_internal_gap_safe, + "topk_order_matches_when_internal_gap_safe": all(order_when_internal_gap_safe), + "topk_set_matches_hf_by_token": set_matches, + "topk_set_matches_hf": all(set_matches), + "rank4_rank5_gap": [float(v) for v in rank_gap.cpu().tolist()], + "min_internal_topk_gap": [float(v) for v in min_internal_gap.cpu().tolist()], + "device_max_abs_error_by_token": [float(v) for v in per_token_error.cpu().tolist()], + "gap_gt_measured_error_by_token": [bool(v) for v in gap_gt_error.cpu().tolist()], + "internal_order_gap_gt_measured_error_by_token": [bool(v) for v in internal_order_gap_gt_error.cpu().tolist()], + "min_gap": float(rank_gap.min().item()), + "min_internal_gap": float(min_internal_gap.min().item()), + "max_device_error": float(per_token_error.max().item()), + "passed": bool(gap_gt_error.all().item()), + } + + +def run_router_gemm(args: argparse.Namespace) -> dict: + mlen = args.mlen + blen = args.blen + build_dir = args.build_dir.expanduser().resolve() + build_dir.mkdir(parents=True, exist_ok=True) + hw = setup_hw(args, build_dir) + + reference = torch.load(args.reference_path.expanduser().resolve(), map_location="cpu") + x = reference["x"].to(torch.bfloat16) + rows, hidden = x.shape + + config = load_json("config.json") + num_experts = int(config["num_local_experts"]) + top_k = int(config["num_experts_per_tok"]) + if hidden != int(config["hidden_size"]): + raise ValueError(f"reference hidden={hidden} does not match config hidden_size={config['hidden_size']}") + if num_experts > mlen: + raise ValueError(f"router experts={num_experts} exceeds MLEN={mlen}; this test covers N