40. Running everywhere: CPUs, Apple, AMD and offloading

One engine, many machines

llama.cpp’s popularity comes as much from where it runs as from how fast: laptops without a discrete GPU, Macs, AMD cards, phones. vLLM and SGLang run on NVIDIA and AMD datacenter GPUs. The engine you’ve built so far is NVIDIA-first: Triton kernels, CUDA graphs, FP8 checks against compute capability. This chapter makes it portable and gives it a fast CPU path. Here’s what changes from platform to platform:

def detect():
    """Inspect the machine. Override the choice with IZH_DEVICE=cpu|cuda|mps|xpu.  (Your engine: Chapter 40)"""
    forced = os.environ.get("IZH_DEVICE")
    if (forced in (None, "cuda")) and torch.cuda.is_available():
        props = torch.cuda.get_device_properties(0)
        if torch.version.hip:                                   # ROCm builds reuse the cuda device API
            arch = getattr(props, "gcnArchName", "gfx").split(":")[0]
            return Platform("rocm", props.name, arch=arch, memory_bytes=props.total_memory, features={
                "bf16": True, "fp8": arch.startswith(("gfx94", "gfx95", "gfx12")), "triton": True, "graphs": True,
                "warp_size": 64})
        cap = (props.major, props.minor)
        return Platform("cuda", props.name, cap, memory_bytes=props.total_memory, features={
            "bf16": cap >= (8, 0), "fp8": cap >= (8, 9), "fp4": cap >= (10, 0), "triton": True, "graphs": True,
            "warp_size": 32})
    if (forced in (None, "mps")) and getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return Platform("mps", "Apple GPU", memory_bytes=_system_memory(), features={
            "bf16": True, "fp8": False, "triton": False, "graphs": False, "unified_memory": True})
    if (forced in (None, "xpu")) and hasattr(torch, "xpu") and torch.xpu.is_available():
        return Platform("xpu", torch.xpu.get_device_name(0), features={"bf16": True, "triton": True, "graphs": False})
    return Platform("cpu", _cpu_name(), memory_bytes=_system_memory(), features={
        "bf16": True, "triton": False, "graphs": False, "simd": _cpu_isa()})

def engine_defaults(platform):
    """EngineConfig fields that suit the hardware; explicit settings always win.  (Your engine: Chapter 40)"""
    if platform.kind in ("cuda", "rocm"):
        return {"attention_backend": "triton", "cuda_graphs": "auto", "kv_cache_dtype": "auto"}
    if platform.kind == "xpu":
        return {"attention_backend": "triton", "cuda_graphs": "off"}
    # MPS and CPU: no Triton (the interpreter is for testing, not serving), no graphs, a modest batch.
    return {"attention_backend": "reference", "cuda_graphs": "off", "max_num_seqs": 16, "max_num_batched_tokens": 512}


def default_dtype(platform):
    if platform.kind in ("cuda", "rocm"):
        return torch.bfloat16 if platform.features.get("bf16") else torch.float16
    if platform.kind == "mps":
        return torch.float16                                   # the fastest MPS matmul path
    return torch.bfloat16 if platform.kind == "cpu" and "avx512" in str(platform.features.get("simd")) else torch.float32

The defaults are a starting point that explicit settings override: Triton attention and CUDA graphs on NVIDIA and AMD GPUs; the reference attention backend, no graphs and a smaller batch on Apple and CPUs; BF16 where the hardware computes it natively. On AMD, nothing else in Parts VIII’s code needs to change: the Triton kernels of Chapters 32, 33 and 38 compile for AMD GPUs as they are. Their tile sizes were chosen for NVIDIA’s 32-thread warps, so autotuning them for 64-thread wavefronts (and MI300’s larger register file) is worth doing; vLLM and SGLang also ship AMD-specific kernels (the AITER library) for the hottest operations.

Decoding on a CPU

A CPU decodes like a GPU: every token reads every weight once, so the time per token is bytes read divided by memory bandwidth (Chapter 10). The difference is the bandwidth: a laptop’s dual-channel DDR5 delivers 60-100 GB/s, an Apple M-series chip 100-800 GB/s (shared with its GPU), a 12-channel server socket 400-600 GB/s, against an H100’s 3,350 GB/s. A CPU is a perfectly good decoder for a quantized model of a few billion parameters, and the same rule as everywhere applies: fewer bytes per weight, more tokens per second.

That makes quantized weights essential, and it creates a kernel problem. Dequantizing a whole weight matrix to float before each matmul (what AffineQuantLinear does without a kernel) reads the compact weights, writes 4 bytes per weight back to memory, and reads them again: slower than not quantizing at all. The weights must be decoded in registers, inside the dot product, as Chapter 20’s Triton kernel does on a GPU.

Integer dot products

CPUs have a faster trick than decoding to float: integer dot products. llama.cpp quantizes the activation vector too, to int8 in blocks of 32 with one float scale each, and then each 32-weight block’s contribution is

$$ y_n \mathrel{+}= \underbrace{s_{n,g}}{\text{weight scale}} \cdot \underbrace{d_b}{\text{activation scale}} \cdot \underbrace{\sum_{k \in b} q_{n,k}, a_k}{\text{integer dot product}} ;+; o{n,g} \sum_{k \in g} x_k , $$

where the integer dot product of 32 4-bit codes with 32 int8 activations is one or two instructions on modern CPUs:

The offsets need only the activations’ per-group sums, computed once per call. Quantizing activations costs a little accuracy (about 1% relative error per matmul here), which is why llama.cpp offers it per type and why AffineQuantLinear uses these kernels only when asked (set_cpu_kernels(True)).

