vllm - 💡(How to fix) Fix [Bug]: chunked prefill can't be disabled. maybe due to race condition in scheduler [2 comments, 2 participants]

import os

// if use VLLM_ENABLE_V1_MULTIPROCESSING will cause chunked prefill can't be disabled
// only add this environment variable can make sure chunked prefill is disabled.

//os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"

import argparse
from datetime import datetime

import torch
from vllm import LLM, SamplingParams


def main():
    parser = argparse.ArgumentParser()

    // some args parse here

    args = parser.parse_args()

    timestamp = datetime.now().strftime("%m%d%H%M%S")
    profile_dir = os.path.abspath(f"profile_{timestamp}")
    os.makedirs(profile_dir, exist_ok=True)

    max_model_len = args.seqlen + args.max_tokens

    llm = LLM(
        model=args.model,
        tensor_parallel_size=args.tp,
        load_format="dummy",
        dtype="bfloat16",
        enforce_eager=True,
        max_model_len=max_model_len,
        max_num_seqs=args.bs,
        max_num_batched_tokens=args.bs * max_model_len,
        enable_chunked_prefill=False,
        enable_prefix_caching=False,
        profiler_config={
            "profiler": "torch",
            "torch_profiler_dir": profile_dir,
            "torch_profiler_with_stack": True,
            "torch_profiler_record_shapes": True,
        },
    )

    vocab_size = llm.llm_engine.model_config.get_vocab_size()
    input_ids = torch.randint(0, vocab_size, (args.seqlen,)).tolist()
    sampling_params = SamplingParams(max_tokens=args.max_tokens)
    prompts = [{"prompt_token_ids": input_ids} for _ in range(args.bs)]

    # Warmup
    for i in range(args.warmup):
        print(f"Warmup {i + 1}/{args.warmup}")
        llm.generate(prompts=prompts, sampling_params=sampling_params)
    torch.cuda.synchronize()

    # Profile
    print(f"Profiling batch size {args.bs}... traces will be saved to {profile_dir}")
    llm.start_profile(profile_prefix=f"mixed_precision_{timestamp}")
    llm.generate(prompts=prompts, sampling_params=sampling_params)
    torch.cuda.synchronize()
    llm.stop_profile()
    print(f"Profile saved to {profile_dir}")


if __name__ == "__main__":
    main()