feat: add CPU benchmark tool with context sweep

- Add benchmark tool for LLaMA/TinyLLaMA models
- Support automated context length sweep (256-4096)
- Add KV cache memory estimation
- Include plotting script for visualization

Tested on WSL2 Ubuntu 24.04 with AMD Ryzen 7 6800H

Author: huangzhenhua111

tools/mllm-llm-benchmark/README.md

@@ -1,217 +1,114 @@
-# MLLM LLM Benchmark Tool
+# CPU Benchmark Tool
 
-## Overview
+Added a benchmark tool for measuring prefill/decode latency across different context lengths on CPU.
 
-This is a benchmark tool for measuring MLLM model performance, including:
-- **TTFT (Time To First Token)**: Time to first token latency
-- **Prefill Speed**: Prefill speed (tokens/s)
-- **Decode Speed**: Decode generation speed (tokens/s)
+## Why
 
-## Build
+There wasn't a unified way to benchmark mllm performance with varying context lengths. The existing benchmark tools had prefill length and decode length settings, but no automated sweep across contexts. So I put together this tool + a bash script to run sweeps automatically.
+
+## Test Setup
 
-Build from the mllm_v2 project root directory:
+Ran on my machine:
+- **System**: WSL2 Ubuntu 24.04  
+- **CPU**: AMD Ryzen 7 6800H with Radeon Graphics
+- **Arch**: x86_64
+- **Threads**: 8
+- **Model**: TinyLLaMA fp32 (not quantized)
 
+## Build
 ```bash
-mkdir -p build && cd build
-cmake ..
-make mllm-llm-benchmark
+cmake -S . -B build -G Ninja \
+  -DCMAKE_BUILD_TYPE=RelWithDebInfo \
+  -DBUILD_TESTING=OFF \
+  -DMLLM_ENABLE_TOOLS=ON \
+  -DCMAKE_CXX_FLAGS="-march=native"
+
+cmake --build build -j$(nproc)
 ```
 
 ## Usage
 
-### Basic Usage
-
+### Single run
 ```bash
-./mllm-llm-benchmark \
-  -n qwen3-w4a32-kai \
-  -m /path/to/model.mllm \
-  -c /path/to/config.json \
-  -t 4 \
-  -pp 64,128,256 \
-  -tg 100,200,300 \
-  -cl 2048
+./build/bin/mllm-llm-benchmark \
+  -n tiny_llama \
+  -m /path/to/tinyllama-fp32.mllm \
+  -c /path/to/config_tiny_llama.json \
+  -pp 254 -tg 2 -cl 256 -t 8
 ```
 
-### Parameters
-
-| Parameter | Long Format | Description | Example |
-|-----------|-------------|-------------|---------|
-| `-n` | `--model_name` | Model name (used to select the correct benchmark implementation) | `qwen3-w4a32-kai` |
-| `-m` | `--model_path` | Model weight file path | `/path/to/model.mllm` |
-| `-c` | `--config_path` | Model configuration file path | `/path/to/config.json` |
-| `-t` | `--threads` | Number of CPU threads | `4` |
-| `-pp` | `--prompt_length` | Prompt length list (comma-separated) | `64,128,256` |
-| `-tg` | `--test_generation_length` | Generation length list (comma-separated, must match pp count) | `100,200,300` |
-| `-cl` | `--cache_length` | Maximum KV cache length | `2048` |
+### Context sweep
 