The kernels live in a header shared by the C++ track and a PyTorch extension, and in the Rust track. First the activations, quantized once per call, with each block’s codes stored as its 16 even positions then its 16 odd ones, to meet the engine’s packing (code $k$ in the low nibble when $k$ is even):

// x [K] -> int8 codes per block of 32, one float scale per block, and each block's float sum
// (the offsets multiply sum(x) over a group). For the q4 kernels each block's codes are stored
// as [the 16 even positions | the 16 odd positions], to meet the low and high nibbles.
struct Q8Activations {
    std::vector<int8_t> q, q_split;
    std::vector<float> scale, sum;
};

inline Q8Activations quantize_activations(const float* x, int K) {
    Q8Activations a;
    int blocks = K / 32;
    a.q.resize(K), a.q_split.resize(K), a.scale.resize(blocks), a.sum.resize(blocks);
    for (int b = 0; b < blocks; ++b) {
        const float* xb = x + 32 * b;
        float amax = 0, total = 0;
        for (int i = 0; i < 32; ++i) amax = std::max(amax, std::fabs(xb[i])), total += xb[i];
        float d = amax / 127.f, inv = d > 0 ? 1.f / d : 0.f;
        a.scale[b] = d, a.sum[b] = total;
        for (int i = 0; i < 32; ++i) {
            int8_t v = int8_t(std::lround(xb[i] * inv));
            a.q[32 * b + i] = v;
            a.q_split[32 * b + (i % 2) * 16 + i / 2] = v;           // evens first, then odds
        }
    }
    return a;
}

Then one block’s integer dot product, for each instruction set. Each function is compiled for exactly its own instruction set with a target attribute, and the best one is chosen at run time, so one binary runs on any x86 CPU:

// One block of 32: integer dot product of 4-bit codes (16 bytes) with 32 int8 activations.
inline int32_t dot_q4_scalar(const uint8_t* c, const int8_t* x) {
    int32_t s = 0;
    for (int i = 0; i < 16; ++i) s += (c[i] & 15) * x[i] + (c[i] >> 4) * x[16 + i];
    return s;
}
inline int32_t dot_q8_scalar(const int8_t* c, const int8_t* x) {
    int32_t s = 0;
    for (int i = 0; i < 32; ++i) s += c[i] * x[i];
    return s;
}

#ifdef IZH_X86
__attribute__((target("avx2"))) inline __m256i nibbles_avx2(const uint8_t* c) {
    __m128i b = _mm_loadu_si128(reinterpret_cast<const __m128i*>(c));
    __m128i m = _mm_set1_epi8(15);
    return _mm256_set_m128i(_mm_and_si128(_mm_srli_epi16(b, 4), m), _mm_and_si128(b, m));   // [low | high]
}
__attribute__((target("avx2"))) inline __m256i dot_q4_avx2(const uint8_t* c, const int8_t* x) {
    __m256i p16 = _mm256_maddubs_epi16(nibbles_avx2(c), _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x)));
    return _mm256_madd_epi16(p16, _mm256_set1_epi16(1));                // 8 lanes of int32 partial sums
}
__attribute__((target("avx2"))) inline __m256i dot_q8_avx2(const int8_t* c, const int8_t* x) {
    __m256i cv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(c));
    __m256i xv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x));
    // maddubs wants an unsigned left operand: |c| x (x * sign(c)) has the same products.
    __m256i p16 = _mm256_maddubs_epi16(_mm256_sign_epi8(cv, cv), _mm256_sign_epi8(xv, cv));
    return _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
}
__attribute__((target("avx512vnni,avx512vl,avx2"))) inline __m256i dot_q4_vnni(const uint8_t* c, const int8_t* x) {
    return _mm256_dpbusd_epi32(_mm256_setzero_si256(), nibbles_avx2(c),
                               _mm256_loadu_si256(reinterpret_cast<const __m256i*>(x)));
}
__attribute__((target("avx2"))) inline float hsum(__m256 v) {
    __m128 s = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
    s = _mm_hadd_ps(s, s);
    return _mm_cvtss_f32(_mm_hadd_ps(s, s));
}
#endif
#ifdef IZH_NEON_DOT
inline int32_t dot_q4_neon(const uint8_t* c, const int8_t* x) {
    uint8x16_t b = vld1q_u8(c);
    int8x16_t lo = vreinterpretq_s8_u8(vandq_u8(b, vdupq_n_u8(15))), hi = vreinterpretq_s8_u8(vshrq_n_u8(b, 4));
    int32x4_t s = vdotq_s32(vdupq_n_s32(0), lo, vld1q_s8(x));
    return vaddvq_s32(vdotq_s32(s, hi, vld1q_s8(x + 16)));
}
#endif

