-
Notifications
You must be signed in to change notification settings - Fork 2
Pillar 3: Performance Optimization
Goal: Achieve 85+ tokens/sec inference (match/exceed llama.cpp)
Current Status: 25% (memory and latency wins achieved, speed needs work)
Last Updated: 14 December 2025
EMBODIOS v0.3 vs llama.c:
| Metric | EMBODIOS v0.3 | llama.c | Status | Analysis |
|---|---|---|---|---|
| Memory | 120 MB | 160 MB | ✅ 25% better | No OS overhead, identity-mapped memory |
| Latency jitter | ±0.5ms | ±5-10ms | ✅ 10x better | Bare-metal, no context switching |
| Speed | ~17 tok/s | 83-86 tok/s | ❌ 5x slower | No SIMD, no KV cache, unoptimized matmul |
Why we're already better at memory:
- No kernel/userspace boundary → zero-copy operations
- Identity-mapped memory → no page table overhead
- Custom AI heap allocator → efficient model memory layout
- No OS background processes → 100% memory for AI
Why we have better latency:
- Bare-metal execution → no syscall overhead
- No scheduler → deterministic timing
- No interrupts during inference → consistent latency
- Direct hardware access → minimal jitter
Why we're slower (for now):
- No KV cache → recomputes all previous tokens every time
- No SIMD → scalar operations instead of vectorized
- Unoptimized matmul → naive O(n³) implementation
- No precomputed embeddings → lookup overhead
- Heap allocator overhead → frequent malloc/free
Quick Wins (~3.6x total):
- KV cache implementation → 2x speedup
- Pre-compute embedding table → 1.2x speedup
- Enable SSE2 → 1.5x speedup
Medium Wins (~1.45x total): 4. Optimize heap allocator → 1.1x speedup 5. Memory layout optimization → 1.2x speedup 6. Cache-friendly data structures → 1.1x speedup
Major Wins (~3.1x total): 7. SIMD matrix multiplication (SSE2/AVX2) → 2x speedup 8. Fixed-point optimizations → 1.3x speedup 9. Custom quantization kernels → 1.2x speedup
Expected Cumulative: 3.6 × 1.45 × 3.1 ≈ 16.2x
From 17 tok/s → 275 tok/s (theoretical max)
Conservative Target: 85 tok/s (5x speedup, very achievable)
Problem: Autoregressive generation recomputes attention for all previous tokens every time
Current (without KV cache):
Generate token 1: compute attention for token 1 → 1 computation
Generate token 2: compute attention for tokens 1-2 → 2 computations
Generate token 3: compute attention for tokens 1-2-3 → 3 computations
...
Generate token 50: compute attention for tokens 1-50 → 50 computations
Total: 1+2+3+...+50 = 1275 computations (O(n²))
With KV cache:
Generate token 1: compute K₁, V₁, cache them → 1 computation
Generate token 2: compute K₂, V₂, reuse K₁, V₁ → 1 computation
Generate token 3: compute K₃, V₃, reuse K₁-K₂, V₁-V₂ → 1 computation
...
Generate token 50: compute K₅₀, V₅₀, reuse all → 1 computation
Total: 50 computations (O(n))
Implementation:
// kernel/ai/kv_cache.c
struct kv_cache {
fixed_t* key_cache; // [max_seq_len][n_heads][head_dim]
fixed_t* value_cache; // [max_seq_len][n_heads][head_dim]
int seq_len; // Current sequence length
int max_seq_len; // Maximum sequence length (2048)
};
// Allocate KV cache
struct kv_cache* kv_cache_create(int max_seq_len, int n_heads, int head_dim) {
struct kv_cache* cache = kmalloc(sizeof(*cache), GFP_KERNEL);
size_t cache_size = max_seq_len * n_heads * head_dim * sizeof(fixed_t);
cache->key_cache = vmalloc(cache_size);
cache->value_cache = vmalloc(cache_size);
cache->seq_len = 0;
cache->max_seq_len = max_seq_len;
return cache;
}
// Cache K, V for current position
void kv_cache_update(struct kv_cache* cache, fixed_t* k, fixed_t* v,
int pos, int n_heads, int head_dim) {
int offset = pos * n_heads * head_dim;
// Copy current K, V into cache
memcpy(&cache->key_cache[offset], k, n_heads * head_dim * sizeof(fixed_t));
memcpy(&cache->value_cache[offset], v, n_heads * head_dim * sizeof(fixed_t));
cache->seq_len = pos + 1;
}
// Retrieve cached K, V for attention computation
void kv_cache_get(struct kv_cache* cache, fixed_t** k_out, fixed_t** v_out) {
*k_out = cache->key_cache;
*v_out = cache->value_cache;
}Memory Cost:
- TinyLlama: 2048 seq × 32 heads × 64 dim × 2 (K+V) × 4 bytes = 32 MB
- Acceptable for <150 MB total budget
Expected Speedup: 2x (verified by llama.cpp benchmarks)
Problem: Embedding lookup does pointer arithmetic every token
Current:
// Lookup embedding for token ID
fixed_t* embed = &model->token_embeddings[token_id * embed_dim];
// → Multiplication every time!Optimized:
// Pre-compute all embedding pointers during model load
fixed_t** embed_table = malloc(vocab_size * sizeof(fixed_t*));
for (int i = 0; i < vocab_size; i++) {
embed_table[i] = &model->token_embeddings[i * embed_dim];
}
// Fast lookup (single array access)
fixed_t* embed = embed_table[token_id];Memory Cost: 50K tokens × 8 bytes/pointer = 400 KB
Expected Speedup: 1.2x (reduces per-token overhead)
Problem: Compiler uses scalar operations instead of SIMD
Current Compiler Flags:
CFLAGS = -O2 -ffreestandingOptimized:
CFLAGS = -O2 -ffreestanding -msse2 -mfpmath=sseBenefit: Auto-vectorization for loops
Example (RMSNorm):
// Before: scalar operations (1 float per cycle)
for (int i = 0; i < dim; i++) {
sum += x[i] * x[i];
}
// After: SSE2 vectorized (4 floats per cycle)
// Compiler auto-generates SIMD instructions:
// movaps, mulps, addps, etc.Expected Speedup: 1.5x (common SIMD win on x86_64)
Problem: Frequent malloc/free during inference
Current Bottlenecks:
// Allocate temporary buffers every forward pass
fixed_t* q = malloc(embed_dim * sizeof(fixed_t)); // ← heap call
fixed_t* k = malloc(embed_dim * sizeof(fixed_t)); // ← heap call
fixed_t* v = malloc(embed_dim * sizeof(fixed_t)); // ← heap call
...
free(q); free(k); free(v); // ← heap callsOptimization: Arena Allocator
// kernel/ai/arena.c
struct arena {
void* base;
size_t size;
size_t offset;
};
void* arena_alloc(struct arena* arena, size_t size) {
if (arena->offset + size > arena->size) {
return NULL; // Out of memory
}
void* ptr = (char*)arena->base + arena->offset;
arena->offset += ALIGN_UP(size, 16);
return ptr;
}
void arena_reset(struct arena* arena) {
arena->offset = 0; // Instant "free all"
}
// Usage:
struct arena inference_arena;
arena_init(&inference_arena, 16 * 1024 * 1024); // 16 MB pool
fixed_t* q = arena_alloc(&inference_arena, embed_dim * sizeof(fixed_t));
fixed_t* k = arena_alloc(&inference_arena, embed_dim * sizeof(fixed_t));
// ... do inference ...
arena_reset(&inference_arena); // Free all in O(1)Benefits:
- No free() overhead → faster
- No fragmentation → predictable
- Cache-friendly → contiguous allocations
Expected Speedup: 1.1x
Problem: Poor cache locality in weight matrices
Current Layout (AoS - Array of Structures):
// Weights stored row-major, but accessed column-wise during matmul
// → Cache misses!
float weights[4096][2048]; // Attention weights
// Matmul: y = weights @ x
for (int i = 0; i < 4096; i++) {
for (int j = 0; j < 2048; j++) {
y[i] += weights[i][j] * x[j]; // ← Strided access, cache miss!
}
}Optimized Layout (Tiled/Blocked):
// Tile weights into cache-friendly blocks
#define TILE_SIZE 64
// Rearrange weights: [4096][2048] → [64][32][64][64]
// rows cols tiles tiles block
void matmul_tiled(float* y, float* weights, float* x, int M, int N) {
for (int i0 = 0; i0 < M; i0 += TILE_SIZE) {
for (int j0 = 0; j0 < N; j0 += TILE_SIZE) {
// Process TILE_SIZE × TILE_SIZE block
// → Fits in L1 cache (64×64×4 = 16 KB)
for (int i = i0; i < min(i0+TILE_SIZE, M); i++) {
for (int j = j0; j < min(j0+TILE_SIZE, N); j++) {
y[i] += weights[i*N + j] * x[j];
}
}
}
}
}Expected Speedup: 1.2x (fewer cache misses)
Goal: Align data to cache line boundaries (64 bytes on x86)
Current:
struct layer_weights {
fixed_t* wq; // Query weights
fixed_t* wk; // Key weights
fixed_t* wv; // Value weights
fixed_t* wo; // Output weights
fixed_t* w1, *w2, *w3; // FFN weights
};Optimized:
struct layer_weights {
fixed_t* wq __attribute__((aligned(64)));
fixed_t* wk __attribute__((aligned(64)));
fixed_t* wv __attribute__((aligned(64)));
fixed_t* wo __attribute__((aligned(64)));
fixed_t* w1 __attribute__((aligned(64)));
fixed_t* w2 __attribute__((aligned(64)));
fixed_t* w3 __attribute__((aligned(64)));
};Padding to avoid false sharing:
struct kv_cache {
fixed_t* key_cache;
char padding1[64 - sizeof(fixed_t*)];
fixed_t* value_cache;
char padding2[64 - sizeof(fixed_t*)];
};Expected Speedup: 1.1x
Problem: Matmul is 80% of inference time, currently scalar
Current (scalar):
// y = W @ x (M×N matrix, N-dim vector)
for (int i = 0; i < M; i++) {
float sum = 0.0f;
for (int j = 0; j < N; j++) {
sum += W[i*N + j] * x[j]; // ← 1 float per cycle
}
y[i] = sum;
}SSE2 (4 floats per cycle):
#include <emmintrin.h> // SSE2
void matmul_sse2(float* y, float* W, float* x, int M, int N) {
for (int i = 0; i < M; i++) {
__m128 sum = _mm_setzero_ps(); // sum = [0, 0, 0, 0]
for (int j = 0; j < N; j += 4) {
__m128 w = _mm_load_ps(&W[i*N + j]); // Load 4 weights
__m128 x_vec = _mm_load_ps(&x[j]); // Load 4 x values
__m128 prod = _mm_mul_ps(w, x_vec); // 4 multiplies
sum = _mm_add_ps(sum, prod); // 4 adds
}
// Horizontal sum: sum[0] + sum[1] + sum[2] + sum[3]
sum = _mm_hadd_ps(sum, sum);
sum = _mm_hadd_ps(sum, sum);
_mm_store_ss(&y[i], sum);
}
}AVX2 (8 floats per cycle):
#include <immintrin.h> // AVX2
void matmul_avx2(float* y, float* W, float* x, int M, int N) {
for (int i = 0; i < M; i++) {
__m256 sum = _mm256_setzero_ps(); // sum = [0×8]
for (int j = 0; j < N; j += 8) {
__m256 w = _mm256_load_ps(&W[i*N + j]); // Load 8 weights
__m256 x_vec = _mm256_load_ps(&x[j]); // Load 8 x values
sum = _mm256_fmadd_ps(w, x_vec, sum); // FMA: sum += w * x
}
// Horizontal sum
__m128 low = _mm256_castps256_ps128(sum);
__m128 high = _mm256_extractf128_ps(sum, 1);
__m128 sum128 = _mm_add_ps(low, high);
sum128 = _mm_hadd_ps(sum128, sum128);
sum128 = _mm_hadd_ps(sum128, sum128);
_mm_store_ss(&y[i], sum128);
}
}Fixed-Point SIMD (Q16.16):
// Use SSE2 for fixed-point (16.16 format)
void matmul_fixed_sse2(fixed_t* y, fixed_t* W, fixed_t* x, int M, int N) {
for (int i = 0; i < M; i++) {
__m128i sum = _mm_setzero_si128();
for (int j = 0; j < N; j += 4) {
__m128i w = _mm_load_si128((__m128i*)&W[i*N + j]);
__m128i x_vec = _mm_load_si128((__m128i*)&x[j]);
// Fixed-point multiply: (a * b) >> 16
__m128i prod_low = _mm_mullo_epi32(w, x_vec);
__m128i prod_high = _mm_mulhi_epi32(w, x_vec);
// Combine high and low parts
__m128i prod = _mm_or_si128(
_mm_srli_epi64(prod_low, 16),
_mm_slli_epi64(prod_high, 16)
);
sum = _mm_add_epi32(sum, prod);
}
// Horizontal sum
// ... (similar to float version)
}
}Expected Speedup: 2x (SSE2) to 4x (AVX2)
Current: Q16.16 fixed-point with generic operations
Optimizations:
- Fast multiply-accumulate:
// Current:
fixed_t mul_fixed(fixed_t a, fixed_t b) {
return ((int64_t)a * b) >> 16; // ← Division is slow!
}
// Optimized (use fused multiply-add):
fixed_t fma_fixed(fixed_t a, fixed_t b, fixed_t c) {
int64_t tmp = (int64_t)a * b;
return (tmp >> 16) + c; // Single operation
}- Approximate reciprocal (for softmax):
// Current:
fixed_t recip = (1 << 16) / x; // ← Slow division!
// Optimized (Newton-Raphson approximation):
fixed_t fast_recip(fixed_t x) {
// Initial guess from lookup table
fixed_t guess = recip_table[x >> 8];
// One Newton iteration: guess = guess * (2 - x * guess)
guess = mul_fixed(guess, (2 << 16) - mul_fixed(x, guess));
return guess;
}- Lookup tables for exp/log (softmax):
// Pre-compute exp(x) for common range
static const fixed_t exp_table[256] = {
/* exp(0) */ 65536,
/* exp(0.01) */ 66191,
...
};
fixed_t fast_exp(fixed_t x) {
int idx = (x >> 8) & 0xFF;
return exp_table[idx];
}Expected Speedup: 1.3x
Problem: Dequantization happens for every matmul
Current:
// Dequantize Q4_K block to float
void dequant_q4k(const block_q4_K* blocks, float* output, int n) {
for (int i = 0; i < n / QK_K; i++) {
const block_q4_K* b = &blocks[i];
for (int j = 0; j < QK_K; j++) {
uint8_t q = (b->qs[j/2] >> ((j%2)*4)) & 0xF;
float scale = b->scales[j / 16] * b->d;
output[i*QK_K + j] = (q - 8) * scale; // ← Slow!
}
}
// Then: matmul with dequantized weights
matmul(output, x, y, M, N);
}Optimized: Fused dequant + matmul
void matmul_q4k_sse2(float* y, const block_q4_K* W, float* x,
int M, int N) {
for (int i = 0; i < M; i++) {
__m128 sum = _mm_setzero_ps();
for (int j = 0; j < N / QK_K; j++) {
const block_q4_K* block = &W[i * (N/QK_K) + j];
// Dequantize and multiply in one pass
for (int k = 0; k < QK_K; k += 8) {
// Load 8 quantized values (4 bytes)
uint32_t qs = *(uint32_t*)&block->qs[k/2];
// Extract 4-bit values to 8×int8
__m128i q8 = _mm_set_epi8(
(qs >> 28) & 0xF, (qs >> 24) & 0xF,
(qs >> 20) & 0xF, (qs >> 16) & 0xF,
(qs >> 12) & 0xF, (qs >> 8) & 0xF,
(qs >> 4) & 0xF, qs & 0xF
);
// Convert to float and apply scale
__m128 q_float = _mm_cvtepi32_ps(_mm_cvtepi8_epi32(q8));
__m128 scale = _mm_set1_ps(block->d * block->scales[k/16]);
__m128 dequant = _mm_mul_ps(q_float, scale);
// Load x values
__m128 x_vec = _mm_load_ps(&x[j*QK_K + k]);
// Accumulate
sum = _mm_add_ps(sum, _mm_mul_ps(dequant, x_vec));
}
}
// Horizontal sum
sum = _mm_hadd_ps(sum, sum);
sum = _mm_hadd_ps(sum, sum);
_mm_store_ss(&y[i], sum);
}
}Expected Speedup: 1.2x (avoids temporary buffer allocation)
Goal: Identify remaining bottlenecks
Tools:
- Cycle Counters (RDTSC):
// kernel/include/embodios/perf.h
static inline uint64_t rdtsc(void) {
uint32_t lo, hi;
__asm__ __volatile__ ("rdtsc" : "=a"(lo), "=d"(hi));
return ((uint64_t)hi << 32) | lo;
}
#define PROFILE_START() uint64_t _t0 = rdtsc()
#define PROFILE_END(name) \
do { \
uint64_t _t1 = rdtsc(); \
printk("%s: %llu cycles\n", name, _t1 - _t0); \
} while (0)
// Usage:
PROFILE_START();
matmul(y, W, x, M, N);
PROFILE_END("matmul");- Function-Level Profiling:
struct profile_entry {
const char* name;
uint64_t total_cycles;
uint64_t call_count;
};
static struct profile_entry profiles[128];
static int n_profiles = 0;
void profile_report(void) {
printk("=== Performance Profile ===\n");
for (int i = 0; i < n_profiles; i++) {
uint64_t avg = profiles[i].total_cycles / profiles[i].call_count;
printk("%s: %llu cycles avg (%llu calls)\n",
profiles[i].name, avg, profiles[i].call_count);
}
}- Memory Bandwidth Measurement:
void measure_bandwidth(void) {
size_t size = 64 * 1024 * 1024; // 64 MB
void* buf = malloc(size);
uint64_t t0 = rdtsc();
memcpy(buf, model_weights, size);
uint64_t t1 = rdtsc();
uint64_t cycles = t1 - t0;
double seconds = cycles / 3.0e9; // Assuming 3 GHz CPU
double bw = size / seconds / 1e9; // GB/s
printk("Memory bandwidth: %.2f GB/s\n", bw);
}Hot Loop Optimization:
// Before: Branch in hot loop
for (int i = 0; i < n; i++) {
if (x[i] > 0) {
y[i] = x[i];
} else {
y[i] = 0;
}
}
// After: Branchless
for (int i = 0; i < n; i++) {
y[i] = (x[i] > 0) * x[i]; // ← No branch!
}Loop Unrolling:
// Before:
for (int i = 0; i < n; i++) {
y[i] = a * x[i] + b;
}
// After: 4× unrolled
for (int i = 0; i < n; i += 4) {
y[i+0] = a * x[i+0] + b;
y[i+1] = a * x[i+1] + b;
y[i+2] = a * x[i+2] + b;
y[i+3] = a * x[i+3] + b;
}Prefetching:
// Before:
for (int i = 0; i < n; i++) {
y[i] = expensive_computation(x[i]);
}
// After: Prefetch next iteration
for (int i = 0; i < n; i++) {
__builtin_prefetch(&x[i+8], 0, 3); // Prefetch 8 ahead
y[i] = expensive_computation(x[i]);
}Benchmark Suite:
# kernel/benchmarks/
# 1. Tokens/sec measurement
./bench_tokps.sh tinyllama-q4km.gguf
# Expected: 85+ tokens/sec
# 2. Memory usage
./bench_memory.sh tinyllama-q4km.gguf
# Expected: <150 MB total
# 3. Latency jitter
./bench_jitter.sh tinyllama-q4km.gguf 1000
# Expected: ±0.5ms or better
# 4. Boot time
./bench_boot.sh
# Expected: <1 second from power-on to readyComparison with llama.cpp:
# compare_performance.py
import subprocess
def benchmark_embodios():
result = subprocess.run(["./bench_tokps.sh", "tinyllama.gguf"],
capture_output=True, text=True)
# Parse output: "85.3 tokens/sec"
return float(result.stdout.split()[0])
def benchmark_llamacpp():
result = subprocess.run(["./llama-cli", "-m", "tinyllama.gguf",
"-p", "test", "-n", 50],
capture_output=True, text=True)
# Parse llama.cpp output
return extract_tokps(result.stdout)
embodios_tokps = benchmark_embodios()
llamacpp_tokps = benchmark_llamacpp()
print(f"EMBODIOS: {embodios_tokps:.2f} tok/s")
print(f"llama.cpp: {llamacpp_tokps:.2f} tok/s")
print(f"Ratio: {embodios_tokps / llamacpp_tokps:.2f}x")
# Target: ratio >= 1.0 (matching or exceeding llama.cpp)
assert embodios_tokps >= 85.0, "Performance target not met"| Metric | Current | v1.0 Target | Status |
|---|---|---|---|
| Speed | 17 tok/s | 85+ tok/s | ❌ In progress |
| Memory | 120 MB | <150 MB | ✅ On track |
| Latency jitter | ±0.5ms | ±0.5ms | ✅ Achieved |
| Boot time | ~1s | <1s | ✅ On track |
| Model size | 15M params | 1.1B params | ❌ In progress |
v1.1:
- Speed: 150+ tok/s (multi-core, AVX2)
- Models: Phi-2 (2.7B), Mistral-7B (7B)
v2.0:
- Speed: 500+ tok/s (GPU acceleration)
- Distributed inference across multiple nodes
Chapter 1 (Quick Wins):
┌────────────────┐
│ KV Cache │ → Independent
├────────────────┤
│ Embeddings │ → Independent
├────────────────┤
│ SSE2 Enable │ → Independent
└────────────────┘
Chapter 2 (Medium Wins):
┌────────────────┐
│ Heap Allocator │ → Depends on Chapter 1 complete
├────────────────┤
│ Memory Layout │ → Independent
├────────────────┤
│ Cache-Friendly │ → Depends on Memory Layout
└────────────────┘
Chapter 3 (Major Wins):
┌────────────────┐
│ SIMD Matmul │ → Depends on SSE2 enabled (Chapter 1)
├────────────────┤
│ Fixed-Point │ → Independent
├────────────────┤
│ Quant Kernels │ → Depends on SIMD Matmul
└────────────────┘
Chapter 4 (Profiling):
┌────────────────┐
│ Profile │ → Depends on Chapters 1-3
├────────────────┤
│ Micro-Opt │ → Depends on Profile results
├────────────────┤
│ Validation │ → Depends on all optimizations
└────────────────┘
// Test SIMD matmul correctness
void test_simd_matmul(void) {
float A[16][16], B[16], C_ref[16], C_simd[16];
// Initialize with known values
init_test_data(A, B);
// Reference implementation (scalar)
matmul_scalar(C_ref, A, B, 16, 16);
// SIMD implementation
matmul_sse2(C_simd, A, B, 16, 16);
// Compare
for (int i = 0; i < 16; i++) {
assert(fabs(C_ref[i] - C_simd[i]) < 1e-5);
}
}# Ensure optimizations don't break correctness
# Generate reference output (llama.cpp)
echo "Once upon a time" | llama-cli -m tinyllama.gguf > ref.txt
# Generate EMBODIOS output
./embodios_test "Once upon a time" > embodios.txt
# Compare (should be identical)
diff ref.txt embodios.txtPillar 3 is complete when:
- ✅ Achieves 85+ tokens/sec on TinyLlama-1.1B Q4_K_M
- ✅ Memory usage <150 MB
- ✅ Latency jitter ±0.5ms
- ✅ Boot time <1 second
- ✅ Performance matches or exceeds llama.cpp
- ✅ All optimizations profiled and validated
#embodios #performance #optimization #simd #kv-cache #profiling #pillar-3