01 Quantization Memory Bandwidth - Findings

Benchmarked qwen2.5-coder-7B at Q4_K_M vs Q8_0 on an Apple M4 Pro (273 GB/s peak memory bandwidth).

Result: the same quantization change has opposite effects on the two inference phases.

Phase	Q4	Q8	Ratio	Bottleneck
Decode	~47 tok/s	~31.5 tok/s	1.49× slower	memory-bound
Prefill (long prompt)	~394 tok/s	~407 tok/s	0.97 (≈equal)	compute-bound

Why:

• Decode is memory-bound - each token reads all weights, so Q8's ~1.7× more bytes → slower. It slows only 1.49× (not 1.72×) because decode also reads the KV cache, which is quant-independent, diluting the weight scaling.
• Prefill is compute-bound - many tokens processed in parallel; the bottleneck is FLOPs, which are identical for Q4/Q8, so quantization barely changes it.

02 - KV Cache

Implemented single-layer causal self-attention (PyTorch), then built the decode loop two ways - naive (recompute all K/V every step) vs KV-cached (reuse past K/V) - and measured the O(N²) → O(N) speedup.

Result: the KV cache turns O(N²) total decode into O(N)

	per-step latency	total (cumulative)
Naive (stateless, recompute all)	O(N) - rises with position	O(N²) - parabola
KV-cached (stateful, reuse cache)	O(1) - flat	O(N) - straight line

Why:

• In cached decode, q/k/v are always [1, d_model] - only the new token is projected each step (constant O(1) work); past K/V are read from the cache, never recomputed. Naive re-projects all t tokens every step (O(t)) - pure waste, since causal masking freezes past K/V. Constant vs growing per-step work → the O(N) vs O(N²) total.

03 - Tiled Attention (FlashAttention core)

Implemented standard attention and a tiled / online-softmax version, and proved they're identical.

Result: tiled output matches standard softmax(QKᵀ/√d)·V to ~1e-7 (float32 rounding) for every block size (8/16/32/64) → FlashAttention is exact (not an approximation) and block-size-independent (tile to fit SRAM, same answer).

Mechanism: per query, maintain a running max, sum, and output; for each K/V block compute safe exp-weights exp(scores - running_max), and when the max grows, rescale the old accumulators by exp(old_max - new_max). Streams softmax without ever materializing the [L,L] matrix (O(L²)→O(L) memory).

04 - Speculative Decoding (draft + verify)

greedy speculative decoding - a small draft model proposes K tokens, the large target verifies them in a single forward pass, and the longest correct prefix is accepted.

Provably lossless: the output is bit-identical to plain target-greedy decoding.

Results (200 new tokens, MPS / M-series, fp16)

Variant	Time	Target passes	Accept rate
No KV cache	31.1s	97	26.5%
+ KV cache	13.1s	105	22.9%

2.37× speedup - and the cached run did more passes at a lower accept rate. The win is not from accepting more tokens; it's from making each verify pass O(K) new tokens (the cache holds the prefix K/V) instead of O(L) re-encoding the whole sequence. Acceptance rate and wall-clock are decoupled.

max_tokens=500 OOM'd macOS unified memory without the cache (O(L²) per pass); with the cache it runs in 47.7s.

Acceptance is entropy-dependent

Prompt	Accept	Passes
"write a poem on how cats rule the world" (high entropy)	22.9%	105
"Write a C++ code to merge two arrays in a sorted fashion" (low)	66.8%	55

05 - Fused & Tiled Softmax Triton Kernels

GPU kernels in Triton (T4, 320 GB/s peak)

• fused_softmax_kernel - one worker per row, whole row in one block, standard 2-pass (max-subtract → exp → sum → divide). Fast, but needs BLOCK_SIZE ≥ n_cols so it can't handle rows wider than a block.
• online_softmax_kernel - tiled/online softmax (FlashAttention's softmax core). One worker per row streams the row in BLOCK_SIZE tiles, carrying running max m and sum l with l = l·exp(m_old−m_new) + Σexp(tile−m_new) rescaling. Constant SRAM footprint → any row width.

Correctness

• Fused vs torch.softmax: max abs diff 2.98e-08.

• Tiled vs torch.softmax (4096×8192): match=True for every block size, lossless and block-size-independent:

  BLOCK_SIZE	max abs err
  64	3.73e-09
  128	3.73e-09
  256	3.73e-09
  512	1.86e-09

  (Error drops at 512 - fewer, wider tiles = fewer rescale steps = less float noise.)

Fused softmax benchmark (4096 rows)

n_cols	naive	torch	triton (fused)	GB/s	speedup vs naive
512	0.349 ms	0.078 ms	0.076 ms	222	4.61×
1024	0.581 ms	0.156 ms	0.146 ms	229	3.97×
2048	1.128 ms	0.314 ms	0.291 ms	230	3.87×

Fused kernel matches/slightly beats torch's native CUDA softmax and runs ~4× faster than the unfused naive baseline (5 HBM round-trips → 1). At ~72% of the T4's 320 GB/s peak it's memory-bound and near the roofline - softmax has ~zero arithmetic intensity, so bandwidth is the only limit.

Tiled vs fused: the crossover

Tiled reads the row twice (reduce pass + normalize pass) vs the fused 1 read + 1 write, so at moderate width the fused strategy wins. But the fused strategy needs the whole row resident - as rows widen it loses occupancy while the constant-footprint tiled kernel stays flat.

Sweep with total elements held ~constant (~256M):

n_cols	torch ms	tiled ms	ratio (tiled/torch)	torch GB/s (1R+1W)
8192	11.652	12.815	1.10	184
16384	12.645	13.085	1.03	170
32768	16.809	13.316	0.79	128
65536	17.009	13.740	0.81	126
131072	17.114	14.319	0.84	125

Ratio crosses 1.0 around 16k–32k columns. Below it, fused wins (at 8192, tiled is 1.72× slower - which factors exactly as 3/2 traffic × 249/216 bandwidth-efficiency). Above it, tiled is ~20% faster: torch's effective bandwidth collapses 184 → 125 GB/s as its resident-row design loses occupancy (fewer warps to hide HBM latency), while the tiled kernel - fixed BLOCK_SIZE, constant SRAM — holds steady.