#![allow(unused)]
fn main() {
/// x [K] -> int8 per block of 32 (evens then odds, to meet low and high nibbles), scales, sums.
pub fn quantize_activations(x: &[f32]) -> (Vec<i8>, Vec<f32>, Vec<f32>) {
    let mut q = vec![0i8; x.len()];
    let (mut scale, mut sum) = (vec![], vec![]);
    for (b, block) in x.chunks(32).enumerate() {
        let amax = block.iter().fold(0.0f32, |m, v| m.max(v.abs()));
        let d = amax / 127.0;
        let inv = if d > 0.0 { 1.0 / d } else { 0.0 };
        for (i, v) in block.iter().enumerate() {
            q[32 * b + (i % 2) * 16 + i / 2] = (v * inv).round() as i8;
        }
        scale.push(d);
        sum.push(block.iter().sum());
    }
    (q, scale, sum)
}

/// One block: 16 bytes of codes against 32 activations (evens then odds).
pub fn dot_q4_scalar(codes: &[u8], x: &[i8]) -> i32 {
    (0..16).map(|i| (codes[i] & 15) as i32 * x[i] as i32 + (codes[i] >> 4) as i32 * x[16 + i] as i32).sum()
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn dot_q4_avx2(codes: &[u8], x: &[i8]) -> i32 {
    use std::arch::x86_64::*;
    let b = _mm_loadu_si128(codes.as_ptr() as *const __m128i);
    let m = _mm_set1_epi8(15);
    let c = _mm256_set_m128i(_mm_and_si128(_mm_srli_epi16::<4>(b), m), _mm_and_si128(b, m));
    let p16 = _mm256_maddubs_epi16(c, _mm256_loadu_si256(x.as_ptr() as *const __m256i));
    let p32 = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
    let s = _mm_add_epi32(_mm256_castsi256_si128(p32), _mm256_extracti128_si256::<1>(p32));
    let s = _mm_hadd_epi32(s, s);
    _mm_cvtsi128_si32(_mm_hadd_epi32(s, s))
}

/// The best dot product this CPU supports, decided once per call site by runtime detection.
pub fn dot_q4(codes: &[u8], x: &[i8]) -> i32 {
    assert!(codes.len() >= 16 && x.len() >= 32);
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") {
            return unsafe { dot_q4_avx2(codes, x) };
        }
    }
    dot_q4_scalar(codes, x)
}

/// y [N] = W x for 4-bit codes [N, K/2], scales and offsets [N, K/G], rows split across threads.
pub fn gemv_q4(codes: &[u8], scales: &[f32], offsets: &[f32], n: usize, k: usize, g: usize, x: &[f32],
               threads: usize) -> Vec<f32> {
    let (q, xs, xsum) = quantize_activations(x);
    let mut y = vec![0.0f32; n];
    let per = n.div_ceil(threads.max(1));
    std::thread::scope(|scope| {
        for (t, chunk) in y.chunks_mut(per).enumerate() {
            let (q, xs, xsum) = (&q, &xs, &xsum);
            scope.spawn(move || {
                for (j, out) in chunk.iter_mut().enumerate() {
                    let row = t * per + j;
                    let (c, s, o) = (&codes[row * k / 2..], &scales[row * (k / g)..], &offsets[row * (k / g)..]);
                    let mut acc = 0.0f32;
                    for b in 0..k / 32 {
                        acc += s[32 * b / g] * xs[b] * dot_q4(&c[16 * b..], &q[32 * b..]) as f32;
                    }
                    for grp in 0..k / g {
                        acc += o[grp] * xsum[grp * g / 32..(grp + 1) * g / 32].iter().sum::<f32>();
                    }
                    *out = acc;
                }
            });
        }
    });
    y
}
}

Then the row loop, and here two details doubled the speed when this chapter’s first version measured only a third of the memory bandwidth:

• No division in the inner loop. The first version looked up each block’s scale as scales[32 * b / G], an integer division by a runtime value, 20-40 cycles per block. Looping over groups outside and blocks inside removes it.
• Independent accumulators, and four rows at a time. A single accumulator makes every fused multiply-add wait for the previous one (4 cycles of latency each). Four accumulators, and four rows sharing each block’s activation load, keep the multiply units busy.

// y[n] = sum over groups g of  scale[n, g] * sum over blocks b in g of xscale[b] * dot_b   +   offset[n, g] * xsum[g]
// Two things decide the speed: no division in the inner loop (groups outside, blocks inside),
// and four independent accumulators, so consecutive fused multiply-adds don't wait for each other.
inline std::vector<float> group_sums(const Q8Activations& a, int K, int G) {
    std::vector<float> s(K / G, 0.f);
    for (int b = 0; b < K / 32; ++b) s[32 * b / G] += a.sum[b];
    return s;
}

#ifdef IZH_X86
#define IZH_ROW_BODY(DOT, CODES_PER_BLOCK, XQ)                                                            \
    __m256 acc[4] = {_mm256_setzero_ps(), _mm256_setzero_ps(), _mm256_setzero_ps(), _mm256_setzero_ps()};  \
    int per = G / 32;                                                                                      \
    for (int g = 0, b = 0; g < K / G; ++g) {                                                               \
        float sg = scales[g];                                                                              \
        for (int i = 0; i < per; ++i, ++b) {                                                               \
            __m256i p = DOT(codes + CODES_PER_BLOCK * b, XQ + 32 * b);                                      \
            acc[b & 3] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(p), _mm256_set1_ps(sg * a.scale[b]), acc[b & 3]); \
        }                                                                                                  \
    }                                                                                                      \
    return hsum(_mm256_add_ps(_mm256_add_ps(acc[0], acc[1]), _mm256_add_ps(acc[2], acc[3])));

// Separate functions, not one template: each must be compiled for exactly its own instruction set,
// or the compiler may use AVX-512 instructions in the path meant for AVX2-only CPUs.
__attribute__((target("avx2,fma"))) inline float row_q4_avx2(const uint8_t* codes, const float* scales, int K, int G,
                                                              const Q8Activations& a) {
    IZH_ROW_BODY(dot_q4_avx2, 16, a.q_split.data())
}
__attribute__((target("avx512vnni,avx512vl,avx2,fma"))) inline float row_q4_vnni(const uint8_t* codes, const float* scales,
                                                                                int K, int G, const Q8Activations& a) {
    IZH_ROW_BODY(dot_q4_vnni, 16, a.q_split.data())
}
__attribute__((target("avx2,fma"))) inline float row_q8_avx2(const int8_t* codes, const float* scales, int K, int G,
                                                              const Q8Activations& a) {
    IZH_ROW_BODY(dot_q8_avx2, 32, a.q.data())
}
#undef IZH_ROW_BODY

