LLM Inference Lab

A from-first-principles LLM inference engine, progressing from naive autoregressive decoding to continuous batching, KV-cache management, profiling, speculative decoding, and DFlash.

Progress

1. Core inference and measurement

• Autoregressive generation
• KV cache
• Prompt/output-length benchmarking
• TTFT and TPOT measurement
• GPU memory benchmarking
• GPU utilization measurement

2. Batching and scheduling

• Static batching

  • Combine multiple sequences into a single model forward pass.
  • Compare aggregate throughput vs. per-user tokens/sec.

• Naive concurrency

  • Run independent requests concurrently against the same GPU.
  • Observe contention, latency growth, memory pressure, and OOM limits.

Continuous batching

• Stage 1: batched decoding

  • Decode multiple active requests in one model forward pass.
  • Introduce a shared batched KV cache.

• Stage 2A: dynamic request removal

  • Support different output lengths.
  • Requests finish at different decode steps.
  • Remove finished requests from the active batch.
  • Remove their corresponding KV-cache rows.

• Stage 2B: dynamic request admission

  • Support more waiting requests than available active slots.
  • When a request finishes, admit a waiting request immediately.
  • Rebuild and pad the active batch when membership changes.
  • Current limitation: admitting a request rebuilds the KV cache for existing requests.

• Stage 2C: variable prompt lengths

  • Support requests with different prompt lengths.
  • Pad shorter sequences to a common batch width.
  • Use attention masks to distinguish real tokens from padding.
  • Support variable prompt lengths and variable output lengths together.

• Stage 2D: persistent KV-cache slots — no full cache rebuild

  • Stop rebuilding the entire KV cache when a request joins.
  • Assign each active request a stable slot.
  • Free a request's slot when it finishes.
  • Reuse freed slots for waiting requests.
  • Keep existing requests in the same cache location.

  Current progress:

  • Slot allocator
  • Find free slots
  • Stable slot_id per request
  • Free finished slots
  • Reuse freed slots
  • Unit tests for slot lifecycle
  • Preallocated KV storage
  • Prefill a new request into only its assigned slot
  • Decode using active slot IDs
  • Eliminate full KV-cache rebuilds
  • Benchmark Stage 2D vs. Stage 2C

• Stage 2E: block/paged KV-cache management

  • Move from one fixed cache slot per request to reusable KV blocks/pages.
  • Allow requests to grow without reserving a maximum contiguous region.
  • Allocate and free KV blocks dynamically.
  • Understand the core idea behind PagedAttention-style memory management.

3. Engineering and optimization

• Code cleanup and argparse CLI

  • Select benchmarks from the command line.
  • Add result-saving options.
  • Add plotting flags.

• PyTorch Profiler integration

  • CPU vs. CUDA time
  • Kernel breakdown
  • Memory allocations
  • CPU/GPU synchronization
  • GPU timeline / idle gaps
  • Prefill vs. decode profiling

• Profiler-guided optimization

4. Production serving comparison

• Continuous-batching serving benchmark
• Compare this engine with vLLM
• Compare KV-cache management with PagedAttention
• Multi-GPU serving
  • Single-GPU baseline
  • Multiple model replicas
  • Tensor-parallel model execution
• TensorRT-LLM comparison

5. Advanced decoding

• Vanilla speculative decoding

  • Draft model
  • Target verification
  • Acceptance rate
  • Accepted tokens per verification pass

• EAGLE-style speculative decoding comparison

• DFlash

  • Block-diffusion drafting
  • Parallel draft-token generation
  • Draft vs. verification time
  • Acceptance rate
  • Effective tokens per target forward pass
  • Compare AR vs. speculative vs. DFlash
  • Find the workload crossover point where speculation stops helping

End Goal

Understand the inference stack by building each bottleneck and its solution:

AR decoding → KV caching → batching → continuous batching → KV-cache memory management → profiling → production serving → speculative decoding → DFlash