Motivation

SGLang's STANDALONE speculative decoding mode supports arbitrary external draft models, but
currently requires the draft model to share the same vocabulary as the target model. This rules
out the large universe of small, publicly available models as draft models for target models with
different tokenizers.

This PR implements TLI (Token-Level Intersection), an algorithm from our ICML 2025 spotlight paper "Accelerating LLM Inference with Lossless Speculative Decoding Algorithms for Heterogeneous Vocabularies" — Timor et al., arXiv 2502.05202. The same idea is described from a practical angle in our HuggingFace blog post Universal Assisted Generation. A reference implementation for HuggingFace Transformers was merged at huggingface/transformers#35029. A parallel implementation for vLLM is in progress at vllm-project/vllm#38174.

TLI is provably lossless: it always produces the same output distribution as greedy or
sampled decoding with the target model alone. It unlocks speculative decoding for any pair of
models whose vocabularies have a meaningful token overlap, regardless of tokenizer family.

Modifications

New files:

• python/sglang/srt/speculative/vocab_mapping.py

  Builds a normalized token intersection between target and draft vocabularies at startup.
  The space-prefix character (Ġ, ▁, or other) is detected dynamically by tokenizing
  a literal space, so unusual tokenizers are handled correctly. Models without unk_token_id
  (e.g. Llama 3, SmolLM2) are supported via an eos_token_id fallback. Exposes:

  • constrain_draft_logits(logits) — sets non-intersection logits to −∞.
  • map_target_to_draft_ids(ids) — remaps a tensor of target token IDs to draft vocab.
  • map_draft_to_target_ids(ids) — inverse mapping.
  • intersection_draft_ids — sorted draft-vocab indices in the intersection (used for LM head pruning).

• python/sglang/srt/speculative/tli_worker.py

  TLIWorker extends StandaloneWorker (which loads a fully independent draft model with no
  shared embeddings / lm_head). It overrides three hooks:

  Hook	What it does
  forward_draft_extend	Maps prompt token IDs target→draft before KV-cache prefill
  forward_draft_extend_after_decode	Maps accepted tokens target→draft, runs update, restores to target vocab
  capture_for_decode	Constrains draft logits to the intersection before top-k sampling

  The rejection sampling step in the target model is not modified — losslessness follows
  directly from the standard speculative decoding acceptance criterion.

  Additionally implements _try_prune_draft_lm_head(): at startup the draft LM head weight
  matrix is replaced with a compact [intersection_size × hidden_dim] version
  (_PrunedReindexLMHead), reducing the dominant draft-decoding matmul from
  O(B × V_draft × H) to O(B × K × H) where K ≪ V_draft. CUDA graphs are
  captured after pruning so the graph encodes the smaller computation.

Modified files:

• python/sglang/srt/speculative/spec_info.py

  • Add TLI = auto() to SpeculativeAlgorithm enum.
  • Add is_tli() predicate.
  • Dispatch to TLIWorker in get_worker_cls(), with a guard that rejects overlap
    scheduling (not yet supported for spec v2).

• python/sglang/srt/server_args.py

  • Add "TLI" to --speculative-algorithm choices.
  • Add TLI validation block in check_server_args:
    • Requires --speculative-draft-model-path.
    • Disables overlap scheduling and mixed chunked prefill (unsupported).
    • Auto-selects speculative_num_steps / speculative_num_draft_tokens via the existing
      auto_choose_speculative_params helper when not explicitly set.
    • Enforces speculative_eagle_topk == 1 (TLI does not support multi-candidate trees).
    • Enforces speculative_num_draft_tokens == speculative_num_steps + 1.
  • Group "TLI" with "STANDALONE" in the cuda-graph-max-bs default path.

Usage example:

python -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --speculative-algorithm TLI \
    --speculative-draft-model-path Qwen/Qwen2.5-0.5B-Instruct \
    --speculative-num-draft-tokens 5

Parameters are auto-chosen if omitted — no manual tuning needed for basic use.

Accuracy Tests

TLI is provably lossless. A formal proof is given in the ICML 2025 paper
Timor et al., arXiv 2502.05202.

We also ran an empirical losslessness check (check_accuracy.py): 50 prompts are sent with
greedy decoding (temperature=0) to both a baseline server and a TLI server sharing the same
target model, and outputs are compared exactly. In our runs with
Qwen2.5-7B-Instruct + Qwen2.5-0.5B-Instruct, 38/50 outputs matched character-for-character.
The 12 mismatches were all minor whitespace/punctuation differences (e.g. " " vs "\n", or
a dropped filler word like "and") that appear identically whether or not TLI is enabled —
they are caused by non-deterministic floating-point accumulation in Flash Attention (non-associative
parallel reductions whose thread scheduling varies between runs), not by TLI. The same mismatches
occur when comparing two independent baseline runs against each other.

Speed Tests and Profiling

All benchmarks use the speed_bench_mixed dataset (long-context summarization, output length ≈ 1024 tokens),
batch_size=8, num_prompts=100, num_draft_tokens=5, speculative_eagle_topk=1. Vocab intersection is computed after
space-prefix normalization by VocabMapping.

Notes

Llama-3.1-8B-Instruct + Qwen2.5-0.5B-Instruct is a genuine cross-family pair (different
tokenizer families, 85.43% vocab overlap). TLI delivers a strong 1.65–2.22× speedup across
devices, with the memory-bandwidth-bound A6000 showing the highest gain (2.22×).

Comparison with Related Implementations

This PR extends the reference implementation in HuggingFace Transformers
(PR #35029, merged ✅) with
several improvements suited to a high-throughput serving engine:

Checklist

• Format your code according to the Format code with pre-commit.
  • All pre-commit hooks pass: trim-trailing-whitespace, isort, ruff, black, codespell.
• Add unit tests according to the Run and add unit tests.
  • test/registered/unit/spec/test_vocab_mapping.py: 36 CPU-only unit tests covering _detect_space_sign, _normalize_token, intersection construction, unk_token_id fallback, bidirectional ID mapping, and logit constraining.
• Update documentation according to Write documentations.
  • docs/advanced_features/speculative_decoding.md: added TLI section (How it works, constraints, usage example), updated quick-guidance bullets, method comparison table, and --speculative-algorithm parameter reference.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

cc: @orenpereg @keyboardAnt @wan-danfeng

CI States

Latest PR Test (Base): ❌ Missing run-ci label -- add it to run CI tests.
Latest PR Test (Extra): ❌ Blocked -- run-ci is required first.