AKP Anthropic Kernel Path

0/21 read · no cycles yet

Month 4 — FA2 / decode kernels, toy pager, Megatron TP costed, weekly KV numbers

03-kernels/curriculum.md Layer 3 (FA2, paged decode) + Layer 5 start (TP volumes). 04-llm-theory/curriculum.md §D parallelism cheat sheet + §F. Roadmap P1 exit / P2 start.

Month goals

  • FA2-style changes in your kernel: parallelize over Q, better occupancy. Still Triton is fine. Backward plan written; implement backward if the forward is solid — do not fake a backward you did not test.
  • Decode kernel: seqlen_q = 1 (or small), KV long. Separate from prefill. FlashDecoding idea.
  • Toy pager: block table, allocate/free, prefix reuse, fragmentation stat. Not vLLM. A pager.
  • Megatron TP drawn and costed: column-parallel + row-parallel, 2× AllReduce (or RS+AG with SP), bytes on the wire for a named MLP.
  • KV cache numbers every week (same config, growing intuition): 70B-class GQA, vary seq and batch and dtype.

Take-home: hold < 2,000 and hunt the next 1.2×. Do not stall the kernels to chase 1487 this month.

Week 13 — FA2 algorithm into the kernel

DayDeep 1Deep 2Paper / bookFri / Sat
MonFA2 paper: the three changes vs FA1. Patch your Triton loop.Work/span of the new parallelization.Dao 2023 FA2.
TueCausal load-imbalance. Measure.Decode-shaped kernel v0 (q=1).FlashDecoding.
WedDebug day.Take-home.None.
ThuKV numbers (week 13): 80L, 8 kv heads, d=128, S=8k, B=16, fp16. Then S=128k. Then fp8. Write GB.Interview whiteboard 25 min: same calculation cold.None.
FriTake-home sport.
SatLong build: FA2 vs FA1 speed table on your GPU. Honest.
SunOff-ish.

Week 14 — Toy pager

DayDeep 1Deep 2Paper / bookFri / Sat
MonKwon et al. 2023 PagedAttention. Draw logical→physical block table.Python pager: pages of tokens, alloc/free, refcount.vLLM paper.
TueDecode attention that gathers via block table (03-kernels L3 ex. 8 energy, even if CPU/numpy first).Copy-on-write for a beam of 2.None.
WedDebug pager fragmentation under a synthetic trace.Take-home.None.
ThuKV numbers (week 14): same model, paged with block=16 vs 128. Internal fragmentation.System-design 45 min: "one GPU, variable-length decode, design the batcher." Record.Orca abstract only (full read month 6).
FriTake-home.
SatLong build: pager + decode demo in the public repo.
SunOff-ish.

Week 14 exit: a pager you can explain, not a vLLM clone.

Week 15 — Megatron TP drawn and costed

DayDeep 1Deep 2Paper / bookFri / Sat
MonShoeybi 2019 + Narayanan 2021 enough to draw TP for one Llama-style MLP.Bytes: AllReduce volume vs RS+AG.Megatron 2019.
TueImplement ColumnParallel + RowParallel in PyTorch (1–2 GPU if you have them; else simulate ranks).Match single-GPU Linear numerically.Korthikanti 2022 SP (why SP rides along).
WedDebug comm vs math.Take-home.None.
ThuKV numbers (week 15): TP=2, what shards, what is not sharded (KV often replicated per TP rank — write it).Cost a 70B, seq 8k, global batch, pick TP/PP/DP on paper (04-llm-theory §D whiteboard).ZeRO paper enough to say why ZeRO-3 ≠ TP.
FriTake-home.
SatLong build: one-page "TP for this block" diagram + comm table.
SunOff-ish.

Week 16 — Close month 4

DayDeep 1Deep 2Paper / bookFri / Sat
MonFA2 backward plan or tiny backward on S=64. Honesty rule: untested backward = not done.Pager polish.None.
TueDecode vs prefill roofline on one plot.Interview drill 90 min.None.
WedDebug / measure.Take-home.None.
ThuKV numbers (week 16): pick a new config (MLA from DeepSeek-V2 §2.1 bytes/token vs GQA-8).Blank-page TP diagram.DeepSeek-V2 MLA section only.
FriTake-home. Hold < 2,000.
SatMonth 4 artifact.
SunOff.

Month 4 exit criteria

Artifact

pager/ + tp-notes.md + FA2 table + notes/kv/week-13..16.md + 1000-word writeup. This is the first month that looks like the Inference Systems req.