phi-triton

Custom Triton kernel for single-token autoregressive decode, integrated end-to-end into Phi-3 Mini inference.

All benchmarks: RTX 4060 · 6 GB GDDR6 · fp16 · batch=1

---

Why this exists

During text generation, LLMs spend the majority of time in decode, not prefill.

Decode has fundamentally different characteristics to prefill:

• One token per forward pass — batch size is effectively 1
• Attention is memory-bound, not compute-bound
• KV-cache access patterns dominate; FLOP count is almost irrelevant
• Python dispatch and kernel launch overhead become visible

This project explores those properties concretely: what a Triton decode kernel looks like, how to integrate it correctly (including RoPE and KV-cache write order), and what the honest benchmark numbers say about where time actually goes.

This matters outside research. Automotive AI systems — ADAS pipelines, in-vehicle assistants, real-time co-pilot inference — run under hard latency budgets on embedded GPUs with constrained HBM bandwidth (Orin NX: 102 GB/s, Xavier: 137 GB/s). The same memory-bound bottlenecks that dominate on a 6 GB desktop GPU appear at every deployment tier. Decode-time systems understanding transfers directly.

---

Hardware context

GPU	RTX 4060 (Ada Lovelace)
VRAM	6 GB GDDR6
Memory bandwidth	288 GB/s
SMs	24
L2 cache	2 MB
Model	Phi-3 Mini (fp16, ~3.8 GB)
Usable headroom	~2 GB for KV cache + activations

The 6 GB constraint is real and shapes every decision:

• KV cache preallocated and capped at max_seq_len=512 to avoid OOM
• BLOCK_T=64 — larger tiles spill registers on 24-SM / 2 MB L2 config
• Single sequence only — no batching at this VRAM budget

---

What is implemented

Triton decode attention kernel (kernels/)

• Single-token GQA decode with online softmax — one pass, O(1) extra memory
• FP32 accumulation for numerical stability; result stored in input dtype
• Hq, D_PAD, BLOCK_T all explicit tl.constexpr — no shape inference from strides
• D_PAD = triton.next_power_of_2(head_dim) — correct for Phi-3's head_dim=96 and any other model (not hardcoded to 128)
• Output dtype matches input — no hidden float32 upcast downstream

RoPE-correct decode path (model/phi3_decode_triton.py)

• Rotary Position Embeddings applied to Q and K before cache writes
• position_ids = kv_cache.cur_pos — correct position for every decode step
• KV write order: append() then get() — eliminates read on uninitialised cache memory (was a real bug in the original code)
• Triton handles attention across all 32 layers; LayerNorm, MLP, and residuals remain in PyTorch

Preallocated KV cache (kv_cache/kv_layout.py)

• Allocated once at prefill — no torch.cat during decode
• In-place positional writes: O(1) per token
• Capacity: 32 layers × 2 × 32 heads × 512 × 96 × fp16 ≈ 192 MB

---

Setup

# 1. Edit config.py — set MODEL_PATH to your Phi-3 Mini checkpoint directory
#    Windows: r"E:\models\phi3_mini"
#    WSL:     "/mnt/e/models/phi3_mini"

pip install -r requirements.txt

---

Run order

# 1. Confirm model loads and shapes are correct
python model/inspect_attention.py

# 2. Validate Triton kernel against PyTorch reference (expect max diff < 0.01)
python model/validate_decode_attention_triton.py

# 3. End-to-end text generation — confirms RoPE is applied correctly
python run_baseline.py

# 4. Microbenchmark — Triton vs PyTorch reference, 1 layer, 1 token
python benchmarks/compare_decode_latency.py

# 5. HuggingFace baseline — full model, 256 tokens
python benchmarks/baseline_latency.py

# 6. Triton path — full model, 256 tokens
python benchmarks/generate_latency.py

# 7. Profiler trace — see which ops dominate per-token latency
python profiling/decode_profile.py

---

Benchmark results (RTX 4060 · fp16 · batch=1)

Kernel validation

Metric	Value
Max |ref − triton|	0.000031
Mean |ref − triton|	0.000001
Result	✓ PASS

End-to-end decode — 32 layers · 256 tokens

Path	ms/token
HuggingFace baseline	775 ms
Triton (this project)	307 ms
Speedup	2.5×

Where the 2.5× comes from

The profiler answers this precisely:

Op	CUDA time share	What it means
aten::mm	93.4%	MLP matmuls — same in both paths
aten::cat	5.1%	HF allocating KV cache every token
aten::bmm	0.65%	Actual attention math

Attention accounts for 0.65% of total CUDA time.

The 2.5× speedup does not come from the Triton kernel arithmetic. It comes from two systems-level changes:

1. Preallocated KV cache — eliminates torch.cat across 32 layers × 256 tokens. The HF path reallocates and copies the entire KV cache every single decode step.
2. Tighter Python decode loop — no HF DynamicCache overhead, no boilerplate per-token dispatch cost.

This is the honest result. The Triton kernel is correct and fast, but it is optimising 0.65% of runtime. The 2.5× gain comes from fixing the memory allocation pattern around it.

What this means

Production inference engines (vLLM, TRT-LLM, SGLang) achieve their speedups through the same insight applied more aggressively: paged KV allocation, CUDA graph capture of the full forward pass, and fused kernels for every layer — not just attention. This project demonstrates why kernel-level work alone has diminishing returns, and where the real leverage is.

---

Key lessons

• Kernel arithmetic is not the bottleneck — at decode time, 93% of CUDA time is MLP matmuls, not attention
• KV-cache allocation pattern matters more than attention math — eliminating torch.cat across 32 layers drives a 2.5× end-to-end speedup
• RoPE is not optional — skipping it silently produces incoherent output
• Decode is a systems problem: memory layout and allocation patterns matter more than kernel FLOP efficiency
• These conclusions match vLLM, TRT-LLM, and NVIDIA's inference engineering guidance

---

Project structure

config.py                              ← set MODEL_PATH here

kernels/
  decode_attention_triton.py           ← Triton JIT kernel
  decode_attention_wrapper.py          ← launch wrapper, RTX 4060 tuning

model/
  load_model.py                        ← model loading + DynamicCache compat patch
  phi3_prefill.py                      ← prefill + KVCache init
  phi3_decode_triton.py                ← full decode with Triton + RoPE
  phi3_decode.py                       ← HF baseline decode
  decode_attention_ref.py              ← PyTorch reference attention
  extract_qkv.py                       ← QKV projection helper
  inspect_attention.py                 ← [step 1] model shape check
  validate_decode_attention_triton.py  ← [step 2] kernel validation

kv_cache/
  kv_layout.py                         ← preallocated KV cache

benchmarks/
  compare_decode_latency.py            ← [step 4] 1-layer microbenchmark
  baseline_latency.py                  ← [step 5] HF end-to-end
  generate_latency.py                  ← [step 6] Triton end-to-end

profiling/
  decode_profile.py                    ← [step 7] torch.profiler trace

run_baseline.py                        ← [step 3] text generation sanity check
requirements.txt
README.md

---

What is not included (by design)

• CUDA graph capture
• Paged / chunked KV cache
• Fused MLP kernels
• C++ runtime

These belong at the runtime engine level. This project covers kernel development and inference loop design — the right scope for understanding where optimisation effort pays off.

---

Requirements

• Python 3.10+
• PyTorch ≥ 2.1 (CUDA build — install from pytorch.org)
• Triton ≥ 2.1
• transformers ≥ 4.44
• NVIDIA GPU (benchmarks on RTX 4060 6 GB)