-### Examples
+The script uses environment variables for configuration:
+```bash
+cd tools/mllm-llm-benchmark
+chmod +x scripts/sweep_context_v2.sh
 
-#### Testing Qwen3-0.6B Model
+export BIN=../../build/bin/mllm-llm-benchmark
+export MODEL=/path/to/tinyllama-fp32.mllm
+export CFG=../../examples/llama/config_tiny_llama.json
 
-```bash
-./mllm-llm-benchmark \
-  -n qwen3-w4a32-kai \
-  -m ../models/Qwen3-0.6B-w4a32kai/model.mllm \
-  -c ../models/Qwen3-0.6B-w4a32kai/config.json \
-  -t 4 \
-  -pp 64,128,256 \
-  -tg 100,100,100 \
-  -cl 2048
+./scripts/sweep_context_v2.sh
 ```
 
-#### Quick Test (Single Configuration)
+Output goes to `bench_context/context_sweep_v2.csv`. You can tweak other params like `THREADS`, `RUNS`, `COOLDOWN` if needed - they have sensible defaults.
 
+### Plot results
 ```bash
-./mllm-llm-benchmark \
-  -n qwen3-w4a32-kai \
-  -m ../models/Qwen3-0.6B-w4a32kai/model.mllm \
-  -c ../models/Qwen3-0.6B-w4a32kai/config.json \
-  -t 8 \
-  -pp 128 \
-  -tg 128 \
-  -cl 2048
+python3 scripts/plot_sweep.py bench_context/context_sweep_v2.csv snapshots/
 ```
 
-## Output Example
+## Benchmark Results
 
-```
-MLLM Build Version : abc123def456
-ARCH               : ARM64
-FP16               : true
-...
-
-Create Benchmark: qwen3-w4a32-kai
-Model Info
-========== Model Information ==========
-Model Type         : Qwen3 W4A32 KAI
-Hidden Size        : 1024
-Num Layers         : 28
-...
-=======================================
-
-Warmup Run
-Warming up with 8 tokens prefill and 4 tokens generation...
-Warmup completed
-
-========================================
-Starting Benchmark Tests
-========================================
-
-----------------------------------------
-Test Configuration:
-  Prompt Length (PP)    : 128
-  Generation Length (TG): 128
-----------------------------------------
-  Run 1 of 3...
-    TTFT         : 902.38605 ms
-    Prefill Speed: 141.84618 tokens/s
-    Decode Speed : 78.11022 tokens/s
-    Cooling down for 5 seconds...
-  Run 2 of 3...
-    TTFT         : 911.94403 ms
-    Prefill Speed: 140.3595 tokens/s
-    Decode Speed : 77.60929 tokens/s
-    Cooling down for 5 seconds...
-  Run 3 of 3...
-    TTFT         : 923.905 ms
-    Prefill Speed: 138.54239 tokens/s
-    Decode Speed : 76.48289 tokens/s
-
-========== Average Results ==========
-Configuration: PP= 128  TG= 128
-Average TTFT         : 912.74506 ms
-Average Prefill Speed: 140.24936 tokens/s
-Average Decode Speed : 77.4008 tokens/s
-=====================================
-
-
-========================================
-Benchmark Tests Completed
-========================================
-```
+Tested TinyLLaMA fp32 on Ryzen 7 6800H (WSL2):
 
-## Test Workflow
-
-Each test configuration executes the following steps:
-
-1. **Clear Cache** - Ensures each test starts from a clean state
-2. **Run 3 Test Rounds** - 5-second sleep between rounds to avoid overheating
-3. **Calculate Average** - Average results from 3 rounds
-4. **Output Results** - Display TTFT, Prefill Speed, Decode Speed
-
-## Adding New Model Support
-
-### 1. Create New Benchmark Class
-
-Create `YourModel_Benchmark.hpp` in the `models/` directory:
-
-```cpp
-#include "BenchmarkTemplate.hpp"
-#include <mllm/models/yourmodel/modeling_yourmodel.hpp>
-
-class YourModel_Benchmark final : public BenchmarkTemplate {
- public:
-  void init(const std::string& cfg_path, const std::string& model_path, int32_t cache_length) override {
-    // Initialize your model
-  }
-  
-  void printModelInfo() override {
-    // Print model information
-  }
-  
-  void warmup() override {
-    // Warmup run
-  }
-  
-  void clear() override {
-    // Clear KV cache
-  }
-  
-  BenchmarkTemplateResult run(int32_t pp, int32_t tg) override {
-    // Run test and return results
-  }
-  
- private:
-  std::unique_ptr<YourModelConfig> config_;
-  std::unique_ptr<YourModel> model_;
-};
-```
+| Context Length | TTFT (s) | Decode Latency (ms/token) | Peak RSS (GB) |
+|----------------|----------|---------------------------|---------------|
+| 256            | 8.5      | 0.68                      | 4.0           |
+| 512            | 18.2     | 0.68                      | 4.2           |
+| 1024           | 50.5     | 0.67                      | 4.6           |
+| 2048           | 149.8    | 0.73                      | 5.9           |
+| 4096           | 523.1    | 0.88                      | 7.0           |
 
-### 2. Register in All.hpp
-
-```cpp
-#include "YourModel_Benchmark.hpp"
-
-std::shared_ptr<BenchmarkTemplate> createBenchmark(const std::string& model_name) {
-  auto normalized_model_name = tolower(model_name);
-  
-  // Add your model check
-  if (normalized_model_name.find("yourmodel") != std::string::npos) {
-    return std::make_shared<YourModel_Benchmark>();
-  }
-  
-  // ... other models
-}
-```
+Prefill time scales roughly quadratically with context length (as expected). Decode latency per token stays pretty stable around 0.7ms. Memory usage grows linearly with context due to KV cache.
 
-## Notes
+Note: This is running in WSL2 so performance is slightly worse than native Linux. Also using unquantized fp32 model which is slower than q4/q8 would be.
 
-- Ensure sufficient memory for model loading and inference
-- 5-second sleep between test rounds to avoid device overheating affecting results
-- TTFT is in milliseconds, speed is in tokens/s
-- Each configuration runs 3 rounds and averages results for more stable measurements
+## Test Modes
 
-## Troubleshooting
+The sweep script runs two modes for each context length:
+- `prefill_ttft`: Measures time to first token (prompt length = CL-2, generates 2 tokens)
+- `decode_heavy`: Measures decode throughput (prompt length = CL-256, generates 256 tokens)
 
-### Error: Model not initialized
-- Check if model path is correct
-- Verify configuration file format is correct
+## Configuration
+
+Environment variables for `sweep_context_v2.sh`:
+- `BIN`: Path to benchmark binary (required)
+- `MODEL`: Path to model file (required)
+- `CFG`: Path to config json (default: ./examples/llama/config_tiny_llama.json)
+- `THREADS`: Number of threads (default: 8)
+- `RUNS`: How many runs to average (default: 1)
+- `COOLDOWN`: Seconds to wait between runs (default: 0)
+- `CLS`: Context lengths to test (default: "256 512 1024 2048 4096")
+- `TG_DH`: Generate length for decode_heavy mode (default: 256)
+- `TG_TTFT`: Generate length for prefill_ttft mode (default: 2)
+- `OUTDIR`: Output directory (default: bench_context)
+
+## Code Structure
+```
+tools/mllm-llm-benchmark/
+├── main.cpp                          # CLI entry point
+├── models/
+│   ├── BenchmarkTemplate.hpp         # Base interface
+│   ├── All.hpp                       # Model factory
+│   └── Llama.hpp                     # LLaMA/TinyLLaMA implementation
+├── scripts/
+│   ├── sweep_context_v2.sh           # Context sweep script
+│   └── plot_sweep.py                 # Plotting script
+```
 
-### Error: Benchmark not found
-- Check if model name is correct (using `-n` parameter)
-- Confirm corresponding model registration exists in `All.hpp`
+The benchmark tool supports multiple models through the `BenchmarkTemplate` interface. Currently implemented:
+- LLaMA/TinyLLaMA
+- Qwen3 (w4a32 with KAI)
 
-### Performance Anomalies
-- Check CPU thread count setting (`-t` parameter)
-- Confirm no other programs are consuming system resources
-- Check if correct quantization scheme is enabled
+Adding new models just requires implementing the template interface.