// Four rows at once: each block's 32 activations are loaded once and used four times, and the four
// rows' independent dot products keep the multiply units busy. y[0..3] receive the four results.
#define IZH_ROWS4(NAME, TARGET, DOT)                                                                         \
    __attribute__((target(TARGET))) inline void NAME(const uint8_t* codes, size_t row_bytes, const float* scales, \
                                                     size_t row_groups, int K, int G, const Q8Activations& a,  \
                                                     float* y) {                                               \
        __m256 acc[4] = {_mm256_setzero_ps(), _mm256_setzero_ps(), _mm256_setzero_ps(), _mm256_setzero_ps()};  \
        int per = G / 32;                                                                                      \
        for (int g = 0, b = 0; g < K / G; ++g) {                                                               \
            __m256 s[4];                                                                                       \
            for (int r = 0; r < 4; ++r) s[r] = _mm256_set1_ps(scales[r * row_groups + g]);                     \
            for (int i = 0; i < per; ++i, ++b) {                                                               \
                __m256i xv = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(a.q_split.data() + 32 * b));   \
                __m256 xs = _mm256_set1_ps(a.scale[b]);                                                        \
                for (int r = 0; r < 4; ++r) {                                                                  \
                    __m256i p = DOT(nibbles_avx2(codes + r * row_bytes + 16 * b), xv);                          \
                    acc[r] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(p), _mm256_mul_ps(s[r], xs), acc[r]);           \
                }                                                                                              \
            }                                                                                                  \
        }                                                                                                      \
        for (int r = 0; r < 4; ++r) y[r] = hsum(acc[r]);                                                       \
    }

__attribute__((target("avx2"))) inline __m256i madd_u8s8_avx2(__m256i c, __m256i x) {
    return _mm256_madd_epi16(_mm256_maddubs_epi16(c, x), _mm256_set1_epi16(1));
}
__attribute__((target("avx512vnni,avx512vl,avx2"))) inline __m256i madd_u8s8_vnni(__m256i c, __m256i x) {
    return _mm256_dpbusd_epi32(_mm256_setzero_si256(), c, x);
}
IZH_ROWS4(rows4_q4_avx2, "avx2,fma", madd_u8s8_avx2)
IZH_ROWS4(rows4_q4_vnni, "avx512vnni,avx512vl,avx2,fma", madd_u8s8_vnni)
#undef IZH_ROWS4
#endif

inline float add_offsets(float y, const float* offsets, const std::vector<float>& gsum) {
    for (size_t g = 0; g < gsum.size(); ++g) y += offsets[g] * gsum[g];
    return y;
}

inline float row_q4(const uint8_t* codes, const float* scales, const float* offsets, int K, int G,
                    const Q8Activations& a, const std::vector<float>& gsum, Isa isa) {
    float y = 0;
    switch (isa) {
#ifdef IZH_X86
        case Isa::avx512_vnni: y = row_q4_vnni(codes, scales, K, G, a); break;
        case Isa::avx2: y = row_q4_avx2(codes, scales, K, G, a); break;
#endif
#ifdef IZH_NEON_DOT
        case Isa::neon_dot:
            for (int b = 0; b < K / 32; ++b)
                y += scales[32 * b / G] * a.scale[b] * float(dot_q4_neon(codes + 16 * b, a.q_split.data() + 32 * b));
            break;
#endif
        default:
            for (int b = 0; b < K / 32; ++b)
                y += scales[32 * b / G] * a.scale[b] * float(dot_q4_scalar(codes + 16 * b, a.q_split.data() + 32 * b));
    }
    return add_offsets(y, offsets, gsum);
}

// Rows n .. n+3 into y[0..3]; falls back to one row at a time where there's no 4-row kernel.
inline void rows4_q4(const uint8_t* codes, const float* scales, const float* offsets, int K, int G,
                     const Q8Activations& a, const std::vector<float>& gsum, Isa isa, float* y) {
    size_t row_bytes = size_t(K) / 2, row_groups = size_t(K / G);
#ifdef IZH_X86
    if (isa == Isa::avx512_vnni || isa == Isa::avx2) {
        if (isa == Isa::avx512_vnni)
            rows4_q4_vnni(codes, row_bytes, scales, row_groups, K, G, a, y);
        else
            rows4_q4_avx2(codes, row_bytes, scales, row_groups, K, G, a, y);
        for (int r = 0; r < 4; ++r) y[r] = add_offsets(y[r], offsets + r * row_groups, gsum);
        return;
    }
#endif
    for (int r = 0; r < 4; ++r)
        y[r] = row_q4(codes + r * row_bytes, scales + r * row_groups, offsets + r * row_groups, K, G, a, gsum, isa);
}

inline float row_q8(const int8_t* codes, const float* scales, const float* offsets, int K, int G,
                    const Q8Activations& a, const std::vector<float>& gsum, Isa isa) {
    float y = 0;
#ifdef IZH_X86
    if (isa == Isa::avx2 || isa == Isa::avx512_vnni)
        y = row_q8_avx2(codes, scales, K, G, a);
    else
#endif
        for (int b = 0; b < K / 32; ++b)
            y += scales[32 * b / G] * a.scale[b] * float(dot_q8_scalar(codes + 32 * b, a.q.data() + 32 * b));
    return add_offsets(y, offsets, gsum);
}

