Script to reproduce 124 tok/sec on a RTX 5090
#6
by vince62s - opened
I am using vllm 0.31 torch 2.12 cuda 13.2
and this script:
# flake8: noqa
from vllm import LLM, SamplingParams
"""
CompilationMode:
NONE = 0
No torch.compile compilation is applied, model runs in fully eager pytorch mode.
The model runs as-is.
STOCK_TORCH_COMPILE = 1
The standard `torch.compile` compilation pipeline.
DYNAMO_TRACE_ONCE = 2
Single Dynamo trace through the model, avoiding recompilation.
VLLM_COMPILE = 3
Custom vLLM Inductor-based backend with caching, piecewise compilation,
shape specialization, and custom passes.
CUDAGraphMode:
NONE = 0
PIECEWISE = 1
FULL = 2
FULL_DECODE_ONLY = (FULL, NONE)
FULL_AND_PIECEWISE = (FULL, PIECEWISE)
"""
compilation_config = {"mode": 3, "cudagraph_mode": "NONE"}
def main():
# Initialize the model with a custom max_model_len
llm = LLM(
model="Frozenlock/Qwen3.8-27B-int4-AutoRound",
max_model_len=2048, # <-- Set the total context length here (prompt + generated)
gpu_memory_utilization=0.85,
compilation_config=compilation_config,
speculative_config={"method": "mtp", "num_speculative_tokens": 3},
)
# Define sampling parameters
sampling_params = SamplingParams(temperature=0.0, top_p=1.0, max_tokens=2048) # how many tokens to generate
# Your prompts
prompts = [
"<|im_start|>user\nGenerate a 500 word text talking about George Orwell.<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
]
# Generate
outputs = llm.generate(prompts, sampling_params)
# Print results
for output in outputs:
print(f"Prompt: {output.prompt!r}")
print(f"Generated: {output.outputs[0].text!r}")
print("-" * 50)
if __name__ == "__main__":
main()
I am getting:
Processed prompts: 100%|ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ| 1/1 [00:10<00:00, 10.65s/it, est. speed input: 2.35 toks/s, output: 73.21 toks/s]