Speculative Decoding#

Source NVIDIA/TensorRT-LLM.

  1from typing import Optional
  2
  3import click
  4
  5from tensorrt_llm import LLM, SamplingParams, logger
  6from tensorrt_llm._utils import get_sm_version
  7from tensorrt_llm.llmapi import (Eagle3DecodingConfig, KvCacheConfig,
  8                                 MTPDecodingConfig, NGramDecodingConfig)
  9
 10prompts = [
 11    "What is the capital of France?",
 12    "What is the future of AI?",
 13]
 14
 15
 16def run_MTP(model: Optional[str] = None):
 17    sm_version = get_sm_version()
 18    if sm_version < 90:
 19        logger.warning(
 20            f"Skipping the MTP example: it requires the DeepSeek MLA "
 21            "generation FMHA kernel, which is only available on "
 22            f"Hopper+ (SM>=90). Detected SM{sm_version}.")
 23        return
 24
 25    spec_config = MTPDecodingConfig(use_relaxed_acceptance_for_thinking=True,
 26                                    relaxed_topk=10,
 27                                    relaxed_delta=0.01)
 28
 29    llm = LLM(
 30        # You can change this to a local model path if you have the model downloaded
 31        model=model or "nvidia/DeepSeek-R1-FP4",
 32        speculative_config=spec_config,
 33    )
 34
 35    for prompt in prompts:
 36        response = llm.generate(prompt, SamplingParams(max_tokens=10))
 37        print(response.outputs[0].text)
 38
 39
 40def run_Eagle3():
 41    spec_config = Eagle3DecodingConfig(
 42        max_draft_len=3,
 43        speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
 44        eagle3_one_model=True)
 45
 46    kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8)
 47
 48    llm = LLM(
 49        model="meta-llama/Llama-3.1-8B-Instruct",
 50        speculative_config=spec_config,
 51        kv_cache_config=kv_cache_config,
 52    )
 53
 54    for prompt in prompts:
 55        response = llm.generate(prompt, SamplingParams(max_tokens=10))
 56        print(response.outputs[0].text)
 57
 58
 59def run_ngram():
 60    spec_config = NGramDecodingConfig(
 61        max_draft_len=3,
 62        max_matching_ngram_size=3,
 63        is_keep_all=True,
 64        is_use_oldest=True,
 65        is_public_pool=True,
 66    )
 67
 68    llm = LLM(
 69        model="meta-llama/Llama-3.1-8B-Instruct",
 70        speculative_config=spec_config,
 71        # ngram doesn't work with overlap_scheduler
 72        disable_overlap_scheduler=True,
 73    )
 74
 75    for prompt in prompts:
 76        response = llm.generate(prompt, SamplingParams(max_tokens=10))
 77        print(response.outputs[0].text)
 78
 79
 80@click.command()
 81@click.argument("algo",
 82                type=click.Choice(["MTP", "EAGLE3", "DRAFT_TARGET", "NGRAM"]))
 83@click.option("--model",
 84              type=str,
 85              default=None,
 86              help="Path to the model or model name.")
 87def main(algo: str, model: Optional[str] = None):
 88    algo = algo.upper()
 89    if algo == "MTP":
 90        run_MTP(model)
 91    elif algo == "EAGLE3":
 92        run_Eagle3()
 93    elif algo == "NGRAM":
 94        run_ngram()
 95    else:
 96        raise ValueError(f"Invalid algorithm: {algo}")
 97
 98
 99if __name__ == "__main__":
100    main()