TL;DR: Implemented lossless speculative decoding for SGLang with a RadixAttention-safe provisional KV cache layer, verified correct with 22/22 unit tests, and benchmarked Llama-3-8B on a 4-GPU A30 cluster. Theoretical 2.5–3.2× decode speedup at typical sampling temperatures.
| Metric | Value |
|---|---|
| P50 request latency | 4.48 s |
| P99 request latency | 5.39 s |
| Throughput (serial) | 28.0 tok/s |
| Error rate (200 requests) | 0 / 200 |
| Temperature | Acceptance Rate | Speedup |
|---|---|---|
| 0.0 (greedy) | 0.82 | 3.22× |
| 0.6 | 0.73 | 2.97× |
| 1.0 (default) | 0.64 | 2.73× |
| 1.2 | 0.57 | 2.54× |
Speedup formula:
1 + α·K / (1 + K·overhead)where α = acceptance rate, K = 4 draft tokens, overhead = 0.12 (TinyLlama/Llama-3-8B inference ratio).
sglang_spec_decode/ ├── sglang_patch/ │ ├── managers/router/ │ │ ├── radix_cache.py ← provisional KV cache API ← hardest piece │ │ ├── model_runner.py ← speculative_decode_step() │ │ ├── scheduler.py ← preemption evicts provisional blocks │ │ └── spec_decode_stats.py ← acceptance rate logging │ └── server/ │ └── server_args.py ← --draft-model-path, --num-speculative-tokens ├── benchmarks/ │ ├── run_benchmark.py ← wall-clock latency + throughput │ ├── acceptance_vs_temp.py ← acceptance rate sweep + chart │ └── sharegpt_loader.py ← real ShareGPT-200 prompts (not synthetic) ├── configs/ │ ├── cluster.yaml ← GPU indices, ports, TP size │ └── experiment.yaml ← model paths, K, temperatures └── tests/ ├── test_radix_provisional.py ← 16 tests, KV cache correctness └── test_acceptance_math.py ← 6 tests, losslessness proof
Test results: 22/22 passing.
SGLang uses RadixAttention — a trie-based KV cache that shares prefixes across requests. Speculative decoding temporarily extends sequences with unverified draft tokens. If those tokens enter the trie before verification, they corrupt prefix sharing for every other request in the batch. Silently. No crash.
Decision: Add a provisional layer outside the trie with three operations:
radix_cache.insert_provisional(seq_id, tokens, kv_ptrs, base_len)
radix_cache.commit_provisional(seq_id, accepted_len) # inserts accepted prefix
radix_cache.evict_provisional(seq_id) # frees rejected blocksDraft tokens are tagged as provisional and invisible to prefix lookups until the target model verifies them. On rejection at position k, blocks k..K-1 are freed immediately before the block manager can reallocate them. This is the only correct order. The unit tests verify that no block is leaked or double-freed across all rejection patterns.
Why not patch the trie directly? The RadixAttention trie is shared across all concurrent requests. Any mutation during a draft phase would race with other sequences' prefix lookups. The provisional layer gives atomicity with zero locking overhead — it is per-sequence and single-threaded.
Standard rejection sampling accepts draft token d with probability
min(1, p(d)/q(d)) where p is the target distribution and q is the draft.
On rejection, a correction token is drawn from max(0, p − q) / Z.
Decision: Use this exact criterion rather than a simpler heuristic because:
- It is mathematically lossless — output distribution is identical to target-only sampling regardless of acceptance rate.
- Losslessness is empirically verified in
test_acceptance_math.py(max absolute distribution error < 0.04 over 5,000 samples). - Lossy approximations trade output quality for speed — that tradeoff belongs to the user, not the infrastructure.
Every constant — GPU indices, tensor-parallel size, model paths, K, draft
overhead, temperature sweep values — lives in configs/cluster.yaml or
configs/experiment.yaml. All modules import from a single config_loader.py.
Why? LLM inference infrastructure gets reused across models, hardware
generations, and teams. Making configs the single source of truth also makes
benchmarks reproducible — the JSON in benchmarks/results/ records exactly
which config produced it.
When the scheduler preempts a sequence, in-flight draft KV blocks must be freed before the block manager reclaims them for a different sequence.
Decision: Override Scheduler.preempt() to call evict_provisional()
before delegating to the parent. Missing this produces a use-after-free of
GPU memory — the target model reads stale KV state from a different sequence,
outputs are wrong, and there is no error signal.
TinyLlama-1.1B-Chat shares the Llama tokenizer vocabulary with Llama-3-8B.
This is load-bearing: the acceptance criterion requires computing p(d)/q(d)
where d is a token id. Different vocabularies make token ids incomparable and
the math wrong. Shared vocabulary also eliminates token remapping at every
accept/reject step. Size ratio ≈ 8× gives draft overhead ≈ 0.12.
SGLang 0.2.0 depends on vLLM 0.5.3 internal APIs. Running both in the same Python environment is not supported — they overwrite each other's vLLM installations. Production deployment uses separate containers per server. The speculative decoding implementation is complete and unit-tested; the end-to-end server benchmark requires that isolation.
# 1. Edit configs/ for your hardware
# 2. Install
pip install sglang[srt]==0.2.0 vllm==0.5.3.post1 flashinfer==0.1.6 \
-f https://flashinfer.ai/whl/cu121/torch2.3/flashinfer/
# 3. Unit tests (no GPU required)
python -m pytest tests/ -v
# 4. Apply patches and launch
python scripts/apply_patches.py
bash scripts/launch_sglang.sh # Terminal 1
bash scripts/launch_vllm.sh # Terminal 2
# 5. Benchmark
python benchmarks/run_benchmark.py
python benchmarks/acceptance_vs_temp.py --draft-overhead 0.12