# SPDX-License-Identifier: MIT
name: dspark-vllm-gx10

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
      - ${DSPARK_MODEL_HOST:-${HOME}/models/dsv4-flash-dspark-abliterated}:/models/dsv4-abliterated:ro
      # A boot-time selector bind-mounts XFlash here when available; otherwise
      # this remains a root-filesystem directory, so /tmp always has a safe fallback.
      - ${DSPARK_TMP_HOST:-/var/lib/dspark-tmp}:/tmp

    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}"
      VLLM_CACHE_ROOT: /cache/huggingface/vllm-cache
      # Long-context crash fix kill-switch (see docs/LONG_CONTEXT_CRASH_FIX.md).
      # 1 (default) = clamp stale DSpark draft-KV slot ids; 0 = detect+log only.
      DSPARK_SLOT_CLAMP: "${DSPARK_SLOT_CLAMP:-1}"
      VLLM_HOST_IP: "${VLLM_HOST_IP:-}"

      VLLM_ALLOW_LONG_MAX_MODEL_LEN: "${VLLM_ALLOW_LONG_MAX_MODEL_LEN:-1}"
      VLLM_TRITON_MLA_SPARSE: "${VLLM_TRITON_MLA_SPARSE:-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}"
      VLLM_SKIP_INIT_MEMORY_CHECK: "${VLLM_SKIP_INIT_MEMORY_CHECK:-1}"

      # FlashInfer sampler is part of the 2026-07-03 DSpark garble fix.
      VLLM_USE_FLASHINFER_SAMPLER: "${VLLM_USE_FLASHINFER_SAMPLER:-1}"

      VLLM_USE_B12X_MOE: "${VLLM_USE_B12X_MOE:-1}"
      VLLM_USE_B12X_WO_PROJECTION: "${VLLM_USE_B12X_WO_PROJECTION:-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:-}"
      B12X_W4A16_TC_DECODE: "${B12X_W4A16_TC_DECODE:-0}"

      VLLM_DSPARK_CONFIDENCE_THRESHOLD: "${VLLM_DSPARK_CONFIDENCE_THRESHOLD:-0.0}"
      VLLM_DSPARK_CONFIDENCE_SCHEDULER: "${VLLM_DSPARK_CONFIDENCE_SCHEDULER:-off}"
      VLLM_DSPARK_LOCAL_ARGMAX: "${VLLM_DSPARK_LOCAL_ARGMAX:-1}"
      VLLM_DSPARK_REPLICATE_MARKOV_W1: "${VLLM_DSPARK_REPLICATE_MARKOV_W1:-1}"
      VLLM_DSPARK_FUSED_MARKOV_ARGMAX: "${VLLM_DSPARK_FUSED_MARKOV_ARGMAX:-0}"
      VLLM_DSPARK_GPU_REJECTED_CONTEXT_MASK: "${VLLM_DSPARK_GPU_REJECTED_CONTEXT_MASK:-1}"
      VLLM_DSPARK_REFERENCE_KV_QUANT_DEQUANT: "${VLLM_DSPARK_REFERENCE_KV_QUANT_DEQUANT:-0}"
      VLLM_DSPARK_HARDWARE_SCHEDULER_EARLY_STOP: "${VLLM_DSPARK_HARDWARE_SCHEDULER_EARLY_STOP:-1}"
      VLLM_DSV4_B12X_COMPRESSED_MLA: "${VLLM_DSV4_B12X_COMPRESSED_MLA:-0}"
      VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE: "${VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE:-0}"
      VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT: "${VLLM_DSV4_DSPARK_DEFER_TARGET_CAPTURE_EXACT:-0}"

      TORCH_CUDA_ARCH_LIST: "${TORCH_CUDA_ARCH_LIST:-12.1a}"
      FLASHINFER_CUDA_ARCH_LIST: "${FLASHINFER_CUDA_ARCH_LIST:-12.1a}"
      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:-enp1s0f0np0}"
      GLOO_SOCKET_IFNAME: "${GLOO_SOCKET_IFNAME:-enp1s0f0np0}"
      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_IB_ADDR_RANGE: "${NCCL_IB_ADDR_RANGE:-}"
      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}"

      NODE_RANK: "${NODE_RANK}"
      HEADLESS: "${HEADLESS:-}"
      MASTER_ADDR: "${MASTER_ADDR}"
      MASTER_PORT: "${MASTER_PORT:-25000}"
      MTP_NUM_TOKENS: "${MTP_NUM_TOKENS:-3}"

    command:
      - bash
      - -lc
      - >
        export PATH="/usr/local/cuda/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:-}";
        SPECULATIVE_CONFIG="{\"method\":\"dspark\",\"num_speculative_tokens\":${MTP_NUM_TOKENS:-3},\"draft_sample_method\":\"probabilistic\"}";
        exec /usr/local/bin/vllm serve ${DSPARK_MODEL:-/models/dsv4-abliterated}
        --served-model-name ${SERVED_MODEL_NAME:-deepseek-v4-flash-dspark-abliterated}
        --host ${VLLM_HOST:-127.0.0.1}
        --port ${VLLM_PORT:-8888}
        --trust-remote-code
        --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:-6}
        --max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS:-8192}
        --max-cudagraph-capture-size $$(( ${MAX_NUM_SEQS:-6} * (${MTP_NUM_TOKENS:-3} + 1) ))
        --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION:-0.80}
        --enable-prefix-caching
        --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}
        --jit-monitor-mode ${JIT_MONITOR_MODE:-warn}
        ${JIT_MONITOR_VERBOSE:+--jit-monitor-verbose}
        ${HEADLESS:+--headless}