Power Consumption Collection

.cargo/config.toml

@@ -1,5 +1,6 @@
 [alias]
 tools = "run --quiet -p cli-tools --"
+collect-metrics = "run --quiet --features collect-metrics -p cli-tools -- collect-metrics"
 
 [env]
 RUST_TEST_THREADS = "1"

Cargo.toml

@@ -66,6 +66,7 @@ half = { version = "2.7", features = ["bytemuck", "num-traits"] }
 ndarray = "0.17"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
+csv = "1.3"
 monostate = "0.1"
 thiserror = "2.0"
 tokenizers = { version = "0.22", features = [

crates/backend-uzu/src/lib.rs

@@ -15,6 +15,7 @@ pub mod data_type;
 
 pub mod engine;
 
+pub use parameters::{HeaderSummary, summarize_header};
 pub use utils::version::{TOOLCHAIN_VERSION, VERSION};
 
 #[cfg(test)]

crates/backend-uzu/src/parameters/mod.rs

@@ -4,4 +4,4 @@ mod loader;
 mod safetensors_metadata;
 
 pub use loader::{ParameterLoader, ParameterLoaderError, ParameterTree};
-pub use safetensors_metadata::HeaderLoadingError;
+pub use safetensors_metadata::{HeaderLoadingError, HeaderSummary, summarize_header};

crates/backend-uzu/src/parameters/safetensors_metadata.rs

@@ -1,6 +1,6 @@
 // This code is based on the safetensors implementation: https://docs.rs/safetensors/latest/src/safetensors/tensor.rs.html
 
-use std::{collections::HashMap, fs::File, str::Utf8Error};
+use std::{collections::HashMap, fs::File, path::Path, str::Utf8Error};
 
 use serde::{Deserialize, Serialize};
 use thiserror::Error;
@@ -109,6 +109,38 @@ impl Dtype {
 
 const MAX_HEADER_SIZE: usize = 100_000_000;
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct HeaderSummary {
+    pub tensor_count: usize,
+    pub logical_payload_bytes: u64,
+}
+
+impl HeaderSummary {
+    pub fn from_metadata(metadata: &HashMetadata) -> Result<Self, HeaderLoadingError> {
+        let mut logical_payload_bytes = 0u64;
+        for (key, tensor) in &metadata.tensors {
+            let (begin, end) = tensor.data_offsets;
+            let size = end.checked_sub(begin).ok_or_else(|| HeaderLoadingError::InvalidTensorOffsets {
+                key: key.clone().into_boxed_str(),
+                begin,
+                end,
+            })?;
+            logical_payload_bytes =
+                logical_payload_bytes.checked_add(size as u64).ok_or(HeaderLoadingError::InvalidHeaderLength)?;
+        }
+        Ok(Self {
+            tensor_count: metadata.tensors.len(),
+            logical_payload_bytes,
+        })
+    }
+}
+
+pub fn summarize_header(path: &Path) -> Result<HeaderSummary, HeaderLoadingError> {
+    let file = File::open(path).map_err(HeaderLoadingError::UnableToReadHeader)?;
+    let (_, metadata) = read_metadata(&file)?;
+    HeaderSummary::from_metadata(&metadata)
+}
+
 pub fn read_metadata(file: &File) -> Result<(usize, HashMetadata), HeaderLoadingError> {
     let mut header_buffer = [0u8; size_of::<u64>()];
     file_read_exact_at(file, &mut header_buffer, 0).map_err(HeaderLoadingError::UnableToReadHeader)?;

crates/backend-uzu/tests/unit/parameters/safetensors_metadata.rs

@@ -3,7 +3,7 @@ use std::fs::File;
 use proc_macros::uzu_test;
 use test_runner::path::get_test_weights_path;
 
-use crate::parameters::safetensors_metadata::read_metadata;
+use crate::parameters::safetensors_metadata::{read_metadata, summarize_header};
 
 #[uzu_test]
 fn test_metadata_loading() {
@@ -12,3 +12,11 @@ fn test_metadata_loading() {
     let (_offset, metadata) = read_metadata(&file).expect("read metadata");
     assert!(!metadata.tensors.is_empty());
 }
+
+#[uzu_test]
+fn test_header_summary() {
+    let path = get_test_weights_path();
+    let summary = summarize_header(&path).expect("summarize header");
+    assert!(summary.tensor_count > 0);
+    assert!(summary.logical_payload_bytes > 0);
+}

crates/cli-tools/Cargo.toml

@@ -26,6 +26,30 @@ minijinja.workspace = true
 anyhow.workspace = true
 indexmap.workspace = true
 colored.workspace = true
+tokio = { workspace = true, features = ["rt-multi-thread", "macros"] }
+
+backend-uzu = { workspace = true, optional = true }
+uzu = { workspace = true, default-features = false, features = ["backend-metal", "capability-grammar"], optional = true }
+keisoku = { workspace = true, optional = true }
+kiban = { workspace = true, optional = true }
+shoji = { workspace = true, optional = true }
+reqwest = { workspace = true, optional = true }
+csv = { workspace = true, optional = true }
+tokio-stream = { workspace = true, optional = true }
+
+[features]
+collect-metrics = [
+    "dep:backend-uzu",
+    "dep:uzu",
+    "dep:keisoku",
+    "dep:kiban",
+    "dep:shoji",
+    "dep:reqwest",
+    "dep:csv",
+    "dep:tokio-stream",
+    "tokio/time",
+    "tokio/fs",
+]
 
 [lints]
 workspace = true

crates/cli-tools/src/collect_metrics/README.md

@@ -0,0 +1,129 @@
+# Collect metrics
+
+`uzu-tools collect-metrics` measures per-model power, energy, and DRAM traffic across a prefill/generate sweep on macOS.
+
+## Running
+
+The command lives behind the non-default `collect-metrics` feature of the `cli-tools` crate. Run it from the repo root with the feature enabled:
+
+```bash
+cargo run --release --features collect-metrics -p cli-tools -- collect-metrics \
+  --output metrics.csv
+```
+
+`uzu-tools` (the CLI's `bin_name`) is used as shorthand in the examples below — substitute the `cargo run --release --features collect-metrics -p cli-tools -- collect-metrics` prefix. The repo's `cargo tools` alias does **not** work here: it doesn't enable the `collect-metrics` feature, and the flag can't be appended through it. For a short command, add an alias to `.cargo/config.toml`:
+
+```toml
+[alias]
+collect-metrics = "run --quiet --features collect-metrics -p cli-tools -- collect-metrics"
+```
+
+then run `cargo collect-metrics --output metrics.csv`.
+
+## Modes
+
+### Registry (default)
+
+Uses the live Mirai registry. Configs are downloaded through `Storage`; safetensors headers are range-fetched and cached beside each config.
+
+```bash
+uzu-tools collect-metrics \
+  --output metrics.csv \
+  --model-id mirai:llama-3-8b-instruct \
+  --storage ~/power-cache
+```
+
+- Omit `--model-id` to benchmark every downloadable chat model.
+- Repeat `--model-id` to select multiple registry models by exact `model.identifier`.
+- `--storage` is optional. When set, registry artifacts are written under that directory so they can be replayed in local mode.
+- When omitted, registry mode uses a temporary cache at `$TMPDIR/uzu-collect-metrics`.
+
+### Local
+
+Replays compatible artifacts from a storage directory. No registry access, `UzuEngine`, or network I/O.
+
+```bash
+uzu-tools collect-metrics \
+  --source local \
+  --storage ~/power-cache \
+  --model-id llama/llama-3-8b-instruct/v1.0
+```
+
+- `--storage` is required.
+- Every compatible model under the storage tree is discovered when no `--model-id` is given.
+- Repeat `--model-id` to select artifacts by exact storage-relative ID (path under `models/`).
+
+## Storage layout
+
+`--storage` is the `StorageConfig::base_path`. Model artifacts live at:
+
+```text
+<storage>/.cache/mirai/models/
+  <reference-name>/<cache-identifier>/<checkpoint-version>/
+    config.json
+    model.header.safetensors
+```
+
+Registry mode writes `config.json` through `Storage` and saves the HTTP range-fetched safetensors header as `model.header.safetensors`.
+
+Local mode accepts:
+
+- `model.header.safetensors` (preferred; header-only file), or
+- `model.safetensors` (full weights file; only the header is read)
+
+Both work with random-weight loading because tensor payloads are synthesized deterministically.
+
+## Artifact requirements
+
+- `config.json` must be the converted Uzu `LanguageModelConfig`, not a Hugging Face config.
+- The safetensors header must contain the exact tensor keys, dtypes, shapes, byte-correct offsets, and `__metadata__` weight specs required by that config.
+- No tensor payload is required in `model.header.safetensors`.
+
+## Common options
+
+```text
+--source registry|local                      # default: registry
+--storage <DIR>                              # optional cache base for registry (default: temp); required for local
+--model-id <ID>                              # repeatable exact model ID selector
+--prefill 1,2,4,6,8,10,12,16,32,64,128,256   # prefill token counts (default: 1,2,4,6,8,10,12,16,32,64,128,256)
+--generate 32,128                            # decode token counts (default: 32,128)
+--iterations 6                               # measured iterations per prefill/generate pair (default: 6)
+```
+
+## Registry-to-local replay
+
+Run once against the registry into a dedicated cache directory:
+
+```bash
+uzu-tools collect-metrics \
+  --storage ~/power-cache \
+  --model-id mirai:tiny \
+  --output registry.csv
+```
+
+Replay the same artifacts offline using the storage-relative IDs printed during discovery:
+
+```bash
+uzu-tools collect-metrics \
+  --source local \
+  --storage ~/power-cache \
+  --model-id tiny/tiny-model/v1.0 \
+  --output local.csv
+```
+
+The CSV includes a `source` column (`registry` or `local`) for each row.
+
+## DRAM columns
+
+Each row also carries DRAM memory-subsystem metrics captured over the same measurement window:
+
+- `dram_read_bytes` / `dram_write_bytes` — total bytes moved (volume). Sourced from the AMC/PMP byte counters; populated on M1/A18-class chips.
+- `dram_read_gbps` / `dram_write_gbps` — residency-weighted average bandwidth (rate). Sourced from the PMP bandwidth histogram; populated on M4-class chips.
+
+Depending on the chip, typically one of the two sources is populated and the other reads `0`.
+
+## Power-user workflow
+
+1. Create a directory tree matching the layout above.
+2. Place a valid Uzu `config.json` and safetensors header (or full `model.safetensors`) in each model directory.
+3. Run with `--source local --storage <DIR>`.

crates/cli-tools/src/collect_metrics/artifacts.rs

@@ -0,0 +1,22 @@
+use std::path::{Path, PathBuf};
+
+pub const CACHE_NAME: &str = "mirai";
+pub const CONFIG_FILE: &str = "config.json";
+pub const HEADER_FILE: &str = "model.header.safetensors";
+pub const WEIGHTS_FILE: &str = "model.safetensors";
+
+pub fn cache_models_path(storage_base: &Path) -> PathBuf {
+    storage_base.join(".cache").join(CACHE_NAME).join("models")
+}
+
+pub fn resolve_weights_path(model_dir: &Path) -> Option<PathBuf> {
+    let header_path = model_dir.join(HEADER_FILE);
+    if header_path.is_file() {
+        return Some(header_path);
+    }
+    let weights_path = model_dir.join(WEIGHTS_FILE);
+    if weights_path.is_file() {
+        return Some(weights_path);
+    }
+    None
+}