# GLM-5.2-NVFP4-REAP (168 experts / ~504B) — vLLM serving on 4× 96 GB Blackwell GPUs
# (e.g. RTX PRO 6000, SM120). Works for both the -504B (code-calib) and -504B-term variants.
#
# Best config (defaults below): DCP4 + MTP5 + global-topk + DCP-sharded draft + use_index_cache
#   → ~1.03M-token KV pool (A16-scale-fix image); 70-94 tok/s codegen, ~50 tok/s @256k (flat at depth),
#     256k prefill no-wedge, ~200 tok/s aggregate @ concurrency 4.
#
# Quickstart:
#   MODEL_DIR=/path/to/GLM-5.2-NVFP4-REAP-504B-term docker compose up -d
#   # serves an OpenAI API on :5001, model id "GLM-5.2"
#
# Tuning via env (no file edits):
#   max decode speed, short ctx :  DCP_SIZE=1 MTP=1 MAX_MODEL_LEN=125000
#   absolute max context        :  raise MAX_MODEL_LEN toward the KV pool (~1.03M @ DCP4)
#   no speculative decode       :  MTP=0
# Sampling: temperature 0.6, top_p 0.95, repetition_penalty 1.0 (do NOT exceed 1.0).
services:
  glm52-reap:
    # A16-scale-fix build: frees ~7 GiB/GPU vs the prior image (no stale dequant scales in VRAM) -> ~1.03M-token KV pool on TP4.
    image: ${IMAGE:-voipmonitor/vllm:glm52-dark-devotion-pr31-pr15-w4a16scale-vllm79f154c-b12xaecc88f-cu132-20260622}
    container_name: ${NAME:-glm52-reap}
    network_mode: host
    ipc: host
    shm_size: 32g
    runtime: nvidia
    gpus: all
    restart: "no"
    ulimits:
      memlock: -1
      stack: 67108864
    volumes:
      - ${MODEL_DIR:?set MODEL_DIR to the model directory}:/model:ro
      - ./cache:/cache
      - ./cache/huggingface:/root/.cache/huggingface
    environment:
      CUDA_VISIBLE_DEVICES: ${CUDA_VISIBLE_DEVICES:-0,1,2,3}
      CUDA_DEVICE_ORDER: PCI_BUS_ID
      CUDA_DEVICE_MAX_CONNECTIONS: "32"
      OMP_NUM_THREADS: "16"
      CUTE_DSL_ARCH: sm_120a
      NCCL_IB_DISABLE: "1"
      NCCL_P2P_LEVEL: SYS
      NCCL_PROTO: LL,LL128,Simple
      PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True
      SAFETENSORS_FAST_GPU: "1"
      VLLM_WORKER_MULTIPROC_METHOD: spawn
      VLLM_USE_AOT_COMPILE: "1"
      VLLM_USE_BREAKABLE_CUDAGRAPH: "0"
      VLLM_USE_MEGA_AOT_ARTIFACT: "1"
      VLLM_USE_FLASHINFER_SAMPLER: "1"
      VLLM_USE_B12X_FP8_GEMM: "1"
      VLLM_USE_B12X_MOE: "1"
      VLLM_USE_B12X_SPARSE_INDEXER: "1"
      VLLM_DCP_GLOBAL_TOPK: "1"
      VLLM_DCP_SHARD_DRAFT: "1"
      # JIT/compile-cache persistence -> artifacts land in the mounted ./cache and survive restarts = much faster reboots
      VLLM_CACHE_DIR: /cache/jit/vllm
      TRITON_CACHE_DIR: /cache/jit/triton
      TORCH_EXTENSIONS_DIR: /cache/jit/torch_extensions
      TORCHINDUCTOR_CACHE_DIR: /cache/jit/torchinductor
      FLASHINFER_WORKSPACE_BASE: /cache/jit/flashinfer
      XDG_CACHE_HOME: /cache/jit
      VLLM_USE_V2_MODEL_RUNNER: "1"
      VLLM_ENABLE_PCIE_ALLREDUCE: "1"
      VLLM_PCIE_ALLREDUCE_BACKEND: b12x
      VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE: 64KB
      USES_B12X: "True"
      B12X_DENSE_SPLITK_TURBO: "1"
      B12X_W4A16_TC_DECODE: "1"
      B12X_MOE_FORCE_A16: "1"
      MODEL: ${MODEL:-/model}
      MTP_MODEL: ${MTP_MODEL:-/model}
      SERVED_MODEL_NAME: ${SERVED_MODEL_NAME:-GLM-5.2}
      PORT: ${PORT:-5001}
      TP_SIZE: ${TP_SIZE:-4}
      DCP_SIZE: ${DCP_SIZE:-4}
      MTP: ${MTP:-1}
      GPU_MEMORY_UTILIZATION: ${GPU_MEMORY_UTILIZATION:-0.95}
      MAX_MODEL_LEN: ${MAX_MODEL_LEN:-300000}
      MAX_NUM_SEQS: ${MAX_NUM_SEQS:-8}
      MAX_NUM_BATCHED_TOKENS: ${MAX_NUM_BATCHED_TOKENS:-8192}
      NUM_SPECULATIVE_TOKENS: ${NUM_SPECULATIVE_TOKENS:-5}
      MAX_CUDAGRAPH_CAPTURE_SIZE: ${MAX_CUDAGRAPH_CAPTURE_SIZE:-}
      SPEC_CONFIG: ${SPEC_CONFIG:-}
      HF_OVERRIDES: '{"index_topk_pattern":"FFFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSSFSSS","use_index_cache":true}'
    entrypoint:
      - bash
      - -lc
      - |
        set -euo pipefail
        unset NCCL_GRAPH_FILE NCCL_GRAPH_DUMP_FILE VLLM_B12X_MLA_EXTEND_MAX_CHUNKS
        if [ -z "$${MAX_CUDAGRAPH_CAPTURE_SIZE:-}" ]; then
          if [ "$${MTP}" = "1" ]; then
            MAX_CUDAGRAPH_CAPTURE_SIZE=$$((MAX_NUM_SEQS * (NUM_SPECULATIVE_TOKENS + 1)))
          else
            MAX_CUDAGRAPH_CAPTURE_SIZE="$${MAX_NUM_SEQS}"
          fi
        fi
        SPEC_ARGS=()
        if [ "$${MTP}" = "1" ]; then
          SPEC_CONFIG="$${SPEC_CONFIG:-$$(printf '{"model":"%s","method":"mtp","num_speculative_tokens":%s,"moe_backend":"b12x","draft_sample_method":"probabilistic"}' "$${MTP_MODEL}" "$${NUM_SPECULATIVE_TOKENS}")}"
          SPEC_ARGS=(--speculative-config "$${SPEC_CONFIG}")
        fi
        echo "[glm52-reap] TP$${TP_SIZE} DCP$${DCP_SIZE} MTP=$${MTP} ctx=$${MAX_MODEL_LEN} util=$${GPU_MEMORY_UTILIZATION}"
        cd /
        exec /opt/venv/bin/python -m vllm.entrypoints.cli.main serve "$${MODEL}" \
          --served-model-name "$${SERVED_MODEL_NAME}" \
          --trust-remote-code \
          --host 0.0.0.0 \
          --port "$${PORT}" \
          --tensor-parallel-size "$${TP_SIZE}" \
          --pipeline-parallel-size 1 \
          --decode-context-parallel-size "$${DCP_SIZE}" \
          --dcp-comm-backend ag_rs \
          --dcp-kv-cache-interleave-size 1 \
          --enable-chunked-prefill \
          --enable-prefix-caching \
          --max-model-len "$${MAX_MODEL_LEN}" \
          --load-format fastsafetensors \
          --async-scheduling \
          -cc.pass_config.fuse_allreduce_rms=True \
          --gpu-memory-utilization "$${GPU_MEMORY_UTILIZATION}" \
          --max-num-batched-tokens "$${MAX_NUM_BATCHED_TOKENS}" \
          --max-num-seqs "$${MAX_NUM_SEQS}" \
          --max-cudagraph-capture-size "$${MAX_CUDAGRAPH_CAPTURE_SIZE}" \
          --quantization modelopt_fp4 \
          --attention-backend B12X_MLA_SPARSE \
          --moe-backend b12x \
          --kv-cache-dtype fp8 \
          --tool-call-parser glm47 \
          --enable-auto-tool-choice \
          --reasoning-parser glm45 \
          --hf-overrides "$${HF_OVERRIDES}" \
          "$${SPEC_ARGS[@]}"