Efficient LLM Inference & Scheduling

This page describes targeted experiments in a lightweight LLM inference engine. The scheduling and quantization changes are presented as engineering modifications and benchmarked prototypes, not as a claim that the full nano-vLLM engine was built from scratch or that the optimizations are production-ready across all workloads.

The engineering question behind this project was simple: when a model is already small enough to run on one GPU, where does inference efficiency actually come from? The answer is not only model math. Request scheduling, KV-cache capacity, weight storage, CUDA graph behavior, and workload shape all decide how the system behaves.

I explored two complementary optimization directions:

• Prefill/decode interleaving. Modify the scheduler so long prompt prefill work and active decode requests can both make progress under mixed traffic.
• W8A16 weight-only quantization. Store decoder linear weights as INT8 with per-output-channel scales while keeping activations in FP16/BF16.

The two experiments target different bottlenecks. Scheduling affects request latency and execution order. Quantization affects persistent model storage and how much GPU memory remains for KV cache and runtime state.

The Challenge: Inference Is a Systems Problem

LLM inference alternates between two phases with different resource profiles:

A scheduler has to decide which work enters the next model forward. A long prompt can consume many chunked prefill steps. Active generations need repeated decode steps to keep producing tokens. GPU memory is split among model weights, KV cache, activations, temporary tensors, and runtime overhead. Optimizing one part can expose a different bottleneck somewhere else.

Experiment 01: Prefill/Decode Interleaving

Problem

Chunked prefill makes long prompts possible under a limited max_num_batched_tokens budget, but it creates a scheduling question. If the scheduler prefers waiting prompt work whenever it exists, a long prompt may run several prefill chunks before active decoders continue. If the scheduler always prioritizes decode, new requests may wait for prompt processing.

The trade-off is between time to first token, decode progress, overall completion time, and fairness across request types.

Implementation

The scheduler change adds a configuration switch, prefill_decode_interleave, and tracks whether the previous engine step was prefill. When interleaving is enabled, the scheduler chooses decode first only when all of these are true:

• the previous scheduled step was prefill
• there is waiting prompt work
• there are active running decode requests
• interleaving is enabled

Otherwise, the scheduler follows the normal prefill-first path, then falls back to decode when no prompt work can be scheduled. This is scheduler-level interleaving, not simultaneous prefill and decode execution in one model forward.

Chunked prefill reuses existing cached-prefix mechanics. The scheduler sets seq.num_scheduled_tokens for the next prompt chunk. The model runner starts from seq.num_cached_tokens, writes new KV slots, and lets later chunks attend to cached prefix blocks through the block table.

Measured Results

The recorded scheduling experiment used Qwen3-0.6B on an RTX 4090 with eager execution, max_num_batched_tokens=128, max_num_seqs=8, four initial short requests, one long prompt of about 1024 tokens, four late short requests, and stochastic sampling with temperature=0.6.

Other observations were important too: preemptions stayed at zero, peak KV-cache block usage stayed at 13 blocks, and decode throughput stayed similar. The clearest qualitative result was the execution trace: with interleaving disabled, a long prompt produced consecutive prefill chunks; with interleaving enabled, active decoders resumed every other step while both queues had work.

Engineering Insight

The result supports the scheduler hypothesis for this workload, but it does not prove a universal win. The absolute TTFT numbers include cold-start overhead, one late request waits because the test has nine requests and max_num_seqs=8, and the benchmark was a single controlled workload on one RTX 4090. The useful engineering result is narrower and stronger: a small scheduler policy change produced the intended execution pattern and improved measured latency metrics under the documented mixed workload.

Experiment 02: W8A16 Weight-Only Quantization

Problem

GPU memory is divided between persistent model state and runtime state. Model weights occupy memory before any request arrives. KV cache grows with context length, batch size, and active sequence count. Reducing model weight storage can make more memory available for serving, but it does not automatically make matrix multiplication faster.

Implementation

The W8A16 path quantizes linear weights at load time. For each output channel, the implementation computes a scale from the maximum absolute weight value, stores the quantized weights as INT8, and stores the scale values as FP32 buffers. During forward execution, weights are dequantized back to the activation dtype and used in the existing linear path.

Measured Results

The quantization benchmark measured persistent model storage and request performance across eager and CUDA graph modes. The most stable result was the storage reduction:

The speed story was more mixed. Some W8A16 runs were faster than their baseline counterpart, but the best interleaved CUDA graph result remained the non-quantized baseline. That is consistent with the implementation: the weights are stored compactly, but each forward pass still pays for dequantization and then runs a higher-precision linear operation.

GPU Memory Architecture

The memory budget is shaped by model size, precision, context length, active sequence count, batch size, and temporary allocations created by the execution backend. W8A16 reduced persistent model storage and increased available KV-cache blocks in the recorded runs, but peak CUDA allocation did not clearly drop because runtime dequantization creates temporary tensors.

Experimental Methodology

The scheduling benchmark records per-step traces, per-request latencies, aggregate summaries, preemption counts, KV-cache block usage, and CUDA memory statistics. The workload shape was chosen to force a long prompt to be chunked while short requests were already decoding.

The quantization benchmark measures model storage by summing parameter and buffer bytes, separates INT8 weights and scale buffers, records KV-cache capacity, and compares request latency across eager and CUDA graph execution. Both experiments used synthetic prompts and controlled request mixes rather than production traffic.

Explore the Implementation

The implementation work is organized around a few focused code paths:

• Scheduler policy and chunked prefill behavior in nanovllm/engine/scheduler.py
• Cached-prefix preparation in nanovllm/engine/model_runner.py
• Configuration switches in nanovllm/config.py
• W8A16 linear layers in nanovllm/layers/linear.py
• Benchmark scripts for latency and quantization profiling in the experiment repository

Future work would move the quantization path from compact storage toward compute-efficient inference: optimized INT8 GEMM kernels, fused dequantization, smaller temporary allocations, kernel-level profiling, and repeated measurements across larger models and heavier workloads.