flashlet 🔥

A pocket-sized LLM inference engine: continuous batching + paged KV-cache, FlashInfer as the reference attention backend, hand-written CuTe DSL kernels where we can beat it. Runs Qwen2.5-0.5B (bf16) end-to-end on an RTX 4050 Laptop (sm_89, 6 GB).

flashlet/
├── flashlet/
│   ├── scheduler.py        continuous batching + chunked prefill + radix prefix reuse
│   ├── kv_cache.py         paged KV pool, free-list allocator, radix cache w/ refcounts
│   ├── model.py            Qwen2.5 forward pass (HF weights, our own loop)
│   ├── engine.py           step(): schedule -> plan -> attention -> MLP -> sample
│   ├── backends/
│   │   ├── base.py         AttentionBackend ABC: plan() / run()
│   │   ├── torch_backend.py    dense SDPA reference (correctness oracle)
│   │   └── flashinfer_backend.py  BatchPrefillWithPagedKVCacheWrapper, persistent bufs
│   └── bench/              harness + workloads + `python -m flashlet.bench`
├── kernels/                standalone CuTe DSL kernels (M4)
├── tests/                  pytest suite
└── scripts/m0_check.py     end-to-end sanity check vs HF generate

pip install torch flashinfer-python transformers fastapi uvicorn
python3 scripts/m0_check.py                     # greedy gen + HF parity + TPS
python3 -m pytest tests/ -q                     # unit tests
python3 -m flashlet.bench --backend flashinfer --concurrency 1 8 32 128 \
    --max-new-tokens 32 --max-num-pages 8192 --out results.json

python3 -m flashlet.server --backend flashinfer --port 8000     # HTTP server
curl -s localhost:8000/health
curl -s localhost:8000/generate -H 'Content-Type: application/json' \
    -d '{"prompt": "The capital of France is", "max_new_tokens": 32}'
curl -s -N localhost:8000/generate_stream -H 'Content-Type: application/json' \
    -d '{"prompt": "Tell me a story", "max_new_tokens": 64}'    # SSE, token by token

Qwen2.5-0.5B-Instruct, bf16, RTX 4050 Laptop (sm_89), greedy, max_new_tokens=32, mixed prompt lengths (32-768 tok), closed-loop, --max-num-pages 8192.

Saturated-batch microbenchmark (128 seqs × 300-tok prompts, pure decode): 11,238 tok/s at 11.4 ms/step with the FlashInfer backend.

Single-stream peak: ~124 tok/s (flashinfer) vs ~101 tok/s (torch reference).

TTFT includes queue wait in closed-loop batches (later requests wait for earlier ones at low concurrency). Logits agree with HF transformers on 99.3% of argmax positions; greedy outputs diverge on near-ties due to bf16 kernel reduction order, as expected.

• One ragged BatchPrefillWithPagedKVCacheWrapper.plan()/run() call per step handles chunked prefills AND decodes together (qo_indptr = [512,1,1,...], causal mask covers both).
• Metadata lives in persistent device buffers updated in-place each step; CUDA-graph-friendly by construction.
• KV pool layout [layers, pages, page_size, kv_heads, head_dim] matches FlashInfer NHD exactly: set_kv is a single index_copy_, zero copies.
• Radix prefix cache with explicit page refcounts; adoption never covers the final prompt token (guarantees a query token exists); LRU-free eviction on pool pressure.
• Chunked prefill shares one token budget across all active prefills; decoders always step every iteration.

A standalone CuTe DSL megakernel for Qwen2.5 (whole-model fused kernel: GEMV, GQA attention, RMSNorm, RoPE and residual ops in one launch graph, see qwen_2.5_megakernal) already measures 140 tok/s single-stream decode on the same RTX 4050. That is close to the practical roofline for this GPU: bf16 weights are about 1 GB, so 140-160 tok/s means streaming roughly 140-160 GB/s of weight reads per token, near the sustained bandwidth ceiling of the laptop's memory bus.

Integration into flashlet as a third attention/model backend is planned: the megakernel would replace the per-layer op sequence for pure-decode steps while flashlet keeps the scheduler, paged KV pool and continuous batching around it. Expect single-stream numbers to land in the 140-160 tok/s band once wired in.

• M0 skeleton: greedy decode, coherent text, HF parity
• M1 paged KV cache: alloc/free cycles, pinned→device metadata, radix reuse
• M2 scheduler: continuous batching + chunked prefill, mixed-batch tests
• M3 FlashInfer backend: plan/run integration, parity vs reference, bench table
• M4 CuTe DSL paged GQA decode kernel (sm_89)
• M5 RoPE+residual fusion kernel
• M6 Blackwell port (stretch)