// y [N] = W x for one activation row, rows split across threads. Each thread streams its own
// slice of the weight matrix: decode is memory-bound, so more threads means more bandwidth,
// up to what the memory system delivers.
inline void gemv_q4(const uint8_t* codes, const float* scales, const float* offsets, int N, int K, int G,
                    const float* x, float* y, int threads = 0, Isa isa = best_isa()) {
    Q8Activations a = quantize_activations(x, K);
    std::vector<float> gsum = group_sums(a, K, G);
    threads = threads > 0 ? threads : int(std::max(1u, std::thread::hardware_concurrency()));
    std::vector<std::thread> pool;
    int per = (N + threads - 1) / threads;
    for (int t = 0; t < threads; ++t) {
        pool.emplace_back([&, t] {
            int n = t * per, end = std::min(N, (t + 1) * per);
            for (; n + 4 <= end; n += 4)
                rows4_q4(codes + size_t(n) * K / 2, scales + size_t(n) * (K / G), offsets + size_t(n) * (K / G),
                         K, G, a, gsum, isa, y + n);
            for (; n < end; ++n)
                y[n] = row_q4(codes + size_t(n) * K / 2, scales + size_t(n) * (K / G), offsets + size_t(n) * (K / G),
                              K, G, a, gsum, isa);
        });
    }
    for (auto& th : pool) th.join();
}

The PyTorch extension (izh/kernels/cpu_ext.cpp) wraps the same functions, splits rows across threads with at::parallel_for (which needs -fopenmp, or it silently runs on one thread, the third lesson of this section’s measurements), and is compiled on first use by torch.utils.cpp_extension.load.

Measured

python run.py cpu

On the 4-core Xeon virtual machine used to write this book, with Qwen3-0.6B’s MLP weights for all 28 layers (528 million weights, far beyond the caches, as in real decoding):

{"kind": "cpu", "name": "Intel(R) Xeon(R) Processor @ 2.10GHz", "capability": null, "memory_gib": 15.7, "bf16": true, "triton": false, "graphs": false, "simd": "avx512_vnni"}
{"measured_copy_bandwidth_GBs": 40.6, "threads": 4}
{"kernel": "torch fp32", "weights_MB": 1057, "ms_per_token": 30.5, "effective_GBs": 34.6, "mlp_tokens_per_s": 32.7}
{"kernel": "torch bf16", "weights_MB": 528, "ms_per_token": 28.3, "effective_GBs": 18.7, "mlp_tokens_per_s": 35.3}
{"kernel": "SIMD int8 codes", "weights_MB": 297, "ms_per_token": 27.2, "effective_GBs": 10.9, "mlp_tokens_per_s": 36.7}
{"kernel": "SIMD int4 codes", "weights_MB": 165, "ms_per_token": 12.7, "effective_GBs": 13.0, "mlp_tokens_per_s": 78.8}

PyTorch’s FP32 matrix-vector product (Intel’s oneDNN underneath) reads its weights at 35 GB/s, close to the 41 GB/s this machine can copy, so it’s bandwidth-bound, as decode should be. BF16 halves the bytes but saves little time: this CPU has no native BF16 arithmetic path for this shape. The int4 kernel reads 6.4 times fewer bytes than FP32 and runs 2.4 times faster. It would be 6 times faster at full bandwidth; what stands in the way is overhead per call (56 Python-level calls per token, each quantizing its activations and launching threads) and the arithmetic per weight. That gap is exactly why llama.cpp runs its whole graph in C++ with a persistent thread pool: stretch exercise 1 moves the decode loop into the Rust track to close it.

Apple GPUs

PyTorch runs on Apple GPUs through the mps device, and every model in this book runs there unchanged with the reference attention backend: Triton doesn’t target Metal. Apple’s GPUs share unified memory with the CPU, so a Mac with 64 GB of RAM can hold a 64 GB model on its GPU without any offloading, which is why Macs became popular for local inference despite modest compute.

For quantized weights, the engine needs a Metal kernel. PyTorch (2.6 and later) compiles Metal shading language at run time with torch.mps.compile_shader; here’s the 4-bit GEMV, one SIMD group of 32 threads per output row, with simd_sum adding their partial sums:

"""A Metal kernel for Apple GPUs (Chapter 40): the 4-bit group-affine GEMV of AffineQuantLinear,
compiled at run time with torch.mps.compile_shader (PyTorch 2.6+).

Apple's GPUs share memory with the CPU (unified memory), so a model's weights are read at the
memory system's full bandwidth by either side; decode is bandwidth-bound exactly as on a discrete
GPU, and keeping weights in 4 bits matters just as much.

One SIMD group (32 threads) per output row: each thread takes every 32nd byte of the row, and
simd_sum adds the 32 partial sums. Not validated in the book's test environment (no Apple GPU);
the test is skipped there and runs on any Mac with MPS.
"""
from functools import lru_cache

SOURCE = r"""
#include <metal_stdlib>
using namespace metal;

kernel void affine_gemv_q4(device const uchar* codes   [[buffer(0)]],
                           device const float* scales  [[buffer(1)]],
                           device const float* offsets [[buffer(2)]],
                           device const float* x       [[buffer(3)]],
                           device float* y             [[buffer(4)]],
                           constant uint& K            [[buffer(5)]],
                           constant uint& G            [[buffer(6)]],
                           uint tid  [[thread_position_in_grid]],
                           uint lane [[thread_index_in_simdgroup]]) {
    uint row = tid / 32;
    uint groups = K / G;
    float acc = 0.0f;
    for (uint j = lane; j < K / 2; j += 32) {           // byte j holds codes 2j (low) and 2j + 1 (high)
        uchar b = codes[row * (K / 2) + j];
        uint k = 2 * j, g = k / G;
        float s = scales[row * groups + g], o = offsets[row * groups + g];
        float x0 = x[k], x1 = x[k + 1];
        acc += (float(b & 15) * s + o) * x0 + (float(b >> 4) * s + o) * x1;
    }
    acc = simd_sum(acc);
    if (lane == 0) y[row] = acc;
}
"""


