Runtime Configuration Examples#

Source NVIDIA/TensorRT-LLM.

 1
 2import argparse
 3
 4from tensorrt_llm import LLM, SamplingParams
 5from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig
 6
 7
 8def example_cuda_graph_config():
 9    """
10    Example demonstrating CUDA graph configuration for performance optimization.
11
12    CUDA graphs help with:
13    - Reduced kernel launch overhead
14    - Better GPU utilization
15    - Improved throughput for repeated operations
16    """
17    print("\n=== CUDA Graph Configuration Example ===")
18
19    cuda_graph_config = CudaGraphConfig(
20        batch_sizes=[1, 2, 4],
21        enable_padding=True,
22    )
23
24    llm = LLM(
25        model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
26        cuda_graph_config=cuda_graph_config,  # Enable CUDA graphs
27        max_batch_size=4,
28        max_seq_len=512,
29        kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.5))
30
31    prompts = [
32        "Hello, my name is",
33        "The capital of France is",
34        "The future of AI is",
35    ]
36
37    sampling_params = SamplingParams(max_tokens=50, temperature=0.8, top_p=0.95)
38
39    # This should benefit from CUDA graphs
40    outputs = llm.generate(prompts, sampling_params)
41    for output in outputs:
42        print(f"Prompt: {output.prompt}")
43        print(f"Generated: {output.outputs[0].text}")
44        print()
45
46
47def example_kv_cache_config():
48    print("\n=== KV Cache Configuration Example ===")
49    print("\n1. KV Cache Configuration:")
50
51    llm_advanced = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
52                       max_batch_size=8,
53                       max_seq_len=1024,
54                       kv_cache_config=KvCacheConfig(
55                           free_gpu_memory_fraction=0.5,
56                           enable_block_reuse=True))
57
58    prompts = [
59        "Hello, my name is",
60        "The capital of France is",
61        "The future of AI is",
62    ]
63
64    outputs = llm_advanced.generate(prompts)
65    for i, output in enumerate(outputs):
66        print(f"Query {i+1}: {output.prompt}")
67        print(f"Answer: {output.outputs[0].text[:100]}...")
68        print()
69
70
71def main():
72    """
73    Main function to run all runtime configuration examples.
74    """
75    parser = argparse.ArgumentParser(
76        description="Runtime Configuration Examples")
77    parser.add_argument("--example",
78                        type=str,
79                        choices=["kv_cache", "cuda_graph", "all"],
80                        default="all",
81                        help="Which example to run")
82
83    args = parser.parse_args()
84
85    if args.example == "kv_cache" or args.example == "all":
86        example_kv_cache_config()
87
88    if args.example == "cuda_graph" or args.example == "all":
89        example_cuda_graph_config()
90
91
92if __name__ == "__main__":
93    main()