Flash Attention: Performance Comparison

Philosophy: Own the Hot Path

Unlike PyTorch/TensorFlow which delegate to MKL/Accelerate, llama.cpp writes attention kernels specifically optimized for transformer inference:

• No BLAS dependency: Direct implementation in ggml_compute_forward_flash_attn_ext_*
• Inference-first design: Optimized for the specific attention patterns in decoder-only models
• Memory-layout aware: Designed for head-major and token-major layouts common in LLM inference

// ggml/src/ggml-cpu/ops.cpp
void ggml_compute_forward_flash_attn_ext(...) {
    // Custom CPU implementation, not MKL/Accelerate
    // Uses ggml_vec_dot/ggml_vec_mad directly
}

Work Partitioning Strategy

ggml uses a sophisticated thread scheduling approach that differs from naive OpenMP:

3.1 Row-Based Distribution

const int nth = params->nth;  // total threads
const int ith = params->ith;  // this thread's index

// Divide rows evenly among threads
const int nr = ggml_nrows(src0);
const int dr = (nr + nth - 1) / nth;  // rows per thread
const int ir0 = dr * ith;              // start row for this thread
const int ir1 = MIN(ir0 + dr, nr);     // end row for this thread

3.2 Chunk-Based Attention

For long sequences, attention is computed in chunks to fit L2/L3 cache:

// Process attention in cache-sized chunks
const int chunk_size = 256;  // Fits in L2 cache
for (int k_start = 0; k_start < T_k; k_start += chunk_size) {
    int k_end = MIN(k_start + chunk_size, T_k);

    // Each thread processes its row range for this chunk
    for (int i = ir0; i < ir1; ++i) {
        // Compute Q[i] × K[k_start:k_end]ᵀ
        compute_chunk(i, k_start, k_end);
    }
}

3.3 Benefits

• Load balancing: Even distribution of work across cores
• Cache locality: Chunks fit in L2/L3, reducing memory bandwidth
• NUMA-aware: Thread binding to specific cores reduces cross-socket traffic

Cache-Friendly Tiling

ggml's flash attention uses multi-level tiling to maximize cache utilization:

5.1 L1 Cache Tiling (32KB)

// Tile size chosen to fit in L1 cache (32KB)
const int TILE_Q = 32;   // Query tokens per tile
const int TILE_K = 32;   // Key tokens per tile

// Process tiles that fit in L1
for (int qt = 0; qt < T_q; qt += TILE_Q) {
    for (int kt = 0; kt < T_k; kt += TILE_K) {
        // Load Q[qt:qt+TILE_Q] into L1
        // Load K[kt:kt+TILE_K] into L1
        compute_tile(qt, kt);
    }
}

5.2 L2/L3 Cache Tiling (256KB/1MB)

// Larger tiles for L2/L3
const int BLOCK_HEAD = 8;  // Heads per block (fits in L2)
const int BLOCK_DIM = 64;  // Head dimension per block

// Block decomposition
for (int h0 = 0; h0 < H; h0 += BLOCK_HEAD) {
    for (int d0 = 0; d0 < D; d0 += BLOCK_DIM) {
        // Each block fits in L2/L3
        process_block(h0, d0);
    }
}

5.3 Benefits

• L1 hit rate > 90%: Frequent data stays in L1
• Reduced memory bandwidth: Each value loaded once, used multiple times
• NUMA locality: Data stays close to computing core

7.3 NUMA-Aware Memory Allocation

// Allocate on local NUMA node
float* data = numa_alloc_onnode(
    size * sizeof(float),
    numa_node_of_cpu(sched_getcpu())
);

// Thread binding
cpu_set_t cpuset;
CPU_ZERO(&cpuset);
CPU_SET(ith, &cpuset);
pthread_setaffinity_np(pthread_self(),
                       sizeof(cpu_set_t), &cpuset);

Benefit: Eliminates cross-socket memory traffic on multi-socket systems

1. 50/50 Split: C-Kernel and llama.cpp each win 3 out of 6 tested workloads
2. Context Length Matters: C-Kernel wins with long contexts (1K-8K), llama.cpp wins with medium (256-512)
3. Head Count Pattern: C-Kernel wins with many heads (32), llama.cpp wins with few heads (4-8)
4. No Universal Winner: Performance depends on workload characteristics, not implementation quality

Where C-Kernel Wins

• Long Contexts: 1K-8K tokens benefit from our OpenMP approach
• Many Heads: 32+ heads show good parallelization
• Simplicity: OpenMP parallel for is easier to maintain
• No Thread Overhead: Works well for appropriate workloads

Where llama.cpp Wins

• Few Heads: 4-8 heads benefit from their threading strategy
• Medium Contexts: 256-512 tokens show ggml's strength
• Large Head Dims: 128-dim heads outperform our approach

Optimization Strategy

• Understand Your Workload: Choose approach based on context length and head count
• Target Weaknesses: Optimize for few-head, medium-context cases
• Learn from ggml: Their threading strategy works well for specific patterns

Phase 1: Address Weaknesses (Priority: Medium)

1. Few Heads (4-8): Investigate why llama.cpp wins here
2. Medium Contexts (256-512): Add specialized handling
3. Large Head Dims (128): Optimize for this configuration

Phase 2: Learn from ggml (Priority: Low)

1. Study their threadpool strategy for few-head cases
2. Implement chunk-based distribution for medium contexts
3. Test fused operations on quantized types

Phase 3: Strengthen Wins (Priority: Very Low)

1. Further optimize for long contexts (1K-8K)
2. Enhance many-head (32+) parallelization
3. NUMA-aware allocation for server deployments