@lru_cache(maxsize=1)
def library():
    import torch
    return torch.mps.compile_shader(SOURCE)


def affine_gemv_q4(x, codes, scales, offsets, group_size):
    """x [K] float32 on "mps" -> y [N]."""
    import torch
    n, k = codes.shape[0], x.shape[0]
    y = torch.empty(n, device=x.device, dtype=torch.float32)
    library().affine_gemv_q4(codes, scales.float(), offsets.float(), x.float(), y, k, group_size,
                             threads=n * 32, group_size=32)
    return y

This kernel is not validated in this book’s test environment, which has no Apple hardware; its test runs on any Mac with MPS and is skipped elsewhere (Appendix F). Apple’s own MLX framework, and llama.cpp’s Metal backend, go further with fused kernels for each quantization type; studying the latter is stretch exercise 3.

What about Vulkan?

llama.cpp’s Vulkan backend runs the same GGUF kernels on any GPU with a Vulkan driver: AMD and Intel consumer cards, older NVIDIA cards, many phones. This engine doesn’t have one, and the reason is architectural rather than a matter of effort: everything here runs through PyTorch, which has no Vulkan device, and through Triton, which has no Vulkan target. A Vulkan backend would be a second kernel set (compute shaders in GLSL or Slang for every operation of Chapters 31-39) plus memory management that PyTorch currently does for you. If Vulkan devices matter to you, llama.cpp’s ggml-vulkan is the reference to read, and the GGUF support of Chapter 38 lets this engine and llama.cpp share model files.

Models bigger than the GPU

A 30-billion-parameter MoE at 4 bits is 17 GB; a consumer GPU has 8-24 GB. Three ways to use host memory, each with a different cost:

Splitting layers

Put the last $n$ layers on the GPU and the rest on the CPU (llama.cpp’s --n-gpu-layers). Each token’s hidden state crosses PCIe twice, a few kilobytes, which costs microseconds. But the CPU layers run at CPU bandwidth: with 40% of the layers on a CPU that’s 10 times slower, the model runs at about a fifth of its full-GPU speed. It’s the right tool when the model almost fits.

class SplitModel(FlatModel):
    """Layers on several devices; the hidden state follows them."""

    def __init__(self, model, layer_devices, head_device=None, backends=None):
        super().__init__(model)
        if len(layer_devices) != len(self.layers):
            raise ValueError("One device per layer")
        self.devices = [torch.device(d) for d in layer_devices]
        self.head_device = torch.device(head_device or layer_devices[-1])
        self.backends = backends or {}
        for layer, device in zip(self.layers, self.devices):
            layer.to(device)
        self.backbone.embed_tokens.to(self.devices[0])
        self.backbone.norm.to(self.head_device)
        self.model.lm_head.to(self.head_device)

    def layer_device(self, i):
        return self.devices[i]

    def forward(self, input_ids, positions, kv_caches, meta, backend, embeds=None, **features):
        """FlatModel.forward, with the hidden state moved to each layer's device.  (Your engine: Chapter 40)"""
        x = self.backbone.embed_tokens(input_ids.to(self.devices[0])) if embeds is None else embeds
        delta, per_device = None, {}
        for i, layer in enumerate(self.layers):
            device = self.devices[i]
            if device not in per_device:                    # positions, RoPE tables, metadata: once per device
                pos = positions.to(device)
                per_device[device] = (pos, self.rope(pos, x.dtype), meta_on(meta, device),
                                      self.backends.get(device.type, backend))
            pos, rope, local_meta, local_backend = per_device[device]
            x = x.to(device, non_blocking=True)
            delta = None if delta is None else delta.to(device, non_blocking=True)
            h, x = self.add_norm(layer.input_layernorm, x, delta)
            delta = self.attention(i, layer.self_attn, h, pos, rope, kv_caches[i], local_meta, local_backend)
            h, x = self.add_norm(layer.post_attention_layernorm, x, delta)
            delta = self.mlp(layer.mlp, h)
        x, delta = x.to(self.head_device), delta.to(self.head_device)
        h, _ = self.add_norm(self.backbone.norm, x, delta)
        return h


def split_by_count(num_layers, gpu_layers, gpu="cuda", cpu="cpu"):
    """llama.cpp's rule: the LAST gpu_layers layers go to the GPU (nearest the output head)."""
    return [cpu] * (num_layers - gpu_layers) + [gpu] * gpu_layers

The model moves the hidden state to each layer’s device, computes the positions, RoPE tables and batch metadata once per device per step, and the runner (Chapter 31) allocates each layer’s KV pool on that layer’s device.

Experts in host memory

A MoE reads only its active experts per token: Qwen3-30B-A3B has 30 B parameters but uses 3 B per token. Keep attention, the router and the shared weights on the GPU, the experts in host memory, and either compute the routed tokens’ experts on the CPU (ktransformers; llama.cpp’s -ot exps=CPU), or copy the experts each step needs into a small GPU cache:

