diff --git a/.gitignore b/.gitignore index b28d259f..4f5b7f48 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,9 @@ build/ doc/figure_gen/* +tools/cost_model/latency/*.png +tools/cost_model/utilisation/*.png + moduleparamid default.svf @@ -61,6 +64,7 @@ ckpts_* *.db *.db-journal co_design/interface/config.toml +target/ # behavioral simulator behavioral_simulator/target diff --git a/behavioral_simulator/Cargo.lock b/behavioral_simulator/Cargo.lock index 5053feac..55b0f1c3 100644 --- a/behavioral_simulator/Cargo.lock +++ b/behavioral_simulator/Cargo.lock @@ -146,6 +146,7 @@ dependencies = [ "tch", "tokio", "toml", + "vector_sram", ] [[package]] @@ -1437,6 +1438,15 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "vector_sram" +version = "0.1.0" +dependencies = [ + "quantize", + "tch", + "tokio", +] + [[package]] name = "version_check" version = "0.9.5" diff --git a/behavioral_simulator/Cargo.toml b/behavioral_simulator/Cargo.toml index 792ad426..d9d0d3d5 100644 --- a/behavioral_simulator/Cargo.toml +++ b/behavioral_simulator/Cargo.toml @@ -9,6 +9,7 @@ memory = { path = "lib/memory" } ramulator = { path = "lib/ramulator" } buddy = { path = "lib/buddy" } quantize = { path = "lib/quantize" } +vector_sram = { path = "lib/vector_sram" } tokio = { version = "1.45.1", features = ["macros", "rt", "rt-multi-thread", "sync"] } anyhow = "1" async-trait = "0.1.88" @@ -26,6 +27,7 @@ members = [ "lib/ramulator", "lib/buddy", "lib/quantize", + "lib/vector_sram", ] [profile.release] diff --git a/behavioral_simulator/lib/memory/src/lib.rs b/behavioral_simulator/lib/memory/src/lib.rs index ee0bba36..a28cc83b 100644 --- a/behavioral_simulator/lib/memory/src/lib.rs +++ b/behavioral_simulator/lib/memory/src/lib.rs @@ -1,6 +1,12 @@ use std::mem::ManuallyDrop; use std::sync::Mutex; +#[derive(Copy, Clone)] +pub struct Statistics { + pub total_bytes_read: u64, + pub total_bytes_written: u64, +} + #[async_trait::async_trait] pub trait MemoryTimingModel: Send + Sync { /// Read 64-bytes of memory. @@ -114,3 +120,49 @@ impl MemoryModel for WithTiming { self.data.write(addr, bytes).await } } + +// Memory model with utilization statistics +pub struct WithStats { + model: T, + statistics: Mutex, +} + +impl WithStats { + pub fn new(model: T) -> Self { + let stats = Statistics { + total_bytes_read: 0, + total_bytes_written: 0, + }; + WithStats { + model, + statistics: Mutex::new(stats), + } + } + + pub fn model(&self) -> &T { + &self.model + } + + pub fn statistics(&self) -> Statistics { + self.statistics.lock().unwrap().clone() + } +} + +#[async_trait::async_trait] +impl MemoryModel for WithStats { + async fn read(&self, addr: u64) -> [u8; 64] { + { + let mut guard = self.statistics.lock().unwrap(); + guard.total_bytes_read += 64; + } + self.model.read(addr).await + } + + async fn write(&self, addr: u64, bytes: [u8; 64]) { + { + let mut guard = self.statistics.lock().unwrap(); + guard.total_bytes_written += 64; + } + self.model.write(addr, bytes).await + } +} diff --git a/behavioral_simulator/lib/quantize/src/dtype.rs b/behavioral_simulator/lib/quantize/src/dtype.rs index eb535767..00084f51 100644 --- a/behavioral_simulator/lib/quantize/src/dtype.rs +++ b/behavioral_simulator/lib/quantize/src/dtype.rs @@ -162,9 +162,44 @@ fn test_f16() { ); } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct IntType { + pub width: u32, +} + +impl IntType { + pub const fn size_in_bits(self) -> u8 { + self.width as u8 + } + + /// Convert f32 to integer bits. Truncates the float to an integer. + pub const fn bits_from_f32(self, float: f32) -> u32 { + let int_val = float as i32; + let mask = if self.width >= 32 { + 0xFFFFFFFFu32 + } else { + ((1u64 << self.width) - 1) as u32 + }; + (int_val as u32) & mask + } + + /// Convert integer bits to f32. Interprets bits as unsigned integer. + pub const fn convert_bits_to_f32(self, bits: u32) -> f32 { + let mask = if self.width >= 32 { + 0xFFFFFFFFu32 + } else { + ((1u64 << self.width) - 1) as u32 + }; + let masked_bits = bits & mask; + masked_bits as f32 + } +} + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DataType { Fp(FpType), + Int(IntType), } impl From for DataType { @@ -173,22 +208,27 @@ impl From for DataType { } } + + impl DataType { pub fn size_in_bits(self) -> u8 { match self { DataType::Fp(fp_type) => fp_type.size_in_bits(), + DataType::Int(int_type) => int_type.size_in_bits(), } } pub const fn bits_from_f32(self, float: f32) -> u32 { match self { DataType::Fp(fp_type) => fp_type.bits_from_f32(float), + DataType::Int(int_type) => int_type.bits_from_f32(float), } } pub const fn convert_bits_to_f32(self, bits: u32) -> f32 { match self { DataType::Fp(fp_type) => fp_type.convert_bits_to_f32(bits), + DataType::Int(int_type) => int_type.convert_bits_to_f32(bits), } } @@ -259,6 +299,15 @@ impl MxDataType { MxDataType::Mx { elem, .. } => elem, } } + + /// Returns the size in bits of the element type + /// Works for both Plain (FP and Int) and Mx variants + pub fn size_in_bits(self) -> u8 { + match self { + MxDataType::Plain(data_type) => data_type.size_in_bits(), + MxDataType::Mx { elem, .. } => elem.size_in_bits(), + } + } } impl From for MxDataType { diff --git a/behavioral_simulator/lib/quantize/src/lib.rs b/behavioral_simulator/lib/quantize/src/lib.rs index 4bc0342a..45967a9d 100644 --- a/behavioral_simulator/lib/quantize/src/lib.rs +++ b/behavioral_simulator/lib/quantize/src/lib.rs @@ -1,5 +1,5 @@ mod dtype; mod tensor; -pub use dtype::{DataType, FpType, MxDataType}; +pub use dtype::{DataType, FpType, IntType, MxDataType}; pub use tensor::QuantTensor; diff --git a/behavioral_simulator/lib/runtime/src/time.rs b/behavioral_simulator/lib/runtime/src/time.rs index 333bc367..7cdb9e2c 100644 --- a/behavioral_simulator/lib/runtime/src/time.rs +++ b/behavioral_simulator/lib/runtime/src/time.rs @@ -106,6 +106,10 @@ impl Instant { /// /// Effectively means "never". pub const ETERNITY: Self = Self(u64::MAX); + + pub const fn to_secs(&self) -> f64 { + self.0 as f64 / (1_000_000_000_000_u64 as f64) + } } pub trait Deadline { diff --git a/behavioral_simulator/lib/vector_sram/Cargo.toml b/behavioral_simulator/lib/vector_sram/Cargo.toml new file mode 100644 index 00000000..ae74dc10 --- /dev/null +++ b/behavioral_simulator/lib/vector_sram/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "vector_sram" +version = "0.1.0" +edition = "2021" + +[dependencies] +tokio = { version = "1.45.1", features = ["macros", "rt", "rt-multi-thread", "sync"] } +quantize = { path = "../quantize" } +tch = { version = "0.20.0", features = ["download-libtorch"] } + diff --git a/behavioral_simulator/lib/vector_sram/src/lib.rs b/behavioral_simulator/lib/vector_sram/src/lib.rs new file mode 100644 index 00000000..b37570a1 --- /dev/null +++ b/behavioral_simulator/lib/vector_sram/src/lib.rs @@ -0,0 +1,428 @@ +use quantize::{DataType, MxDataType, QuantTensor}; +use tokio::sync::Mutex; +use tokio::sync::oneshot::Receiver; +use tch::Tensor; + +/// Vector SRAM that stores data in pure binary format with row-based storage. +/// +/// The SRAM supports two data types: +/// - FP (Floating Point): Stored as binary representation of the FP type +/// - INT (Integer): Stored as binary representation of integers +/// +/// The SRAM is organized as rows, where each row has a width of `vlen * element_size_in_bytes`. +/// During read/write operations, data is clipped to VLEN-sized vectors. +pub struct VectorSram { + /// Vector length (VLEN) - determines the size of each vector operation + vlen: u32, + /// Number of rows in the SRAM + depth: usize, + /// Data type for FP storage (used when writing QuantTensor) + fp_type: DataType, + /// Size of integer in bytes (used when writing integer vectors) + int_size_bytes: usize, + /// Raw binary storage: each row is stored as bytes + /// Row width = vlen * element_size_in_bytes + rows: Vec>, +} + +/// Represents a row of data, either ready or pending from a delayed write +enum RowData { + Ready(Vec), + Pending(Receiver), +} + +impl VectorSram { + /// Create a new Vector SRAM with given vector length, depth, and data types. + /// + /// # Arguments + /// * `vlen` - Vector length (VLEN) + /// * `depth` - Number of rows in the SRAM + /// * `fp_type` - Floating point data type for FP operations + /// * `int_size_bytes` - Size of integer in bytes (typically 4 for i32) + pub fn new( + vlen: u32, + depth: usize, + fp_type: DataType, + int_size_bytes: usize, + ) -> Self { + // Use FP type size for row width (can be changed if needed) + let element_size = fp_type.size_in_bits() as usize / 8; + let row_width = vlen as usize * element_size; + + let rows = (0..depth) + .map(|_| { + Mutex::new(RowData::Ready(vec![0u8; row_width])) + }) + .collect(); + + Self { + vlen, + depth, + fp_type, + int_size_bytes, + rows, + } + } + + /// Create a new Vector SRAM from MxDataType (for backward compatibility). + /// + /// Extracts the Plain DataType from MxDataType and uses it for FP storage. + pub fn from_mx_type( + vlen: u32, + depth: usize, + mx_type: MxDataType, + ) -> Self { + let fp_type = match mx_type { + MxDataType::Plain(dt) => dt, + MxDataType::Mx { elem, .. } => elem, + }; + Self::new(vlen, depth, fp_type, 4) // Default to 4 bytes for int (i32) + } + + /// Get the vector length (VLEN) + pub fn tile_size(&self) -> u32 { + self.vlen + } + + /// Get the data type (for backward compatibility) + pub fn ty(&self) -> MxDataType { + MxDataType::Plain(self.fp_type) + } + + /// Get the size of the SRAM in bytes + pub fn size_in_bytes(&self) -> usize { + let element_size = self.fp_type.size_in_bits() as usize / 8; + let row_width = self.vlen as usize * element_size; + row_width * self.depth + } + + /// Read a vector from the SRAM at the given address as FP (QuantTensor). + /// + /// The address must be a multiple of vlen (in element units). + /// Data is read from binary storage and converted to QuantTensor. + pub async fn read(&self, addr: u32) -> QuantTensor { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + let mut guard = self.rows[row_idx].lock().await; + + // Handle pending writes + if let RowData::Pending(ref mut receiver) = *guard { + let tensor = receiver.await.unwrap(); + let row_bytes = self.quant_tensor_to_bytes(&tensor); + *guard = RowData::Ready(row_bytes); + } + + // Read the row data + let row_bytes = match &*guard { + RowData::Ready(bytes) => bytes.clone(), + RowData::Pending(_) => unreachable!(), + }; + + // Convert from binary to QuantTensor + self.bytes_to_quant_tensor(&row_bytes, self.vlen) + } + + /// Read a vector from the SRAM at the given address as integers. + /// + /// The address must be a multiple of vlen (in element units). + /// Returns a vector of i32 values. + pub async fn read_int(&self, addr: u32) -> Vec { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + let mut guard = self.rows[row_idx].lock().await; + + // Handle pending writes (convert to bytes first) + if let RowData::Pending(ref mut receiver) = *guard { + let tensor = receiver.await.unwrap(); + let row_bytes = self.quant_tensor_to_bytes(&tensor); + *guard = RowData::Ready(row_bytes); + } + + // Read the row data + let row_bytes = match &*guard { + RowData::Ready(bytes) => bytes.clone(), + RowData::Pending(_) => unreachable!(), + }; + + // Convert from binary to integers + self.bytes_to_int_vec(&row_bytes, self.vlen) + } + + /// Write a vector to the SRAM at the given address as FP (QuantTensor). + /// + /// The address must be a multiple of vlen (in element units). + /// Data is converted from QuantTensor to binary storage. + pub async fn write(&self, addr: u32, tensor: QuantTensor) { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + // Clip to VLEN + let clipped = self.clip_to_vlen(&tensor); + + // Convert to bytes + let row_bytes = self.quant_tensor_to_bytes(&clipped); + + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + + /// Write a vector to the SRAM at the given address as integers. + /// + /// The address must be a multiple of vlen (in element units). + /// Data is converted from integers to binary storage. + pub async fn write_int(&self, addr: u32, int_vec: &[i32]) { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + assert!(int_vec.len() <= self.vlen as usize, "Vector too long"); + + // Convert integers to bytes + let row_bytes = self.int_vec_to_bytes(int_vec, self.vlen); + + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + + /// Write a vector with delayed delivery (from a channel). + pub async fn write_delayed(&self, addr: u32, tensor: Receiver) { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + *self.rows[row_idx].lock().await = RowData::Pending(tensor); + } + + /// Continuous write delayed - writes multiple rows from a single tensor. + pub async fn continous_write_delayed( + &self, + addr: u32, + write_amount: u32, + tensor: Receiver, + ) { + let start_row_idx = self.addr_to_row_idx(addr); + + // Await the tensor from the channel and extract data immediately to make it Send + let tensor = tensor.await.unwrap(); + 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 chunk_size = self.vlen as usize; + let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); + + for i in 0..num_chunks { + let row_idx = start_row_idx + i as usize; + if row_idx >= self.depth { + break; + } + + let start = (i as usize) * chunk_size; + let end = (start + chunk_size).min(total_elements); + let chunk_data = &data_vec[start..end]; + + // Pad to VLEN if needed + let mut padded_data = vec![0.0f32; chunk_size]; + let chunk_len = end - start; + 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 chunk_qt = QuantTensor::quantize(padded_tensor, MxDataType::Plain(self.fp_type)); + let row_bytes = self.quant_tensor_to_bytes(&chunk_qt); + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + } + + /// Continuous write delayed for integers - writes multiple rows from a single integer vector. + pub async fn continous_write_delayed_int( + &self, + addr: u32, + write_amount: u32, + int_vec: Receiver>, + ) { + let start_row_idx = self.addr_to_row_idx(addr); + + // Await the integer vector from the channel + let int_vec = int_vec.await.unwrap(); + // println!("addr = {:?}", addr); + // println!("in write int_vec = {:?}", int_vec); + let total_elements = int_vec.len(); + let chunk_size = self.vlen as usize; + let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); + + for i in 0..num_chunks { + let row_idx = start_row_idx + i as usize; + if row_idx >= self.depth { + break; + } + + let start = (i as usize) * chunk_size; + let end = (start + chunk_size).min(total_elements); + let chunk = &int_vec[start..end]; + + // Convert to bytes (will pad to VLEN if needed) + let row_bytes = self.int_vec_to_bytes(chunk, self.vlen); + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + } + + /// Load data from bytes into the SRAM. + /// + /// This is used for preloading the SRAM with test data. + pub async fn load_from_bytes(&self, bytes: &[u8]) { + let element_size = self.fp_type.size_in_bits() as usize / 8; + let bytes_per_element = element_size; + let total_elements = bytes.len() / bytes_per_element; + let num_rows = (total_elements + self.vlen as usize - 1) / self.vlen as usize; + + for row_idx in 0..num_rows.min(self.depth) { + let start_element = row_idx * self.vlen as usize; + let end_element = (start_element + self.vlen as usize).min(total_elements); + let elements_in_row = end_element - start_element; + + let start_byte = start_element * bytes_per_element; + let end_byte = end_element * bytes_per_element; + + // Convert bytes to f32 values + let mut vec = vec![0f32; elements_in_row]; + self.fp_type.convert_bytes_to_f32_vec(&bytes[start_byte..end_byte], &mut vec); + + // Pad with zeros if needed + if elements_in_row < self.vlen as usize { + vec.resize(self.vlen as usize, 0.0f32); + } + + // Create QuantTensor and convert to bytes + let tensor = Tensor::from_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 = RowData::Ready(row_bytes); + } + } + + /// Dump the entire SRAM content as bytes. + /// + /// This returns the raw binary representation of all stored data. + pub async fn as_bytes(&self) -> Vec { + let mut result = Vec::new(); + let mut row_idx = 0; + + for row_mutex in &self.rows { + let mut guard = row_mutex.lock().await; + + // Handle pending writes + if let RowData::Pending(ref mut receiver) = *guard { + let tensor = receiver.await.unwrap(); + let row_bytes = self.quant_tensor_to_bytes(&tensor); + *guard = RowData::Ready(row_bytes); + } + + // Read the row data + let row_bytes = match &*guard { + RowData::Ready(bytes) => bytes.clone(), + RowData::Pending(_) => unreachable!(), + }; + row_idx += 1; + result.extend_from_slice(&row_bytes); + } + + result + } + + // Helper methods + + /// Convert address (in element units) to row index + fn addr_to_row_idx(&self, addr: u32) -> usize { + assert!(addr % self.vlen == 0, "Address must be multiple of vlen"); + (addr / self.vlen) as usize + } + + /// Clip a tensor to VLEN size + fn clip_to_vlen(&self, tensor: &QuantTensor) -> QuantTensor { + let tensor_data = tensor.as_tensor(); + let len = tensor_data.size1().unwrap() as i64; + + if len <= self.vlen as i64 { + tensor.clone() + } else { + let clipped = tensor_data.narrow(0, 0, self.vlen as i64); + QuantTensor::quantize(clipped, tensor.data_type()) + } + } + + /// 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 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); + bytes + } + + /// Convert bytes to QuantTensor (FP format) + fn bytes_to_quant_tensor(&self, bytes: &[u8], expected_len: u32) -> QuantTensor { + let bytes_per_element = self.fp_type.size_in_bits() as usize / 8; + let num_elements = bytes.len() / bytes_per_element; + let actual_len = num_elements.min(expected_len as usize); + + let mut vec = vec![0f32; actual_len]; + self.fp_type.convert_bytes_to_f32_vec( + &bytes[..actual_len * bytes_per_element], + &mut vec, + ); + + // Pad to expected_len if needed + if actual_len < expected_len as usize { + vec.resize(expected_len as usize, 0.0f32); + } + + let tensor = Tensor::from_slice(&vec); + QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type)) + } + + /// Convert integer vector to bytes + fn int_vec_to_bytes(&self, int_vec: &[i32], expected_len: u32) -> Vec { + let mut bytes = Vec::with_capacity(expected_len as usize * self.int_size_bytes); + + // Write the actual integers + for &val in int_vec.iter() { + bytes.extend_from_slice(&val.to_le_bytes()); + } + + // Pad with zeros to expected_len + bytes.resize(expected_len as usize * self.int_size_bytes, 0); + + bytes + } + + /// Convert bytes to integer vector + fn bytes_to_int_vec(&self, bytes: &[u8], expected_len: u32) -> Vec { + let mut result = Vec::with_capacity(expected_len as usize); + let mut offset = 0; + + for _ in 0..expected_len as usize { + if offset + self.int_size_bytes <= bytes.len() { + let mut int_bytes = [0u8; 4]; + let copy_len = self.int_size_bytes.min(4); + int_bytes[..copy_len].copy_from_slice(&bytes[offset..offset + copy_len]); + let val = i32::from_le_bytes(int_bytes); + result.push(val); + offset += self.int_size_bytes; + } else { + result.push(0); + } + } + + result + } +} diff --git a/behavioral_simulator/readme.md b/behavioral_simulator/readme.md index bc2dccc8..72d6dc3f 100644 --- a/behavioral_simulator/readme.md +++ b/behavioral_simulator/readme.md @@ -72,3 +72,4 @@ Writes a (BLEN, BLEN) accumulator matrix (`m_accum`) to the Vector SRAM. This op - **Linear Projection Testing** (`linear`) - **RMSNorm Testing** (`rms`) - **Attention Testing** (`attn`) +- **FFN Testing** (`ffn`) diff --git a/behavioral_simulator/src/load_config.rs b/behavioral_simulator/src/load_config.rs index 4791835d..dcaf09bc 100644 --- a/behavioral_simulator/src/load_config.rs +++ b/behavioral_simulator/src/load_config.rs @@ -4,7 +4,7 @@ use std::time::Duration; use std::{fs, sync::LazyLock}; // Import the types from your main module -use quantize::{DataType, FpType, MxDataType}; +use quantize::{DataType, FpType, IntType, MxDataType}; #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ConfigValue { @@ -29,10 +29,16 @@ pub struct FpTypeConfig { pub mantissa: u8, } +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct IntTypeConfig { + pub width: u32, +} + #[derive(Debug, Serialize, Deserialize, Clone)] #[serde(tag = "type")] pub enum DataTypeConfig { Fp(FpTypeConfig), + Int(IntTypeConfig), } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -110,6 +116,8 @@ pub struct PrecisionSection { pub hbm_v_act_type: MxDataTypeConfig, #[serde(rename = "HBM_V_KV_TYPE")] pub hbm_v_kv_type: MxDataTypeConfig, + #[serde(rename = "HBM_V_INT_TYPE")] + pub hbm_v_int_type: MxDataTypeConfig, #[serde(rename = "SCALAR_FP")] pub scalar_fp: DataTypeConfig, } @@ -250,6 +258,14 @@ impl Default for AcceleratorConfig { }), }, }, + hbm_v_int_type: MxDataTypeConfig { + format: "Plain".to_string(), + data: MxDataTypeData::Plain { + data_type: DataTypeConfig::Int(IntTypeConfig { + width: 32, + }), + }, + }, scalar_fp: DataTypeConfig::Fp(FpTypeConfig { sign: true, exponent: 8, @@ -333,10 +349,19 @@ impl From for FpType { } } +impl From for IntType { + fn from(config: IntTypeConfig) -> Self { + IntType { + width: config.width, + } + } +} + impl From for DataType { fn from(config: DataTypeConfig) -> Self { match config { DataTypeConfig::Fp(fp_config) => DataType::Fp(fp_config.into()), + DataTypeConfig::Int(int_config) => DataType::Int(int_config.into()), } } } @@ -446,6 +471,10 @@ pub fn vector_kv_type() -> MxDataType { CONFIG.precision.hbm_v_kv_type.clone().into() } +pub fn vector_int_type() -> MxDataType { + CONFIG.precision.hbm_v_int_type.clone().into() +} + // Additional accessor functions for new parameters pub fn mlen() -> u32 { CONFIG.config.mlen.value diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 80546f5c..ecebeb7f 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -14,9 +14,10 @@ use futures::StreamExt; use futures::stream::FuturesUnordered; use half::f16; use memory::MemoryModel; -use quantize::{MxDataType, QuantTensor}; +use quantize::{DataType, MxDataType, QuantTensor}; use runtime::{Duration, Executor, Instant}; use tch::{IndexOp, Tensor}; +use vector_sram::VectorSram; use tokio::sync::Mutex; use tokio::sync::oneshot::{self, Receiver}; @@ -56,6 +57,7 @@ static MATRIX_WEIGHT_TYPE: LazyLock = LazyLock::new(|| matrix_weight static MATRIX_KV_TYPE: LazyLock = LazyLock::new(|| matrix_kv_type()); static VECTOR_ACTIVATION_TYPE: LazyLock = LazyLock::new(|| vector_activation_type()); static VECTOR_KV_TYPE: LazyLock = LazyLock::new(|| vector_kv_type()); +static VECTOR_INT_TYPE: LazyLock = LazyLock::new(|| vector_int_type()); static PREFETCH_M_AMOUNT: LazyLock = LazyLock::new(|| hbm_m_prefetch_amount()); static PREFETCH_V_AMOUNT: LazyLock = LazyLock::new(|| hbm_v_prefetch_amount()); static WRITEBACK_V_AMOUNT: LazyLock = LazyLock::new(|| hbm_v_writeback_amount()); @@ -199,143 +201,7 @@ impl MatrixSram { } } -/// Behaviour modelling of vector SRAM. -/// -/// The timing aspect is to be considered by the matrix and vector machines themselves. -struct VectorSram { - tile_size: u32, - tiles: Vec>>>, - ty: MxDataType, -} - -impl VectorSram { - /// Creata a matrix SRAM with given tile size and depth. - fn new(tile_size: u32, depth: usize, ty: MxDataType) -> Self { - let tiles = (0..depth) - .map(|_| Mutex::new(Ok(QuantTensor::zeros(tile_size as usize, ty)))) - .collect(); - Self { - tile_size, - tiles, - ty, - } - } - - fn size_in_bytes(&self) -> usize { - self.tile_size as usize * self.tiles.len() - } - - async fn read(&self, addr: u32) -> QuantTensor { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - - let mut guard = self.tiles[addr_in_tiles as usize].lock().await; - if let Err(ref mut fut) = *guard { - *guard = Ok(fut.await.unwrap()); - } - - guard.as_ref().map_err(|_| ()).unwrap().clone() - } - - async fn write(&self, addr: u32, tensor: QuantTensor) { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - - assert_eq!(tensor.data_type(), self.ty); - *self.tiles[addr_in_tiles as usize].lock().await = Ok(tensor); - } - - async fn write_delayed(&self, addr: u32, tensor: Receiver) { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - - *self.tiles[addr_in_tiles as usize].lock().await = Err(tensor); - } - - async fn continous_write_delayed( - &self, - addr: u32, - write_amount: u32, - tensor: Receiver, - ) { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - // Await the tensor from the channel (blocks until data arrives) - if let Ok(tensor) = tensor.await { - let dims = tensor.as_tensor().size(); - let chunk_size = self.tile_size as i64; - let total = dims[0]; - - // Split the tensor into chunks of self.tile_size and store each in self.tiles. - for i in 0..write_amount.min((total as u32 + self.tile_size - 1) / self.tile_size) { - let start = (i as i64) * chunk_size; - let end = ((i as i64 + 1) * chunk_size).min(total); - let chunk = tensor - .as_tensor() - .narrow(0, start, end - start) - .shallow_clone(); - let chunk_qt = QuantTensor::quantize(chunk, self.ty); - *self.tiles[(addr_in_tiles + i) as usize].lock().await = Ok(chunk_qt); - } - } - } - - /// TODO: used to preload Vector SRAM to facilitate testing. - async fn load_from_bytes(&self, bytes: &[u8]) { - let element_ty = self.ty.element_type(); - let element_bits = element_ty.size_in_bits(); - let bytes_per_element = (element_bits / 8) as usize; - - // Total number of elements that can be loaded - let total_elements = bytes.len() / bytes_per_element; - let tile_size = self.tile_size as usize; - let num_tiles = (total_elements + tile_size - 1) / tile_size; // Round up - - for tile_idx in 0..num_tiles.min(self.tiles.len()) { - let start_element = tile_idx * tile_size; - let end_element = (start_element + tile_size).min(total_elements); - let elements_in_tile = end_element - start_element; - - let start_byte = start_element * bytes_per_element; - let end_byte = end_element * bytes_per_element; - - // Convert bytes to f32 values - let mut vec = vec![0f32; elements_in_tile]; - element_ty.convert_bytes_to_f32_vec(&bytes[start_byte..end_byte], &mut vec); - - // Pad with zeros if needed - if elements_in_tile < tile_size { - vec.resize(tile_size, 0.0f32); - } - - // Create QuantTensor and store it - let tensor = tch::Tensor::from_slice(&vec); - let quant_tensor = QuantTensor::quantize(tensor, self.ty); - *self.tiles[tile_idx].lock().await = Ok(quant_tensor); - } - } - - async fn as_bytes(&self) -> Vec { - let element_ty = self.ty.element_type(); - let mut result = Vec::new(); - - for tile_mutex in &self.tiles { - let mut guard = tile_mutex.lock().await; - if let Err(ref mut fut) = *guard { - *guard = Ok(fut.await.unwrap()); - } - let tensor = guard.as_ref().map_err(|_| ()).unwrap(); - 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) }; - // 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); - result.extend_from_slice(&tile_bytes); - } - - result - } -} +// VectorSram is now imported from the vector_sram library struct MatrixMachine { mram: Arc, @@ -357,9 +223,9 @@ impl MatrixMachine { let (mat_base, mat_offset) = m_addr.multiple_and_offset(self.mlen * self.mlen); // println!("mat_offset = {:?}", mat_offset); // println!("mat_base = {:?}", mat_base); - assert!(mat_offset.is_multiple_of(self.mlen)); - let mat_row_offset = mat_offset as i64 / self.mlen as i64; - + assert!(mat_offset.is_multiple_of(self.blen)); + assert!(mat_offset <= self.mlen); + let mat_row_offset = mat_offset as i64; let full_mat = self.mram.read(mat_base).await; // Slice columns instead of rows: [mlen, blen] let mat = full_mat @@ -379,19 +245,21 @@ impl MatrixMachine { } // Stack along dimension 0 to get [blen, mlen] let vec = tch::Tensor::stack(&tensors, 0); + // println!("vec = {}", vec); + // println!("mat = {}", mat); // Now vec @ mat: [blen, mlen] @ [mlen, blen] = [blen, blen] self.m_accum += vec.matmul(&mat); } async fn bmm(&mut self, m_addr: u32, v_addr: u32, stride_len: u32, bmm_scale: f32) { - println!("m_addr = {:?}", m_addr); - println!("v_addr = {:?}", v_addr); + // println!("m_addr = {:?}", m_addr); + // println!("v_addr = {:?}", v_addr); assert!(self.broadcast_amount * self.hlen == self.mlen); // Load matrix from matrix SRAM. let (mat_base, mat_offset) = m_addr.multiple_and_offset(self.mlen * self.blen); let (mat_offset, head_offset) = mat_offset.multiple_and_offset(self.mlen); - println!("mat_offset = {:?}", mat_offset); + // println!("mat_offset = {:?}", mat_offset); assert!(mat_offset.is_multiple_of(self.blen)); assert!(head_offset.is_multiple_of(self.hlen)); let full_mat = self.mram.read(mat_base).await; @@ -534,17 +402,17 @@ impl MatrixMachine { let (vec_base, vec_offset) = v_addr.multiple_and_offset(self.mlen); assert!(vec_offset.is_multiple_of(self.blen)); cycle!(1); - println!("======================== MM_WO =========================="); - println!("m accum = {}", self.m_accum); - println!("vec_base = {}, vec_offset = {}, stride_len = {}", vec_base, vec_offset, stride_len); + // println!("======================== MM_WO =========================="); + // println!("m accum = {}", self.m_accum); + // println!("vec_base = {}, vec_offset = {}, stride_len = {}", vec_base, vec_offset, stride_len); for i in 0..self.blen { let tensor = self.m_accum.i((i as i64, ..)); let old = self.vram.read(vec_base + i * self.mlen * stride_len).await; - println!("old = {}", old.as_tensor()); + // println!("old = {}", old.as_tensor()); let new = old.as_tensor().copy(); new.i(vec_offset as i64..(vec_offset + self.blen) as i64) .copy_(&tensor); - println!("new = {}", new); + // println!("new = {}", new); self.vram .write( vec_base + i * self.mlen * stride_len, @@ -561,13 +429,13 @@ impl MatrixMachine { async fn bmm_wo(&mut self, v_addr: u32) { let (vec_base, vec_offset) = v_addr.multiple_and_offset(self.mlen); - println!("======================== BMM_WO =========================="); + // println!("======================== BMM_WO =========================="); assert!(vec_offset.is_multiple_of(self.mlen)); cycle!(1); for j in 0..self.broadcast_amount { for i in 0..self.mlen { let tensor = self.h_accum.i((j as i64, i as i64, ..)); - self.vram.write(vec_base + (j * self.mlen + i) * self.mlen, QuantTensor::quantize(tensor, self.vram.ty)).await; + self.vram.write(vec_base + (j * self.mlen + i) * self.mlen, QuantTensor::quantize(tensor, self.vram.ty())).await; } } self.h_accum = Tensor::zeros( @@ -611,7 +479,7 @@ impl MatrixMachine { } async fn mv_wo(&mut self, v_addr: u32) { - let quant = QuantTensor::quantize(self.v_accum.shallow_clone(), self.vram.ty); + let quant = QuantTensor::quantize(self.v_accum.shallow_clone(), self.vram.ty()); self.vram.write(v_addr, quant).await; cycle!(1); self.v_accum = Tensor::zeros([self.mlen as i64], (tch::Kind::Float, tch::Device::Cpu)); @@ -653,6 +521,44 @@ impl VectorMachine { } } + async fn sub_scalar(&mut self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32, rorder: op::VectorOrder) { + let a = self.vram.read(vs1).await; + if rmask == 0 { + if matches!(rorder, op::VectorOrder::Normal) { + let c = QuantTensor::quantize(a.as_tensor() - (f as f64), a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } else { + let c = QuantTensor::quantize((f as f64) - a.as_tensor(), a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } + } else { + // mask is a bitmask; each bit controls whether to apply 'f' to corresponding mask_unit-section + let mut 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 { + // Mask is set for this head + 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 = if matches!(rorder, op::VectorOrder::Normal) { + &sliced - (f as f64) + } else { + (f as f64) - &sliced + }; + // Overwrite this section with calculated values + result.narrow(0, start, end - start).copy_(&updated); + } + // else leave unchanged + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } + } + async fn mul_scalar(&mut self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; if rmask == 0 { @@ -660,10 +566,10 @@ impl VectorMachine { cycle!(*VECTOR_MUL_CYCLES); self.vram.write(vd, c).await; } else { - println!("======================== V_ADD_VF =========================="); - println!("add: mask = {:?}", mask); - println!("a = {}", a.as_tensor()); - println!("f = {}", f); + // println!("======================== V_ADD_VF =========================="); + // println!("add: mask = {:?}", mask); + // println!("a = {}", a.as_tensor()); + // println!("f = {}", f); let mut result = a.as_tensor().shallow_clone(); let total_heads = self.tile_size / self.mask_unit; for head in 0..total_heads { @@ -689,10 +595,10 @@ impl VectorMachine { cycle!(*VECTOR_ADD_CYCLES); self.vram.write(vd, c).await; } else { - println!("======================== V_ADD =========================="); - println!("add: mask = {:?}", mask); - println!("a = {}", a.as_tensor()); - println!("b = {}", b.as_tensor()); + // println!("======================== V_ADD =========================="); + // println!("add: mask = {:?}", mask); + // println!("a = {}", a.as_tensor()); + // println!("b = {}", b.as_tensor()); let mut result = a.as_tensor().shallow_clone(); let total_heads = self.tile_size / self.mask_unit; for head in 0..total_heads { @@ -713,44 +619,111 @@ impl VectorMachine { async fn sub(&mut self, vd: u32, vs1: u32, vs2: u32, rmask: u8, mask: u32) { let (a, b) = tokio::join!(self.vram.read(vs1), self.vram.read(vs2)); - let c = QuantTensor::quantize(a.as_tensor() - b.as_tensor(), a.data_type()); - cycle!(*VECTOR_ADD_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor() - b.as_tensor(), a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut 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 - b.as_tensor().narrow(0, start, end - start); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } } async fn mul(&mut self, vd: u32, vs1: u32, vs2: u32, rmask: u8, mask: u32) { let (a, b) = tokio::join!(self.vram.read(vs1), self.vram.read(vs2)); - let c = QuantTensor::quantize(a.as_tensor() * b.as_tensor(), a.data_type()); - cycle!(*VECTOR_MUL_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor() * b.as_tensor(), a.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut 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 * b.as_tensor().narrow(0, start, end - start); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } } async fn exp(&mut self, vd: u32, vs1: u32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; - let c = QuantTensor::quantize(a.as_tensor().exp(), a.data_type()); - cycle!(*VECTOR_EXP_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor().exp(), a.data_type()); + cycle!(*VECTOR_EXP_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut 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.exp(); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_EXP_CYCLES); + self.vram.write(vd, c).await; + } } async fn reciprocal(&mut self, vd: u32, vs1: u32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; - let c = QuantTensor::quantize(a.as_tensor().reciprocal(), a.data_type()); - cycle!(*VECTOR_RECI_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor().reciprocal(), a.data_type()); + cycle!(*VECTOR_RECI_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut 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.reciprocal(); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_RECI_CYCLES); + self.vram.write(vd, c).await; + } } - // async fn broadcast(&mut self, vd: u32, f: f32) { - // let c = QuantTensor::quantize( - // Tensor::full( - // [self.vram.tile_size as i64], - // f as f64, - // (tch::Kind::Float, tch::Device::Cpu), - // ), - // self.vram.ty, - // ); - // cycle!(*VECTOR_BASIC_CYCLES); - // self.vram.write(vd, c).await; - // } + async fn vector_transfer_fp(&mut self, vd: u32, f: &[f16]) { + assert_eq!(f.len(), self.vram.tile_size() as usize, "Input vector length must match tile_size"); + // Convert f16 slice to f32 vector + let f32_vec: Vec = f.iter().map(|x| f32::from(*x)).collect(); + // Create tensor from f32 vector + let tensor = tch::Tensor::from_slice(&f32_vec); + // Quantize the tensor according to vram data type + let c = QuantTensor::quantize(tensor, self.vram.ty()); + cycle!(*VLEN); + self.vram.write(vd, c).await; + } async fn reduce_sum(&mut self, vs1: u32, f: f32, rmask: u8, mask: u32) -> f32 { let a = self.vram.read(vs1).await; @@ -799,7 +772,7 @@ impl Accelerator { /// - stride: Byte offset between consecutive loads /// - load_dim: Number of elements to load per iteration /// - load_amount: Number of strided loads to perform - fn transfer_from_hbm( + fn transfer_mx_from_hbm( &mut self, index: u64, scale_index: u64, @@ -826,10 +799,6 @@ impl Accelerator { } else { load_dim }; - println!("Call transfer_from_hbm"); - println!("stride = {:?}", stride); - println!("index = {:?}", index); - println!("scale_index = {:?}", scale_index); Executor::current().spawn(async move { let element_ty = hbm_type.element_type(); @@ -896,7 +865,7 @@ impl Accelerator { let load_iter = write_idx * write_amount + block_idx; let element_addr = index + (load_iter * stride) as u64; let scale_addr = scale_index + (load_iter as f32 * stride_scale) as u64; - println!("element_addr = {:?}, scale_addr = {:?}", element_addr, scale_addr); + // println!("element_addr = {:?}, scale_addr = {:?}", element_addr, scale_addr); let byte_offset = (write_idx * write_amount * len_in_bytes_per_load) as usize + block_idx as usize * len_in_bytes_per_load as usize; let scale_byte_offset = (write_idx * write_amount * scale_len_in_bytes_per_load) @@ -1019,6 +988,112 @@ impl Accelerator { receiver } + fn transfer_int_from_hbm( + &mut self, + index: u64, + hbm_type: MxDataType, + sram_type: MxDataType, + rstride: u8, + load_dim: u32, + load_amount: u32, + ) -> Receiver> { + // Transfer integer data from HBM: loads elements based on index and stride + // Converts bytes directly to Vec for writing to vector SRAM + + let (sender, receiver) = oneshot::channel(); + + let hbm_clone = self.hbm.clone(); + let stride = if rstride == 1 { + self.reg_file.stride + } else { + load_dim + }; + + Executor::current().spawn(async move { + let element_bits = hbm_type.size_in_bits(); + assert!(element_bits.is_power_of_two()); + + let len_in_bits_per_load = element_bits as u32 * load_dim; + assert!(len_in_bits_per_load.is_multiple_of(8 * 64)); + let len_in_bytes_per_load = len_in_bits_per_load / 8; + + // Total elements/bytes for all loads: + let total_elements = (load_dim * load_amount) as usize; + let total_bytes = (len_in_bytes_per_load * load_amount) as usize; + + let mut bytes = vec![0u8; total_bytes]; + let hbm_clone = &hbm_clone; + + enum ChunkType { + Element(usize, [u8; 64], usize), // (offset, data, size) + } + let mut futures = + FuturesUnordered:: + Send>>>::new(); + + // Gather all HBM reads + for load_idx in 0..load_amount { + let element_addr = index + (load_idx * stride) as u64; + let byte_offset = (load_idx * len_in_bytes_per_load) as usize; + + // Element chunks: + for i in 0..(len_in_bytes_per_load as usize + 63) / 64 { + let chunk_offset = byte_offset + i * 64; + let chunk_size = std::cmp::min(64, total_bytes - chunk_offset); + let addr = element_addr + (i * 64) as u64; + assert!(addr.is_multiple_of(64)); + futures.push(Box::pin(async move { + let data = hbm_clone.read(addr).await; + ChunkType::Element(chunk_offset, data, chunk_size) + })); + } + } + + // Collect all HBM reads + while let Some(chunk_result) = futures.next().await { + match chunk_result { + ChunkType::Element(offset, data, size) => { + bytes[offset..offset + size].copy_from_slice(&data[..size]); + } + } + } + + // Convert bytes to Vec + // For integer types, interpret bytes as little-endian integer values + let element_size_bytes = element_bits as usize / 8; + let mut int_vec = Vec::with_capacity(total_elements); + + for i in 0..total_elements { + let byte_offset = i * element_size_bytes; + if byte_offset + element_size_bytes <= bytes.len() { + // Read bytes as little-endian and convert to i32 + let int_value = match element_size_bytes { + 1 => bytes[byte_offset] as i8 as i32, + 2 => { + let bytes_slice = &bytes[byte_offset..byte_offset + 2]; + i16::from_le_bytes([bytes_slice[1], bytes_slice[0]]) as i32 + } + 4 => { + let bytes_slice = &bytes[byte_offset..byte_offset + 4]; + // println!("bytes_slice = {:?}", bytes_slice.iter().map(|b| format!("{:02x}", b)).collect::>()); + i32::from_le_bytes([ + bytes_slice[3], + bytes_slice[2], + bytes_slice[1], + bytes_slice[0], + ]) + } + _ => panic!("Unsupported integer size: {} bytes (must be 1, 2, or 4)", element_size_bytes), + }; + int_vec.push(int_value); + } + } + let _ = sender.send(int_vec); + }); + + receiver + } + + async fn do_ops(&mut self, ops: &[op::Opcode]) { for op in ops { println!("execute op = {:?}", op); @@ -1140,15 +1215,16 @@ impl Accelerator { ) .await; } - op::Opcode::V_SUB_VF { rd, rs1, rs2, rmask } => { + op::Opcode::V_SUB_VF { rd, rs1, rs2, rmask, rorder} => { let mask = if rmask == 0 { (1 << *HLEN as u32) - 1 } else {self.reg_file.v_mask }; self.v_machine - .add_scalar( + .sub_scalar( self.reg_file.gp_reg[rd as usize], self.reg_file.gp_reg[rs1 as usize], - (-self.reg_file.fp_reg[rs2 as usize]).into(), + self.reg_file.fp_reg[rs2 as usize].into(), rmask, mask, + rorder, ) .await; } @@ -1285,8 +1361,18 @@ impl Accelerator { self.reg_file.fp_reg[rd as usize]; cycle!(1); } - op::Opcode::S_MAP_V_FP { rd, rs1, imm } => todo!(), - + op::Opcode::S_MAP_V_FP { rd, rs1, imm } => { + let start_idx = (self.reg_file.gp_reg[rs1 as usize] + imm) as usize; + let end_idx = start_idx + *VLEN as usize; + let f = &self.fpsram[start_idx..end_idx]; + self.v_machine + .vector_transfer_fp( + self.reg_file.gp_reg[rd as usize], + f, + ) + .await; + cycle!(*VLEN); + } op::Opcode::S_ADD_INT { rd, rs1, rs2 } => { self.reg_file.gp_reg[rd as usize] = self.reg_file.gp_reg[rs1 as usize] .wrapping_add(self.reg_file.gp_reg[rs2 as usize]); @@ -1343,7 +1429,7 @@ impl Accelerator { / (elem.size_in_bits() as u32 * block / scale.size_in_bits() as u32) } // Element addr shifted by (element to scale ratio) }; - let xfer = self.transfer_from_hbm( + let xfer = self.transfer_mx_from_hbm( addr + offset as u64, addr + self.reg_file.scale as u64 + scale as u64, dtype, @@ -1369,39 +1455,64 @@ impl Accelerator { rs2, rstride, precision, + loadtype, } => { // TODO: rstride support to be added let offset = self.reg_file.gp_reg[rs1 as usize]; let addr = self.reg_file.hbm_addr_reg[rs2 as usize]; + if matches! (loadtype, op::HBM_LOAD_TYPE::MX) { + let dtype = match precision { + op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, + op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, + op::VectorPrecision::INT => *VECTOR_INT_TYPE, + }; + + let scale = match dtype { + MxDataType::Plain(_) => 0, + MxDataType::Mx { elem, scale, block } => { + offset + / (elem.size_in_bits() as u32 * block / scale.size_in_bits() as u32) + } + }; + let xfer = self.transfer_mx_from_hbm( + addr + offset as u64, + addr + self.reg_file.scale as u64 + scale as u64, + dtype, + self.v_machine.vram.ty(), + rstride, + *VLEN, + *PREFETCH_V_AMOUNT, + 1, + ); + + let dest = self.reg_file.gp_reg[rd as usize]; + self.v_machine + .vram + .continous_write_delayed(dest, *PREFETCH_V_AMOUNT, xfer) + .await; + } else { + let dtype = match precision { + // TODO: Left for future support FP load. + op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, + op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, + op::VectorPrecision::INT => *VECTOR_INT_TYPE, + }; + let xfer = self.transfer_int_from_hbm( + addr + offset as u64, + dtype, + self.v_machine.vram.ty(), + rstride, + *VLEN, + *PREFETCH_V_AMOUNT, + ); + let dest = self.reg_file.gp_reg[rd as usize]; + self.v_machine + .vram + .continous_write_delayed_int(dest, *PREFETCH_V_AMOUNT, xfer) + .await; + } - let dtype = match precision { - op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, - op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, - }; - - let scale = match dtype { - MxDataType::Plain(_) => 0, - MxDataType::Mx { elem, scale, block } => { - offset - / (elem.size_in_bits() as u32 * block / scale.size_in_bits() as u32) - } - }; - let xfer = self.transfer_from_hbm( - addr + offset as u64, - addr + self.reg_file.scale as u64 + scale as u64, - dtype, - self.v_machine.vram.ty, - rstride, - *VLEN, - *PREFETCH_V_AMOUNT, - 1, - ); - - let dest = self.reg_file.gp_reg[rd as usize]; - self.v_machine - .vram - .continous_write_delayed(dest, *PREFETCH_V_AMOUNT, xfer) - .await; + } op::Opcode::H_STORE_V { rd, @@ -1460,7 +1571,7 @@ struct Opts { async fn start() { let opts = Opts::parse(); let mram = Arc::new(MatrixSram::new(*MLEN, *MATRIX_SRAM_SIZE, *MATRIX_SRAM_TYPE)); // Matrix SRAM - let vram = Arc::new(VectorSram::new(*VLEN, *VECTOR_SRAM_SIZE, *VECTOR_SRAM_TYPE)); // Vector SRAM + let vram = Arc::new(VectorSram::from_mx_type(*VLEN, *VECTOR_SRAM_SIZE, *VECTOR_SRAM_TYPE)); // Vector SRAM let machine = MatrixMachine { mram, vram: vram.clone(), @@ -1480,10 +1591,10 @@ async fn start() { }; let v_machine = VectorMachine { vram, tile_size: *VLEN, mask_unit: *HLEN }; // Share same dim with VSRAM - let hbm = Arc::new(memory::WithTiming::new( + let hbm = Arc::new(memory::WithStats::new(memory::WithTiming::new( ManuallyDrop::new(ramulator::Ramulator::hbm2_preset(8).unwrap()), memory::MemoryBacked::with_capacity(*HBM_SIZE), - )); + ))); let mut accelerator = Accelerator { m_machine: machine, @@ -1506,7 +1617,7 @@ async fn start() { use std::fs; let op_file = fs::read_to_string(opts.opcode).unwrap(); - eprintln!("Loaded opcode file: {:?}", op_file); + // eprintln!("Loaded opcode file: {:?}", op_file); let op: Vec = op_file .split_whitespace() // split by spaces/newlines @@ -1521,8 +1632,7 @@ async fn start() { // Memory Initialization // - HBM Preload let hbm_data = std::fs::read(opts.hbm).unwrap(); - - hbm.data().with_data(|f| { + hbm.model().data().with_data(|f| { f[..hbm_data.len()].copy_from_slice(&hbm_data); }); @@ -1558,11 +1668,13 @@ async fn start() { } // - Execute Instructions - accelerator - .do_ops(&dbg!( - op.into_iter().map(op::Opcode::decode).collect::>() - )) - .await; + // accelerator + // .do_ops(&dbg!( + // op.into_iter().map(op::Opcode::decode).collect::>() + // )) + // .await; + let decoded_ops = op.into_iter().map(op::Opcode::decode).collect::>(); + accelerator.do_ops(&decoded_ops).await; println!("gp1 = {:x}", accelerator.reg_file.gp_reg[1]); println!("scale = {}", accelerator.reg_file.scale); @@ -1574,7 +1686,7 @@ async fn start() { "Matrix SRAM Contents: \n {}", accelerator.m_machine.mram.read(0x0000).await.as_tensor() ); - println!("FP SRAM Contents: \n {:?}", accelerator.fpsram); + // println!("FP SRAM Contents: \n {:?}", accelerator.fpsram); // Dump MRAM let mram_dump_path = "mram_dump.bin"; @@ -1589,6 +1701,14 @@ async fn start() { let mut vram_file = std::fs::File::create(vram_dump_path).unwrap(); vram_file.write_all(&vram_bytes).unwrap(); eprintln!("Dumped VRAM content to: {:?}", vram_dump_path); + + let memory_stats = hbm.statistics(); + let utilization = (memory_stats.total_bytes_read + memory_stats.total_bytes_written) as f64 + / Executor::current().now().to_secs(); + eprintln!( + "HBM Statistics - Bytes read: {:?} | Bytes written: {:?} | Utilization: {:.2e} bytes/sec", + memory_stats.total_bytes_read, memory_stats.total_bytes_written, utilization + ); } #[tokio::main] diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index 81446b54..cc49ccd8 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -8,8 +8,22 @@ pub enum MatrixPrecision { pub enum VectorPrecision { Activation, KeyValue, + INT, } +#[derive(Debug, Clone, Copy)] +pub enum VectorOrder { + Normal, + Reverse, +} + +#[derive(Debug, Clone, Copy)] +pub enum HBM_LOAD_TYPE { + MX, + Normal, +} + + #[allow(non_camel_case_types)] #[derive(Debug)] pub enum Opcode { @@ -90,6 +104,7 @@ pub enum Opcode { rs1: u8, rs2: u8, rmask: u8, + rorder: VectorOrder, }, V_MUL_VV { rd: u8, @@ -224,6 +239,7 @@ pub enum Opcode { rs2: u8, rstride: u8, precision: VectorPrecision, + loadtype: HBM_LOAD_TYPE, }, H_STORE_V { rd: u8, @@ -260,6 +276,44 @@ const fn mask(width: u32) -> u32 { } impl Opcode { + #[inline] + fn matrix_precision_from(funct1: u8) -> MatrixPrecision { + if funct1 == 0 { + MatrixPrecision::Weights + } else { + MatrixPrecision::KeyValue + } + } + + #[inline] + fn vector_precision_from(funct1: u8) -> VectorPrecision { + if funct1 == 0 { + VectorPrecision::Activation + } else if funct1 == 1 { + VectorPrecision::KeyValue + } else { + VectorPrecision::INT + } + } + + #[inline] + fn vector_order_from(funct1: u8) -> VectorOrder { + if funct1 == 0 { + VectorOrder::Normal + } else { + VectorOrder::Reverse + } + } + + #[inline] + fn hbm_load_type_from(funct2: u8) -> HBM_LOAD_TYPE { + if funct2 == 0 { + HBM_LOAD_TYPE::MX + } else { + HBM_LOAD_TYPE::Normal + } + } + pub fn decode(instr: u32) -> Self { // eprintln!( // "decode(): instr = 0x{instr:08X} ({instr:032b})" @@ -269,6 +323,8 @@ impl Opcode { let rs1 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH)) & mask(OPERAND_WIDTH)) as u8; let rs2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 2)) & mask(OPERAND_WIDTH)) as u8; let rs3 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 3)) & mask(OPERAND_WIDTH)) as u8; + let funct1 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 4)) & mask(OPERAND_WIDTH)) as u8; + let funct2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 5)) & mask(OPERAND_WIDTH)) as u8; let imm = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH)) & mask(IMM_WIDTH)) as u32; let imm2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 2)) & mask(IMM_2_WIDTH)) as u32; @@ -292,7 +348,7 @@ impl Opcode { 0x0D => Self::V_ADD_VV { rd, rs1, rs2, rmask: rs3 }, 0x0E => Self::V_ADD_VF { rd, rs1, rs2, rmask: rs3 }, 0x0F => Self::V_SUB_VV { rd, rs1, rs2, rmask: rs3 }, - 0x10 => Self::V_SUB_VF { rd, rs1, rs2, rmask: rs3 }, + 0x10 => Self::V_SUB_VF { rd, rs1, rs2, rmask: rs3, rorder: Self::vector_order_from(funct1) }, 0x11 => Self::V_MUL_VV { rd, rs1, rs2, rmask: rs3 }, 0x12 => Self::V_MUL_VF { rd, rs1, rs2, rmask: rs3 }, 0x13 => Self::V_EXP_V { rd, rs1, rmask: rs3 }, @@ -326,7 +382,7 @@ impl Opcode { rs1, rs2, rstride: rs3, - precision: MatrixPrecision::Weights, + precision: Self::matrix_precision_from(funct1), }, // 0x29 => Self::H_PREFETCH_M { rd, rs1, rs2, rstride: rs3, precision: MatrixPrecision::KeyValue }, 0x29 => Self::H_PREFETCH_V { @@ -334,7 +390,8 @@ impl Opcode { rs1, rs2, rstride: rs3, - precision: VectorPrecision::KeyValue, + precision: Self::vector_precision_from(funct1), + loadtype: Self::hbm_load_type_from(funct2), }, // 0x2A => Self::H_PREFETCH_V { rd, rs1, rs2, rstride: rs3, precision: VectorPrecision::KeyValue }, 0x2A => Self::H_STORE_V { @@ -342,7 +399,7 @@ impl Opcode { rs1, rs2, rstride: rs3, - precision: VectorPrecision::Activation, + precision: Self::vector_precision_from(funct1), }, // 0x2B => Self::H_STORE_V { rd, rs1, rs2, rstride: rs3, precision: VectorPrecision::KeyValue }, 0x2B => Self::C_SET_ADDR_REG { rd, rs1, rs2 }, @@ -356,4 +413,4 @@ impl Opcode { } } } -} +} \ No newline at end of file diff --git a/behavioral_simulator/testbench/attn_test.py b/behavioral_simulator/testbench/attn_test.py index 3fc54f38..d05b5630 100644 --- a/behavioral_simulator/testbench/attn_test.py +++ b/behavioral_simulator/testbench/attn_test.py @@ -14,7 +14,7 @@ from create_sim_env import create_sim_env from transformers.modeling_flash_attention_utils import _flash_attention_forward as _flash_attention_forward_ref from aria_lm_ops.models.llama import flash_attn2_gemv -from sim_env_utils import build_fake_sim_env +from sim_env_utils import create_mem_for_sim # Reshape q, k, v so that each batch's data stays together and rows are of length mem_row_size @@ -145,4 +145,4 @@ ) create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = ["q", "k", "v"]) \ No newline at end of file + create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = ["q", "k", "v"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/bmm_test.py b/behavioral_simulator/testbench/bmm_test.py index af19f0e4..0a123b06 100644 --- a/behavioral_simulator/testbench/bmm_test.py +++ b/behavioral_simulator/testbench/bmm_test.py @@ -6,9 +6,9 @@ from torch import Tensor, nn # from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ from test_data_gen import get_weights_path, generate_and_save_random_weights -from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm +from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm, reset_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import create_mem_for_sim if __name__ == "__main__": # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) @@ -44,6 +44,10 @@ addr_reg_val=[int(m * k * batch_size * real_data_ratio)] ) + gen_assembly_code += reset_reg_asm( + alive_registers=[1] + ) + gen_assembly_code += batched_matmul_asm( mlen=mlen, blen=blen, @@ -61,7 +65,7 @@ create_sim_env(input_tensor, weight_2_tensor, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/create_sim_env.py b/behavioral_simulator/testbench/create_sim_env.py index c13a157d..f87da7d7 100644 --- a/behavioral_simulator/testbench/create_sim_env.py +++ b/behavioral_simulator/testbench/create_sim_env.py @@ -14,7 +14,7 @@ def np_array_to_str_2f(arr): # For higher dimensions, default to numpy's print (rare for this context) return np.array2string(arr, formatter={'float_kind':lambda x: "%.2f" % x}) -def create_sim_env(input_tensor, input_weight, generated_code, golden_result, fp_preload = None, int_preload = None): +def create_sim_env(input_tensor, generated_code, golden_result, fp_preload = None, int_preload = None): build_dir = os.path.join(os.path.dirname(__file__), "build") os.makedirs(build_dir, exist_ok=True) if isinstance(input_tensor, dict): @@ -24,8 +24,6 @@ def create_sim_env(input_tensor, input_weight, generated_code, golden_result, fp else: with open(os.path.join(build_dir, "input_tensor.pt"), "wb") as f: torch.save(input_tensor, f) - with open(os.path.join(build_dir, "model_weights.pt"), "wb") as f: - torch.save(input_weight, f) with open(os.path.join(build_dir, "generated_asm_code.asm"), "w") as f: f.write(generated_code) # Store golden_result in a readable format, including tensor contents. @@ -46,15 +44,6 @@ def create_sim_env(input_tensor, input_weight, generated_code, golden_result, fp else: value_np = input_tensor.detach().cpu().float().numpy() f.write(np_array_to_str_2f(value_np)) - f.write("\n\nWeights (state_dict):\n") - if isinstance(golden_result["weights"], dict): - for key, value in golden_result["weights"].items(): - # Convert BFloat16 to float32 before converting to numpy - value_np = value.detach().cpu().float().numpy() - f.write(f"{key}:\n{np_array_to_str_2f(value_np)}\n") - else: - value_np = golden_result["weights"].detach().cpu().float().numpy() - f.write(np_array_to_str_2f(value_np)) f.write("\n\nOriginal Output:\n") # Convert BFloat16 to float32 before converting to numpy output_np = golden_result["original_output"].detach().cpu().float().numpy() diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py new file mode 100644 index 00000000..0c466729 --- /dev/null +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -0,0 +1,93 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import torch +from torch import Tensor, nn +# from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ +from test_data_gen import get_weights_path, generate_and_save_random_weights +from compiler.asm_templates import preload_act_asm, reset_reg_asm, preload_addr_reg_asm +from create_sim_env import create_sim_env +from sim_env_utils import create_mem_for_sim +import torch.nn.functional as F + +from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth +from tools.memory_mapping.addr_align import align_addr_up +from transformers import AutoTokenizer + + +if __name__ == "__main__": + + + # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) + vocal_size = 2 + hidden_size = 128 + vlen = 64 + batch_size = 4 + preload_amount = 4 + real_data_ratio = (8*8 + 8) / (8 * 8) + hbm_data_width = 64 + fp_preload = [0.0, 1e-6] + + torch.manual_seed(42) + logits = torch.randn(batch_size, hidden_size) + weights = logits + original_output = logits + + # Generate vlen random int32 data + int_preload = torch.randint(low=0, high=10, size=(batch_size, hidden_size), dtype=torch.int32) + print("int_preload", int_preload) + + + print('logits.shape= ', logits.shape) + print('original_output.shape= ', original_output.shape) + + + input_tensor = { + "logits": logits, + "int": int_preload, + } + + golden_result = { + "input_tensor": input_tensor, + "weights": weights, + "original_output": original_output + } + print('original_output.shape = ',original_output.shape) + print('original_output = ',original_output) + + gen_assembly_code = "; DLLM Test Generation \n" + + # Set the addr offset for mask + + # gen_assembly_code += preload_addr_reg_asm( + # addr_reg_to_set=[1,2], + # available_registers=[1,2], + # addr_reg_val=[int(align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width)),int(2*align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width))] + # ) + + + # Reset the registers + # gen_assembly_code += reset_reg_asm( + # alive_registers=[1,2,3,4,5,6] + # ) + + # Gen logtis Preload (B,L,V) + gen_assembly_code += preload_act_asm( + vlen=vlen, + preload_len=preload_amount, + batch=batch_size, + hidden_size=hidden_size, + alive_registers=[1,2,3], # [a_actual_register, set_stride_register, result_register] + act_vram_offset=0, + activation_offset_reg=0, + stride_size=hidden_size + ) + + # # Preload Integer Activation to the later section. + # gen_assembly_code += f"S_ADDI_INT gp1, gp0, {hidden_size * batch_size} \n" + # gen_assembly_code += f"S_ADDI_INT gp2, gp0, {align_addr_up(hidden_size * batch_size * real_data_ratio, hbm_data_width)} \n" + # gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 2, 1 \n" + + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits", "int"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py new file mode 100644 index 00000000..791cb583 --- /dev/null +++ b/behavioral_simulator/testbench/ffn_test.py @@ -0,0 +1,180 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import torch +from torch import Tensor, nn +# from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ +from test_data_gen import get_weights_path, generate_and_save_random_weights +from compiler.asm_templates import ffn_asm, preload_addr_reg_asm, reset_reg_asm, preload_act_asm +from create_sim_env import create_sim_env +from sim_env_utils import create_mem_for_sim + + +class LlamaFeedForward(nn.Module): + """ + Standard FeedForward layer used in Llama architectures: + y = W2(activation(W1(x))) + where activation is typically SwiGLU in Llama2. + + Args: + dim (int): input and output dimension (hidden size) + inter_dim (int): intermediate/fc dimension + activation (callable): nonlinearity to use (default: SwiGLU) + """ + def __init__(self, dim: int, inter_dim: int, activation: str = "silu"): + super().__init__() + # Llama uses SwiGLU: x * silu(x) + self.w1 = nn.Linear(dim, inter_dim, bias=False) + print("w1:", self.w1.weight.shape) + self.w2 = nn.Linear(dim, inter_dim, bias=False) + print("w2:", self.w2.weight.shape) + self.w3 = nn.Linear(inter_dim, dim, bias=False) + print("w3:", self.w3.weight.shape) + self.act = torch.nn.SiLU() if activation == "silu" else getattr(torch.nn, activation)() + + def forward(self, x: Tensor) -> Tensor: + # SwiGLU: (x @ w1) * silu(x @ w3) + # print("input_tensor:\n", x) + # print("up weight:\n", self.w1.weight.t()) + # up_proj = self.w1(x) + print("up projection (all elements):") + w1_out = self.w1(x) + w1_out = w1_out.reshape(w1_out.shape[0] * w1_out.shape[1], w1_out.shape[-1]) + print("self.w1(x) (upper all elements):") + print(w1_out[:, :8]) + print("self.w1(x) (mid upper all elements):") + print(w1_out[:, 64:72]) + print("self.w1(x) (mid lower all elements):") + print(w1_out[:, 128:135]) + print("self.w1(x) (lower all elements):") + print(w1_out[:, 192:199]) + + # w2_out = self.w2(x) + # w2_out = w2_out.reshape(w2_out.shape[0] * w2_out.shape[1], w2_out.shape[-1]) + # print("self.w2(x) (upper all elements):") + # print(w2_out[:, :8]) + # print("self.w2(x) (mid upper all elements):") + # print(w2_out[:, 64:72]) + # print("self.w2(x) (mid lower all elements):") + # print(w2_out[:, 128:135]) + # print("self.w2(x) (lower all elements):") + # print("gate projection:\n", self.w2(x)) + + # print("silu activation:\n", self.act(self.w1(x))) + # print("product of silu activation and gate projection:\n", self.act(self.w1(x)) * self.w2(x)) + silu_mixed_out = self.act(self.w1(x)) * self.w2(x) + # silu_mixed_out = silu_mixed_out.reshape(silu_mixed_out.shape[0] * silu_mixed_out.shape[1], silu_mixed_out.shape[-1]) + # print(f"silu mixed out of shape {silu_mixed_out.shape}: \n") + # print("silu mixed out (upper all elements):") + # print(silu_mixed_out[:, :8]) + # print("silu mixed out (mid upper all elements):") + # print(silu_mixed_out[:, 64:72]) + # print("silu mixed out (mid lower all elements):") + # print(silu_mixed_out[:, 128:135]) + # print("silu mixed out (lower all elements):") + # print(silu_mixed_out[:, 192:199]) + outcome = self.w3(silu_mixed_out) + print("final output (upper all elements):") + print(outcome[:, :8]) + print("final output (mid upper all elements):") + print(outcome[:, 64:72]) + # print("final output:\n", self.w3(self.act(self.w1(x)) * self.w2(x))) + return self.w3(self.act(self.w1(x)) * self.w2(x)) + +if __name__ == "__main__": + # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) + hidden_size = 128 + inter_dim = 256 + batch_size = 4 + real_data_ratio = (8*8 + 8) / (8 * 8) + fp_preload = [0.0, 1] + mlen = 64 + blen = 4 + vlen = 64 + seq_len = 2 + + torch.manual_seed(42) + act_tensor = torch.randn(batch_size, seq_len, hidden_size) + + original_layer = LlamaFeedForward(dim=hidden_size, inter_dim=inter_dim) + weight_up_layer = torch.randn(inter_dim, hidden_size) + weight_gate_layer = torch.randn(inter_dim, hidden_size) + weight_down_layer = torch.ones(hidden_size, inter_dim) + print(f"weight_down_layer of shape {weight_down_layer.t().shape}: \n") + print("upper all elements:") + print(weight_down_layer.t()[:, :8]) + print("lower all elements:") + print(weight_down_layer.t()[:, 64:72]) + + # Set weights for w1, w2, w3 to the generated tensors + with torch.no_grad(): + original_layer.w1.weight.copy_(weight_up_layer) + original_layer.w2.weight.copy_(weight_gate_layer) + original_layer.w3.weight.copy_(weight_down_layer) + + original_output = original_layer(act_tensor) + + input_tensor = { + "act_tensor": act_tensor.reshape(batch_size * seq_len, hidden_size), + "weight_up_layer": weight_up_layer.t(), + "weight_gate_layer": weight_gate_layer.t(), + "weight_down_layer": weight_down_layer.t(), + } + + golden_result = { + "input_tensor": input_tensor, + "original_output": original_output + } + + gen_assembly_code = "; FFN Test Generation \n" + + # Set the addr offset for weight and bias + gen_assembly_code += preload_addr_reg_asm( + addr_reg_to_set=[1, 2, 3], + available_registers=[1, 2, 3], + addr_reg_val=[int(hidden_size * batch_size * seq_len * real_data_ratio), + int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio), + int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio) + int(inter_dim * hidden_size * real_data_ratio)] + ) + + print("up_addr_hbm_val:", int(hidden_size * batch_size * seq_len * real_data_ratio)) + print("gate_addr_hbm_val:", int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio)) + print("down_addr_hbm_val:", int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio) + int(inter_dim * hidden_size * real_data_ratio)) + + # Reset the registers + gen_assembly_code += reset_reg_asm( + alive_registers=[1,2,3] + ) + + # Preload Activation + gen_assembly_code += preload_act_asm( + vlen=vlen, + preload_len=4, + batch=batch_size * seq_len, + hidden_size=hidden_size, + alive_registers=[1,2,3], + act_vram_offset=0, + activation_offset_reg=0, + stride_size=hidden_size + ) + + # FFN Generation + gen_assembly_code += ffn_asm( + mlen=mlen, + vlen=vlen, + blen=blen, + batch=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + intermediate_size=inter_dim, + alive_registers=[1,2,3,4,5,6,7], + up_weight_hbm_offset_reg=1, + gate_weight_hbm_offset_reg=2, + down_weight_hbm_offset_reg=3, + const_one_fp_address=1, + activation_base_address=0 + ) + + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + create_mem_for_sim(data_size=256, mode="behave_sim", asm=None, data=None, specified_data_order = ["act_tensor", "weight_up_layer", "weight_gate_layer", "weight_down_layer"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/linear_test.py b/behavioral_simulator/testbench/linear_test.py index 990de79f..a9bc1aff 100644 --- a/behavioral_simulator/testbench/linear_test.py +++ b/behavioral_simulator/testbench/linear_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import create_mem_for_sim # TODOs: Need to integrate the MX quantizer here. @@ -34,13 +34,11 @@ # generate_and_save_random_weights(hidden_size, hidden_size, get_weights_path('model_weights.pt')) torch.manual_seed(42) - input_tensor = torch.randn(batch_size, hidden_size) - # Print input_tensor split in half along columns, as two (4, 64) tensors - print("input_tensor lhs (4, 64):\n", input_tensor[:, :64]) - print("input_tensor rhs (4, 64):\n", input_tensor[:, 64:]) - + act_tensor = torch.randn(batch_size, hidden_size) original_layer = nn.Linear(in_features=hidden_size, out_features=hidden_size, bias=False) weights = original_layer.state_dict() + + original_output = original_layer(act_tensor) # Print weight k, (128, 128) -> print 4 quadrants of (64, 64) each # Quadrant indices: @@ -48,36 +46,41 @@ # 1: [0:64, 64:128] # 2: [64:128, 0:64] # 3: [64:128, 64:128] - w_k = weights['weight'] if isinstance(weights, dict) else weights + w_k = weights['weight'].t() if isinstance(weights, dict) else weights print("Weight k shape:", w_k.shape) - for idx, (r_slice, c_slice) in enumerate([ - (slice(0, 64), slice(0, 64)), - (slice(0, 64), slice(64, 128)), - (slice(64, 128), slice(0, 64)), - (slice(64, 128), slice(64, 128)), - ]): - print(f"---------- Quadrant {idx}: Rows {r_slice}, Cols {c_slice} ----------") - print(w_k[r_slice, c_slice]) + # for idx, (r_slice, c_slice) in enumerate([ + # (slice(0, 64), slice(0, 64)), + # (slice(0, 64), slice(64, 128)), + # (slice(64, 128), slice(0, 64)), + # (slice(64, 128), slice(64, 128)), + # ]): + # print(f"---------- Quadrant {idx}: Rows {r_slice}, Cols {c_slice} ----------") + # print(w_k[r_slice, c_slice]) # Print the matmul result of input_tensor[:, :64] and weight[0:64, 0:4] - matmul_result_11 = input_tensor[:, :64] @ w_k[:64, :4] - print("Matmul result of input_tensor[:, :64] @ weight[0:64, 0:4]:") - print(matmul_result_11) + matmul_result_11 = act_tensor[:, :64] @ w_k[:64, :4] + # print("act_tensor[:, :64]: \n", act_tensor[:, :64]) + # print("w_k[:64, :4]: \n", w_k[:64, :4]) + # print("Matmul result of input_tensor[:, :64] @ weight[0:64, 0:4]:") + # print(matmul_result_11) - matmul_result_22 = input_tensor[:, 64:] @ w_k[64:, :4] + matmul_result_22 = act_tensor[:, 64:] @ w_k[64:, :4] print("Matmul result of input_tensor[:, 64:] @ weight[64:, :4]:") + print("act_tensor[:, 64:]: \n", act_tensor[:, 64:]) + print("w_k[64:, :4]: \n", w_k[64:, :4]) print(matmul_result_22) print ("sum of two matmul results:", matmul_result_11 + matmul_result_22) - original_output = original_layer(input_tensor) - print ("original_output:", original_output) + input_tensor = { + "act_tensor": act_tensor, + "weights": weights['weight'].t(), + } golden_result = { "input_tensor": input_tensor, - "weights": weights, "original_output": original_output } @@ -90,8 +93,8 @@ addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] ) - print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) - print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) + # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) + # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) # Reset the registers gen_assembly_code += reset_reg_asm( @@ -106,7 +109,8 @@ hidden_size=128, alive_registers=[1,2,3], act_vram_offset=0, - activation_offset_reg=0 + activation_offset_reg=0, + stride_size=hidden_size ) # Reset the registers @@ -120,17 +124,14 @@ batch=4, hidden_size=128, alive_registers=[1,2,3,4], - head_dim=128, w_base_hbm_offset_reg=1, - rope_hbm_offset_reg=0, - rope_on_chip_address=0, activation_base_address=0, result_base_address=hidden_size * batch_size, rope_enabled=False ) - create_sim_env(input_tensor, weights['weight'].t(), gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["act_tensor", "weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/rms_test.py b/behavioral_simulator/testbench/rms_test.py index f1e77e6e..8a9a8644 100644 --- a/behavioral_simulator/testbench/rms_test.py +++ b/behavioral_simulator/testbench/rms_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import create_mem_for_sim # Taken from LLAMA RMSNorm implementation @@ -73,33 +73,30 @@ def forward(self, x): # generate_and_save_random_weights(hidden_size, hidden_size, get_weights_path('model_weights.pt')) torch.manual_seed(42) - input_tensor = torch.randn(batch_size, hidden_size) + act_tensor = torch.randn(batch_size, hidden_size) # Print input_tensor split in half along columns, as two (4, 64) tensors - print("input_tensor lhs (4, 64):\n", input_tensor[:, :64]) - print("input_tensor rhs (4, 64):\n", input_tensor[:, 64:]) + print("act_tensor lhs (4, 64):\n", act_tensor[:, :64]) + print("act_tensor rhs (4, 64):\n", act_tensor[:, 64:]) original_layer = RMSNorm(dim=hidden_size) weights = original_layer.state_dict() - original_output = original_layer(input_tensor) + input_tensor = { + "act_tensor": act_tensor, + "weights": weights['weight'].t(), + } + + original_output = original_layer(act_tensor) golden_result = { "input_tensor": input_tensor, - "weights": weights, "original_output": original_output } gen_assembly_code = "; RMSNorm Test Generation \n" - - # Set the addr offset for weight and bias - gen_assembly_code += preload_addr_reg_asm( - addr_reg_to_set=[1, 2], - available_registers=[1, 2], - addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] - ) - print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) - print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) + # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) + # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) # Reset the registers @@ -115,12 +112,13 @@ def forward(self, x): hidden_size=128, alive_registers=[1,2,3], act_vram_offset=0, - activation_offset_reg=0 + activation_offset_reg=0, + stride_size=hidden_size ) # Reset the registers gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3,4] + alive_registers=[1,2,3] ) gen_assembly_code += rms_norm_asm( @@ -134,5 +132,9 @@ def forward(self, x): hidden_dim=128 ) - create_sim_env(input_tensor, weights['weight'].t(), gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["input_tensor", "model_weights"]) \ No newline at end of file + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["act_tensor", "weights"]) + + print("================================================") + print("Finished generating assembly code") + print("================================================") diff --git a/behavioral_simulator/testbench/s_map_v_test.py b/behavioral_simulator/testbench/s_map_v_test.py new file mode 100644 index 00000000..51bdec7f --- /dev/null +++ b/behavioral_simulator/testbench/s_map_v_test.py @@ -0,0 +1,68 @@ +from re import I +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import torch +from torch import Tensor, nn +# from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ +from test_data_gen import get_weights_path, generate_and_save_random_weights +from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm +from create_sim_env import create_sim_env +from sim_env_utils import create_mem_for_sim +from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth +import torch.nn.functional as F + +if __name__ == "__main__": + + + vlen = 64 + load_len = vlen * vlen + batch_size = 1 + real_data_ratio = (8*8 + 8) / (8 * 8) + fp_preload = [0.0, 1e-6] + preload_amount = 4 + hbm_data_width = 64 + + torch.manual_seed(42) + fp_preload = [0, 1] + # fp_sram = [fp_preload[0]] * vlen + + gen_assembly_code = "; S MAP V Test Generation \n" + torch.manual_seed(42) + + input_tensor1 = torch.randn(batch_size, load_len) + input_tensor = { + "input_tensor1": input_tensor1, + "input_tensor2": input_tensor1 + } + weights = input_tensor1 + original_output = torch.max(input_tensor1.reshape(batch_size, vlen, vlen), dim=-1).values + + golden_result = { + "input_tensor": input_tensor, + "weights": weights, + "original_output": original_output + } + # Gen Activation Preload for fp0 + gen_assembly_code += preload_act_asm( + vlen=vlen, + preload_len=preload_amount, + batch=batch_size, + hidden_size=load_len, + alive_registers=[1,2,3], # [a_actual_register, set_stride_register, result_register] + act_vram_offset=0, + activation_offset_reg=0 + ) + gen_assembly_code += f"S_ADDI_INT gp1, gp0, 0 \n" + for i in range(vlen): + gen_assembly_code += f"V_RED_MAX f1, gp1, 0 \n" + gen_assembly_code += f"S_ADDI_INT gp1, gp1, {vlen} \n" + gen_assembly_code += f"S_ST_FP f1, gp0, {i} \n" + + # gen_assembly_code += "; V_RED_MAX f1, gp0, 0 \n" + gen_assembly_code += f"S_ADDI_INT gp1, gp0, 0 \n" + gen_assembly_code += f"S_MAP_V_FP gp1, gp0, 0 \n" + + create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/two_input_test.py b/behavioral_simulator/testbench/two_input_test.py index f4d9e978..704f488f 100644 --- a/behavioral_simulator/testbench/two_input_test.py +++ b/behavioral_simulator/testbench/two_input_test.py @@ -9,7 +9,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import create_mem_for_sim from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth import torch.nn.functional as F @@ -76,10 +76,6 @@ activation_offset_reg=1 ) - # Reset the registers - gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3,4] - ) - create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 5674434a..c1dbedad 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -3,12 +3,105 @@ import os import struct -def view_bin_file_by_row(bin_file, +def view_bin_file_by_row_int(bin_file, + int_width=32, + num_bytes_per_val=None, + row_size = 64 * 2, + start_row_idx=0, + load_row_size=None, + signed=True): + """ + Reads a binary file and parses each value as an integer. + + - bin_file: Path to the binary file + - row_dim: Number of values per row to group/print + - int_width: Bit width of the integer (default 32 for i32) + Common values: 8 (i8), 16 (i16), 32 (i32) + - num_bytes_per_val:Number of bytes per value (auto-calculated from int_width if None) + - start_row_idx: Starting row index + - load_row_size: If set, limit number of rows to print + - signed: If True, interpret as signed integer (default True) + + Note: The Rust code uses little-endian byte order. + If int_width is less than num_bytes_per_val * 8, the value will be masked. + """ + # Auto-calculate bytes per value if not specified + if num_bytes_per_val is None: + num_bytes_per_val = (int_width + 7) // 8 # Round up to nearest byte + + # Validate that num_bytes_per_val is sufficient + if int_width > num_bytes_per_val * 8: + raise ValueError(f"int_width ({int_width}) exceeds num_bytes_per_val * 8 ({num_bytes_per_val * 8})") + + # Create mask for the specified bit width + if int_width >= 64: + mask = (1 << 64) - 1 + else: + mask = (1 << int_width) - 1 + + with open(bin_file, "rb") as f: + data = f.read() + + row_dim = row_size // num_bytes_per_val + num_vals = len(data) // num_bytes_per_val + total_rows = (num_vals + row_dim - 1) // row_dim + + for row_idx in range(total_rows): + if row_idx < start_row_idx: + continue + if load_row_size is not None and row_idx >= load_row_size + start_row_idx: + print("... (truncated)") + break + vals = [] + for col_idx in range(row_dim): + val_idx = row_idx * row_size + col_idx * num_bytes_per_val + if val_idx >= num_vals: + break + chunk = data[val_idx : val_idx + num_bytes_per_val] + # print("chunk = {:?}", chunk) + if not chunk or len(chunk) < num_bytes_per_val: + vals.append(None) + continue + # Use little-endian byte order to match Rust's byte packing + if signed: + int_val = int.from_bytes(chunk, byteorder='little', signed=True) + else: + int_val = int.from_bytes(chunk, byteorder='little', signed=False) + + # Apply bit width mask + int_val = int_val & mask + + # Sign extend if signed and the sign bit is set + if signed and int_width < num_bytes_per_val * 8: + sign_bit = (int_val >> (int_width - 1)) & 1 + if sign_bit == 1: + # Sign extend: set all upper bits to 1 + sign_extend_mask = ((1 << (num_bytes_per_val * 8 - int_width)) - 1) << int_width + int_val = int_val | sign_extend_mask + # Convert to signed integer + if num_bytes_per_val == 1: + int_val = int_val if int_val < 128 else int_val - 256 + elif num_bytes_per_val == 2: + int_val = int_val if int_val < 32768 else int_val - 65536 + elif num_bytes_per_val == 4: + int_val = int_val if int_val < 2147483648 else int_val - 4294967296 + + vals.append(int_val) + print(f"Row {row_idx:3d}: ", end="") + for v in vals: + if v is not None: + print(f"{v:8d}", end=" ") + else: + print(" ", end=" ") + print() + +def view_bin_file_by_row_fp(bin_file, exp_width, man_width, row_dim, num_bytes_per_val=2, start_row_idx=0, + offset=0, load_row_size=None): """ Reads a binary file generated by Rust behavioral simulator, parses each value as custom FP. @@ -103,22 +196,30 @@ def raw_to_fp(bits_val): # VRAM uses BF16 format by default: sign=1, exponent=8, mantissa=7 (16 bits total = 2 bytes) print("Viewing VRAM dump from 0 Base Address") - view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) + view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) + + # print("Viewing VRAM dump from 8 Base Address") + # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. + view_bin_file_by_row_int(vram_file, row_size=64 * 2, int_width=32, start_row_idx=8, num_bytes_per_val=4, load_row_size=4) + + + # print("Viewing VRAM dump from 48 Base Address") + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=48, load_row_size=32) # print("Viewing VRAM dump from Q Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) # print("\nViewing VRAM dump from S Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=16) # print("\nViewing VRAM dump from PV Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=320, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=320, load_row_size=16) # print("\nViewing VRAM dump from O_Old Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=576, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=576, load_row_size=16) # print("Viewing MRAM dump 0 to 7 rows (BF16 format)") - # view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) + # view_bin_file_by_row_fp(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) # print("Viewing MRAM dump 64 to 71 rows (BF16 format)") - # view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=8) \ No newline at end of file + # view_bin_file_by_row_fp(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=8) \ No newline at end of file diff --git a/compiler/asm_templates/batched_matmul_asm.py b/compiler/asm_templates/batched_matmul_asm.py index ff41bf17..8083538d 100644 --- a/compiler/asm_templates/batched_matmul_asm.py +++ b/compiler/asm_templates/batched_matmul_asm.py @@ -40,6 +40,8 @@ def batched_matmul_asm( assert k % mlen == 0, "k must be divisible by mlen" assert m % blen == 0, "m must be divisible by blen" assert n % blen == 0, "n must be divisible by blen" + print(f"b = {b}, m = {m}, k = {k}, n = {n}") + print(f"mlen = {mlen}, blen = {blen}") a_actual_register = alive_registers[0] w_actual_register = alive_registers[1] @@ -47,7 +49,7 @@ def batched_matmul_asm( generated_code += f"S_ADDI_INT gp{result_actual_register}, gp0, {result_base_address} \n" - for batch in range(b): + for batch in range(1, b + 1): # preload the activation matrix for i in range(math.ceil(n // blen)): assert w_prefetch_amount >= k, "w_prefetch_amount must be greater than or equal to k" @@ -56,7 +58,7 @@ def batched_matmul_asm( generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {n} \n" generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{w_base_hbm_offset_reg}, {w_prefetch_amount // k}, 0 \n" + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{w_base_hbm_offset_reg}, 1, 0 \n" for j in range(math.ceil(m // blen)): if j == 0: @@ -68,7 +70,7 @@ def batched_matmul_asm( if j % math.ceil((a_prefetch_amount * mlen) // k) == 0: generated_code += f"H_PREFETCH_V gp{a_actual_register}, gp{a_actual_register}, a{a_base_hbm_offset_reg}, 1, 0 \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, 0 \n" - for k in range(math.ceil(k // mlen)): + for g in range(math.ceil(k // mlen)): generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {blen * mlen} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {blen * mlen} \n" diff --git a/compiler/asm_templates/ffn_asm.py b/compiler/asm_templates/ffn_asm.py index b56fcc61..626f41da 100644 --- a/compiler/asm_templates/ffn_asm.py +++ b/compiler/asm_templates/ffn_asm.py @@ -1,102 +1,192 @@ import os from typing import Dict, List, Any, Optional from pathlib import Path - - +import math +IMM2_BOUND = 2**18 def ffn_asm( mlen: int, vlen: int, blen: int, batch: int, + seq_len: int, hidden_size: int, - alive_registers: List[int], - weight_hbm_offset_reg: int, intermediate_size: int, - const_address: int, - activation_base_address: int, - result_base_address: int + + alive_registers: List[int], + gate_weight_hbm_offset_reg: int, + up_weight_hbm_offset_reg: int, + down_weight_hbm_offset_reg: int, + const_one_fp_address: int, + + activation_base_address: int ) -> str: """ - Generates assembly code for a general matrix multiplication operation. + Generates assembly code for a FFN operation. Args: - mlen (int): The number of rows in the first matrix. + mlen (int): The number of rows in the matrix. + vlen (int): The number of columns in the matrix. blen (int): The number of columns in the second matrix. + batch (int): The number of batches. + hidden_size (int): The number of rows in the hidden size. + intermediate_size (int): The number of rows in the intermediate size. alive_registers (List[int]): List of registers that are alive. - rope_base_address (int): index for the address mapper pointing to the base addr of the rope matrix. + gate_weight_hbm_offset_reg (int): index for the address mapper pointing to the base addr of the gate weight matrix. + up_weight_hbm_offset_reg (int): index for the address mapper pointing to the base addr of the up weight matrix. + down_weight_hbm_offset_reg (int): index for the address mapper pointing to the base addr of the down weight matrix. activation_base_address (int): index for the address mapper pointing to the base addr of the activation matrix. + result_base_address (int): index for the address mapper pointing to the base addr of the result matrix. Functionality: Upsize linear (b, s, hidden_size) @ (hidden_size, intermediate_size) - > (b, s, intermediate_size) - Activation (b, s, intermediate_size) -> (b, s, intermediate_size) + Gate Projection (b, s, hidden_size) @ (hidden_size, intermediate_size) -> (b, s, intermediate_size) + SILU Activation (b, s, intermediate_size) -> (b, s, intermediate_size) Downsize linear (b, s, intermediate_size) @ (intermediate_size, hidden_size) -> (b, s, hidden_size) """ - generated_code = " ; FFN Upsize Linear Generation \n" - # Dot product of weight (Hidden Size, Hidden Size) and activation (Batch, 1, Hidden Size) - assert batch < blen, "Batch size must be greater than blen" - # get two registers from alive_registers, 1 as w address, 1 as a address - w_base_register = alive_registers[0] - a_base_register = alive_registers[1] - result_register = alive_registers[2] - w_actual_register = alive_registers[3] - a_actual_register = alive_registers[4] - intermediate_register = alive_registers[5] - # reset the registers - set_w_base_register = f"S_ADDI_INT gp{w_base_register}, gp0, 0 \n" - set_a_base_address = f"S_ADDI_INT gp{a_base_register}, gp0, {activation_base_address} \n" - set_result_address = f"S_ADDI_INT gp{result_register}, gp0, {result_base_address} \n" + - set_w_actual_address = f"S_ADD_INT gp{w_actual_register}, gp0, {w_base_register} \n" - set_a_actual_address = f"S_ADD_INT gp{a_actual_register}, gp0, {a_base_register} \n" + # memory assignment + # 0 -> activation + # b * s * hidden_size -> upsize intermediate results + # b * s * (hidden_size + intermediate_size) -> gate projection results - increment_result_actual_address = f"S_ADDI_INT gp{result_register}, gp{result_register}, {mlen} \n" + w_actual_register = alive_registers[0] + w_temp_register = alive_registers[1] + a_actual_register = alive_registers[2] + up_result_register = alive_registers[3] + intermediate_register = alive_registers[4] + gate_result_register = alive_registers[5] + w_hbm_offset_register = alive_registers[6] - row_loop_over_hid = intermediate_size // blen - vect_loop_over_hid = intermediate_size // vlen - col_loop_over_hid = hidden_size // mlen - generated_code += set_w_base_register - generated_code += set_a_base_address - generated_code += set_result_address - - generated_code += f"PREFETECH_M {w_actual_register}, gp0, a{weight_hbm_offset_reg}, 1, 0 \n" - for i in range(row_loop_over_hid): - generated_code += f"; <---- Generating New Row Tile at index {i} ----> \n" - for j in range(col_loop_over_hid): - generated_code += f"; <---- Generating New Column Tile at row {i} col {j} ----> \n" - generated_code += f"M_MM 0, {w_actual_register}, {a_actual_register} \n" - generated_code += set_w_actual_address - generated_code += set_a_actual_address - generated_code += f"M_MM_WO {result_register}, 0, 0 \n" - if (i % blen) == 0: - generated_code += increment_result_actual_address - - generated_code += "; SILU Generation \n" - fp_const_reg = "f1" - generated_code += f"S_LD_FP {fp_const_reg}, gp0, {const_address} \n" - for i in range(vect_loop_over_hid): - generated_code += f"; <---- per VLEN block {i} ----> \n" - generated_code += f"V_SUB_VV {intermediate_register}, {intermediate_register}, {result_register} \n" - generated_code += f"V_EXP_V {intermediate_register}, {intermediate_register} \n" - generated_code += f"V_ADD_VF {intermediate_register}, {intermediate_register}, {fp_const_reg} \n" - generated_code += f"V_ADD_VF {intermediate_register}, {intermediate_register}, {fp_const_reg} \n" - generated_code += f"V_REC_V {intermediate_register}, {intermediate_register} \n" - generated_code += f"V_MUL_VV {result_register}, {result_register}, {intermediate_register} \n" + # reset the registers + generated_code = "; FFN Generation \n" + # Settings for up and gate weight matrices prefetching + assert hidden_size * intermediate_size < IMM2_BOUND, f"hidden_size * hidden_size must be less than {IMM2_BOUND}" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" + generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {intermediate_size} \n" + generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + assert hidden_size * batch * seq_len < IMM2_BOUND, f"hidden_size * batch * seq_len must be less than {IMM2_BOUND}" + # Set the address for on-chip sram + generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{up_result_register}, {intermediate_size * batch * seq_len} \n" - generated_code += "; FFN Downsize Linear Generation \n" + generated_code += " ; FFN Upsize Linear Generation \n" + for weight_row in range (intermediate_size // blen): + if weight_row % (mlen // blen) == 0: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{up_result_register}, 0 \n" + + for weight_col in range (hidden_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{up_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * intermediate_size} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + else: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{up_result_register}, {(weight_row % (mlen // blen)) * blen} \n" + for act_col in range ((batch * seq_len) // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address + act_col * mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n" + for inner_loop_index in range (hidden_size // mlen): + generated_code += f"M_MM 0, gp{w_temp_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_temp_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n" # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address} \n" + if (weight_row + 1) % (mlen // blen) == 0 and weight_row != intermediate_size // blen - 1: + generated_code += f"S_ADDI_INT gp{up_result_register}, gp{up_result_register}, {mlen * batch * seq_len} \n" - row_loop_over_hid = hidden_size // blen - col_loop_over_hid = intermediate_size // mlen + generated_code += " ; FFN Gate Projection Generation \n" + for weight_row in range (intermediate_size // blen): + if weight_row % (mlen // blen) == 0: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" + + for weight_col in range (hidden_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{gate_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * intermediate_size} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + else: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, {(weight_row % (mlen // blen)) * blen} \n" + for act_col in range ((batch * seq_len) // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address + act_col * mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n" + for inner_loop_index in range (hidden_size // mlen): + generated_code += f"M_MM 0, gp{w_temp_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_temp_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n" # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address} \n" + if (weight_row + 1) % (mlen // blen) == 0 and weight_row != intermediate_size // blen - 1: + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {mlen * batch * seq_len} \n" - for i in range(row_loop_over_hid): - generated_code += f"; <---- Generating New Row Tile at index {i} ----> \n" - for j in range(col_loop_over_hid): - generated_code += f"; <---- Generating New Column Tile at row {i} col {j} ----> \n" - generated_code += f"M_MM 0, {w_actual_register}, {result_register} \n" - generated_code += set_w_actual_address - generated_code += set_result_address - generated_code += f"M_MM_WO {a_actual_register}, 0, 0 \n" - if (i % blen) == 0: - generated_code += set_a_actual_address + # Intermediate Dim SILU Activation Generation, now x in shape of (b, s, intermediate_size) + generated_code += "; SILU Generation \n" + generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" + # Reset the addr for up and gate result + generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{up_result_register}, {intermediate_size * batch * seq_len} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp0, {activation_base_address} \n" + + # Treat the original activation region as the place for scratchpad. + for b in range(batch * seq_len): + for i in range(intermediate_size // vlen): + # 0 : -x + generated_code += f"V_SUB_VF gp{intermediate_register}, gp{up_result_register}, f0, 0, 1 \n" + # 1 : exp(-x) + generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # 2 : 1 + exp(-x) + generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1, 0 \n" + # 3 : 1 / (1 + exp(-x)) + generated_code += f"V_RECI_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # 4 : (1 / (1 + exp(-x))) * gate_result + generated_code += f"V_MUL_VV gp{intermediate_register}, gp{intermediate_register}, gp{up_result_register}, 0 \n" + # 5: multiply by gate result and store to the up result region + generated_code += f"V_MUL_VV gp{up_result_register}, gp{intermediate_register}, gp{gate_result_register}, 0 \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {vlen} \n" + generated_code += f"S_ADDI_INT gp{up_result_register}, gp{up_result_register}, {vlen} \n" + + generated_code += "; FFN Downsize Linear Generation \n" + # Reset the addr for up and gate result + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" + generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size} \n" + generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + # Storing the results to the activation base region + act_result_register = gate_result_register + generated_code += f"S_ADDI_INT gp{act_result_register}, gp0, {activation_base_address} \n" + generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" + for weight_row in range (hidden_size // blen): + if weight_row % (mlen // blen) == 0: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{act_result_register}, 0 \n" + for weight_col in range (intermediate_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{down_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * hidden_size} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + else: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{act_result_register}, {(weight_row % (mlen // blen)) * blen} \n" + for act_col in range (batch * seq_len // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{up_result_register}, {act_col * mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n" + for inner_loop_index in range (intermediate_size // mlen): + generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n" # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address} \n" + if (weight_row + 1) % (mlen // blen) == 0 and weight_row != intermediate_size // blen - 1: + generated_code += f"S_ADDI_INT gp{act_result_register}, gp{act_result_register}, {mlen * batch * seq_len} \n" return generated_code diff --git a/compiler/asm_templates/flash_attn_asm.py b/compiler/asm_templates/flash_attn_asm.py index 17dc6ae9..134ea75e 100644 --- a/compiler/asm_templates/flash_attn_asm.py +++ b/compiler/asm_templates/flash_attn_asm.py @@ -199,7 +199,7 @@ def _computing_pv_code( # Update v_base_register and reset p_base_register and pv_base_register generated_code += f"S_ADDI_INT gp{p_base_register}, gp0, {p_base_address + q_head_index * mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{pv_base_register}, gp0, {pv_base_address + q_head_index * head_dim + (i+1) * blen} \n" - generated_code += f"S_ADDI_INT gp{v_base_register}, gp{v_base_register}, {blen} \n" + generated_code += f"S_ADDI_INT gp{v_base_register}, gp{v_base_register}, {mlen} \n" return generated_code @@ -296,7 +296,6 @@ def _computing_row_wise_scaling_code( generated_code += f"S_RECI_FP f{l_old_fp_register}, f{l_old_fp_register}, 0 \n" # multiply o_old with the inverse of l_old generated_code += f"V_MUL_VF gp{o_old_vector_address_register}, gp{o_old_vector_address_register}, f{l_old_fp_register} \n" - # update o_old base address generated_code += f"S_ADDI_INT gp{o_old_vector_address_register}, gp{o_old_vector_address_register}, {mlen} \n" diff --git a/compiler/asm_templates/load_int.py b/compiler/asm_templates/load_int.py new file mode 100644 index 00000000..e69de29b diff --git a/compiler/asm_templates/normalization_asm.py b/compiler/asm_templates/normalization_asm.py index da489889..dec889bd 100644 --- a/compiler/asm_templates/normalization_asm.py +++ b/compiler/asm_templates/normalization_asm.py @@ -32,7 +32,7 @@ def rms_norm_asm( for batch in range(batch_size): for i in range(hidden_dim // vlen): # Compute square of the activation vector and summation - generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{act_addr}, gp{act_addr} \n" + generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{act_addr}, gp{act_addr}, 0 \n" generated_code += f"V_RED_SUM f2, gp{scratchpad_addr} \n" # Move to next vector @@ -41,23 +41,22 @@ def rms_norm_asm( # Taking the avg generated_code += f"S_MUL_FP f2, f2, f3 \n" - # # Plus epsilon + # Plus epsilon generated_code += f"S_ADD_FP f2, f2, f1 \n" - # # Compute square root + # Compute square root generated_code += "S_SQRT_FP f2, f2 \n" - # # Compute reciprocal + # Compute reciprocal generated_code += "S_RECI_FP f2, f2 \n" for i in range(hidden_dim // vlen): # Normalize the activation vector - generated_code += f"V_MUL_VF gp{act_addr}, gp{act_addr}, f2 \n" + generated_code += f"V_MUL_VF gp{act_addr}, gp{act_addr}, f2, 0 \n" # Move to next vector generated_code += f"S_ADDI_INT gp{act_addr}, gp{act_addr}, {vlen * batch_size} \n" generated_code += "S_ADD_FP f2, f0, f0 \n" generated_code += f"S_ADDI_INT gp{act_addr}, gp0, {activation_base_address + vlen * batch} \n" - return generated_code \ No newline at end of file diff --git a/compiler/asm_templates/preload_act.py b/compiler/asm_templates/preload_act.py index fa2c1bfc..cf5851ae 100644 --- a/compiler/asm_templates/preload_act.py +++ b/compiler/asm_templates/preload_act.py @@ -1,6 +1,9 @@ import os from typing import Dict, List, Any, Optional from pathlib import Path +import math + +IMM2_BOUND = 2**18 def preload_act_asm( vlen: int, @@ -14,6 +17,7 @@ def preload_act_asm( ) -> str: """ Generates assembly code for preloading activation. + Memory Layout: Here we assume the activation is stored in (Hidden // MLEN , Batch (Integrate with Seq Len), MLEN) """ generated_code = "; Preload Activation Generation \n" # get two registers from alive_registers, 1 as a address @@ -27,23 +31,23 @@ def preload_act_asm( generated_code += f"C_SET_SCALE_REG gp{a_actual_register} \n" # reset the registers - set_a_base_address = f"S_ADDI_INT gp{a_actual_register}, gp0, 0 \n" - set_act_vram_base_address = f"S_ADDI_INT gp{result_register}, gp0, {act_vram_offset} \n" - generated_code += set_a_base_address - generated_code += set_act_vram_base_address + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp0, {act_vram_offset} \n" + load_amount_per_hidden = math.ceil(hidden_size / vlen) if batch == 1: - generated_code += f"S_ADDI_INT gp{set_stride_register}, gp0, {stride_len} \n" - generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" - for i in range((hidden_size + (vlen * preload_len) - 1) // (vlen * preload_len)): - generated_code += f"H_PREFETCH_V gp{a_actual_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0 \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen * preload_len} \n" + for i in range(math.ceil(hidden_size / (vlen * preload_len))): + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0, 0 \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {load_amount_per_hidden} \n" else: generated_code += f"S_ADDI_INT gp{set_stride_register}, gp0, {stride_len} \n" generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" - for i in range((batch * hidden_size) // (vlen * preload_len)): - generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0 \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen} \n" - generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" - + assert batch * hidden_size <= IMM2_BOUND, "batch * hidden_size must be less than {IMM2_BOUND}" + for i in range(load_amount_per_hidden): + for j in range(math.ceil(batch / preload_len)): + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0, 0 \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {hidden_size * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {(i + 1) * vlen} \n" return generated_code \ No newline at end of file diff --git a/compiler/asm_templates/projection_asm.py b/compiler/asm_templates/projection_asm.py index 3483cc5c..e534547e 100644 --- a/compiler/asm_templates/projection_asm.py +++ b/compiler/asm_templates/projection_asm.py @@ -1,27 +1,27 @@ import os from typing import Dict, List, Any, Optional from pathlib import Path - +import math IMM2_BOUND = 2**18 + def projection_asm( mlen: int, blen: int, batch: int, hidden_size: int, alive_registers: List[int], - head_dim: int, w_base_hbm_offset_reg: int, - rope_hbm_offset_reg: int, - rope_on_chip_address: int, activation_base_address: int, result_base_address: int, - rope_enabled: bool = True + rope_enabled: bool = False, + rope_hbm_offset_reg: int = 0, + rope_on_chip_address: int = 0 ) -> str: """ Generates assembly code for a general matrix multiplication operation. (Batch, Hidden Size) @ (Hidden Size, Hidden Size) -> (Batch, Hidden Size) - + assume Batch = BLEN for this version. Args: mlen (int): The number of rows in the first matrix. blen (int): The number of columns in the second matrix. @@ -32,16 +32,14 @@ def projection_asm( Returns: str: Generated assembly code for projection, including dot product and RoPE(cond) """ - generated_code = "" - assert batch <= blen, "Batch size must be less than blen" - # get two registers from alive_registers, 1 as w address, 1 as a address - result_register = alive_registers[0] - w_actual_register = alive_registers[1] - w_hbm_offset_register = alive_registers[2] - a_actual_register = alive_registers[3] + generated_code = "; Projection Generation \n" + + result_register = alive_registers[0] + w_actual_register = alive_registers[1] + w_hbm_offset_register = alive_registers[2] + a_actual_register = alive_registers[3] # Set scale offset - #TODO: when hidden is large, cannot use addi command. assert hidden_size * hidden_size < IMM2_BOUND, f"hidden_size * hidden_size must be less than {IMM2_BOUND}" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {hidden_size * hidden_size} \n" generated_code += f"C_SET_SCALE_REG gp{a_actual_register} \n" @@ -49,7 +47,6 @@ def projection_asm( generated_code += f"C_SET_STRIDE_REG gp{a_actual_register} \n" # reset the registers - row_loop_over_hid = hidden_size // blen col_loop_over_hid = hidden_size // mlen generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address} \n" @@ -60,56 +57,21 @@ def projection_asm( # Load a complete col of hidden size into on-chip memory (hidden_size, mlen) generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {i * blen} \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp0, {result_base_address + (math.floor(i / (mlen // blen))) * mlen * blen} \n" for k in range (hidden_size // mlen): - generated_code += f"; <---- Generating New Row Tile at index {i} col {k} ----> \n" generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{w_base_hbm_offset_reg}, 1, 0 \n" generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * hidden_size} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" else: generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(i % (mlen // blen)) * blen} \n" - + generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {blen} \n" for j in range(col_loop_over_hid): # Loop over the hidden size dimension - generated_code += f"; <---- Generating New Column Tile at row {i} col {j} \n" generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" generated_code += f"M_MM_WO {result_register}, gp0, 0 \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address} \n" - generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {blen} \n" - - - # RoPE - if rope_enabled: - generated_code += "; Generating RoPE code here \n" - generated_code += f"S_ADDI_INT gp{result_register}, gp0, {result_base_address} \n" - upper_base_register = alive_registers[0] - lower_base_register = alive_registers[1] - roped_upper_base_register = alive_registers[2] - roped_lower_base_register = alive_registers[3] - cos_base_register = alive_registers[4] - sin_base_register = alive_registers[5] - intermediate_1_register = alive_registers[6] - intermediate_2_register = alive_registers[7] - - generated_code += f"S_ADDI_INT gp{cos_base_register}, gp0, {rope_on_chip_address} \n" - generated_code += f"H_PREFETCH_V gp{cos_base_register}, gp{rope_on_chip_address}, a{rope_hbm_offset_reg}, 0, 0 \n" - generated_code += f"S_ADDI_INT gp{sin_base_register}, gp{cos_base_register}, {head_dim} \n" - generated_code += f"H_PREFETCH_V gp{sin_base_register}, gp{rope_on_chip_address}, a{rope_hbm_offset_reg}, 0, 0 \n" - - for i in range(batch * head_dim): - generated_code += f"; <---- Generating RoPE code for batch {i // head_dim} head {i % head_dim} ----> \n" - generated_code += f"V_MUL_VV gp{intermediate_1_register}, gp{upper_base_register}, gp{cos_base_register} \n" - generated_code += f"V_MUL_VV gp{intermediate_2_register}, gp{lower_base_register}, gp{sin_base_register} \n" - generated_code += f"V_SUB_VV gp{roped_upper_base_register}, gp{intermediate_1_register}, gp{intermediate_2_register} \n" - generated_code += f"V_MUL_VV gp{intermediate_1_register}, gp{upper_base_register}, gp{sin_base_register} \n" - generated_code += f"V_MUL_VV gp{intermediate_2_register}, gp{lower_base_register}, gp{cos_base_register} \n" - generated_code += f"V_ADD_VV gp{roped_lower_base_register}, gp{intermediate_1_register}, gp{intermediate_2_register} \n" - generated_code += f"S_ADDI_INT gp{upper_base_register}, gp{upper_base_register}, {mlen} \n" - generated_code += f"S_ADDI_INT gp{lower_base_register}, gp{lower_base_register}, {mlen} \n" - generated_code += f"S_ADDI_INT gp{roped_upper_base_register}, gp{roped_upper_base_register}, {mlen} \n" - generated_code += f"S_ADDI_INT gp{roped_lower_base_register}, gp{roped_lower_base_register}, {mlen} \n" - return generated_code diff --git a/compiler/doc/plena_isa_spec.md b/compiler/doc/plena_isa_spec.md new file mode 100644 index 00000000..92e24c2d --- /dev/null +++ b/compiler/doc/plena_isa_spec.md @@ -0,0 +1,542 @@ +# PLENA Instruction Set Architecture (ISA) Specification + +## Register Types + +The PLENA architecture supports four types of registers: + +- **gp_reg** (`gp0` to `gp15`): General-purpose integer registers +- **fp_reg** (`f0` to `f7`): Floating-point registers +- **hbm_addr_reg** (`a0` to `a7`): HBM address registers + +## Instruction Format + +Instructions follow one of the following formats: + +- `opcode, rd, rs1, rs2, rstride, precision` +- `opcode, rd, rs1, rs2, rmask` (V two sources instructions) +- `opcode, rd, rs1, rmask` (V one source instructions) +- `opcode, 0, rs1, rs2` (M_MM, M_TMM, M_MV, M_TMV) +- `opcode, rd, rs1, rs2` +- `opcode, rd, imm` (M_BMM_WO, M_MV_WO) +- `opcode, rd, rs1, imm` + +## Parameters +Refer to `plena_settings.toml` for the detailed parameters. +- **MLEN**: Tile size used in matrix machine +- **BLEN**: Tile size used in systolic array +- **HLEN**: Tile size used in partitioned systolic array +- **VLEN**: Tile size used in vector machine +- **HBM_M_Prefetch_Amount**: Number of MLEN rows fetched from HBM +- **HBM_V_Prefetch_Amount**: Number of VLEN rows fetched from HBM + + +## Matrix (M-Type) Instructions + +### Notation + +| Notation | Description | +|----------|-------------| +| **Matrix[i]** | i-th entry of the Matrix SRAM | +| **Vector[i]** | i-th entry of the Vector SRAM | + +### M_MM + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `Systolic Array = Vector_SRAM[gp_reg] @ Matrix_SRAM[gp_reg]` + +**Description:** + +Fetch an (BLEN,MLEN) vector from the Vector SRAM using the address provided by `rs1` and an (MLEN, BLEN) matrix from the Matrix SRAM using the address provided by `rs2`. Then, perform an array of dot products. The result matrix (MLEN, BLEN) is internally accumulated in every PE of the systolic array. + +### M_TMM + +**Format:** `opcode, x, rs1, rs2` + +**Operation:** `Systolic Array = Vector[gp_reg] @ Matrix[gp_reg]^T` + +**Description:** + +Similar to `M_MM`, but transposes the matrix. + +### M_BMM + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `Systolic Array = Per Head (Vector_SRAM[gp_reg] @ Matrix_SRAM[gp_reg + gp_reg])` + +`[MLEN // HLEN, MLEN, HLEN] @ [HLEN, MLEN] = [MLEN // HLEN, MLEN, MLEN]` + +**Description:** + +Only take the sliced (HLEN, MLEN) matrix from the Matrix SRAM using the address provided by `gp_reg + gp_reg`, and the vector of shape (MLEN // HLEN, MLEN, HLEN) from the Vector SRAM using the address provided by `gp_reg`. Then, perform an array of dot products. The result matrix [MLEN // HLEN, MLEN, MLEN] is internally accumulated in every PE of the systolic array. + +### M_BTMM + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `Systolic Array = Per Head (Vector_SRAM[gp_reg] @ Matrix_SRAM[gp_reg + gp_reg])^T` + +**Description:** + +Similar to `M_BMM`, but the matrix from Matrix SRAM is transposed before the operation. + +### M_BMM_WO + +**Format:** `opcode, rd, imm` + +**Description:** + +Store the accumulated result [MLEN // HLEN, MLEN, MLEN] to the Vector SRAM at the address specified by `gp_reg + imm` with stride `MLEN // HLEN` and precision `Weights` or `KeyValue` depending on the precision of the MXFP data. + +### M_MM_WO + +**Format:** `opcode, rd, rstride, imm` + +**Description:** + +Output the accumulated result (BLEN, BLEN) stored in the first row of the systolic array to the Vector SRAM at the address specified by `gp_reg` with stride `rstride`. + +### M_MV + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `First Row of Sys Array = Vector[gp_reg] @ Matrix[gp_reg]` +**Description:** + +Fetch an (MLEN, MLEN) matrix from the Matrix SRAM using the address provided by `gp_reg`, and an (MLEN, 1) vector from the Vector SRAM using the address provided by `gp_reg`. Then, perform a dot product and store the resulting (MLEN, 1) vector in the **First Row of Sys Array**. + +### M_TMV + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `First Row of Sys Array = Vector[gp_reg] @ Matrix[gp_reg]^T` + +**Description:** + +This instruction is similar to `M_MV`, but transposes the Matrix when fetching from the Matrix SRAM at the address set by `rs2`. + +### M_BMV (TODO: Implement) + +### M_BTMV (TODO: Implement) + +### M_MV_WO + +**Format:** `opcode, rd, imm` + +**Description:** + +Store the accumulated result (MLEN, 1) stored in the first row of the systolic array to the Vector SRAM at the address specified by `gp_reg + imm` + +### M_BMV_WO (TODO: Implement) + +--- + +## Vector (V-Type) Instructions + +### Notation + +| Notation | Description | +|----------|-------------| +| **Vector[i]** | i-th entry of the Vector SRAM | + +`rmask` is a binary flag indicating whether to apply the mask to the result. The mask is set by the `C_SET_V_MASK_REG` instruction. + +### V_ADD_VV + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) + (Vector[gp_reg]) & gp_rmask` + +**Description:** + +Fetch two (MLEN, 1) vectors from the Vector SRAM using the addresses provided by `rs2` and `rs1`, and then perform element-wise addition. Store the resulting vector back to the Vector SRAM at the address provided by `rd`. + +### V_ADD_VF + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) + Broadcast(fp_reg) & gp_reg` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, then fetch a single floating-point value from the FP register file using the index provided by `rs2`. Broadcast this value by duplicating it to form an (MLEN, 1) vector, and then perform element-wise addition. Store the resulting vector back to Vector SRAM at the address provided by `rd`. + +### V_SUB_VV + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) - (Vector[gp_reg] & gp_reg)` + +**Description:** + +Similar to `V_ADD_VV`, but performs element-wise subtraction. + +### V_SUB_VF + +**Format:** `opcode, rd, rs1, fp2, rmask, rorder` + +**Operation:** +- If `rorder = Normal`: `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) - Broadcast(fp_reg) & gp_reg` +- If `rorder = Reverse`: `Vector[gp_reg] & gp_rmask = Broadcast(fp_reg) & gp_reg - (Vector[gp_reg] & gp_reg)` + +**Description:** + +Similar to `V_ADD_VF`, but performs element-wise subtraction. The `rorder` parameter (decoded from `funct1`) controls the order of subtraction: +- `0` (Normal): vector - scalar +- `>0` (Reverse): scalar - vector + +### V_MUL_VV + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) * (Vector[gp_reg] & gp_reg)` + +**Description:** + +Similar to `V_ADD_VV`, but performs element-wise multiplication. + +### V_MUL_VF + +**Format:** `opcode, rd, rs1, fp2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) * Broadcast(fp_reg) & gp_reg` + +**Description:** + +Similar to `V_ADD_VF`, but performs element-wise multiplication. + +### V_EXP_V + +**Format:** `opcode, rd, rs1, x, rmask` + +**Operation:** `Vector[gp_reg] = exp(Vector[gp_reg])` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, perform element-wise exponentiation, and store the resulting vector back into the Vector SRAM at the address specified by `rd`. + +### V_RECI_V + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `Vector[gp_reg] = reciprocal(Vector[gp_reg])` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, perform element-wise reciprocal, and store the resulting vector back into the Vector SRAM at the address specified by `rd`. + +### V_RED_SUM + +**Format:** `opcode, rd, rs1, 0` + +**Operation:** `fp_reg = sum(Vector[gp_reg], fp_reg)` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, and a single floating-point value from the FP register file using the index specified by `rd`. Perform element-wise addition on the combined vector, and store the resulting sum back into the FP register file at the index specified by `rd`. This instruction is designed to facilitate continuous summation along a high-dimension vector. + +### V_RED_MAX + +**Format:** `opcode, rd, rs1, 0` + +**Operation:** `fp_reg = max(Vector[gp_reg], fp_reg)` + +**Description:** + +Similar to `V_RED_SUM` but performs the max value selection operation. + +--- + +## Scalar (S-Type) Instructions + +### Integer Operations + +#### Notation + +| Notation | Description | +|----------|-------------| +| **INT_MEM[i]** | i-th entry of the SRAM within the scalar machine specifically designed for integer operations | + +#### S_ADD_INT + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `gp_reg = gp_reg + gp_reg` + +#### S_ADDI_INT + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `gp_reg = gp_reg + imm2` + +#### S_SUB_INT + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `gp_reg = gp_reg - gp_reg` + +#### S_MUL_INT + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `gp_reg = gp_reg * gp_reg` + +#### S_LUI_INT + +**Format:** `opcode, rd, imm` + +**Operation:** `gp_reg = imm << 12` + +**Description:** + +Load upper immediate value into the integer register. + +#### S_LD_INT + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `gp_reg = INT_MEM[gp_reg + imm2]` + +#### S_ST_INT + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `INT_MEM[gp_reg + imm2] = gp_reg` + +### Floating-Point Operations + +#### Notation + +| Notation | Description | +|----------|-------------| +| **FP_MEM[i]** | i-th entry of the SRAM within the scalar machine specifically designed for floating-point operations | + +#### S_ADD_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = fp_reg + fp_reg` + +#### S_SUB_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = fp_reg - fp_reg` + +#### S_MAX_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = max(fp_reg, fp_reg)` + +#### S_MUL_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = fp_reg * fp_reg` + +#### S_EXP_FP + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `fp_reg = exp(fp_reg)` + +#### S_RECI_FP + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `fp_reg = reciprocal(fp_reg)` + +#### S_SQRT_FP + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `fp_reg = sqrt(fp_reg)` + +#### S_LD_FP + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `fp_reg = FP_MEM[gp_reg + imm2]` + +#### S_ST_FP + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `FP_MEM[gp_reg + imm2] = fp_reg` + +#### S_MAP_V_FP + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `Vector[gp_reg :+ VLEN] = FP_MEM[gp_reg + imm2 :+ VLEN]` + +**Description:** + +Copy a vector of length VLEN from FP_MEM to Vector SRAM. + +--- + +## Memory (H-Type) Instructions + +### Notation + +| Notation | Description | +|----------|-------------| +| **Matrix[i]** | The i-th entry of the Matrix SRAM | +| **Vector[i]** | The i-th entry of the Vector SRAM | +| **HBM[i]** | The i-th entry of the HBM | + +### H_PREFETCH_M + +**Format:** `opcode, rd, rs1, rs2, rstride, precision` + +**Operation:** `Matrix[gp_reg] = HBM[gp_reg + hbm_addr_reg_]` + +**Description:** + +Prefetch a matrix of size **HBM_M_Prefetch_Amount × MLEN** from the HBM to the Matrix SRAM, with a stride width specified by **hbm_stride_reg[rstride]**. + +The `precision` field (decoded from `funct1`) determines the data precision: +- `0` (Weights): High precision weights +- `>0` (KeyValue): Lower precision key-value data + +### H_PREFETCH_V + +**Format:** `opcode, rd, rs1, rs2, rstride, precision` + +**Operation:** `Vector[gp_reg] = HBM[gp_reg + hbm_addr_reg_]` + +**Description:** + +Prefetch a matrix of size **HBM_V_Prefetch_Amount × VLEN** from the HBM to the Vector SRAM, with a stride width specified by **hbm_stride_reg[rstride]**. + +The `precision` field (decoded from `funct1`) determines the data precision: +- `0` (Activation): High precision activation data +- `>0` (KeyValue): Lower precision key-value data + +### H_STORE_V + +**Format:** `opcode, rd, rs1, rs2, rstride, precision` + +**Operation:** `HBM[gp_reg + hbm_addr_reg_] = Vector[gp_reg]` + +**Description:** + +Store a matrix of size **HBM_V_Writeback_Amount × VLEN** from Vector SRAM to HBM, with a stride width specified by **hbm_stride_reg[rstride]**. + +The `precision` field (decoded from `funct1`) determines the data precision: +- `0` (Activation): High precision activation data +- `>0` (KeyValue): Lower precision key-value data + +--- + +## Control and Status Register (C-Type) Instructions + +### C_SET_ADDR_REG + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `hbm_addr_reg_ = {gp_reg, gp_reg}` + +**Description:** + +This instruction is used to set the value of `hbm_addr_reg[rd]`, assuming it has double the bit width of `fix_reg`, by concatenating two `fix_reg` entries and storing the result in `hbm_addr_reg[rd]`. The concatenation order is `{rs2, rs1}`. + +### C_SET_SCALE_REG + +**Format:** `opcode, rd` + +**Operation:** `SCALE_OFFSET = rd` + +**Description:** + +This instruction is used to set the scale offset. The blocks and scales of the MXFP data are stored separately in HBM for memory alignment purposes. Their distance is set by the scale offset. This value differs depending on the precision of the MXFP and the data size. For example, for Q(B, S, H, D), the scales are stored after the blocks, and the offset is `B × S × H × D × (EXP_WIDTH + MANT_WIDTH + 1) / 8`. + +### C_SET_STRIDE_REG + +**Format:** `opcode, rd` + +**Operation:** `STRIDE_SIZE = rd` + +**Description:** + +This instruction is used to set the stride size for the prefetch instructions. + +### C_SET_V_MASK_REG + +**Format:** `opcode, rd` + +**Operation:** `V_MASK = rd` + +**Description:** + +This instruction is used to set the vector mask register for masked vector operations. + +### C_BREAK + +**Format:** `opcode, 0, 0, 0` + +**Operation:** Breakpoint exception + +**Description:** + +Triggers a breakpoint exception, typically used for debugging purposes. + +--- + +## Instruction Encoding Summary + +### Opcode Map + +| Opcode | Instruction | Type | +|--------|-------------|------| +| 0x00 | Invalid | - | +| 0x01 | M_MM | M-Type | +| 0x02 | M_TMM | M-Type | +| 0x03 | M_BMM | M-Type | +| 0x04 | M_BTMM | M-Type | +| 0x05 | M_BMM_WO | M-Type | +| 0x06 | M_MM_WO | M-Type | +| 0x07 | M_MV | M-Type | +| 0x08 | M_TMV | M-Type | +| 0x09 | M_BMV | M-Type | +| 0x0A | M_BTMV | M-Type | +| 0x0B | M_MV_WO | M-Type | +| 0x0C | M_BMV_WO | M-Type | +| 0x0D | V_ADD_VV | V-Type | +| 0x0E | V_ADD_VF | V-Type | +| 0x0F | V_SUB_VV | V-Type | +| 0x10 | V_SUB_VF | V-Type | +| 0x11 | V_MUL_VV | V-Type | +| 0x12 | V_MUL_VF | V-Type | +| 0x13 | V_EXP_V | V-Type | +| 0x14 | V_RECI_V | V-Type | +| 0x15 | V_RED_SUM | V-Type | +| 0x16 | V_RED_MAX | V-Type | +| 0x17 | S_ADD_FP | S-Type | +| 0x18 | S_SUB_FP | S-Type | +| 0x19 | S_MAX_FP | S-Type | +| 0x1A | S_MUL_FP | S-Type | +| 0x1B | S_EXP_FP | S-Type | +| 0x1C | S_RECI_FP | S-Type | +| 0x1D | S_SQRT_FP | S-Type | +| 0x1E | S_LD_FP | S-Type | +| 0x1F | S_ST_FP | S-Type | +| 0x20 | S_MAP_V_FP | S-Type | +| 0x21 | S_ADD_INT | S-Type | +| 0x22 | S_ADDI_INT | S-Type | +| 0x23 | S_SUB_INT | S-Type | +| 0x24 | S_MUL_INT | S-Type | +| 0x25 | S_LUI_INT | S-Type | +| 0x26 | S_LD_INT | S-Type | +| 0x27 | S_ST_INT | S-Type | +| 0x28 | H_PREFETCH_M | H-Type | +| 0x29 | H_PREFETCH_V | H-Type | +| 0x2A | H_STORE_V | H-Type | +| 0x2B | C_SET_ADDR_REG | C-Type | +| 0x2C | C_SET_SCALE_REG | C-Type | +| 0x2D | C_SET_STRIDE_REG | C-Type | +| 0x2E | C_SET_V_MASK_REG | C-Type | +| 0x2F | C_BREAK | C-Type | \ No newline at end of file diff --git a/justfile b/justfile index 2231845a..2e4aa42e 100644 --- a/justfile +++ b/justfile @@ -44,7 +44,7 @@ build-behave-sim-debug arg: cd behavioral_simulator && \ RUST_BACKTRACE=1 cargo run --release -- --opcode "$asm_path" --hbm "$data_path" --fpsram "$fp_sram_path" python3 behavioral_simulator/testbench/view_mem.py - python3 tools/utils/compare_match.py + # python3 tools/utils/compare_match.py build-rtl-sim arg: rm -rf test/Instr_Level_Benchmark/build/{{arg}} diff --git a/src/definitions/operation.svh b/src/definitions/operation.svh index 19dfb6bd..1a974fce 100644 --- a/src/definitions/operation.svh +++ b/src/definitions/operation.svh @@ -218,4 +218,4 @@ typedef struct { logic update_v_waddr; } OP_BUNDLE; -`endif +`endif \ No newline at end of file diff --git a/src/definitions/plena_settings.toml b/src/definitions/plena_settings.toml index 3f27e552..f355eecf 100644 --- a/src/definitions/plena_settings.toml +++ b/src/definitions/plena_settings.toml @@ -19,6 +19,9 @@ value = 4 [CONFIG.HBM_SIZE] value = 1073741824 +[CONFIG.HBM_WIDTH] +value = 512 + [CONFIG.MATRIX_SRAM_SIZE] value = 1024 @@ -106,6 +109,12 @@ sign = false exponent = 8 mantissa = 0 +[PRECISION.HBM_V_INT_TYPE] +format = "Plain" +[PRECISION.HBM_V_INT_TYPE.DATA_TYPE] +type = "Int" +width = 32 + [PRECISION.SCALAR_FP] type = "Fp" sign = true diff --git a/tools/assembler/assembly_to_binary.py b/tools/assembler/assembly_to_binary.py index 1c38e77e..36b12aab 100644 --- a/tools/assembler/assembly_to_binary.py +++ b/tools/assembler/assembly_to_binary.py @@ -38,6 +38,7 @@ def _convert_to_binary(self, instruction): rs2 = instruction.rs2 rstride = instruction.rstride funct1 = instruction.funct1 + funct2 = instruction.funct2 imm = instruction.imm rmask = instruction.rmask binary_instruction = 0 @@ -69,7 +70,7 @@ def _convert_to_binary(self, instruction): (rd << opw) + opcode ) - elif instruction.opcode in [ "H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V"]: + elif instruction.opcode in ["H_PREFETCH_M", "V_SUB_VF"]: binary_instruction = ( (funct1 << (opw + 4 * ow)) + (rstride << (opw + 3 * ow)) + @@ -78,7 +79,17 @@ def _convert_to_binary(self, instruction): (rd << opw) + opcode ) - elif instruction.opcode in ["V_ADD_VV", "V_ADD_VF", "V_SUB_VV", "V_SUB_VF", "V_MUL_VV", "V_MUL_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"]: + elif instruction.opcode in ["H_PREFETCH_V", "H_STORE_V"]: + binary_instruction = ( + (funct2 << (opw + 5 * ow)) + + (funct1 << (opw + 4 * ow)) + + (rstride << (opw + 3 * ow)) + + (rs2 << (opw + 2 * ow)) + + (rs1 << (opw + ow)) + + (rd << opw) + + opcode + ) + elif instruction.opcode in ["V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"]: binary_instruction = ( (rmask << (opw + 3 * ow)) + (rs2 << (opw + 2 * ow)) + diff --git a/tools/assembler/parser.py b/tools/assembler/parser.py index 5c31533e..e81d6d55 100644 --- a/tools/assembler/parser.py +++ b/tools/assembler/parser.py @@ -62,7 +62,7 @@ def load_isa_settings(file_path: str) -> dict: class Instruction: - def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], rstride: Optional[str], funct1: Optional[int], imm: Optional[int] = None): + def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], rstride: Optional[str], funct1: Optional[int], funct2: Optional[int], imm: Optional[int] = None, rflag: Optional[int] = None): self.opcode = opcode self.rd = rd @@ -70,11 +70,12 @@ def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], self.rs2 = rs2 self.rstride = rstride self.funct1 = funct1 + self.funct2 = funct2 self.imm = imm self.rmask = rstride def __repr__(self): - return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, imm={self.imm})" + return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, funct2={self.funct2}, imm={self.imm}, rflag={self.rflag})" def parse_asm_file(file_path: str) -> List[Instruction]: @@ -117,6 +118,7 @@ def parse_asm_file(file_path: str) -> List[Instruction]: rs2 = None rstride = None funct1 = None + funct2 = None imm = None # Helper to parse a register or int operand @@ -221,9 +223,44 @@ def parse_reg_or_int(operand): funct1 = int(funct1_raw) except ValueError: funct1 = funct1_raw # fallback, if not int, keep as string + elif len(operands) == 6: + operand_0, operand_1, operand_2, operand_3, operand_4, operand_5 = operands + rd = parse_reg_or_int(operand_0) + if operand_1.strip().startswith(('gp','f','a')): + rs1 = parse_reg_or_int(operand_1) + else: + try: + imm = int(operand_1) + except ValueError: + imm = None + if operand_2.strip().startswith(('gp','f','a')): + rs2 = parse_reg_or_int(operand_2) + else: + try: + imm = int(operand_2) + except ValueError: + pass + try: + rstride = int(operand_3) + except ValueError: + rstride = None + funct1_raw = operand_4.strip() + if funct1_raw.endswith(';'): + funct1_raw = funct1_raw[:-1] + try: + funct1 = int(funct1_raw) + except ValueError: + funct1 = funct1_raw # fallback, if not int, keep as string + funct2_raw = operand_5.strip() + if funct2_raw.endswith(';'): + funct2_raw = funct2_raw[:-1] + try: + funct2 = int(funct2_raw) + except ValueError: + funct2 = funct2_raw # fallback, if not int, keep as string - instructions.append(Instruction(opcode, rd, rs1, rs2, rstride, funct1, imm)) + instructions.append(Instruction(opcode, rd, rs1, rs2, rstride, funct1, funct2, imm)) return instructions diff --git a/tools/cost_model/latency/ablation_study.py b/tools/cost_model/latency/ablation_study.py new file mode 100644 index 00000000..e8fb334e --- /dev/null +++ b/tools/cost_model/latency/ablation_study.py @@ -0,0 +1,334 @@ +import json +import os +from pathlib import Path +import matplotlib.pyplot as plt +import numpy as np +from overall_inference_estimation import model_config + +batch_size = 4 + +# Experiment configurations +# sample_1: baseline, MLEN = BLEN = 64, VLEN = 1024, using selfattention +# sample_2: MLEN = BLEN = 64, VLEN = 1024, using flash attention +# sample_3: MLEN = 1024, BLEN = 4, VLEN = 1024, using flash attention, not partitioned optimised +# sample_4: MLEN = 1024, BLEN = 4, VLEN = 1024, using flash attention, partitioned optimised + + +def compute_attention_ffn_distribution(model_param_path, hardware_config, batch_size, input_seq, decode_seq_lengths, + use_flash_attention=True, partitioned_optimized=False, device_num=1): + """ + Compute attention and FFN computation distribution for prefill and decode stages at different sequence lengths. + + Args: + model_param_path: Path to model parameter JSON file + hardware_config: Dictionary with MLEN, BLEN, VLEN + batch_size: Batch size + input_seq: Input sequence length (for prefill) + decode_seq_lengths: List of total sequence lengths for decode (e.g., [10k, 40k, 80k]) + use_flash_attention: Whether to use flash attention (True) or self attention (False) + partitioned_optimized: Whether to use partitioned optimization (for sample 4) + device_num: Number of devices + + Returns: + dict: Dictionary with 'prefill' and 'decode' keys + 'prefill' contains 'attention' and 'ffn' for input_seq + 'decode' contains list of dicts, each with 'attention' and 'ffn' for each decode_seq_length + """ + print("<", "="*10, "starting computation", "="*10, ">") + print("hardware config: ", hardware_config) + print("use flash attention: ", use_flash_attention) + print("partitioned optimized: ", partitioned_optimized) + + print("starting computation") + # Use max decode length for model initialization + max_decode_len = max(decode_seq_lengths) if decode_seq_lengths else 0 + model = model_config( + model_param_path=model_param_path, + hardware_config=hardware_config, + batch_size=batch_size, + seq_len=int(input_seq), + output_token=1000, + device_num=device_num + ) + + results = {} + + # Prefill stage + mode = "prefill" + if use_flash_attention: + attention_inst = model.flash_attention(mode, partitioned_optimized) + else: + attention_inst = model.self_attention(mode) + + ffn_inst = model.feed_forward(mode) + + results['prefill'] = { + 'attention': attention_inst, + 'ffn': ffn_inst, + 'total': attention_inst + ffn_inst + } + print("completed prefill computation") + print("prefill attention: ", attention_inst) + print("prefill ffn: ", ffn_inst) + + # Decode stage - compute for each sequence length + mode = "decode" + results['decode'] = [] + + for total_seq_len in decode_seq_lengths: + model.kv_size = total_seq_len # Reset KV size to initial value + + # Accumulate instruction counts as KV cache grows + attention_inst_decode = 0 + ffn_inst_decode = 0 + + for token_idx in range(total_seq_len): + if use_flash_attention: + attention_inst_per_token = model.flash_attention(mode, partitioned_optimized) + else: + attention_inst_per_token = model.self_attention(mode) + ffn_inst_per_token = model.feed_forward(mode) + + attention_inst_decode += attention_inst_per_token + ffn_inst_decode += ffn_inst_per_token + + # KV cache grows by 1 for each token + model.kv_size += 1 + + results['decode'].append({ + 'attention': attention_inst_decode, + 'ffn': ffn_inst_decode, + 'total': attention_inst_decode + ffn_inst_decode + }) + print(f"completed decode computation for seq_len={total_seq_len}") + print(f" decode attention: {attention_inst_decode:.2e}, ffn: {ffn_inst_decode:.2e}") + + return results + + +def plot_ablation_study(): + """ + Conduct ablation study experiments and plot bar chart showing + FFN + Attention computation distribution across sequence lengths. + Three plots: Left (Prefill 5.6k), Middle (Decode 5k), Right (Decode 8k) + Y-axis: Execution Time (% relative to S1) + """ + colors = [ + "#762a83", # deep purple + "#af8dc3", # light violet + "#e7d4e8", # very light purple-pink + "#d9f0d3", # pale green + "#7fbf7b", # medium green + # Optional sixth if needed: "#1b7837" # dark green + ] + + # Get model parameter path + current_dir = Path(__file__).resolve().parents[3] + model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + + # Sequence lengths: prefill at 5.6k, decode at 5k and 8k total + input_seq = 5.6 * 1000 + decode_seq_lengths = [5 * 1000, 8 * 10000] # Total sequence lengths after decode (5k and 8k) + + # Define experiment samples + experiments = [ + { + 'name': 'Sample 1', + 'hardware_config': {'MLEN': 64, 'BLEN': 64, 'VLEN': 1024}, + 'use_flash_attention': False, + 'partitioned_optimized': False, + 'label': 'Sample 1 (Baseline): MLEN=BLEN=64, VLEN=1024' + }, + { + 'name': 'Sample 2', + 'hardware_config': {'MLEN': 64, 'BLEN': 64, 'VLEN': 1024}, + 'use_flash_attention': True, + 'partitioned_optimized': False, + 'label': 'Sample 2 (Flash Attention): MLEN=BLEN=64, VLEN=1024' + }, + { + 'name': 'Sample 3', + 'hardware_config': {'MLEN': 1024, 'BLEN': 4, 'VLEN': 1024}, + 'use_flash_attention': True, + 'partitioned_optimized': True, + 'label': 'Sample 3 (Flash Attention + Flattened Systolic Array): MLEN=1024, BLEN=4, VLEN=1024' + } + ] + + # Collect data for all sequence lengths + # Structure: data[sample_idx][seq_idx] = {'attention': ..., 'ffn': ...} + all_data = [] + experiment_names = [] + + for exp in experiments: + results = compute_attention_ffn_distribution( + model_param_path=model_param_path, + hardware_config=exp['hardware_config'], + batch_size=batch_size, + input_seq=input_seq, + decode_seq_lengths=decode_seq_lengths, + use_flash_attention=exp['use_flash_attention'], + partitioned_optimized=exp['partitioned_optimized'] + ) + + experiment_names.append(exp['name']) + sample_data = [] + + # Add prefill data (at 5.6k) + sample_data.append({ + 'attention': results['prefill']['attention'], + 'ffn': results['prefill']['ffn'] + }) + + # Add decode data (at 10k, 40k, 80k) + for decode_result in results['decode']: + sample_data.append({ + 'attention': decode_result['attention'], + 'ffn': decode_result['ffn'] + }) + + all_data.append(sample_data) + + print("completed computation") + + # Prepare data for plotting + num_samples = len(experiments) + + # Extract data for each plot: prefill (index 0), decode 5k (index 1), decode 8k (index 2) + prefill_attention = np.array([all_data[s][0]['attention'] for s in range(num_samples)]) + prefill_ffn = np.array([all_data[s][0]['ffn'] for s in range(num_samples)]) + + decode_5k_attention = np.array([all_data[s][1]['attention'] for s in range(num_samples)]) + decode_5k_ffn = np.array([all_data[s][1]['ffn'] for s in range(num_samples)]) + + decode_8k_attention = np.array([all_data[s][2]['attention'] for s in range(num_samples)]) + decode_8k_ffn = np.array([all_data[s][2]['ffn'] for s in range(num_samples)]) + + # Normalize all values relative to S1 (S1 = 100%) + s1_prefill_total = prefill_attention[0] + prefill_ffn[0] + s1_decode_5k_total = decode_5k_attention[0] + decode_5k_ffn[0] + s1_decode_8k_total = decode_8k_attention[0] + decode_8k_ffn[0] + + # Normalize prefill + prefill_attention_norm = (prefill_attention / s1_prefill_total) * 100 + prefill_ffn_norm = (prefill_ffn / s1_prefill_total) * 100 + + # Normalize decode 5k + decode_5k_attention_norm = (decode_5k_attention / s1_decode_5k_total) * 100 + decode_5k_ffn_norm = (decode_5k_ffn / s1_decode_5k_total) * 100 + + # Normalize decode 8k + decode_8k_attention_norm = (decode_8k_attention / s1_decode_8k_total) * 100 + decode_8k_ffn_norm = (decode_8k_ffn / s1_decode_8k_total) * 100 + + # Create figure with 3 subplots, shared y-axis - reduced height for compactness + fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(10, 3), sharey=True) + + # Colors - using the color palette defined above + attention_color = colors[0] # dark blue + ffn_color = colors[4] # teal + + # Bar positioning - make more compact + x = np.arange(num_samples) + bar_width = 0.5 # Slightly wider bars for better visibility + + # Find maximum value across all plots for consistent y-axis + max_total = max( + (prefill_attention_norm + prefill_ffn_norm).max(), + (decode_5k_attention_norm + decode_5k_ffn_norm).max(), + (decode_8k_attention_norm + decode_8k_ffn_norm).max() + ) + y_max = max_total * 1.05 # Reduced padding to 5% for more compact y-axis + + # Data for each plot (normalized) + plot_data = [ + { + 'ax': ax1, + 'xlabel': 'Prefill (5.6k)', + 'attention': prefill_attention_norm, + 'ffn': prefill_ffn_norm + }, + { + 'ax': ax2, + 'xlabel': 'Decode (10k)', + 'attention': decode_5k_attention_norm, + 'ffn': decode_5k_ffn_norm + }, + { + 'ax': ax3, + 'xlabel': 'Decode (80k)', + 'attention': decode_8k_attention_norm, + 'ffn': decode_8k_ffn_norm + } + ] + + # Plot each subplot + for plot_info in plot_data: + ax = plot_info['ax'] + attn_data = plot_info['attention'] + ffn_data = plot_info['ffn'] + + # Plot stacked bars + ax.bar(x, attn_data, bar_width, label='Attention', color=attention_color, edgecolor='black', alpha=0.8) + ax.bar(x, ffn_data, bar_width, bottom=attn_data, label='FFN', color=ffn_color, edgecolor='black', alpha=0.8) + + # X-axis - make more compact + ax.set_xticks(x) + ax.set_xticklabels([f'S{i+1}' for i in range(num_samples)], fontsize=13) + # Set title at the top instead of x-axis label + ax.set_title(plot_info['xlabel'], fontsize=14) + + # Set x-axis limits to reduce space between bars + ax.set_xlim(-0.5, num_samples - 0.5) + + # Set shared y-axis limits with minimal padding + ax.set_ylim(0, y_max) + ax.grid(axis='y', alpha=0.3, linestyle='--') + + # Reduce tick label sizes for more compact appearance + ax.tick_params(axis='both', labelsize=12) + + # Add legend only to the first plot + if ax == ax1: + ax.legend(loc='upper right', fontsize=13, framealpha=0.9) + # Y-axis label only on the leftmost plot - positioned lower + ax.set_ylabel('Exec Time (% relative to S1)', fontsize=12) + # Move y-axis label down by adjusting its position + ax.yaxis.set_label_coords(-0.15, 0.50) + else: + # Remove y-axis labels from middle and right plots + ax.set_ylabel('') + + # Add overall title + # fig.suptitle('Ablation Study for LLaMA 3.1 8B', fontsize=18, y=1.02) + + # Add sample legend at the bottom in three rows (one per sample) + sample_legend = [exp["label"] for exp in experiments] + # Arrange in three rows: one sample per row + legend_text = '\n'.join(sample_legend) + + # More compact vertical spacing + plt.subplots_adjust(wspace=0.15, left=0.12, right=0.95, top=0.88, bottom=0.25) + + # Add sample legend at the bottom in three rows + # fig.text(0.5, 0.12, legend_text, + # fontsize=12, va='top', ha='center', + # bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=4)) + + plt.savefig('ablation_study_attention_ffn.png', dpi=300, bbox_inches='tight') + print("Plot saved as 'ablation_study_attention_ffn.png'") + + # Print summary + print("\n" + "="*70) + print("Ablation Study Results Summary") + print("="*70) + for i, exp in enumerate(experiments): + print(f"\n{exp['label']}:") + print(f" Prefill (5.6k) - Attention: {prefill_attention[i]:.2e}, FFN: {prefill_ffn[i]:.2e}") + print(f" Decode (5k) - Attention: {decode_5k_attention[i]:.2e}, FFN: {decode_5k_ffn[i]:.2e}") + print(f" Decode (8k) - Attention: {decode_8k_attention[i]:.2e}, FFN: {decode_8k_ffn[i]:.2e}") + print("="*70) + + +if __name__ == "__main__": + plot_ablation_study() diff --git a/tools/cost_model/latency/latency_model.py b/tools/cost_model/latency/latency_model.py index 520a16e7..769004e3 100644 --- a/tools/cost_model/latency/latency_model.py +++ b/tools/cost_model/latency/latency_model.py @@ -78,10 +78,10 @@ def obtain_per_instr_alone_latency(self, output_file: str = "instr_alone_latency print(f"Alone latency model saved to {output_file}") def obtain_overall_latency(self, updated_config): - batch_size = 1024 - input_seq_len = 1024 - output_seq_len = 128 - device_num = 4 + batch_size = 4 + input_seq_len = 2048 + output_seq_len = 1024 + device_num = 1 hardware_settings = self.hardware_config for key, value in updated_config.items(): hardware_settings[key] = value diff --git a/tools/cost_model/latency/overall_inference_estimation.py b/tools/cost_model/latency/overall_inference_estimation.py index 793ce487..ff80547f 100644 --- a/tools/cost_model/latency/overall_inference_estimation.py +++ b/tools/cost_model/latency/overall_inference_estimation.py @@ -24,6 +24,7 @@ def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = self.theoratical_frequency = 10**9 # 1 GHz self.hardware_config = hardware_config self.batch_size = batch_size + self.device_batch_size = batch_size // device_num self.kv_size = seq_len self.device_num = device_num print("=" * 15, "Model Settings","=" * 15) @@ -31,6 +32,7 @@ def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = print("batch size: ", self.batch_size) print("input token: ", self.input_token) print("output token: ", self.output_token) + print("hidden size: ", self.hidden_size) print("head dim: ", self.head_dim) print("num key value heads: ", self.num_key_value_heads) print("num attention heads: ", self.num_attention_heads) @@ -54,66 +56,129 @@ def rms_layer(self, mode = "prefill"): instruction_num = 0 instruction_num += setting_inst_num instruction_num += loop_num * loop_inst_num - return instruction_num * self.batch_size + return instruction_num * self.device_batch_size def projection(self, mode = "prefill"): # Compute Q, K, V Projection + RoPE if mode == "prefill": overall_inst_num = 0 # Q, K Projection + RoPE - overall_inst_num += self.batch_size * (math.ceil(self.hidden_size / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) - overall_inst_num += self.batch_size * (self.num_attention_heads * (self.input_token // self.hardware_config["VLEN"])) * 3 + overall_inst_num += self.device_batch_size * (math.ceil(self.hidden_size / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) + overall_inst_num += self.device_batch_size * (self.num_attention_heads * math.ceil(self.input_token / self.hardware_config["VLEN"])) * 3 - overall_inst_num += self.batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) - overall_inst_num += self.batch_size * (self.num_key_value_heads * (self.input_token // self.hardware_config["VLEN"])) * 3 + overall_inst_num += self.device_batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) + overall_inst_num += self.device_batch_size * (self.num_key_value_heads * math.ceil(self.input_token / self.hardware_config["VLEN"])) * 3 # V - overall_inst_num += self.batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) + overall_inst_num += self.device_batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) elif mode == "decode": overall_inst_num = 0 # Q, K Projection + RoPE overall_inst_num += (math.ceil(self.hidden_size / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * 2 + 10)) - overall_inst_num += self.batch_size * (self.num_attention_heads) * 4 + overall_inst_num += self.device_batch_size * (self.num_attention_heads) * 4 overall_inst_num += (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * 2 + 10)) - overall_inst_num += self.batch_size * (self.num_key_value_heads) * 4 + overall_inst_num += self.device_batch_size * (self.num_key_value_heads) * 4 # V overall_inst_num += (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * 2 + 10)) return overall_inst_num - def flash_attention(self, mode = "prefill"): + def flash_attention(self, mode = "prefill", partitoned_optimised = True): overall_inst_num = 0 mlen = self.hardware_config["MLEN"] blen = self.hardware_config["BLEN"] - tile_in_atten = min(self.head_dim, mlen) + prefill_len = min(self.input_token, mlen) + decode_len = min(self.kv_size, mlen) + max_head_per_mlen = math.ceil(mlen / self.head_dim) + + per_head_iter = math.ceil(self.num_attention_heads / max_head_per_mlen) + if mode == "prefill": # Outer loop - for kv_head_index in range(self.num_key_value_heads): - for i in range(math.ceil(self.input_token // self.hardware_config["MLEN"])): - overall_inst_num += mlen * 2 # Reset - for j in range(math.ceil(self.input_token // self.hardware_config["MLEN"])): - overall_inst_num += mlen * 2 # QKT - overall_inst_num += 2 + mlen * 14 # Softmax - overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(mlen / blen) * 4) #PV - overall_inst_num += mlen * 5 + 4 #Compute O + if partitoned_optimised: + for i in range(math.ceil(self.input_token / mlen)): + overall_inst_num += mlen * max_head_per_mlen # Reset + for j in range(math.ceil(self.input_token / mlen)): + # overall_inst_num += math.ceil(prefill_len / blen) * blen * math.ceil(prefill_len / blen) * 5 # QKT + overall_inst_num += (2 + prefill_len * 5) * max_head_per_mlen # Softmax + overall_inst_num += math.ceil(mlen / blen) * blen * math.ceil(self.head_dim / blen) * max_head_per_mlen + overall_inst_num += (prefill_len * 5 + 4) * max_head_per_mlen# Compute O + overall_inst_num += 8 * max_head_per_mlen + self.kv_size = self.input_token + return overall_inst_num * per_head_iter * self.device_batch_size + else: + for i in range(math.ceil(self.input_token / mlen)): + # overall_inst_num += mlen # Reset + for j in range(math.ceil(self.input_token / mlen)): + # (mlen, head_dim) @ (head_dim, mlen) + overall_inst_num += math.ceil(prefill_len / blen) * (math.ceil(self.head_dim / mlen)) * blen * math.ceil(prefill_len / blen) + overall_inst_num += 2 + prefill_len * 5 # Softmax + overall_inst_num += math.ceil(mlen / blen) * blen * math.ceil(self.head_dim / blen) + overall_inst_num += prefill_len * 5 + 4 # Compute O overall_inst_num += 8 - self.kv_size = self.input_token + return overall_inst_num * self.device_batch_size * self.num_attention_heads elif mode == "decode": - for kv_head_index in range(self.num_key_value_heads): - for i in range(math.ceil(self.kv_size // self.hardware_config["MLEN"])): - # overall_inst_num += mlen * 2 # Reset - overall_inst_num += mlen * 2 # QKT - overall_inst_num += 2 + mlen * 14 # Softmax - overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(mlen / blen) * 4) #PV - overall_inst_num += mlen * 5 + 4 #Compute O - overall_inst_num += 8 - self.kv_size = self.kv_size + 1 - return overall_inst_num * self.batch_size + if partitoned_optimised: + overall_inst_num = decode_len + overall_inst_num += (2 + decode_len * 3) * max_head_per_mlen # Softmax + # PV (1, decode_len) @ (decode_len, head_dim) + overall_inst_num += math.ceil(decode_len / mlen) * blen * math.ceil(self.head_dim / blen) * max_head_per_mlen #PV + overall_inst_num += (decode_len * 3 + 4) * max_head_per_mlen #Compute O + overall_inst_num += 8 * max_head_per_mlen + return overall_inst_num * self.device_batch_size * math.ceil(self.kv_size / self.hardware_config["MLEN"]) * (per_head_iter) + else: + overall_inst_num = math.ceil(mlen / blen) * (math.ceil(self.head_dim / mlen) * blen) + overall_inst_num += (2 + decode_len * 3) + overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(decode_len / blen) * 4) + overall_inst_num += (decode_len * 3 + 4) + overall_inst_num += 8 + return overall_inst_num * self.device_batch_size * math.ceil(self.kv_size / self.hardware_config["MLEN"]) * self.num_attention_heads + def self_attention(self, mode = "prefill"): + # Note: this it the latency estimation model for self-attention without using the flash attention algorithm + overall_inst_num = 0 + mlen = self.hardware_config["MLEN"] + blen = self.hardware_config["BLEN"] + vlen = self.hardware_config["VLEN"] + if mode == "prefill": + for i in range(self.num_attention_heads): + # QKT (batch, s, h, d) @ (batch, s, h, d) + overall_inst_num += (math.ceil(self.input_token / blen)) * math.ceil(self.head_dim / mlen) * blen * math.ceil(self.input_token / blen) * 4 + # Store QKT (h, s, s) + overall_inst_num += self.input_token * 4 * (math.ceil(self.input_token / mlen)) * self.num_attention_heads * 2 + # Softmax (batch, s, h, d) + overall_inst_num += self.input_token * (math.ceil(self.input_token / mlen)) * 30 * self.num_attention_heads + # PV (h, s, s) @ (h, s, d) + for i in range(self.num_attention_heads): + overall_inst_num += math.ceil(self.input_token / blen) * math.ceil(self.input_token / mlen) * blen * math.ceil(self.head_dim / blen) + # Store PV (h, s, d) + overall_inst_num += self.head_dim * 4 * (math.ceil(self.input_token / mlen)) * self.num_attention_heads* 2 + # Compute O (batch, s, h, d) + for i in range(self.num_attention_heads): + overall_inst_num += (math.ceil(self.input_token / vlen)) * 5 + 4 + overall_inst_num += 8 + overall_inst_num = overall_inst_num * self.device_batch_size + elif mode == "decode": + for i in range(self.num_attention_heads): + # QKT (batch, 1, h, d) @ (batch, kv, h, d) + overall_inst_num += 4 * math.ceil(self.head_dim / mlen) * math.ceil(self.kv_size / blen) * blen + # Store QKT (h, s, s) + overall_inst_num += 4 * math.ceil(self.kv_size / mlen) * self.num_attention_heads * 36 + # Softmax (batch, s, h, d) + overall_inst_num += (math.ceil(self.kv_size / mlen)) * 30 * self.num_attention_heads + # PV (h, 1, s) @ (h, s, d) + for i in range(self.num_attention_heads): + overall_inst_num += math.ceil(self.kv_size / mlen) * blen * (self.head_dim // blen) * 2 + # Compute O (batch, s, h, d) + for i in range(self.num_attention_heads): + overall_inst_num += (math.ceil(self.kv_size / mlen)) * 5 + 4 + overall_inst_num = overall_inst_num * self.device_batch_size + return overall_inst_num + def residual (self, mode = "prefill"): overall_inst_num = 0 # -- Residual @@ -123,7 +188,7 @@ def residual (self, mode = "prefill"): elif mode == "decode": iteration = self.hidden_size // self.hardware_config["VLEN"] overall_inst_num = 5 * iteration + 3 - return overall_inst_num + return overall_inst_num * self.device_batch_size def feed_forward(self, mode = "prefill"): mlen = self.hardware_config["MLEN"] @@ -133,14 +198,16 @@ def feed_forward(self, mode = "prefill"): overall_inst_num = 0 # -- MLP if mode == "prefill": - overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * 4 * (self.input_token // blen) - overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 - overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * 4 * (self.input_token // blen) - overall_inst_num = (overall_inst_num) * self.batch_size + # Upprojection, and Gate (seq, hidden) @ (hidden, intermediate) + overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.input_token / blen) * math.ceil(self.hidden_size / mlen) * blen + overall_inst_num += math.ceil(self.intermediate_size / vlen) * 3 * self.input_token + # Downprojection (seq, intermediate) @ (intermediate, hidden) + overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * math.ceil(self.input_token / blen) * blen + overall_inst_num = (overall_inst_num) * self.device_batch_size elif mode == "decode": - overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.batch_size // blen) ) * 4 - overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 - overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * 4 + overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * blen * math.ceil(self.device_batch_size / blen) + overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 * self.device_batch_size + overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * blen * math.ceil(self.device_batch_size / blen) return overall_inst_num def embeddings(self, mode = "prefill"): @@ -173,7 +240,7 @@ def compute_prefill_time(self): overall_inst_num += self.residual(mode) overall_inst_num += self.rms_layer(mode) overall_inst_num += self.feed_forward(mode) - # overall_inst_num += self.residual(mode) + # overall_inst_num += self.residual(mode) # overall_inst_num += self.rms_layer() overall_inst_num += self.lm_head() overall_exe_cycle = overall_inst_num * 2 @@ -204,6 +271,7 @@ def compute_decode_time(self, output_token_size): residual_count += self.residual(mode) rms_count += self.rms_layer(mode) feed_forward_count += self.feed_forward(mode) + self.kv_size = self.kv_size + 1 overall_inst_num = rms_count + projection_count + flash_attention_count + residual_count + feed_forward_count overall_exe_cycle = overall_inst_num * 2 # avg 2 execution cycles theoratical_execution_time = overall_exe_cycle / self.theoratical_frequency @@ -216,16 +284,8 @@ def compute_decode_time(self, output_token_size): print(f"Feed Forward: {feed_forward_count / overall_inst_num * 100}%") return theoratical_execution_time - def compute_overall_perf(self): + # Per batch performance. ttft = (self.compute_prefill_time() + self.compute_decode_time(1)) / self.device_num - tps = (self.device_num * (self.batch_size * self.output_token)) / self.compute_decode_time(self.output_token // self.device_num) - return ttft, tps - - - - - - - - + tps = (self.batch_size * self.output_token) / self.compute_decode_time(self.output_token) + return ttft, tps \ No newline at end of file diff --git a/tools/cost_model/utilisation/attainable.py b/tools/cost_model/utilisation/attainable.py index f3392ad9..82993e0f 100644 --- a/tools/cost_model/utilisation/attainable.py +++ b/tools/cost_model/utilisation/attainable.py @@ -1,9 +1,11 @@ import os from typing import Dict, List, Any, Optional import json +import math class attn_model_config: - def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = 2048, output_token = 128, device_num = 1): + def __init__(self, MLEN, BLEN, VLEN,partitioned_matrix, model_param_path, batch_size = 1, seq_len = 2048, output_token = 128, device_num = 1): + print(f"Model param path: {model_param_path}") model_param = json.load(open(model_param_path)) self.hidden_size = model_param["hidden_size"] self.num_attention_heads = model_param["num_attention_heads"] @@ -17,15 +19,20 @@ def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = self.vocab_size = model_param["vocab_size"] self.DataTypeSize = 2 self.theoratical_frequency = 10**9 # 1 GHz - self.hardware_config = hardware_config self.output_token = output_token self.batch_size = batch_size + print("=" * 25) + print(f"MLEN: {MLEN}, BLEN: {BLEN}, VLEN: {VLEN}") print(f"Batch size: {self.batch_size}") + print(f"num_attention_heads: {self.num_attention_heads}, num_key_value_heads: {self.num_key_value_heads}") + print(f"head_dim: {self.head_dim}, num_head_groups: {self.num_head_groups}") self.kv_size = seq_len self.device_num = device_num - self.M = hardware_config["BLEN"] - self.K = hardware_config["MLEN"] - self.N = self.M + self.MLEN = MLEN + self.BLEN = BLEN + self.VLEN = VLEN + self.partitioned_matrix = partitioned_matrix + self.max_head_per_mlen = math.ceil(self.MLEN / self.head_dim) def _report_flash_attn_utilization(self, mode = "prefill") -> None: """ @@ -40,47 +47,220 @@ def _report_flash_attn_utilization(self, mode = "prefill") -> None: input_token_size = self.input_seq_len theoretical_operation = 0 attainable_operation = 0 + attainable_gemm_operation_amount = 0 overall_operation_amount = 0 + bubble_ratio = 1.2 # Decoding if mode == "prefill": # Projection - operation_amount = ((head_dim * num_attn_heads) // self.M) * ( hidden_size // self.K) * (self.input_seq_len // self.M) + ((head_dim * num_kv_heads) // self.M) * ( hidden_size// self.K) * (self.input_seq_len // self.M) * 2 - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - # QKT - operation_amount = batch_size * num_attn_heads * (head_dim // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - - # PV - operation_amount = batch_size * num_attn_heads * (input_token_size // self.K) * (head_dim // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + gemm_operation_amount = (self.input_seq_len // self.BLEN) * ( hidden_size / self.MLEN) * self.BLEN * ((head_dim * num_attn_heads) // self.BLEN) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size / self.MLEN) * self.BLEN * (self.input_seq_len // self.BLEN) * 2 + attainable_gemm_operation_amount = gemm_operation_amount * batch_size + attention_operation_amount = gemm_operation_amount * bubble_ratio * batch_size + + # Flash Attention + if self.partitioned_matrix: + for b in range(batch_size): + for i in range(math.ceil(num_attn_heads / self.max_head_per_mlen)): + for j in range(math.ceil(self.input_seq_len / self.MLEN)): + for k in range(math.ceil(self.input_seq_len / self.MLEN)): + # QKT (MLEN, Hq) * (Hq, MLEN), (num_attn_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = self.head_dim # need to clarify this, coz the actual implementation is not like this + attention_operation_amount += gemm_operation_amount * bubble_ratio # Prefetch Data + if num_attn_heads > self.max_head_per_mlen: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim * num_attn_heads / self.MLEN) + # Softmax + attention_operation_amount += self.MLEN * 10 + # PV (MLEN, MLEN) @ (MLEN, head_dim), (num_attn_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = math.ceil(self.MLEN / self.BLEN) * self.BLEN * math.ceil(head_dim / self.BLEN) * self.num_head_groups + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + attention_operation_amount += self.MLEN * self.num_head_groups * (head_dim // self.BLEN) * bubble_ratio + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + else: + for b in range(batch_size): + for i in range(math.ceil(num_attn_heads)): + for j in range(math.ceil(self.input_seq_len / self.MLEN)): + for k in range(math.ceil(self.input_seq_len / self.MLEN)): + # QKT (MLEN, head_dim) @ (head_dim, MLEN) Full Utilization + gemm_operation_amount = (self.MLEN // self.BLEN) * math.ceil(head_dim / self.MLEN) * self.BLEN * (self.MLEN // self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio # Prefetch Data + if head_dim > self.MLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.MLEN) + # Softmax + attention_operation_amount += self.MLEN * 10 + # PV (MLEN, MLEN) @ (MLEN, head_dim) + gemm_operation_amount = (self.MLEN // self.BLEN) * math.ceil(head_dim / self.BLEN) * self.BLEN + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount + elif mode == "decode": + # Projection Q, K, V + gemm_operation_amount = ((head_dim * num_attn_heads) // self.BLEN) * ( hidden_size // self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size// self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) * 2 + attention_operation_amount = gemm_operation_amount * bubble_ratio + if batch_size > self.BLEN: + attainable_gemm_operation_amount = gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount = gemm_operation_amount * (batch_size / self.BLEN) + + # Flash Attention + if self.partitioned_matrix: + for b in range(batch_size): + for i in range(math.ceil(num_kv_heads / self.max_head_per_mlen)): + for j in range(math.ceil(self.kv_size / self.MLEN)): + # QKT (self.num_head_groups, 1, Hq) * (Hq, MLEN), (num_kv_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = math.ceil(self.num_head_groups / self.BLEN) * self.head_dim * math.ceil(self.MLEN / self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if num_kv_heads > (self.MLEN // self.head_dim): + if self.num_head_groups > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (self.num_head_groups / self.BLEN) + else: + if self.num_head_groups > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 * (num_kv_heads / (self.MLEN // self.head_dim)) + else: + attainable_gemm_operation_amount += gemm_operation_amount * (self.num_head_groups / self.BLEN) * (num_kv_heads / (self.MLEN // self.head_dim)) + # Softmax + attention_operation_amount += self.MLEN * 10 + # PV (self.num_head_groups, 1, MLEN) @ (MLEN, head_dim), (num_kv_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = math.ceil(self.num_head_groups / self.BLEN) * math.ceil(head_dim / self.BLEN) * self.num_head_groups + attention_operation_amount += gemm_operation_amount * bubble_ratio + if self.num_head_groups > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (self.num_head_groups / self.BLEN) + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + + else: + for b in range(batch_size): + for i in range(math.ceil(num_attn_heads)): + for j in range(math.ceil(self.kv_size / self.MLEN)): + # QKT (1, head_dim) * (head_dim, MLEN), (num_kv_heads / (self.MLEN // self.head_dim)) + gemm_operation_amount = math.ceil(self.head_dim / self.MLEN) * math.ceil(self.MLEN / self.BLEN) * self.BLEN + attention_operation_amount += gemm_operation_amount * bubble_ratio # Prefetch Data + attainable_gemm_operation_amount += gemm_operation_amount * 1 * (1 / self.BLEN) * (self.head_dim / self.MLEN) + + # Softmax + attention_operation_amount += self.MLEN * 10 + + # PV (1, MLEN) * (MLEN, head_dim), + gemm_operation_amount = (head_dim // self.BLEN) * self.BLEN + attention_operation_amount += gemm_operation_amount * bubble_ratio + attainable_gemm_operation_amount += gemm_operation_amount * 1 * (1 / self.BLEN) + + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount + + return [attainable_operation, theoretical_operation] + + def _report_self_attention_utilization(self, mode = "prefill") -> None: + """ + Report the utilization of self attention without using flash attention. + """ + batch_size = self.batch_size + hidden_size = self.hidden_size + num_attn_heads = self.num_attention_heads + num_kv_heads = self.num_key_value_heads + + head_dim = self.head_dim + input_token_size = self.input_seq_len + theoretical_operation = 0 + attainable_operation = 0 + attainable_gemm_operation_amount = 0 + overall_operation_amount = 0 + bubble_ratio = 1.2 + + if mode == "prefill": # Projection - operation_amount = ((head_dim * num_attn_heads) // self.M) * ( hidden_size // self.K) + ((head_dim * num_kv_heads) // self.M) * ( hidden_size// self.K) * 2 - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - - # QKT - operation_amount = batch_size * num_attn_heads * (head_dim // self.K) * (self.kv_size // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - - # PV - operation_amount = batch_size * num_attn_heads * (self.kv_size // self.K) * (head_dim // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - self.kv_size = self.kv_size + 1 + gemm_operation_amount = ((head_dim * num_attn_heads) // self.BLEN) * ( hidden_size // self.MLEN) * self.BLEN * (self.input_seq_len // self.BLEN) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size// self.MLEN) * self.BLEN * (self.input_seq_len // self.BLEN) * 2 + attainable_gemm_operation_amount = gemm_operation_amount * batch_size + attention_operation_amount = gemm_operation_amount * bubble_ratio * batch_size + + # QKT (batch, input_token_size, num_attn_heads, head_dim) @ (batch, input_token_size, num_attn_heads, head_dim) + for b in range(batch_size): + for i in range(num_attn_heads): + gemm_operation_amount = (self.input_token_size // self.BLEN) * (head_dim // self.MLEN) * (self.input_token_size // self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.MLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.MLEN) + # Storing and fetching QKT (input_token_size, input_token_size) + attention_operation_amount += self.input_token_size * 4 * (self.input_token_size / self.VLEN) + # Softmax + attention_operation_amount += self.input_token_size * (self.input_token_size / self.VLEN) * 30 + # Storing and fetching O (input_token_size, head_dim) + attention_operation_amount += self.input_token_size * (self.input_token_size / self.VLEN) * 5 + + # PV (self.input_token_size, self.input_token_size) @ (self.input_token_size, head_dim) + gemm_operation_amount = math.ceil(self.input_token_size / self.BLEN) * math.ceil(head_dim / self.BLEN) * self.num_head_groups + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + # Compute O + attention_operation_amount += self.input_token_size * (self.input_token_size / self.VLEN) * 5 + 4 + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount - return [operation_amount, attainable_operation, theoretical_operation] + elif mode == "decode": + # Projection Q, K, V + gemm_operation_amount = ((head_dim * num_attn_heads) // self.BLEN) * ( hidden_size // self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size// self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) * 2 + attention_operation_amount = gemm_operation_amount * bubble_ratio + if batch_size > self.BLEN: + attainable_gemm_operation_amount = gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount = gemm_operation_amount * (batch_size / self.BLEN) + # QKT (batch, 1, num_attn_heads, head_dim) @ (batch, kv_size, num_attn_heads, head_dim) + for b in range(batch_size): + for i in range(num_attn_heads): + gemm_operation_amount = (self.kv_size // self.BLEN) * (head_dim // self.MLEN) * (self.kv_size // self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.MLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.MLEN) + # Storing and fetching QKT (1, kv_size) + attention_operation_amount += 4 * (self.kv_size / self.VLEN) + # Softmax + attention_operation_amount += (self.kv_size / self.VLEN) * 30 + # Storing and fetching O (input_token_size, head_dim) + attention_operation_amount += (self.input_token_size / self.VLEN) * 5 + + # PV (1 self.kv_size) @ (self.kv_size, head_dim) + gemm_operation_amount = math.ceil(self.kv_size / self.MLEN) * math.ceil(head_dim / self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + # Compute O + attention_operation_amount += (self.kv_size / self.VLEN) * 15 + 4 + + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount + return [attainable_operation, theoretical_operation] def _report_embedding_utilization(self, mode = "prefill") -> None: """ @@ -95,13 +275,13 @@ def _report_embedding_utilization(self, mode = "prefill") -> None: if mode == "prefill": # Assuming Decoding only - operation_amount = (hidden_size // self.M) * (hidden_size // self.K) * (self.input_seq_len // self.N) - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (hidden_size // self.BLEN) * (hidden_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount * (self.BLEN * self.MLEN * self.BLEN) + theoretical_operation += operation_amount * (self.BLEN * self.MLEN * self.BLEN) elif mode == "decode": - operation_amount = (hidden_size // self.M) * (hidden_size // self.K) - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (hidden_size // self.BLEN) * (hidden_size // self.MLEN) + attainable_operation += operation_amount * (self.BLEN * self.MLEN * min(self.batch_size, self.BLEN)) + theoretical_operation += operation_amount * (self.BLEN * self.MLEN * self.BLEN) return [operation_amount, attainable_operation, theoretical_operation] @@ -109,115 +289,84 @@ def _report_ffn_utilization(self, mode = "prefill") -> None: """ Report the utilization of flash attention for a given node. """ - batch_size = self.batch_size hidden_size = self.hidden_size intermediate_size = self.intermediate_size - overall_operation_amount = 0 theoretical_operation = 0 attainable_operation = 0 + bubble_ratio = 1.2 if mode == "prefill": # Up Projection - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (intermediate_size // self.BLEN) * (hidden_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount + theoretical_operation += operation_amount * bubble_ratio # Gate Projection - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (intermediate_size // self.BLEN) * (hidden_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount + theoretical_operation += operation_amount * bubble_ratio + + # SiLU + theoretical_operation += (intermediate_size // self.VLEN) * 5 # Down Projection - operation_amount = (hidden_size // self.M) * (intermediate_size // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (hidden_size // self.BLEN) * (intermediate_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount + theoretical_operation += operation_amount * bubble_ratio elif mode == "decode": - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + + # Up projection + fill_bubble_time = self.BLEN * (intermediate_size // self.BLEN) * math.ceil(self.batch_size / self.BLEN) + operation_amount = (intermediate_size // self.BLEN) * self.BLEN * (hidden_size // self.MLEN) * math.ceil(self.batch_size / self.BLEN) + attainable_operation += operation_amount * (min(self.batch_size, self.BLEN) / self.BLEN) + theoretical_operation += operation_amount + fill_bubble_time # Gate Projection - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + fill_bubble_time = self.BLEN * (intermediate_size // self.BLEN) * math.ceil(self.batch_size / self.BLEN) + operation_amount = (intermediate_size // self.BLEN) * self.BLEN * (hidden_size // self.MLEN) * math.ceil(self.batch_size / self.BLEN) + attainable_operation += operation_amount * (min(self.batch_size, self.BLEN) / self.BLEN) + theoretical_operation += operation_amount + fill_bubble_time - # Down Projection - operation_amount = (hidden_size // self.M) * (intermediate_size // self.K) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + # SiLU + theoretical_operation += (intermediate_size // self.VLEN) * 5 + # Down Projection + fill_bubble_time = self.BLEN * (hidden_size // self.BLEN) * math.ceil(self.batch_size / self.BLEN) + operation_amount = (hidden_size // self.BLEN) * self.BLEN * (intermediate_size // self.MLEN) * math.ceil(self.batch_size / self.BLEN) + attainable_operation += operation_amount * (min(self.batch_size, self.BLEN) / self.BLEN) + theoretical_operation += operation_amount+ fill_bubble_time - return [overall_operation_amount, attainable_operation, theoretical_operation] + return [attainable_operation, theoretical_operation] def _report_prefill_utilization(self): - overall_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_attainable_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_theoretical_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - single_op_operation = self._report_embedding_utilization("prefill") - overall_operations["embedding"] += (single_op_operation[0] * 2) / (10 ** 9) - overall_attainable_FLOPS["embedding"] += (single_op_operation[1] * 2) / (10 ** 9) - overall_theoretical_FLOPS["embedding"] += (single_op_operation[2] * 2) / (10 ** 9) - for i in range(self.num_hidden_layers): - single_op_operation = self._report_flash_attn_utilization("prefill") - overall_operations["attention"] += single_op_operation[0] / (10 ** 9) - overall_attainable_FLOPS["attention"] += single_op_operation[1] / (10 ** 9) - overall_theoretical_FLOPS["attention"] += single_op_operation[2] / (10 ** 9) - - single_op_operation = self._report_ffn_utilization("prefill") - overall_operations["ffn"] += single_op_operation[0] / (10 ** 9) - overall_attainable_FLOPS["ffn"] += single_op_operation[1] / (10 ** 9) - overall_theoretical_FLOPS["ffn"] += single_op_operation[2] / (10 ** 9) + overall_attainable_GEMM = {"attention": 0, "ffn": 0} + overall_theoretical_GEMM = {"attention": 0, "ffn": 0} + overall_attainable_GEMM["attention"] = self._report_flash_attn_utilization("prefill")[0] + overall_theoretical_GEMM["attention"] = self._report_flash_attn_utilization("prefill")[1] + overall_attainable_GEMM["ffn"] = self._report_ffn_utilization("prefill")[0] + overall_theoretical_GEMM["ffn"] = self._report_ffn_utilization("prefill")[1] return { - "operations": overall_operations, - "attainable_FLOPS": overall_attainable_FLOPS, - "theoretical_FLOPS": overall_theoretical_FLOPS + "attainable_FLOPS": overall_attainable_GEMM, + "theoretical_FLOPS": overall_theoretical_GEMM } def _report_decode_utilization(self): - overall_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_attainable_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_theoretical_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - for j in range (self.output_token): - per_token_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - per_token_attainable_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - per_token_theoretical_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - - for i in range(self.num_hidden_layers): - single_op_operation = self._report_flash_attn_utilization("decode") - per_token_operations["attention"] += single_op_operation[0] - per_token_attainable_FLOPS["attention"] += single_op_operation[1] - per_token_theoretical_FLOPS["attention"] += single_op_operation[2] - - single_op_operation = self._report_ffn_utilization("decode") - per_token_operations["ffn"] += single_op_operation[0] - per_token_attainable_FLOPS["ffn"] += single_op_operation[1] - per_token_theoretical_FLOPS["ffn"] += single_op_operation[2] - - overall_operations["attention"] += per_token_operations["attention"] / (10**9) - overall_operations["ffn"] += per_token_operations["ffn"] / (10**9) - - overall_attainable_FLOPS["attention"] += per_token_attainable_FLOPS["attention"] / (10**9) - overall_theoretical_FLOPS["attention"] += per_token_theoretical_FLOPS["attention"] / (10**9) - - overall_attainable_FLOPS["ffn"] += per_token_attainable_FLOPS["ffn"] / (10**9) - overall_theoretical_FLOPS["ffn"] += per_token_theoretical_FLOPS["ffn"] / (10**9) + overall_attainable_GEMM = {"attention": 0, "ffn": 0} + overall_theoretical_GEMM = {"attention": 0, "ffn": 0} + overall_attainable_GEMM["attention"] = self._report_flash_attn_utilization("decode")[0] + overall_theoretical_GEMM["attention"] = self._report_flash_attn_utilization("decode")[1] + overall_attainable_GEMM["ffn"] = self._report_ffn_utilization("decode")[0] + overall_theoretical_GEMM["ffn"] = self._report_ffn_utilization("decode")[1] return { - "operations": overall_operations, - "attainable_FLOPS": overall_attainable_FLOPS, - "theoretical_FLOPS": overall_theoretical_FLOPS + "attainable_FLOPS": overall_attainable_GEMM, + "theoretical_FLOPS": overall_theoretical_GEMM } def compute_overall_perf(self): @@ -226,4 +375,17 @@ def compute_overall_perf(self): decode_perf = self._report_decode_utilization() print(f"Decode Performance: {decode_perf}") utilization = (prefill_perf["attainable_FLOPS"]["ffn"] + prefill_perf["attainable_FLOPS"]["attention"] + decode_perf["attainable_FLOPS"]["ffn"] + decode_perf["attainable_FLOPS"]["attention"]) / (prefill_perf["theoretical_FLOPS"]["ffn"] + prefill_perf["theoretical_FLOPS"]["attention"] + decode_perf["theoretical_FLOPS"]["ffn"] + decode_perf["theoretical_FLOPS"]["attention"]) - return utilization \ No newline at end of file + return utilization + + +if __name__ == "__main__": + import os + from pathlib import Path + # Get the absolute path of the current file's directory + current_dir = Path(__file__).resolve().parents[3] + model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + model = attn_model_config(MLEN=1024, BLEN=4, VLEN=1024, partitioned_matrix=False, model_param_path=model_param_path, batch_size=1, seq_len=5600, output_token=1000, device_num=1) + model.kv_size = 80000 + actual, theoretical = model._report_flash_attn_utilization("decode") + print(f"Actual: {actual}, Theoretical: {theoretical}") + print(f"Utilization: {actual / theoretical}") \ No newline at end of file diff --git a/tools/cost_model/utilisation/attainable_gemm_test.py b/tools/cost_model/utilisation/attainable_gemm_test.py new file mode 100644 index 00000000..4dec52c2 --- /dev/null +++ b/tools/cost_model/utilisation/attainable_gemm_test.py @@ -0,0 +1,104 @@ +from attainable import attn_model_config + +import os +from pathlib import Path +# Get the absolute path of the current file's directory +current_dir = Path(__file__).resolve().parents[3] +model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + +import matplotlib.pyplot as plt +import numpy as np + +# Define MLEN and BLEN configurations and their labels +configs = [ + {"MLEN": 64, "BLEN": 64, "label": "MLEN=64, BLEN=64"}, + {"MLEN": 128, "BLEN": 32, "label": "MLEN=128, BLEN=32"}, + {"MLEN": 256, "BLEN": 16, "label": "MLEN=256, BLEN=16"}, + {"MLEN": 512, "BLEN": 8, "label": "MLEN=512, BLEN=8"}, + {"MLEN": 1024, "BLEN": 4, "label": "MLEN=1024, BLEN=4"} +] + +batch_sizes = [1, 2, 4, 8, 16] + +bar_width = 0.06 # Further reduced bar width +bar_gap = 0.02 # Gap between bars in same group +group_gap = 0 # Gap between groups of bars + +# Create x positions with spacing between groups (reduced distance) +x = np.arange(len(batch_sizes)) * (0.5 + group_gap) + +# Create two subplots side by side: utilization (left) and execution time (right) +fig, (ax1_ffn, ax2_time) = plt.subplots(1, 2, figsize=(16, 4)) + +colors = [ + (240/255, 249/255, 232/255), # light green + (186/255, 228/255, 188/255), # green + (123/255, 204/255, 196/255), # teal + (67/255, 162/255, 202/255), # blue + (8/255, 104/255, 172/255) # dark blue +] + +# Process each configuration +for idx, cfg in enumerate(configs): + # FFN data + ffn_utilizations = [] + execution_times = [] # Execution time in microseconds (cycles / 1GHz) + + for bs in batch_sizes: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=True, + model_param_path=model_param_path, batch_size=bs, seq_len=8024, output_token=100, device_num=1 + ) + + # FFN stage + ffn_attainable, ffn_theoretical = model._report_ffn_utilization("decode")[0:2] + ffn_utilization = ffn_attainable / ffn_theoretical if ffn_theoretical != 0 else 0 + ffn_utilizations.append(ffn_utilization) + + # Execution time: theoretical cycles / 1GHz = time in seconds, convert to microseconds + execution_time_us = (ffn_theoretical / 1e9) * 1e6 # Convert to microseconds + execution_times.append(execution_time_us) + + # Bar positions for configs - with spacing between groups + pos = x + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot FFN utilization (left) + rects_ffn = ax1_ffn.bar(pos, ffn_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + + # Plot execution time (right) + rects_time = ax2_time.bar(pos, execution_times, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Configure FFN utilization plot (left) +ax1_ffn.set_xlabel("Batch Size", fontsize=18) +ax1_ffn.set_ylabel("Utilization", fontsize=18) +ax1_ffn.tick_params(axis='both', which='major', labelsize=14) +ax1_ffn.set_title("FFN Decode Systolic Array Utilization vs. Batch Size", fontsize=20) +ax1_ffn.set_xticks(x) +ax1_ffn.set_xticklabels([str(bs) for bs in batch_sizes]) +ax1_ffn.set_ylim(0, 1) + +# Set x-axis limits to show all bars with spacing +total_config_width = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_ffn.set_xlim(x[0] - total_config_width/2 - group_gap/2, x[-1] + total_config_width/2 + group_gap/2) + +# Configure execution time plot (right) +ax2_time.set_xlabel("Batch Size", fontsize=18) +ax2_time.set_ylabel("Execution Time (μs)", fontsize=18) +ax2_time.tick_params(axis='both', which='major', labelsize=14) +ax2_time.set_title("FFN Decode Execution Time vs. Batch Size (1GHz)", fontsize=20) +ax2_time.set_xticks(x) +ax2_time.set_xticklabels([str(bs) for bs in batch_sizes]) + +# Set x-axis limits to show all bars with spacing (same as left plot) +ax2_time.set_xlim(x[0] - total_config_width/2 - group_gap/2, x[-1] + total_config_width/2 + group_gap/2) + +# Create shared legend at the bottom in two rows (placed lower) +handles1, labels1 = ax1_ffn.get_legend_handles_labels() +# Create legend with all config labels - transparent with no frame, arranged in 2 rows +# With 5 configs, use ncol=3 to get 2 rows (3 in first row, 2 in second row) +# Push legend lower using bbox_to_anchor +fig.legend(handles1, labels1, loc='upper center', bbox_to_anchor=(0.5, 0.1), ncol=5, fontsize=15, frameon=False) + +plt.tight_layout(rect=[0, 0.08, 1, 0.98]) +plt.savefig("decode_FFN_GEMM_Utilization_and_Latency.png", dpi=300, bbox_inches='tight') diff --git a/tools/cost_model/utilisation/combined_ffn_fa_utilization.py b/tools/cost_model/utilisation/combined_ffn_fa_utilization.py new file mode 100644 index 00000000..59258b22 --- /dev/null +++ b/tools/cost_model/utilisation/combined_ffn_fa_utilization.py @@ -0,0 +1,178 @@ +from attainable import attn_model_config + +import os +from pathlib import Path +# Get the absolute path of the current file's directory +current_dir = Path(__file__).resolve().parents[3] +model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + +import matplotlib.pyplot as plt +import matplotlib.ticker as ticker +import numpy as np + +# Define MLEN and BLEN configurations and their labels +configs = [ + {"MLEN": 64, "BLEN": 64, "label": "MLEN=64, BLEN=64"}, + {"MLEN": 128, "BLEN": 32, "label": "MLEN=128, BLEN=32"}, + {"MLEN": 256, "BLEN": 16, "label": "MLEN=256, BLEN=16"}, + {"MLEN": 512, "BLEN": 8, "label": "MLEN=512, BLEN=8"}, + {"MLEN": 1024, "BLEN": 4, "label": "MLEN=1024, BLEN=4"} +] + +batch_sizes = [1, 2, 4, 8, 16] +sequence_lengths = [128, 1024, 5120, 20480, 81920] # 128, 1k, 5k, 20k, 80k +seq_labels = ["128", "1k", "5k", "20k", "80k"] + +bar_width = 0.06 # Bar width +bar_gap = 0.02 # Gap between bars in same group +group_gap = 0 # Gap between groups of bars + +# Create x positions with spacing between groups +x_ffn = np.arange(len(batch_sizes)) * (0.5 + group_gap) +x_fa = np.arange(len(sequence_lengths)) * (0.5 + group_gap) + +# Create 1x4 grid: FFN util, FFN time, FA util, FA time +fig, (ax1_ffn_util, ax2_ffn_time, ax1_fa_util, ax2_fa_time) = plt.subplots(1, 4, figsize=(20, 4)) + +colors = [ + "#762a83", # deep purple + "#af8dc3", # light violet + "#e7d4e8", # very light purple-pink + "#d9f0d3", # pale green + "#7fbf7b", # medium green + # Optional sixth if needed: "#1b7837" # dark green +] + +# Process FFN data +for idx, cfg in enumerate(configs): + # FFN data + ffn_utilizations = [] + ffn_execution_times = [] # Execution time in microseconds (cycles / 1GHz) + + for bs in batch_sizes: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=True, + model_param_path=model_param_path, batch_size=bs, seq_len=8024, output_token=100, device_num=1 + ) + + # FFN stage + ffn_attainable, ffn_theoretical = model._report_ffn_utilization("decode")[0:2] + ffn_utilization = (ffn_attainable / ffn_theoretical if ffn_theoretical != 0 else 0) * 100 # Convert to percentage + ffn_utilizations.append(ffn_utilization) + + # Execution time: theoretical cycles / 1GHz = time in seconds, convert to milliseconds + execution_time_ms = (ffn_theoretical / 1e9) * 1e3 # Convert to milliseconds + ffn_execution_times.append(execution_time_ms) + + # Bar positions for configs - with spacing between groups + pos_ffn = x_ffn + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot FFN utilization + rects_ffn_util = ax1_ffn_util.bar(pos_ffn, ffn_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + + # Plot FFN execution time + rects_ffn_time = ax2_ffn_time.bar(pos_ffn, ffn_execution_times, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Process Flash Attention data +for idx, cfg in enumerate(configs): + # Flash Attention data + fa_utilizations = [] + fa_execution_times = [] # Execution time in microseconds (cycles / 1GHz) + + for seq_len in sequence_lengths: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=True, + model_param_path=model_param_path, batch_size=1, seq_len=seq_len, output_token=100, device_num=1 + ) + + # Flash Attention stage + fa_attainable, fa_theoretical = model._report_flash_attn_utilization("decode")[0:2] + fa_utilization = (fa_attainable / fa_theoretical if fa_theoretical != 0 else 0) * 100 # Convert to percentage + fa_utilizations.append(fa_utilization) + + # Execution time: theoretical cycles / 1GHz = time in seconds, convert to milliseconds + execution_time_ms = (fa_theoretical / 1e9) * 1e3 # Convert to milliseconds + fa_execution_times.append(execution_time_ms) + + # Bar positions for configs - with spacing between groups + pos_fa = x_fa + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot Flash Attention utilization + rects_fa_util = ax1_fa_util.bar(pos_fa, fa_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + + # Plot Flash Attention execution time + rects_fa_time = ax2_fa_time.bar(pos_fa, fa_execution_times, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Configure FFN utilization plot +ax1_ffn_util.set_xlabel("Batch Size", fontsize=18) +ax1_ffn_util.set_ylabel("Utilization (%)", fontsize=18) +ax1_ffn_util.tick_params(axis='both', which='major', labelsize=14) +ax1_ffn_util.set_title("FFN Decode SA Utilization", fontsize=20) +ax1_ffn_util.set_xticks(x_ffn) +ax1_ffn_util.set_xticklabels([str(bs) for bs in batch_sizes]) +ax1_ffn_util.set_ylim(0, 100) +# Format y-axis to show only integer ticks +ax1_ffn_util.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax1_ffn_util.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing +total_config_width_ffn = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_ffn_util.set_xlim(x_ffn[0] - total_config_width_ffn/2 - group_gap/2, x_ffn[-1] + total_config_width_ffn/2 + group_gap/2) + +# Configure FFN execution time plot +ax2_ffn_time.set_xlabel("Batch Size", fontsize=18) +ax2_ffn_time.set_ylabel("Exec Time (ms)", fontsize=18) +ax2_ffn_time.tick_params(axis='both', which='major', labelsize=14) +ax2_ffn_time.set_title("FFN Decode Exec Time", fontsize=20) +ax2_ffn_time.set_xticks(x_ffn) +ax2_ffn_time.set_xticklabels([str(bs) for bs in batch_sizes]) +# Format y-axis to show only integer ticks +ax2_ffn_time.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax2_ffn_time.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing (same as FFN util) +ax2_ffn_time.set_xlim(x_ffn[0] - total_config_width_ffn/2 - group_gap/2, x_ffn[-1] + total_config_width_ffn/2 + group_gap/2) + +# Configure Flash Attention utilization plot +ax1_fa_util.set_xlabel("Number of Tokens", fontsize=18) +ax1_fa_util.set_ylabel("Utilization (%)", fontsize=18) +ax1_fa_util.tick_params(axis='both', which='major', labelsize=14) +ax1_fa_util.set_title("FA Decode SA Utilization", fontsize=20) +ax1_fa_util.set_xticks(x_fa) +ax1_fa_util.set_xticklabels(seq_labels) +ax1_fa_util.set_ylim(0, 100) +# Format y-axis to show only integer ticks +ax1_fa_util.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax1_fa_util.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing +total_config_width_fa = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_fa_util.set_xlim(x_fa[0] - total_config_width_fa/2 - group_gap/2, x_fa[-1] + total_config_width_fa/2 + group_gap/2) + +# Configure Flash Attention execution time plot +ax2_fa_time.set_xlabel("Number of Tokens", fontsize=18) +ax2_fa_time.set_ylabel("Exec Time (ms)", fontsize=18) +ax2_fa_time.tick_params(axis='both', which='major', labelsize=14) +ax2_fa_time.set_title("FA Decode Exec Time", fontsize=20) +ax2_fa_time.set_xticks(x_fa) +ax2_fa_time.set_xticklabels(seq_labels) +# Format y-axis to show only integer ticks +ax2_fa_time.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax2_fa_time.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing (same as FA util) +ax2_fa_time.set_xlim(x_fa[0] - total_config_width_fa/2 - group_gap/2, x_fa[-1] + total_config_width_fa/2 + group_gap/2) + +# Create shared legend at the bottom in two rows (placed lower) +handles1, labels1 = ax1_ffn_util.get_legend_handles_labels() +# Create legend with all config labels - transparent with no frame, arranged in 2 rows +# With 5 configs, use ncol=3 to get 2 rows (3 in first row, 2 in second row) +# Push legend lower using bbox_to_anchor +fig.legend(handles1, labels1, loc='upper center', bbox_to_anchor=(0.5, 0.1), ncol=5, fontsize=15, frameon=False) + +plt.tight_layout(rect=[0, 0.1, 1, 0.98]) +plt.savefig("combined_FFN_FA_Utilization.png", dpi=300, bbox_inches='tight') +print("Plot saved as 'combined_FFN_FA_Utilization.png'") + diff --git a/tools/cost_model/utilisation/fa_utilizaiton_test.py b/tools/cost_model/utilisation/fa_utilizaiton_test.py new file mode 100644 index 00000000..e6511d42 --- /dev/null +++ b/tools/cost_model/utilisation/fa_utilizaiton_test.py @@ -0,0 +1,85 @@ +from attainable import attn_model_config + +import os +from pathlib import Path +# Get the absolute path of the current file's directory +current_dir = Path(__file__).resolve().parents[3] +model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + +import matplotlib.pyplot as plt +import numpy as np + +# Define MLEN and BLEN configurations and their labels +configs = [ + {"MLEN": 64, "BLEN": 64, "label": "MLEN=64, BLEN=64"}, + {"MLEN": 128, "BLEN": 32, "label": "MLEN=128, BLEN=32"}, + {"MLEN": 256, "BLEN": 16, "label": "MLEN=256, BLEN=16"}, + {"MLEN": 512, "BLEN": 8, "label": "MLEN=512, BLEN=8"}, + {"MLEN": 1024, "BLEN": 4, "label": "MLEN=1024, BLEN=4"} +] + +sequence_lengths = [128, 1024, 5120, 20480, 81920] # 128, 1k, 5k, 20k, 80k +seq_labels = ["128", "1k", "5k", "20k", "80k"] + +bar_width = 0.1 # Thinner bars +bar_gap = 0.02 # Gap between bars in same group +group_gap = 0 # Gap between groups of bars + +# Create x positions with spacing between groups +x = np.arange(len(sequence_lengths)) * (0.7 + group_gap) + +# Create single plot for Flash Attention +fig, ax1_fa = plt.subplots(1, 1, figsize=(12, 4)) + +colors = [ + (240/255, 249/255, 232/255), # light green + (186/255, 228/255, 188/255), # green + (123/255, 204/255, 196/255), # teal + (67/255, 162/255, 202/255), # blue + (8/255, 104/255, 172/255) # dark blue +] + +# Process each configuration +for idx, cfg in enumerate(configs): + # Flash Attention data + fa_utilizations = [] + + for seq_len in sequence_lengths: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=cfg["partitioned_matrix"], + model_param_path=model_param_path, batch_size=1, seq_len=seq_len, output_token=100, device_num=1 + ) + + # Flash Attention stage + fa_attainable, fa_theoretical = model._report_flash_attn_utilization("decode")[0:2] + fa_utilization = fa_attainable / fa_theoretical if fa_theoretical != 0 else 0 + fa_utilizations.append(fa_utilization) + + # Bar positions for configs - with spacing between groups + pos = x + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot Flash Attention + rects_fa = ax1_fa.bar(pos, fa_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Configure Flash Attention plot +ax1_fa.set_xlabel("Sequence Length", fontsize=18) +ax1_fa.set_ylabel("Utilization", fontsize=18) +ax1_fa.tick_params(axis='both', which='major', labelsize=14) +ax1_fa.set_title("Flash Attention Decode Systolic Array Utilization vs. Sequence Length", fontsize=20) +ax1_fa.set_xticks(x) +ax1_fa.set_xticklabels(seq_labels) +ax1_fa.set_ylim(0, 1) + +# Set x-axis limits to show all bars with spacing +total_config_width = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_fa.set_xlim(x[0] - total_config_width/2 - group_gap/2, x[-1] + total_config_width/2 + group_gap/2) + +# Create shared legend at the bottom in two rows +handles1, labels1 = ax1_fa.get_legend_handles_labels() +# Create legend with all config labels - transparent with no frame, arranged in 2 rows +# With 5 configs, use ncol=3 to get 2 rows (3 in first row, 2 in second row) +fig.legend(handles1, labels1, loc='lower center', ncol=3, fontsize=15, frameon=False) + +plt.tight_layout(rect=[0, 0.05, 1, 0.98]) +plt.savefig("decode_FA_GEMM_Utilization.png", dpi=300, bbox_inches='tight') diff --git a/tools/cost_model/utilisation/individual_units_lib.json b/tools/cost_model/utilisation/individual_units_lib.json index 6de0e56d..69cdd9b1 100644 --- a/tools/cost_model/utilisation/individual_units_lib.json +++ b/tools/cost_model/utilisation/individual_units_lib.json @@ -2,10 +2,10 @@ "MatrixMachine": { "Coefficients": { "P1": 50, - "P2": 500, + "P2": 100, "P3": 10, "P4": 10, - "P5": 400 + "P5": 200 }, "Relationship": "P1 * MLEN * BLEN * (FP_EXP_WIDTH + FP_MANT_WIDTH + 1) + P2 * MLEN * BLEN * (WT_MX_MANT_WIDTH + WT_MX_EXP_WIDTH + 1 + MX_SCALE_WIDTH) + P3 * (MLEN // BLOCK_DIM) * BLEN * (BLOCK_DIM * ACT_ELEMENT_WIDTH + MX_SCALE_WIDTH) + P4 * MLEN * (FP_EXP_WIDTH + FP_MANT_WIDTH + 1) + P5" @@ -45,7 +45,7 @@ "Coefficients": { "P1": 300, "P2": 300, - "P3": 800, + "P3": 100, "P4": 100, "P5": 100, "P6": 100 diff --git a/tools/cost_model/utilisation/memory_utilization_plot.py b/tools/cost_model/utilisation/memory_utilization_plot.py new file mode 100644 index 00000000..2200e138 --- /dev/null +++ b/tools/cost_model/utilisation/memory_utilization_plot.py @@ -0,0 +1,179 @@ +import json +import re + +import os +from pathlib import Path + +def compute_hbm_storage(batch_size, model_param_path, kv_size, kv_precision=2, act_precision=2, wt_precision=2): + """ + Compute HBM storage requirements for model weights, kv_cache, and other elements. + + Args: + batch_size (int): The batch size. + model_param_path (str or Path): Path to the model parameter JSON file. + kv_size (int): The kv cache sequence length. + kv_precision (int, optional): Byte size of KV cache data type (default: 2 for fp16). + act_precision (int, optional): Byte size of activation data type (unused; default: 2). + wt_precision (int, optional): Byte size of weight data type (default: 2 for fp16). + + Returns: + dict: Dictionary with keys 'weights', 'kv_cache', and 'other' with size in GB. + """ + model_param = json.load(open(model_param_path)) + hidden_size = model_param["hidden_size"] + num_attention_heads = model_param["num_attention_heads"] + num_hidden_layers = model_param["num_hidden_layers"] + intermediate_size = model_param["intermediate_size"] + num_key_value_heads = model_param["num_key_value_heads"] + repeat_layer = model_param["num_hidden_layers"] + vocab_size = model_param["vocab_size"] + head_dim = hidden_size // num_attention_heads + + # Extract model size in Billion parameters from the file name (e.g., ".../llama-3.1-8b.json" -> 8) + model_filename = str(model_param_path).split('/')[-1] + match = re.search(r'(\d+)[bB]', model_filename) + if match: + model_size_b = int(match.group(1)) + else: + model_size_b = None # Optionally raise error + + if model_size_b is not None: + weight_size_gb = model_size_b * 1e9 * wt_precision / (1024 ** 3) + else: + weight_size_gb = 0 # Or raise error + + kv_cache_gb = (2 * head_dim * num_key_value_heads * kv_size * batch_size * kv_precision * repeat_layer) / 1024 / 1024 / 1024 + + other_gb = (hidden_size * 4 + intermediate_size * 4 + vocab_size * 4) / 1024 / 1024 / 1024 + + hbm_storage = { + "weights": weight_size_gb, + "kv_cache": kv_cache_gb, + "other": other_gb + } + return hbm_storage + + + +if __name__ == "__main__": + current_dir = Path(__file__).resolve().parents[3] + model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + batch_size = 1 + kv_size = 8024 + kv_precision = 2 + act_precision = 2 + wt_precision = 2 + hbm_storage = compute_hbm_storage(batch_size, model_param_path, kv_size, kv_precision, act_precision, wt_precision) + print(hbm_storage) + + import matplotlib.pyplot as plt + import numpy as np + + # kv_size values + kv_sizes = [1000, 10000, 50000, 100000] # [1k, 10k, 50k, 100k] + kv_size_labels = ['1k', '10k', '50k', '100k'] + + colors = [ + (240/255, 249/255, 232/255), # light green + (186/255, 228/255, 188/255), # green + (123/255, 204/255, 196/255), # teal + (67/255, 162/255, 202/255), # blue + (8/255, 104/255, 172/255) # dark blue + ] + + + # Define the four precision settings (kv, act, wt) + # Format: (kv_precision, act_precision, wt_precision) + precision_settings = [ + (2, 2, 2), + (1.125, 1.125, 1.125), + (1.125, 1.125, 0.625), + (0.625, 1.125, 0.625), + ] + precision_labels = [ + "KV=BF16 ,ACT=BF16, Weight=BF16", + "KV=MXINT E4M3 ,ACT=MXINT E4M3, Weight= MXINT E4M3", + "KV=MXINT E4M3 ,ACT=MXINT E4M3, Weight= MXINT E2M1", + "KV=MXINT E2M1 ,ACT=MXINT E4M3, Weight= MXINT E2M1", + ] + num_settings = len(precision_settings) + + # Colors for each memory type + color_dict = { + "weights": colors[2], + "kv_cache": colors[4], + } + mem_keys = ["weights", "kv_cache"] + + # Collect data: result[setting_idx][kv_idx][mem_type] + results = [] # (settings, kv_size, mem_type) + for setting in precision_settings: + setting_results = [] + for kv in kv_sizes: + kv_precision, act_precision, wt_precision = setting + hbm_storage = compute_hbm_storage(batch_size, model_param_path, kv, kv_precision, act_precision, wt_precision) + vals = [hbm_storage[k] for k in mem_keys] + setting_results.append(vals) + results.append(setting_results) + results = np.array(results) # shape: (num_settings, num_kv, 3) + + # Plotting + fig, ax = plt.subplots(figsize=(12, 6)) + + bar_width = 0.1 # Thinner bars + spacing = 0.02 # Space between each precision group + group_width = num_settings * bar_width + (num_settings - 1) * spacing + x = np.arange(len(kv_sizes)) * 0.6 # Reduce spacing between bar groups + + # For bar placement (center group at each kv_size) + for sidx, setting_label in enumerate(precision_labels): + for kidx in range(len(kv_sizes)): + left = x[kidx] - group_width/2 + sidx*(bar_width + spacing) + bar_width/2 + bottom = 0 + prev_bottom = 0 + # Plot stacked: weights, kv_cache, other + for midx, mem in enumerate(mem_keys): + val = results[sidx, kidx, midx] + color = color_dict[mem] + bar = ax.bar(left, val, bar_width, bottom=prev_bottom, color=color, edgecolor="black" if sidx==0 else None) + prev_bottom += val + # Put legend entries only for the "weights" bar for nice stacking + if sidx == 0: + # Just for the legend + handles = [plt.Rectangle((0,0),1,1, color=color_dict[m]) for m in mem_keys] + ax.legend(handles, mem_keys, title="Memory Type", loc='upper right', fontsize=13, title_fontsize=14) + + # X-axis + ax.set_xticks(x) + xtick_labels = [] + for k in kv_size_labels: + label = f"{k}" + xtick_labels.append(label) + ax.set_xticklabels(xtick_labels, fontsize=12) + + # Put precision labels (P1,P2,...) centered on each bar + for kidx, xpos in enumerate(x): + for sidx in range(num_settings): + left = xpos - group_width/2 + sidx*(bar_width + spacing) + bar_width/2 + height = results[sidx, kidx].sum() + # Center P-label on top part of each stacked bar + ax.text( + left, height + 0.03*max(results.flatten()), + f"P{sidx+1}", + ha='center', + va='bottom', + fontsize=10, + rotation=0 + ) + + # Additional legend for precision index -> actual precision setting (top left) + precision_legend = [f"P{i+1}: {lbl}" for i, lbl in enumerate(precision_labels)] + ax.text(0.02, 0.98, "Precisions:\n" + "\n".join(precision_legend), + transform=ax.transAxes, fontsize=13, va='top', ha='left', bbox=dict(facecolor='white', alpha=0.7, edgecolor='gray')) + + ax.set_xlabel("KV Size", fontsize=15) + ax.set_ylabel("Memory Usage (GB, Stacked)", fontsize=15) + ax.set_title("Single Batch Memory Distribution vs KV Size", fontsize=17) + ax.set_ylim(bottom=0, top=40) + plt.tight_layout() + plt.savefig("memory_utilization.png", dpi=300, bbox_inches='tight') diff --git a/tools/cost_model/utilisation/utilisation_model.py b/tools/cost_model/utilisation/utilisation_model.py index ba31ca6d..5377c2b2 100644 --- a/tools/cost_model/utilisation/utilisation_model.py +++ b/tools/cost_model/utilisation/utilisation_model.py @@ -3,27 +3,52 @@ import os import json from pathlib import Path -from ...utils import load_toml_config, load_json, load_svh_settings +from utils import load_toml_config, load_json, load_svh_settings from attainable import attn_model_config class utilisation_model: def __init__(self, hardware_settings_file: str = "plena_settings.toml", precision_settings_file: str = "precision.toml", unit_info_file: str = "unit_info.json"): self.unit_info = load_json(unit_info_file) + config_settings = load_svh_settings(hardware_settings_file) precision_settings = load_svh_settings(precision_settings_file) self.hardware_settings = {**config_settings, **precision_settings} def obtain_resource_utilisation(self, updated_config): resource_utilisation = 0 hardware_settings = self.hardware_settings - for key, value in updated_config.items(): - hardware_settings[key] = value + #TODOs + # for key, value in updated_config.items(): + # hardware_settings[key] = value + custom_config= { + "MLEN": 1024, + "BLEN": 64, + "VLEN": 1024, + "WT_MX_MANT_WIDTH": 1, + "WT_MX_EXP_WIDTH": 2, + "KV_ELEMENT_WIDTH": 8, + "BLOCK_DIM": 8, + "MX_SCALE_WIDTH" : 8, + # "ACT_MXFP_MANT_WIDTH": 3, + # "ACT_MXFP_EXP_WIDTH": 4, + "ACT_ELEMENT_WIDTH" : 8, + "FP_EXP_WIDTH": 3, + "FP_MANT_WIDTH": 4, + "HBM_ELE_WIDTH": 512, + "HBM_SCALE_WIDTH": 512, + "MATRIX_SRAM_DEPTH": 4096, + "VECTOR_SRAM_DEPTH": 4096, + "INT_DATA_WIDTH": 32, + "INT_SRAM_DEPTH": 256, + "FP_SRAM_DEPTH": 256, + + } for unit, info in self.unit_info.items(): - if "Coefficients" in info and "Relationship" in info: - hardware_settings.update(info["Coefficients"]) + if "Coefficients" in info and "Relationship" in info: + custom_config.update(info["Coefficients"]) relationship = info["Relationship"] - resource_utilisation += eval(relationship, {}, hardware_settings) + resource_utilisation += eval(relationship, {}, custom_config) return resource_utilisation diff --git a/tools/memory_mapping/addr_align.py b/tools/memory_mapping/addr_align.py new file mode 100644 index 00000000..db6f61b9 --- /dev/null +++ b/tools/memory_mapping/addr_align.py @@ -0,0 +1,43 @@ +""" +Address alignment utilities. +""" + + +def align_addr_up(addr: int, multiple: int) -> int: + """ + Round an address up to the next multiple of the given value. + + Args: + addr: The input address to align + multiple: The alignment multiple (e.g., 64 for 64-byte alignment) + + Returns: + The address rounded up to the next multiple + """ + if multiple <= 0: + raise ValueError("multiple must be positive") + + # Round up: (addr + multiple - 1) // multiple * multiple + return int(((addr + multiple - 1) // multiple) * multiple) + + +if __name__ == "__main__": + # Test cases + test_cases = [ + (32, 64, 64), + (16, 64, 64), + (64, 64, 64), + (65, 64, 128), + (100, 64, 128), + (0, 64, 0), + (1, 64, 64), + (63, 64, 64), + (128, 64, 128), + ] + + print("Testing align_addr_up:") + for addr, multiple, expected in test_cases: + result = align_addr_up(addr, multiple) + status = "✓" if result == expected else "✗" + print(f"{status} align_addr_up({addr}, {multiple}) = {result} (expected {expected})") + diff --git a/tools/memory_mapping/memory_map.py b/tools/memory_mapping/memory_map.py index 3b0e9282..b6831801 100644 --- a/tools/memory_mapping/memory_map.py +++ b/tools/memory_mapping/memory_map.py @@ -111,7 +111,7 @@ def map_data_to_fake_hbm_for_rtl_sim(blocks, element_width, block_width, bias, b f.write("0x" + insert_bias_row + "\n") -def map_data_to_fake_hbm_for_behave_sim(blocks, element_width, block_width, bias, bias_width, directory, append=True, hbm_row_width=64): +def map_mx_data_to_hbm_for_behave_sim(blocks, element_width, block_width, bias, bias_width, directory, append=True, hbm_row_width=64): """ Maps the quantized blocks and bias to binary memory file for fake HBM memory. Writes raw bytes instead of ASCII hex text. @@ -124,6 +124,7 @@ def map_data_to_fake_hbm_for_behave_sim(blocks, element_width, block_width, bias output_file = os.path.join(directory, "hbm_for_behave_sim.bin") mode = 'ab' if append else 'wb' + for row_idx, row in enumerate(blocks): hex_row = " ".join(f"0x{val:02X}" for val in row) @@ -179,6 +180,29 @@ def map_data_to_fake_hbm_for_behave_sim(blocks, element_width, block_width, bias print_outputfile_contents(output_file) +def map_normal_data_to_hbm_for_behave_sim(data, data_width, directory, append=True, hbm_row_width=64): + """ + Maps the normal data to binary memory file for fake HBM memory. + """ + if not os.path.exists(directory): + os.makedirs(directory) + output_file = os.path.join(directory, "hbm_for_behave_sim.bin") + mode = 'ab' if append else 'wb' + data = data.flatten() + with open(output_file, mode) as f: + row_buffer = bytearray() + for i, element in enumerate(data): + hex_str = map_scale_to_value(element, data_width) + data_bytes = hex_to_bytes(hex_str) + row_buffer.extend(data_bytes) + if len(row_buffer) >= hbm_row_width: + f.write(row_buffer[:hbm_row_width]) + row_buffer = bytearray() + if len(row_buffer) > 0: + f.write(row_buffer) + print_outputfile_contents(output_file) + + if __name__ == "__main__": directory = "../../test/weight" fake_hbm_dir = "../../test/load_mem" diff --git a/tools/memory_mapping/rand_gen.py b/tools/memory_mapping/rand_gen.py index bc91174b..f392ea0d 100644 --- a/tools/memory_mapping/rand_gen.py +++ b/tools/memory_mapping/rand_gen.py @@ -139,58 +139,41 @@ def quantize_tensor(self, tensor): # INSERT_YOUR_CODE # Accept tensor as either a torch.Tensor or an OrderedDict of tensors # If it's an OrderedDict, quantize each value (assuming each is a tensor), otherwise just quantize the single tensor - - block_list = [] - scaling_list = [] - tensors = [] - - if isinstance(tensor, collections.OrderedDict): - tensors = list(tensor.values()) - elif isinstance(tensor, dict): # in case sometimes dict is used, not strictly OrderedDict - tensors = list(tensor.values()) - elif tensor == None: - return block_list, scaling_list - else: - tensors = [tensor] - - for t in tensors: - # print("quantizing tensor", t.shape) - # If the input is 1D, add a dimension to make it 2D (row vector) - if t.ndim == 1: - t = t.unsqueeze(0) - print("reshaped to", t.shape) - - bm_x, per_block_exponent, per_block_mantissa, per_block_scaling = _mx_fp_quantize_hardware( - t, - width = self.quant_config["exp_width"] + self.quant_config["man_width"] + 1, - exponent_width = self.quant_config["exp_width"], - exponent_bias_width = self.quant_config["exp_bias_width"], - block_size = self.quant_config["block_size"], - skip_first_dim = self.quant_config["skip_first_dim"], + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + print("reshaped to", tensor.shape) + + bm_x, per_block_exponent, per_block_mantissa, per_block_scaling = _mx_fp_quantize_hardware( + tensor, + width = self.quant_config["exp_width"] + self.quant_config["man_width"] + 1, + exponent_width = self.quant_config["exp_width"], + exponent_bias_width = self.quant_config["exp_bias_width"], + block_size = self.quant_config["block_size"], + skip_first_dim = self.quant_config["skip_first_dim"], + ) + + logger.debug(f"per_block_mantissa: {per_block_mantissa.shape}") + logger.debug(f"per_block_exponent: {per_block_exponent.shape}") + logger.debug(f"per_block_quant_bias: {per_block_scaling.shape}") + + inner_block_list = [] + inner_scaling_list = [] + + for i in range(per_block_mantissa.shape[0]): + bin_block = pack_fp_to_bin( + per_block_exponent[i], + per_block_mantissa[i], + self.quant_config["exp_width"], + self.quant_config["man_width"], ) + inner_block_list.append(bin_block.tolist()) + inner_scaling_list.append(int(per_block_scaling[i])) + # note here the block_mantissa was represented as unsigned integer + # the exponent was represented as signed integer - logger.debug(f"per_block_mantissa: {per_block_mantissa.shape}") - logger.debug(f"per_block_exponent: {per_block_exponent.shape}") - logger.debug(f"per_block_quant_bias: {per_block_scaling.shape}") - - inner_block_list = [] - inner_scaling_list = [] - - for i in range(per_block_mantissa.shape[0]): - bin_block = pack_fp_to_bin( - per_block_exponent[i], - per_block_mantissa[i], - self.quant_config["exp_width"], - self.quant_config["man_width"], - ) - inner_block_list.append(bin_block.tolist()) - inner_scaling_list.append(int(per_block_scaling[i])) - # note here the block_mantissa was represented as unsigned integer - # the exponent was represented as signed integer - - block_list.append(inner_block_list) - scaling_list.append(inner_scaling_list) - return block_list, scaling_list + # block_list.append(inner_block_list) + # scaling_list.append(inner_scaling_list) + return inner_block_list, inner_scaling_list if __name__ == "__main__": diff --git a/tools/sim_env_utils/__init__.py b/tools/sim_env_utils/__init__.py index 54f36c84..3b52183e 100644 --- a/tools/sim_env_utils/__init__.py +++ b/tools/sim_env_utils/__init__.py @@ -1 +1 @@ -from .build_env import build_fake_sim_env \ No newline at end of file +from .build_env import create_mem_for_sim \ No newline at end of file diff --git a/tools/sim_env_utils/build_env.py b/tools/sim_env_utils/build_env.py index 63d473f3..09e6fab8 100644 --- a/tools/sim_env_utils/build_env.py +++ b/tools/sim_env_utils/build_env.py @@ -3,16 +3,70 @@ from cfl_cocotb import SRC_PATH from cfl_tools.logger import get_logger from memory_mapping.rand_gen import Random_MXFP_Tensor_Generator -from utils.load_config import load_svh_settings +from utils.load_config import load_toml_config from pathlib import Path +import torch logger = get_logger("testbench") logger.setLevel(logging.DEBUG) -def build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): +class MemoryDataManager: + """Manages memory data from pt files, supporting multiple mx and int entries.""" + def __init__(self): + self.mx_entries = [] # Can have multiple mx entries + self.int_entries = [] # Can have multiple int entries - config_settings = load_svh_settings(str(SRC_PATH / "definitions" / "configuration.svh")) - precision_settings = load_svh_settings(str(SRC_PATH / "definitions" / "precision.svh")) + def add_mx_file(self, filename, blocks, bias): + """Add an mx type data entry.""" + self.mx_entries.append({ + "filename": filename, + "type": "mx", + "blocks": blocks, + "bias": bias + }) + + def add_int_file(self, filename, data): + """Add an int type data entry.""" + self.int_entries.append({ + "filename": filename, + "type": "int", + "data": data + }) + + def get_all_entries(self): + """Get all entries as a list for iteration.""" + entries = [] + entries.extend(self.mx_entries) + entries.extend(self.int_entries) + return entries + + def to_dict(self): + """Convert to dictionary format for backward compatibility if needed.""" + result = {} + if self.mx_entries: + result["mx"] = { + "blocks": [entry["blocks"] for entry in self.mx_entries], + "bias": [entry["bias"] for entry in self.mx_entries] + } + if self.int_entries: + # For backward compatibility, use "normal" key + # If multiple int entries, combine them or use the last one + if len(self.int_entries) == 1: + result["normal"] = { + "data": self.int_entries[0]["data"] + } + else: + # If multiple int entries, use the last one (or could combine) + result["normal"] = { + "data": self.int_entries[-1]["data"] + } + return result + +def create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): + + plena_toml_path = str(SRC_PATH / "definitions" / "plena_settings.toml") + config_settings = load_toml_config(plena_toml_path, "CONFIG") + precision_settings = load_toml_config(plena_toml_path, "PRECISION") if mode == "behave_sim": asm_file = Path(PROJECT_PATH / "behavioral_simulator" / "testbench" / "build" / "generated_asm_code.asm") else: @@ -22,14 +76,15 @@ def build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, data_config = { "tensor_size": [1, data_size], - "block_size" : [1, precision_settings["BLOCK_DIM"]], + "block_size" : [1, precision_settings["HBM_M_WEIGHT_TYPE"]["block"]], } quant_config = { - "exp_width": precision_settings["ACT_MXFP_EXP_WIDTH"], - "man_width": precision_settings["ACT_MXFP_MANT_WIDTH"], - "exp_bias_width": precision_settings["MX_SCALE_WIDTH"], + "exp_width": precision_settings["HBM_V_ACT_TYPE"]["ELEM"]["exponent"], + "man_width": precision_settings["HBM_V_ACT_TYPE"]["ELEM"]["mantissa"], + "exp_bias_width": precision_settings["HBM_V_ACT_TYPE"]["SCALE"]["exponent"], "block_size": data_config["block_size"], + "int_width": precision_settings["HBM_V_INT_TYPE"]["DATA_TYPE"]["width"], "skip_first_dim": False, } @@ -50,32 +105,35 @@ def build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, grp_bias.append(bias) else: # The provided path (args.data) is a directory. Enumerate all .pt and .pth files within, - # then load and quantize all of them. Collect the results in dictionaries keyed by filename. + # then load and quantize all of them. Collect the results in a MemoryDataManager. target_dir = PROJECT_PATH / "behavioral_simulator" / "testbench" / "build" if specified_data_order is not None: pt_files = [target_dir / f"{data}.pt" for data in specified_data_order] else: pt_files = list(target_dir.glob("*.pt")) + list(target_dir.glob("*.pth")) - grp_blocks = [] - grp_bias = [] + memory_data_manager = MemoryDataManager() for pt_file in pt_files: - print("loading file", pt_file) - file_raw_data = Random_MXFP_Tensor_Generator( - shape = tuple(data_config["tensor_size"]), - quant_config = quant_config, - config_settings = config_settings, - directory = Path(asm_file).parent, - filename = pt_file - ) - file_tensor = file_raw_data.tensor_load() - blocks, bias = file_raw_data.quantize_tensor(file_tensor) - for block, b in zip(blocks, bias): - grp_blocks.append(block) - grp_bias.append(b) + if pt_file.stem != "int": + print("loading file", pt_file) + file_raw_data = Random_MXFP_Tensor_Generator( + shape = tuple(data_config["tensor_size"]), + quant_config = quant_config, + config_settings = config_settings, + directory = Path(asm_file).parent, + filename = pt_file + ) + file_tensor = file_raw_data.tensor_load() + blocks, bias = file_raw_data.quantize_tensor(file_tensor) + # Multiple mx files are all kept + memory_data_manager.add_mx_file(pt_file.name, blocks, bias) + else: + print("loading file", pt_file) + int_data = torch.load(pt_file) + memory_data_manager.add_int_file(pt_file.name, int_data) # generate_golden_result(data, logger, precision_settings, data_config) - env_setup(grp_blocks, grp_bias, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]) + env_setup(memory_data_manager, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]["value"]) if __name__ == "__main__": - build_fake_sim_env() + create_mem_for_sim() pass \ No newline at end of file diff --git a/tools/sim_env_utils/build_sys_tools.py b/tools/sim_env_utils/build_sys_tools.py index 81b6db14..cec85d2e 100644 --- a/tools/sim_env_utils/build_sys_tools.py +++ b/tools/sim_env_utils/build_sys_tools.py @@ -8,7 +8,7 @@ from cfl_cocotb.torch_fp_conversion import pack_fp_to_bin, fp_2_bin from cfl_tools import PROJECT_PATH -from memory_mapping.memory_map import map_fp_data_to_fake_hbm, map_data_to_fake_hbm_for_rtl_sim, map_data_to_fake_hbm_for_behave_sim +from memory_mapping.memory_map import map_fp_data_to_fake_hbm, map_data_to_fake_hbm_for_rtl_sim, map_mx_data_to_hbm_for_behave_sim, map_normal_data_to_hbm_for_behave_sim from assembler.assembly_to_binary import AssemblyToBinary def generate_golden_result(data, logger, precision_settings, data_config): @@ -51,7 +51,19 @@ def generate_golden_result(data, logger, precision_settings, data_config): return qdata -def env_setup(grp_blocks, grp_bias, build_path: str, data_config, quant_config, hbm_row_width=256, test_file_name=None): +def env_setup(memory_data_manager, build_path: str, data_config, quant_config, hbm_row_width=256, test_file_name=None): + """ + Setup environment for simulation using MemoryDataManager. + Each pt file entry is processed based on its type (mx or int). + + Args: + memory_data_manager: MemoryDataManager instance or dict (for backward compatibility) + build_path: Path to build directory + data_config: Data configuration dictionary + quant_config: Quantization configuration dictionary + hbm_row_width: HBM row width + test_file_name: Optional test file name + """ isa_file_path = PROJECT_PATH / 'src' / 'definitions' / 'operation.svh' config_file_path = PROJECT_PATH / 'src' / 'definitions' / 'configuration.svh' @@ -62,29 +74,45 @@ def env_setup(grp_blocks, grp_bias, build_path: str, data_config, quant_config, assembler = AssemblyToBinary(str(isa_file_path), str(config_file_path)) assembler.generate_binary(build_path / f'{test_file_name}.asm', build_path / f'{test_file_name}.mem') - for blocks, bias in zip(grp_blocks, grp_bias): - print("blocks", blocks) - print("bias", bias) - map_data_to_fake_hbm_for_rtl_sim( - blocks =blocks, - element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, - block_width =data_config["block_size"][1], - bias =bias, - bias_width =quant_config["exp_bias_width"], - combined_blk_dim=hbm_row_width // data_config["block_size"][1], - directory =build_path, - append =True, - hbm_row_width =hbm_row_width) - - map_data_to_fake_hbm_for_behave_sim( - blocks =blocks, - element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, - block_width =data_config["block_size"][1], - bias =bias, - bias_width =quant_config["exp_bias_width"], - directory =build_path, - append =True, - hbm_row_width =hbm_row_width) + entries = memory_data_manager.get_all_entries() + + # Process each entry based on its type + for entry in entries: + if entry["type"] == "mx": + blocks = entry["blocks"] + bias = entry["bias"] + print("blocks", blocks) + print("bias", bias) + print(f"Processing mx file: {entry.get('filename', 'unknown')}") + # map_data_to_fake_hbm_for_rtl_sim( + # blocks =blocks, + # element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, + # block_width =data_config["block_size"][1], + # bias =bias, + # bias_width =quant_config["exp_bias_width"], + # combined_blk_dim=hbm_row_width // data_config["block_size"][1], + # directory =build_path, + # append =True, + # hbm_row_width =hbm_row_width) + + map_mx_data_to_hbm_for_behave_sim( + blocks =blocks, + element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, + block_width =data_config["block_size"][1], + bias =bias, + bias_width =quant_config["exp_bias_width"], + directory =build_path, + append =True, + hbm_row_width =hbm_row_width) + elif entry["type"] == "int": + data = entry["data"] + print(f"Processing int file: {entry.get('filename', 'unknown')}") + map_normal_data_to_hbm_for_behave_sim( + data =data, + data_width =quant_config["int_width"], + directory =build_path, + append =True, + hbm_row_width =hbm_row_width) def parse_args(): diff --git a/tools/utils/load_config.py b/tools/utils/load_config.py index f97490de..e812d234 100644 --- a/tools/utils/load_config.py +++ b/tools/utils/load_config.py @@ -34,29 +34,7 @@ def load_json(file_path): return ml_config -def load_toml_config(file_path, mode=None): - section_to_load = ["CONFIG", "PRECISION", "INSTR"] - config = {} - +def load_toml_config(file_path, section_to_load=None): with open(file_path, "r") as f: full_toml = toml.load(f) - for section in section_to_load: - toml_config = full_toml.get(section, {}) - if not isinstance(toml_config, dict): - continue - - for param, values in toml_config.items(): - if mode is None: - config[param] = values - continue - - if isinstance(values, dict): - if mode in values: - config[param] = values[mode] - elif "value" in values: - config[param] = values["value"] - else: - config[param] = values - else: - config[param] = values - return config \ No newline at end of file + return full_toml.get(section_to_load, {}) \ No newline at end of file