runtime/docker-compose.yml

services:
  vllm-dspark:
    image: ${DSPARK_VLLM_IMAGE:-ghcr.io/anemll/dspark-vllm-gx10:0.1.1}
    pull_policy: if_not_present
    entrypoint: []
    network_mode: host
    ipc: host
    shm_size: "64gb"

    ulimits:
      memlock: -1
      stack: 67108864

    gpus: all

    devices:
      - /dev/infiniband:/dev/infiniband

    volumes:
      - ${HF_CACHE:-${HOME}/.cache/huggingface}:/cache/huggingface
      - ${MODEL_PATH:?MODEL_PATH must point to the verified 0731 checkpoint}:/model:ro
      # Optional Stage-C overlay: host proposer only when path exists on the
      # Anemll layout (vLLM ships DSpark under v1/worker/gpu/spec_decode).
      # GB10 hybrid plugin tree for ENABLE_VLLM_GB10_PATCH=1 only.
      - ${VLLM_GB10_PATCH_DIR:-./vllm_patch_gb10}:/opt/vllm-gb10-hybrid-nvfp4
      - ${DSPARK_TMP_HOST:-${HOME}/.cache/dspark-tmp}:/tmp

    # Env surface is split by runtime image:
    #   - This file: knobs registered (or non-VLLM_) on Anemll 0.1.1
    #   - docker-compose.stage-c.override.yml: Stage-C overlay kill-switches
    # See docs/ENVS.md before adding VLLM_* keys back to the default path.
    environment:
      HF_HOME: /cache/huggingface
      HF_HUB_OFFLINE: "${HF_HUB_OFFLINE:-0}"
      TRANSFORMERS_OFFLINE: "${TRANSFORMERS_OFFLINE:-0}"
      HF_HUB_DISABLE_XET: "${HF_HUB_DISABLE_XET:-1}"
      DSPARK_ENCODING_FILE: "${DSPARK_ENCODING_FILE:-}"
      VLLM_CACHE_ROOT: /cache/huggingface/vllm-cache
      VLLM_HOST_IP: "${VLLM_HOST_IP:-}"

      VLLM_ALLOW_LONG_MAX_MODEL_LEN: "${VLLM_ALLOW_LONG_MAX_MODEL_LEN:-1}"
      VLLM_SPARSE_INDEXER_MAX_LOGITS_MB: "${VLLM_SPARSE_INDEXER_MAX_LOGITS_MB:-256}"
      VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS: "${VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS:-0}"
      ENABLE_VLLM_GB10_PATCH: "${ENABLE_VLLM_GB10_PATCH:-0}"
      GB10_HYBRID_NVFP4_M_THRESHOLD: "${GB10_HYBRID_NVFP4_M_THRESHOLD:-128}"

      VLLM_USE_FLASHINFER_SAMPLER: "${VLLM_USE_FLASHINFER_SAMPLER:-1}"
      VLLM_USE_BREAKABLE_CUDAGRAPH: "${VLLM_USE_BREAKABLE_CUDAGRAPH:-0}"

      VLLM_USE_B12X_MOE: "${VLLM_USE_B12X_MOE:-1}"
      VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM: "${VLLM_B12X_W4A16_FORCE_BLOCKS_PER_SM:-0}"
      VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M: "${VLLM_B12X_W4A16_FORCE_BLOCKS_MAX_M:-16}"
      VLLM_B12X_W4A16_FORCE_TILE_CONFIG: "${VLLM_B12X_W4A16_FORCE_TILE_CONFIG:-}"

      TORCH_CUDA_ARCH_LIST: "${TORCH_CUDA_ARCH_LIST:-12.1a}"
      FLASHINFER_CUDA_ARCH_LIST: "${FLASHINFER_CUDA_ARCH_LIST:-12.1a}"
      # Targets CuTeDSL kernel compilation at GB10 (sm_121a). Without this,
      # b12x/CuTeDSL may JIT a slower W4A16 fused MoE path mid-inference
      # instead of GB10-native kernels — confirmed via live boot logs.
      # Not a VLLM_* registry key (no unknown-env warning).
      CUTE_DSL_ARCH: "${CUTE_DSL_ARCH:-sm_121a}"
      FLASHINFER_DISABLE_VERSION_CHECK: "${FLASHINFER_DISABLE_VERSION_CHECK:-1}"
      FLASHINFER_WORKSPACE_BASE: "${FLASHINFER_WORKSPACE_BASE:-/cache/huggingface/flashinfer}"
      TILELANG_CLEANUP_TEMP_FILES: "${TILELANG_CLEANUP_TEMP_FILES:-1}"
      DG_JIT_USE_NVRTC: "${DG_JIT_USE_NVRTC:-0}"
      DG_JIT_NVCC_COMPILER: "${DG_JIT_NVCC_COMPILER:-/usr/local/cuda/bin/nvcc}"
      PYTORCH_CUDA_ALLOC_CONF: "${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"

      NCCL_NET: "${NCCL_NET:-IB}"
      NCCL_IB_DISABLE: "${NCCL_IB_DISABLE:-0}"
      NCCL_IB_HCA: "${NCCL_IB_HCA}"
      NCCL_SOCKET_IFNAME: "${NCCL_SOCKET_IFNAME}"
      TP_SOCKET_IFNAME: "${TP_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}"
      GLOO_SOCKET_IFNAME: "${GLOO_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}"
      NCCL_IB_GID_INDEX: "${NCCL_IB_GID_INDEX:-}"
      NCCL_IB_ADDR_FAMILY: "${NCCL_IB_ADDR_FAMILY:-AF_INET}"
      NCCL_IB_ROCE_VERSION_NUM: "${NCCL_IB_ROCE_VERSION_NUM:-2}"
      NCCL_CROSS_NIC: "${NCCL_CROSS_NIC:-1}"
      NCCL_CUMEM_ENABLE: "${NCCL_CUMEM_ENABLE:-0}"
      NCCL_IGNORE_CPU_AFFINITY: "${NCCL_IGNORE_CPU_AFFINITY:-1}"
      NCCL_DEBUG: "${NCCL_DEBUG:-WARN}"
      NCCL_NVLS_ENABLE: "${NCCL_NVLS_ENABLE:-0}"
      UCX_MEM_MMAP_HOOK_MODE: "${UCX_MEM_MMAP_HOOK_MODE:-none}"
      UCX_RCACHE_MAX_UNRELEASED: "${UCX_RCACHE_MAX_UNRELEASED:-1024}"

      NODE_RANK: "${NODE_RANK}"
      HEADLESS: "${HEADLESS:-}"
      MASTER_ADDR: "${MASTER_ADDR}"
      MASTER_PORT: "${MASTER_PORT:-25000}"
      # Checkpoint dspark_block_size is 5; k<5 truncates draft blocks (silent on 0.25.2 image).
      MTP_NUM_TOKENS: "${MTP_NUM_TOKENS:-5}"

    command:
      - bash
      - -lc
      - >
        export PATH="/usr/local/cuda/bin:/usr/local/bin:$${PATH:-}";
        export CUDA_HOME="$${CUDA_HOME:-/usr/local/cuda}";
        export CUDA_PATH="$${CUDA_PATH:-$${CUDA_HOME}}";
        export CUDAToolkit_ROOT="$${CUDAToolkit_ROOT:-$${CUDA_HOME}}";
        export LD_LIBRARY_PATH="/usr/local/cuda/lib64:$${LD_LIBRARY_PATH:-}";
        ENCODING_SOURCE="$${DSPARK_ENCODING_FILE:-}";
        if [ -z "$${ENCODING_SOURCE}" ]; then for candidate in /cache/huggingface/hub/models--deepseek-ai--DeepSeek-V4-Flash-0731/snapshots/*/encoding/encoding_dsv4.py; do if [ -f "$${candidate}" ]; then ENCODING_SOURCE="$${candidate}"; break; fi; done; fi;
        if [ -z "$${ENCODING_SOURCE}" ] && [ -f /models/deepseek-ai/DeepSeek-V4-Flash-0731/encoding/encoding_dsv4.py ]; then ENCODING_SOURCE=/models/deepseek-ai/DeepSeek-V4-Flash-0731/encoding/encoding_dsv4.py; fi;
        if [ -n "$${ENCODING_SOURCE}" ] && [ -f "$${ENCODING_SOURCE}" ]; then cp "$${ENCODING_SOURCE}" /usr/local/lib/python3.12/dist-packages/vllm/tokenizers/deepseek_v4_encoding.py; python3 -c 'from pathlib import Path; p=Path("/usr/local/lib/python3.12/dist-packages/vllm/tokenizers/deepseek_v4.py"); s=p.read_text(); old="elif reasoning_effort in (\"max\", \"xhigh\"):\n                reasoning_effort = \"max\"\n            else:\n                reasoning_effort = \"high\""; new="elif reasoning_effort in (\"max\", \"xhigh\"):\n                reasoning_effort = \"max\"\n            elif reasoning_effort == \"high\":\n                reasoning_effort = \"high\"\n            else:\n                reasoning_effort = \"low\""; updated=s.replace(old,new); assert new in updated; p.write_text(updated)'; fi;
        VLLM_QUANTIZATION_ARGS="";
        if [ "$${ENABLE_VLLM_GB10_PATCH:-0}" = "1" ]; then
          python3 -m pip install -e /opt/vllm-gb10-hybrid-nvfp4 --no-deps;
          export VLLM_PLUGINS="$${VLLM_PLUGINS:-gb10_hybrid_nvfp4}";
          VLLM_QUANTIZATION_ARGS="--quantization modelopt_gb10_hybrid";
        fi;
        SPECULATIVE_CONFIG="{\"method\":\"dspark\",\"num_speculative_tokens\":$${MTP_NUM_TOKENS:-5},\"draft_sample_method\":\"probabilistic\"}";
        exec /usr/local/bin/vllm serve ${DSPARK_MODEL:-deepseek-ai/DeepSeek-V4-Flash-DSpark}
        --served-model-name ${SERVED_MODEL_NAME:-deepseek-v4-flash-dspark} deepseek-v4-flash DeepSeek-V4-Flash
        --host ${VLLM_HOST:-127.0.0.1}
        --port ${VLLM_PORT:-8888}
        --trust-remote-code
        $${VLLM_QUANTIZATION_ARGS}
        --tensor-parallel-size 2
        --pipeline-parallel-size 1
        --kv-cache-dtype nvfp4_ds_mla
        --block-size 256
        --max-model-len ${MAX_MODEL_LEN:-1048576}
        --max-num-seqs ${MAX_NUM_SEQS:-2}
        --max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS:-4096}
        --max-cudagraph-capture-size ${MAX_CUDAGRAPH_CAPTURE_SIZE:-8}
        --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION:-0.75}
        --enable-prefix-caching
        --enable-prompt-tokens-details
        --async-scheduling
        --enable-chunked-prefill
        --speculative-config "$${SPECULATIVE_CONFIG}"
        --tokenizer-mode deepseek_v4
        --distributed-executor-backend mp
        --moe-backend flashinfer_b12x
        --tool-call-parser deepseek_v4
        --enable-auto-tool-choice
        --reasoning-parser deepseek_v4
        --reasoning-config '{"reasoning_parser":"deepseek_v4","reasoning_start_str":"<think>","reasoning_end_str":"</think>"}'
        --default-chat-template-kwargs '{"thinking":false}'
        --generation-config vllm
        --enable-flashinfer-autotune
        --nnodes 2
        --node-rank ${NODE_RANK}
        --master-addr ${MASTER_ADDR}
        --master-port ${MASTER_PORT:-25000}
        ${HEADLESS:+--headless}