class OffloadedExperts(nn.Module):
    """Stacked expert weights kept in (pinned) host memory.

    mode="cpu":    send the step's routed rows to the CPU, run the experts there, send results back.
    mode="stream": copy the experts this step needs into a GPU cache of `capacity` slots (least
                   recently used evicted), then compute on the GPU.
    """

    def __init__(self, experts, device, mode="stream", capacity=8):
        super().__init__()
        if mode not in ("cpu", "stream"):
            raise ValueError("mode is cpu or stream")
        self.mode, self.device, self.capacity = mode, torch.device(device), capacity
        pin = self.device.type == "cuda"
        host = lambda t: t.detach().cpu().pin_memory() if pin else t.detach().cpu().clone()     # noqa: E731
        self.host_gate_up, self.host_down = host(experts.gate_up_proj), host(experts.down_proj)
        self.slots_gate_up = torch.empty((capacity, *self.host_gate_up.shape[1:]), dtype=self.host_gate_up.dtype,
                                         device=self.device)
        self.slots_down = torch.empty((capacity, *self.host_down.shape[1:]), dtype=self.host_down.dtype, device=self.device)
        self.slot_of, self.lru = {}, []
        self.stats = {"hits": 0, "misses": 0, "bytes_copied": 0}

    @property
    def num_experts(self):
        return self.host_gate_up.shape[0]

    def fetch(self, e):
        """Make expert e resident in a GPU slot (LRU eviction); returns the slot.  (Your engine: Chapter 40)"""
        if e in self.slot_of:
            self.stats["hits"] += 1
            self.lru.remove(e)
            self.lru.append(e)
            return self.slot_of[e]
        self.stats["misses"] += 1
        if len(self.slot_of) < self.capacity:
            slot = len(self.slot_of)
        else:
            victim = self.lru.pop(0)
            slot = self.slot_of.pop(victim)
        self.slots_gate_up[slot].copy_(self.host_gate_up[e], non_blocking=True)
        self.slots_down[slot].copy_(self.host_down[e], non_blocking=True)
        self.stats["bytes_copied"] += self.host_gate_up[e].nbytes + self.host_down[e].nbytes
        self.slot_of[e] = slot
        self.lru.append(e)
        return slot

    def expert_weights(self, e):
        if self.mode == "cpu":
            return self.host_gate_up[e], self.host_down[e]
        slot = self.fetch(e)
        return self.slots_gate_up[slot], self.slots_down[slot]

    def forward_grouped(self, x, weights, experts):
        device = x.device
        if self.mode == "cpu":
            x, weights, experts = x.cpu(), weights.cpu(), experts.cpu()
        out = torch.zeros_like(x)
        # Experts run one after another, each fetched just before it's used, so a step may route to
        # more experts than the cache holds: a slot is reused once its last expert has run (on one
        # CUDA stream, the copy into it is ordered after that computation).
        for e in experts.unique().tolist():
            token, slot = torch.where(experts == e)
            gate_up, down = self.expert_weights(e)
            gate, up = nn.functional.linear(x[token], gate_up.to(x.dtype)).chunk(2, dim=-1)
            y = nn.functional.linear(nn.functional.silu(gate) * up, down.to(x.dtype))
            out.index_add_(0, token, y * weights[token, slot, None])
        return out.to(device, non_blocking=True)

    forward_loop = forward_grouped


def offload_experts(model, device, mode="stream", capacity=8):
    """Replace every MoE layer's experts; returns the new modules (for their stats)."""
    from .moe import SparseMoeBlock
    replaced = []
    for module in model.modules():
        if isinstance(module, SparseMoeBlock):
            module.experts = OffloadedExperts(module.experts, device, mode, capacity)
            replaced.append(module.experts)
    return replaced

Whether copying pays depends on reuse:

python run.py offload --new-tokens 24

{"requests": 1, "gpu_slots_per_layer": "4 of 32", "hit_rate": 0.076, "MB_copied_per_token": 3.167, "all_experts_MB": 50.3}
{"requests": 1, "gpu_slots_per_layer": "8 of 32", "hit_rate": 0.174, "MB_copied_per_token": 2.83, "all_experts_MB": 50.3}
{"requests": 1, "gpu_slots_per_layer": "16 of 32", "hit_rate": 0.377, "MB_copied_per_token": 2.135, "all_experts_MB": 50.3}
{"requests": 1, "gpu_slots_per_layer": "32 of 32", "hit_rate": 0.74, "MB_copied_per_token": 0.892, "all_experts_MB": 50.3}
{"requests": 4, "gpu_slots_per_layer": "4 of 32", "hit_rate": 0.0, "MB_copied_per_token": 2.38, "all_experts_MB": 50.3}
{"requests": 4, "gpu_slots_per_layer": "8 of 32", "hit_rate": 0.0, "MB_copied_per_token": 2.38, "all_experts_MB": 50.3}
{"requests": 4, "gpu_slots_per_layer": "16 of 32", "hit_rate": 0.363, "MB_copied_per_token": 1.517, "all_experts_MB": 50.3}
{"requests": 4, "gpu_slots_per_layer": "32 of 32", "hit_rate": 0.906, "MB_copied_per_token": 0.225, "all_experts_MB": 50.3}

The test model’s router is random, so a single request’s hit rate is about the fraction of experts cached: there’s no locality to exploit. With four requests in a batch, each step routes to most experts, and a small LRU cache thrashes: every expert is evicted just before it’s needed again, and the hit rate drops to zero. Real MoEs route with more locality than a random router, but the lesson holds: streaming experts suits single-user decoding with a cache that holds a good fraction of the experts; for batched serving, computing experts where they live (on the CPU, or on other GPUs with expert parallelism, Chapter 41) is better.

Streaming layers

When nothing else fits, keep every layer in pinned host memory and stream them through two GPU slots, copying layer $i + 1$ while layer $i$ computes:

