x-common: &common
  build:
    context: .
    dockerfile: docker/vllm/Dockerfile
    args:
      VLLM_IMAGE: "${VLLM_IMAGE:-vllm/vllm-openai:gemma4-0505-arm64-cu130}"
      APPLY_VLLM_PR_PATCH: "${APPLY_VLLM_PR_PATCH:-1}"
      VLLM_PR_NUMBER: "${VLLM_PR_NUMBER:-41745}"
      VLLM_PATCH_REPO: "${VLLM_PATCH_REPO:-https://github.com/vllm-project/vllm.git}"
      PATCH_TRANSFORMERS_FOR_MTP: "${PATCH_TRANSFORMERS_FOR_MTP:-1}"
      HF_HUB_VERSION: "${HF_HUB_VERSION:-1.13.0}"
      REGEX_VERSION: "${REGEX_VERSION:-2025.10.23}"
      TOKENIZERS_VERSION: "${TOKENIZERS_VERSION:-0.22.1}"
      TRANSFORMERS_SOURCE: "${TRANSFORMERS_SOURCE:-git}"
  image: gemma4/gemma4-official-mtp-pr41745:latest
  network_mode: host
  ipc: host
  shm_size: "32gb"
  env_file:
    - path: .env
      required: false
  volumes:
    - ./models/hf:/root/.cache/huggingface
    - ./logs:/logs
  deploy:
    resources:
      reservations:
        devices:
          - driver: nvidia
            count: all
            capabilities: [gpu]

services:
  gemma4-26b:
    <<: *common
    container_name: gemma4-26b
    restart: unless-stopped
    environment:
      PROFILE_NAME: gemma4-26b
      MODEL_ID: "${GEMMA26_MODEL:-google/gemma-4-26B-A4B-it}"
      SERVED_MODEL_NAME: "${GEMMA26_MODEL_NAME:-gemma4-26b-a4b}"
      PORT: "${GEMMA26_PORT:-8100}"
      ENABLE_MTP: "0"
      EXTRA_ARGS_PROFILE: "${GEMMA26_EXTRA_ARGS:-}"
      NVIDIA_VISIBLE_DEVICES: all
      NVIDIA_DRIVER_CAPABILITIES: compute,utility

  gemma4-26b-mtp:
    <<: *common
    container_name: gemma4-26b-mtp
    restart: unless-stopped
    environment:
      PROFILE_NAME: gemma4-26b-mtp
      MODEL_ID: "${GEMMA26_MODEL:-google/gemma-4-26B-A4B-it}"
      SERVED_MODEL_NAME: "${GEMMA26_MTP_MODEL_NAME:-gemma4-26b-a4b-mtp}"
      PORT: "${GEMMA26_MTP_PORT:-8101}"
      ENABLE_MTP: "1"
      ASSISTANT_MODEL: "${GEMMA26_ASSISTANT_MODEL:-google/gemma-4-26B-A4B-it-assistant}"
      NUM_SPECULATIVE_TOKENS: "${GEMMA26_MTP_NUM_SPECULATIVE_TOKENS:-4}"
      EXTRA_ARGS_PROFILE: "${GEMMA26_MTP_EXTRA_ARGS:-}"
      NVIDIA_VISIBLE_DEVICES: all
      NVIDIA_DRIVER_CAPABILITIES: compute,utility

  gemma4-31b:
    <<: *common
    container_name: gemma4-31b
    restart: unless-stopped
    environment:
      PROFILE_NAME: gemma4-31b
      MODEL_ID: "${GEMMA31_MODEL:-google/gemma-4-31B-it}"
      SERVED_MODEL_NAME: "${GEMMA31_MODEL_NAME:-gemma4-31b}"
      PORT: "${GEMMA31_PORT:-8102}"
      ENABLE_MTP: "0"
      EXTRA_ARGS_PROFILE: "${GEMMA31_EXTRA_ARGS:-}"
      NVIDIA_VISIBLE_DEVICES: all
      NVIDIA_DRIVER_CAPABILITIES: compute,utility

  gemma4-31b-mtp:
    <<: *common
    container_name: gemma4-31b-mtp
    restart: "${GEMMA31_MTP_RESTART:-no}"
    mem_limit: "${GEMMA31_MTP_MEM_LIMIT:-96g}"
    memswap_limit: "${GEMMA31_MTP_MEMSWAP_LIMIT:-96g}"
    environment:
      PROFILE_NAME: gemma4-31b-mtp
      MODEL_ID: "${GEMMA31_MODEL:-google/gemma-4-31B-it}"
      SERVED_MODEL_NAME: "${GEMMA31_MTP_MODEL_NAME:-gemma4-31b-mtp}"
      PORT: "${GEMMA31_MTP_PORT:-8103}"
      ENABLE_MTP: "1"
      ASSISTANT_MODEL: "${GEMMA31_ASSISTANT_MODEL:-google/gemma-4-31B-it-assistant}"
      NUM_SPECULATIVE_TOKENS: "${GEMMA31_MTP_NUM_SPECULATIVE_TOKENS:-4}"
      MAX_MODEL_LEN: "${GEMMA31_MTP_MAX_MODEL_LEN:-8192}"
      MAX_NUM_BATCHED_TOKENS: "${GEMMA31_MTP_MAX_NUM_BATCHED_TOKENS:-12288}"
      MAX_NUM_SEQS: "${GEMMA31_MTP_MAX_NUM_SEQS:-1}"
      GPU_MEMORY_UTILIZATION: "${GEMMA31_MTP_GPU_MEMORY_UTILIZATION:-0.72}"
      LOAD_FORMAT: "${GEMMA31_MTP_LOAD_FORMAT:-safetensors}"
      ATTENTION_BACKEND: "${GEMMA31_MTP_ATTENTION_BACKEND:-auto}"
      EXTRA_ARGS_PROFILE: "${GEMMA31_MTP_EXTRA_ARGS:-}"
      NVIDIA_VISIBLE_DEVICES: all
      NVIDIA_DRIVER_CAPABILITIES: compute,utility