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
45    kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8)
46
47    llm = LLM(
48        model="meta-llama/Llama-3.1-8B-Instruct",
49        speculative_config=spec_config,
50        kv_cache_config=kv_cache_config,
51    )
52
53    for prompt in prompts:
54        response = llm.generate(prompt, SamplingParams(max_tokens=10))
55        print(response.outputs[0].text)
56
57
58def run_ngram():
59    spec_config = NGramDecodingConfig(
60        max_draft_len=3,
61        max_matching_ngram_size=3,
62        is_keep_all=True,
63        is_use_oldest=True,
64        is_public_pool=True,
65    )
66
67    llm = LLM(
68        model="meta-llama/Llama-3.1-8B-Instruct",
69        speculative_config=spec_config,
70        # ngram doesn't work with overlap_scheduler
71        disable_overlap_scheduler=True,
72    )
73
74    for prompt in prompts:
75        response = llm.generate(prompt, SamplingParams(max_tokens=10))
76        print(response.outputs[0].text)
77
78
79@click.command()
80@click.argument("algo",
81                type=click.Choice(["MTP", "EAGLE3", "DRAFT_TARGET", "NGRAM"]))
82@click.option("--model",
83              type=str,
84              default=None,
85              help="Path to the model or model name.")
86def main(algo: str, model: Optional[str] = None):
87    algo = algo.upper()
88    if algo == "MTP":
89        run_MTP(model)
90    elif algo == "EAGLE3":
91        run_Eagle3()
92    elif algo == "NGRAM":
93        run_ngram()
94    else:
95        raise ValueError(f"Invalid algorithm: {algo}")
96
97
98if __name__ == "__main__":
99    main()