# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

# default model parameters can be overwritten from command line
llm_model_default_parameters = {
    "meta-llama/Llama-3.1-3B-Instruct": {
        "name": "3B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 256,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "llama3_json",
    },
    "meta-llama/Llama-3.1-8B-Instruct": {
        "name": "8B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "llama3_json",
    },
    "meta-llama/Llama-3.1-70B-Instruct": {
        "name": "70B",
        "tensor-parallel-size": 4,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max-model-len": 30960,
        "gpu-memory-utilization": 0.8,
        "max_ongoing_requests": 100,
        "tool-call-parser": "llama3_json",
    },
    "meta-llama/Llama-3.1-405B-Instruct-FP8": {
        "name": "405B-FP8",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max-model-len": 10240,
        "gpu-memory-utilization": 0.8,
        "max_ongoing_requests": 50,
        "tool-call-parser": "llama3_json",
    },
    "meta-llama/Llama-3.1-405B-Instruct": {
        "name": "405B",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 2,
        "enable-prefix-caching": True,
        "max-model-len": 10240,  # 30960 (4 node), 61440 (6 node), 128000 (10 nodes)
        "gpu-memory-utilization": 0.8,
        "max_ongoing_requests": 50,
        "tool-call-parser": "llama3_json",
    },
    "meta-llama/Llama-3.3-70B-Instruct": {
        "name": "3_3_70B",
        "tensor-parallel-size": 4,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
        "max_ongoing_requests": 100,
        "tool-call-parser": "llama3_json",
    },
    "deepseek-ai/DeepSeek-R1": {
        "name": "deepseek-r1",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 3,
        "enable-prefix-caching": True,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.9,
        "max_ongoing_requests": 80,
        "trust-remote-code": True,
    },
    "meta-llama/Llama-4-Scout-17B-16E-Instruct": {
        "name": "scout",
        "tensor-parallel-size": 4,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.9,
        "max_ongoing_requests": 100,
        "use_v1_engine": "true",
        "tool-call-parser": "llama3_json",
    },
    "meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": {
        "name": "maverick-fp8",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max-model-len": 128000,
        "gpu-memory-utilization": 0.9,
        "max_ongoing_requests": 100,
        "dtype": "auto",
        "kv-cache-dtype": "auto",
        "quantization": "compressed-tensors",
        "use_v1_engine": "true",
        "tool-call-parser": "llama3_json",
    },
    "unsloth/mistral-7b-instruct-v0.2-bnb-4bit": {
        "name": "unsloth-mistral-7B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 256,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.4,
        "enable-lora": True,
        "quantization": "bitsandbytes",
        "load-format": "bitsandbytes",
        "max_lora_rank": 32,
    },
    "google/gemma-3-12b-it": {
        "name": "gemma3-12B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
    },
    "mistralai/Ministral-8B-Instruct-2410": {
        "name": "mistral-8B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
        "disable-sliding-window": True,
        "tool-call-parser": "minimax",
    },
    "mistralai/Mistral-7B-Instruct-v0.1": {
        "name": "mistralai-7B",
        "tensor-parallel-size": 2,
        "pipeline-parallel-size": 1,
        "disable-sliding-window": True,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 4096,
        "gpu-memory-utilization": 0.8,
        "dtype": "half",
        "tool-call-parser": "minimax",
    },
    "Qwen/Qwen2.5-7B-Instruct": {
        "name": "Qwen2.5-7B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "hermes",
    },
    "Qwen/Qwen2.5-32B-Instruct": {
        "name": "Qwen2.5-32B",
        "tensor-parallel-size": 2,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 100,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "hermes",
    },
    "Qwen/Qwen2.5-72B-Instruct": {
        "name": "Qwen2.5-72B",
        "tensor-parallel-size": 4,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 100,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
    },
    "Qwen/Qwen3-8B": {
        "name": "Qwen3-8B",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "hermes",
    },
    "Qwen/Qwen3-32B": {
        "name": "Qwen3-32B",
        "tensor-parallel-size": 2,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 100,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "hermes",
    },
    "Qwen/Qwen3-235B-A22B": {
        "name": "Qwen3-235B-A22B",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 2,
        "enable-prefix-caching": True,
        "max-model-len": 32768,
        "gpu-memory-utilization": 0.8,
        "max_ongoing_requests": 50,
        "tool-call-parser": "hermes",
    },
    # >= vllm 0.11.0
    "Qwen/Qwen3-VL-30B-A3B-Instruct": {
        "name": "Qwen3-VL-30B-A3B-Instruct",
        "tensor-parallel-size": 2,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 128,
        "max-model-len": 262000,
        "gpu-memory-utilization": 0.8,
        "tool-call-parser": "hermes",
        "mm-encoder-tp-mode": "data",
        # "enable-expert-parallel": True,
        "no-async-scheduling": True,  # async does not work with ray backend
        "use_v1_engine": "true",
    },
    # >= vllm 0.10.2
    "openai/gpt-oss-20b": {
        "name": "gpt-oss-20b",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.85,
        "max_ongoing_requests": 64,
        "use_v1_engine": "true",
        "tool-call-parser": "openai",
    },
    # >= vllm 0.10.2
    "openai/gpt-oss-120b": {
        "name": "gpt-oss-120b",
        "tensor-parallel-size": 2,
        "pipeline-parallel-size": 1,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.85,
        "max_ongoing_requests": 64,
        "use_v1_engine": "true",
        "tool-call-parser": "openai",
    },
    "google/gemma-3-27b-it": {
        "name": "gemma-3-27b-it",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 1,
        "max-model-len": 131072,
        "max_ongoing_requests": 64,
        "limit_mm_per_prompt": "image=128",
    },
    "google/gemma-3-12b-it": {
        "name": "gemma-3-12b-it",
        "tensor-parallel-size": 4,
        "pipeline-parallel-size": 1,
        "max-model-len": 131072,
        "max_ongoing_requests": 64,
        "limit_mm_per_prompt": "image=128",
    },
    "google/gemma-3-4b-it": {
        "name": "gemma-3-4b-it",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "max-model-len": 131072,
        "max_ongoing_requests": 64,
        "limit_mm_per_prompt": "image=128",
    },
    "google/gemma-3-270m-it": {
        "name": "gemma-3-270m-it",
        "tensor-parallel-size": 1,
        "pipeline-parallel-size": 1,
        "max-model-len": 131072,
        "max_ongoing_requests": 64,
        "limit_mm_per_prompt": "image=128",
    },
    "moonshotai/Kimi-K2-Instruct-0905": {
        "name": "kimi-k2",
        "tensor-parallel-size": 8,
        "pipeline-parallel-size": 2,
        "enable-prefix-caching": True,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.85,
        "max_ongoing_requests": 64,
        "trust-remote-code": True,
        "use_v1_engine": "true",
        "tool-call-parser": "kimi_k2",
    },
    "facebook/cwm": {
        "name": "cwm",
        "tensor-parallel-size": 2,
        "pipeline-parallel-size": 1,
        "enable-prefix-caching": True,
        "max_ongoing_requests": 150,
        "max-model-len": 131072,
        "gpu-memory-utilization": 0.85,
        "use_v1_engine": "true",
        "tool-call-parser": "llama3_json",
    },
}
