1from tensorrt_llm import LLM, SamplingParams
- 2
+ 1from tensorrt_llm import BuildConfig, SamplingParams
+ 2from tensorrt_llm._tensorrt_engine import LLM # NOTE the change
3
- 4def main():
- 5
- 6 # Model could accept HF model name, a path to local HF model,
- 7 # or TensorRT Model Optimizer's quantized checkpoints like nvidia/Llama-3.1-8B-Instruct-FP8 on HF.
- 8 llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
- 9
-10 # Sample prompts.
-11 prompts = [
-12 "Hello, my name is",
-13 "The capital of France is",
-14 "The future of AI is",
-15 ]
-16
-17 # Create a sampling params.
-18 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
-19
-20 for output in llm.generate(prompts, sampling_params):
-21 print(
-22 f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
-23 )
-24
-25 # Got output like
-26 # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming'
-27 # Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the'
-28 # Prompt: 'The capital of France is', Generated text: 'Paris.'
-29 # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are'
+ 4
+ 5def main():
+ 6
+ 7 build_config = BuildConfig()
+ 8 build_config.max_batch_size = 256
+ 9 build_config.max_num_tokens = 1024
+10
+11 # Model could accept HF model name, a path to local HF model,
+12 # or TensorRT Model Optimizer's quantized checkpoints like nvidia/Llama-3.1-8B-Instruct-FP8 on HF.
+13 llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+14 build_config=build_config)
+15
+16 # Sample prompts.
+17 prompts = [
+18 "Hello, my name is",
+19 "The capital of France is",
+20 "The future of AI is",
+21 ]
+22
+23 # Create a sampling params.
+24 sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
+25
+26 for output in llm.generate(prompts, sampling_params):
+27 print(
+28 f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
+29 )
30
-31
-32if __name__ == '__main__':
-33 main()
+31 # Got output like
+32 # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming'
+33 # Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the'
+34 # Prompt: 'The capital of France is', Generated text: 'Paris.'
+35 # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are'
+36
+37
+38if __name__ == '__main__':
+39 main()
Scaling Expert Parallelism in TensorRT-LLM (Part 2: Performance Status and Optimization)
Running a High Performance GPT-OSS-120B Inference Server with TensorRT-LLM
+
Use TensorRT Engine
+
+ 
+
+
+