consciousness/research/sparse-kernel-napkin-math.md

# Sparse Kernel Compilation: Napkin Math for Qwen3.5-27B

## Architecture recap

| Parameter | Value |
|-----------|-------|
| Layers | 64 (48 linear attention, 16 full attention) |
| Hidden dim (H) | 5,120 |
| Query heads | 24 |
| KV heads | 4 (GQA) |
| Head dim | 256 |
| FFN intermediate | 17,408 |
| Full attention interval | every 4th layer |
| Total params | ~27.5B |

**Key discovery**: Qwen3.5 already uses sparse attention — 48/64 layers use
linear attention (O(N)), only 16 use full O(N^2) attention. The attention
sparsity question is partially answered by the architecture itself. The
remaining opportunity is **weight sparsity** in the projection and FFN matrices.

## Measured weight sparsity (from B200, all 11 shards)

| Component | Count | Total params | <0.001 | <0.005 | <0.01 | <0.02 |
|-----------|-------|-------------|--------|--------|-------|-------|
| down_proj | 65 | 5,793,382,400 | 7.4% | 35.8% | 64.3% | 92.8% |
| gate_proj | 65 | 5,793,382,400 | 7.2% | 35.0% | 63.3% | 92.3% |
| up_proj | 65 | 5,793,382,400 | 7.3% | 35.2% | 63.5% | 92.5% |
| q_proj | 17 | 1,069,547,520 | 5.7% | 27.9% | 51.9% | 82.8% |
| k_proj | 17 | 89,128,960 | 6.0% | 29.4% | 54.1% | 84.1% |
| v_proj | 17 | 89,128,960 | 4.8% | 23.7% | 44.7% | 74.2% |
| o_proj | 17 | 534,773,760 | 5.3% | 25.9% | 48.7% | 79.6% |
| other | 605 | 6,070,652,272 | 5.4% | 26.4% | 49.6% | 80.9% |

**Key findings**:
- FFN layers (gate/up/down) are remarkably sparse: ~64% of weights below 0.01
- At threshold 0.02, FFN sparsity exceeds 92%
- Attention projections are less sparse but still significant: 45-52% below 0.01
- v_proj is the least sparse component (44.7% below 0.01)

## Per-layer parameter breakdown

- **QKV projection**: H × (Q_dim + K_dim + V_dim) ≈ 5120 × 13312 ≈ 68M
- **Output projection**: ~26M
- **FFN (gate + up + down)**: 5120 × 17408 × 3 ≈ 267M
- **Total per layer**: ~361M
- **64 layers**: ~23.1B (rest is embeddings, norms, etc.)

## Dense baseline: FLOPs per token

Each weight parameter contributes 2 FLOPs per token (multiply + accumulate).

- Per layer: ~722M FLOPs/token
- 64 layers: **~46.2B FLOPs/token**

On a B200 (theoretical ~4.5 PFLOPS FP8, ~1.1 PFLOPS BF16):
- BF16 throughput: 46.2B / 1.1e15 ≈ 0.042ms per token (compute-bound limit)
- But inference is usually **memory-bandwidth-bound** for small batch sizes

B200 HBM bandwidth: ~8 TB/s. At BF16 (2 bytes/param):
- Loading all weights once: 27.5B × 2 = 55GB → 55/8000 ≈ **6.9ms per token** (batch=1)
- This is the real bottleneck. Compute is cheap; loading weights is expensive.

## What sparsity buys you

At **X% sparsity** (X% of weights are zero), you need to load (1-X)% of the weights:

| Sparsity | Params loaded | Time (batch=1) | Speedup |
|----------|--------------|-----------------|---------|
| 0% (dense) | 27.5B | 6.9ms | 1.0x |
| 50% | 13.8B | 3.4ms | 2.0x |
| 75% | 6.9B | 1.7ms | 4.0x |
| 90% | 2.8B | 0.7ms | 10x |

**This is the key insight**: inference at small batch sizes is memory-bandwidth-bound.
Sparse weights = fewer bytes to load = directly proportional speedup,
**IF** the sparse kernel can avoid the gather/scatter overhead.

## The compilation problem

Dense GEMM is fast because:
1. Weights are contiguous in memory → sequential reads
2. Tiling fits perfectly in SRAM → high reuse
3. Hardware tensor cores expect dense blocks

Naive sparse matmul kills this:
- Irregular memory access → low bandwidth utilization
- Poor SRAM tiling → cache thrashing
- Tensor cores can't help

### The FlashAttention analogy

FlashAttention's insight: the N×N attention matrix doesn't fit in SRAM,
but you can tile it so each tile does. You recompute instead of materializing.

**Sparse kernel compilation insight**: the sparsity pattern is **known at compile time**.
A compiler can:

1. **Analyze the sparsity graph** of each weight matrix
2. **Find blocks** of non-zero weights that are close in memory
3. **Generate a tiling schedule** that loads these blocks into SRAM efficiently
4. **Emit fused kernels** where the memory access pattern is baked in as constants

