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]

Sign up or log in to comment