class StreamedModel(FlatModel):
    """All layers in host memory; two GPU slots, the next layer copied while this one computes."""

    def __init__(self, model, device):
        super().__init__(model)
        self.device = torch.device(device)
        pin = self.device.type == "cuda"
        self.host_layers = self.layers
        for layer in self.host_layers:
            layer.to("cpu")
            if pin:
                for p in layer.parameters():
                    p.data = p.data.pin_memory()
        self.slots = [copy.deepcopy(self.host_layers[0]).to(self.device) for _ in range(2)]
        self.backbone.embed_tokens.to(self.device)
        self.backbone.norm.to(self.device)
        self.model.lm_head.to(self.device)
        self.copy_stream = torch.cuda.Stream(self.device) if self.device.type == "cuda" else None
        self.ready = [None, None]
        self.bytes_copied = 0

    def layer_device(self, i):
        return self.device

    def load(self, i):
        """Copy layer i's weights into slot i % 2, on the copy stream when there is one."""
        slot = self.slots[i % 2]
        if self.copy_stream is not None:
            self.copy_stream.wait_stream(torch.cuda.current_stream(self.device))   # the slot's last reader is done
            with torch.cuda.stream(self.copy_stream):
                for dst, src in zip(slot.parameters(), self.host_layers[i].parameters()):
                    dst.data.copy_(src.data, non_blocking=True)
                self.ready[i % 2] = torch.cuda.Event()
                self.ready[i % 2].record(self.copy_stream)
        else:
            for dst, src in zip(slot.parameters(), self.host_layers[i].parameters()):
                dst.data.copy_(src.data)
        self.bytes_copied += sum(p.nbytes for p in self.host_layers[i].parameters())

    def forward(self, input_ids, positions, kv_caches, meta, backend, embeds=None, **features):
        """Run layer i from slot i % 2 while layer i + 1 loads into the other.  (Your engine: Chapter 40)"""
        x = self.backbone.embed_tokens(input_ids) if embeds is None else embeds
        rope = self.rope(positions, x.dtype)
        delta = None
        self.load(0)
        for i in range(len(self.host_layers)):
            if self.ready[i % 2] is not None:
                torch.cuda.current_stream(self.device).wait_event(self.ready[i % 2])
            if i + 1 < len(self.host_layers):
                self.load(i + 1)                          # overlaps with this layer's compute
            layer = self.slots[i % 2]
            h, x = self.add_norm(layer.input_layernorm, x, delta)
            delta = self.attention(i, layer.self_attn, h, positions, rope, kv_caches[i], meta, backend)
            h, x = self.add_norm(layer.post_attention_layernorm, x, delta)
            delta = self.mlp(layer.mlp, h)
        h, _ = self.add_norm(self.backbone.norm, x, delta)
        return h

Each step moves the whole model over PCIe: a 70 B model at 4 bits is 35 GB, a bit over a second per step at 25-30 GB/s. That’s hopeless for interactive decode and fine for batch work: a 4,096-token prefill, or 256 sequences decoded together, amortize the same 35 GB over thousands of tokens. FlexGen (Sheng et al., 2023) built a whole throughput-oriented engine on this observation.

Build it

pytest tests/test_ch40_everywhere.py
python run.py cpu --impl engine
python run.py offload --impl engine

Stretch exercises

1. ★★★ Move decode into the Rust track: give rust/src/qwen3.rs a path that keeps its linear layers as 4-bit codes and uses simd::gemv_q4 with a persistent thread pool. Compare tokens per second with run.py cpu’s Python-driven kernel, and with llama.cpp on the same GGUF file. Where: rust/src/qwen3.rs and rust/src/simd.rs.
2. ★★ Prefill on a CPU is compute-bound: write a tiled int8 GEMM (4 × 4 output tiles, VNNI) for batches of 64 tokens, or use Intel AMX tiles on Sapphire Rapids. Measure prompt tokens per second. Where: CPU SIMD kernels in cpp/izh_cpu.hpp or rust/src/simd.rs; expose the C++ path through engine/kernels/cpu_ext.cpp and engine/kernels/cpu.py.
3. ★★ Read llama.cpp’s Metal kernel for Q4_K (ggml-metal.metal) and port its structure to kernels/metal.py. What does it do that the simple SIMD-group kernel doesn’t? Where: the Metal shader and launcher in engine/kernels/metal.py.
4. ★★ Prefetch experts: run the next layer’s router on the current layer’s output (an approximation some MoEs make exact with a pre-gating design) and start copying its likely experts early, on a second CUDA stream. Measure the stall time saved. Where: OffloadedExperts.fetch and the cache lifecycle in engine/offload.py, with scheduling hooks in the model’s layer loop.

Check your understanding

1. Why does a CPU decode at a rate set by memory bandwidth, and what does that imply about quantization on CPUs?
2. Why is dequantizing the whole weight matrix before a matmul slower than not quantizing at all?
3. What does quantizing the activations to int8 buy, and what does it cost?
4. Why are the SIMD functions compiled with per-function target attributes instead of -march=native?
5. Why does an LRU expert cache with 8 of 32 slots thrash when four requests share a step?
6. When does streaming every layer through the GPU make sense, and when doesn’t it?

Going deeper

• llama.cpp’s ggml/src/ggml-cpu/ (the vec_dot_q4_0_q8_0 family for AVX2, AVX-512, NEON and others) and ggml-metal/, ggml-vulkan/ for its other backends.
• Intel’s Intel® 64 and IA-32 Architectures Optimization Reference Manual (VNNI, AMX) and Arm’s Neon Programmer’s Guide (sdot); Agner Fog’s Optimizing software in C++ for the latency and throughput reasoning of this chapter.
• Sheng et al., FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (ICML 2023); Chen et al., KTransformers (SOSP 2025), for CPU/GPU hybrid MoE inference.
• AMD’s ROCm documentation for PyTorch and Triton, and vLLM’s ROCm notes (AITER, FP8 formats on MI300 and MI350); Apple’s Metal Shading Language Specification and the MLX documentation.