The resulting kernel looks like a dense kernel to the hardware —
sequential reads, high SRAM reuse, maybe even tensor core compatible
(if the compiler finds dense sub-blocks within the sparse matrix).

## Block sparsity vs unstructured sparsity

**Block sparse** (e.g., 4×4 or 16×16 blocks zeroed out):
- GPU-friendly: blocks map to tensor core operations
- Less flexible: coarser pruning granularity → less achievable sparsity
- NVIDIA's 2:4 structured sparsity gets ~50% sparsity with tensor core support
- Real-world: typically 50-70% sparsity achievable without quality loss

**Unstructured sparse** (individual weights zeroed):
- Maximally flexible: fine-grained pruning → higher achievable sparsity
- GPU-hostile: the gather/scatter problem
- Real-world: 80-95% sparsity achievable in many layers without quality loss

**The compiled kernel approach bridges this**: take unstructured sparsity
(maximally flexible, high compression) and compile it into a kernel that
runs as efficiently as block-sparse. Best of both worlds.

## Recurrent depth composability

From our April 10 discussion: middle transformer layers are doing
open-coded simulated annealing — similar weights, similar computation.

If layers 8-24 have cosine similarity > 0.95:
- Replace 16 layers with 1 layer × 16 iterations
- **Parameter reduction**: 16 × 361M = 5.8B → 361M (16x reduction for those layers)
- **Memory bandwidth**: load one layer's weights, iterate in SRAM
- Combined with 50% sparsity on the remaining unique layers:
  - Unique layers (48): 48 × 361M × 0.5 = 8.7B params
  - Recurrent layer: 361M × 0.5 = 180M params (but iterated 16x in SRAM)
  - Total loaded per token: ~8.9B × 2 bytes = 17.8GB
  - Time: 17.8/8000 ≈ **2.2ms per token** (vs 6.9ms dense) — **3.1x speedup**

With higher sparsity (75%) + recurrence:
  - Unique layers: 48 × 361M × 0.25 = 4.3B
  - Recurrent: 90M
  - Total: ~4.4B × 2 = 8.8GB → **1.1ms per token** — **6.3x speedup**

## What needs to happen

### Phase 1: Measure (can do now with B200 access)
1. Extract all weight matrices from Qwen3.5-27B
2. For each matrix, compute:
   - Magnitude distribution (what % of weights are near-zero?)
   - Achievable sparsity at various thresholds (L1 magnitude pruning)
   - Dense sub-block statistics (how many 4×4, 16×16 blocks are all-zero?)
3. Layer similarity: pairwise cosine similarity of weight matrices across layers
   - Which layers are nearly identical? (recurrence candidates)
4. Validate quality: run perplexity eval at various sparsity levels

### Phase 2: Compile (research project)
1. For a single sparse weight matrix, generate an optimized Triton kernel
2. Benchmark vs dense GEMM and vs NVIDIA's 2:4 sparse
3. Iterate on the tiling strategy

### Phase 3: End-to-end
1. Full model with compiled sparse kernels
2. Perplexity + latency benchmarks
3. Compare: dense, 2:4 structured, compiled unstructured

## Related work to read

- **SparseGPT** (Frantar & Alistarh, 2023): one-shot pruning to 50-60% unstructured sparsity
  with minimal quality loss. Key result: large models are more prunable than small ones.
- **Wanda** (Sun et al., 2023): pruning by weight magnitude × input activation.
  Simpler than SparseGPT, comparable results.
- **NVIDIA 2:4 sparsity**: hardware-supported structured sparsity on Ampere+.
  50% sparsity, ~2x speedup on tensor cores. The existence proof that sparse can be fast.
- **Triton** (Tillet et al.): Python DSL for GPU kernel generation.
  The right compilation target — can express arbitrary tiling strategies.
- **TACO** (Kjolstad et al.): tensor algebra compiler. Generates kernels for
  specific sparse tensor formats. Academic but the ideas are right.
- **FlashAttention** (Dao et al.): the tiling strategy to learn from.
- **DejaVu** (Liu et al., 2023): contextual sparsity — predicting which neurons
  to activate per input. Dynamic sparsity, complementary to weight sparsity.

## The bigger picture

Current state of the art: dense models with FlashAttention for the N×N attention part.
Weight sparsity is known to work (SparseGPT, Wanda) but isn't deployed because
the GPU kernels don't exist to exploit it efficiently.

The gap: nobody has built a compiler that takes a specific sparse weight matrix
and emits a kernel optimized for that exact pattern. FlashAttention proved
that custom kernels for specific computational patterns beat general-purpose ones.
The same should hold for sparse weight patterns.

**The bet**: a compiled sparse kernel for Qwen3.5-27B's actual sparsity pattern
would be within 80% of the theoretical bandwidth-bound speedup. If true,
50% sparsity → 1.6x real speedup, 75% → 3.2x, composing with recurrent depth
for potentially 5-6x total.

That would make 27B inference as fast as a 5B dense model, with 27B quality.