# GLM-5.2 QuantTrio Int4-Int8Mix — 4× DGX Spark / GB10 (sm_121), TP=4.
# MTP k=4, cudagraph FULL, prefix caching, fp8_ds_mla KV, 350k ctx.
#
# Model: QuantTrio/GLM-5.2-Int4-Int8Mix (mixed Int4/Int8, 10% expert-pruned)
#   → https://huggingface.co/QuantTrio/GLM-5.2-Int4-Int8Mix
# MTP draft: CosmicRaisins/GLM-5.2-MTP-INT4-aligned
#   → https://huggingface.co/CosmicRaisins/GLM-5.2-MTP-INT4-aligned
#
# BUILD THE IMAGE FIRST:
#   Reference: https://github.com/CosmicRaisins/glm-5.2-gb10
#   The custom image vllm-node-tf5-glm52-b12x:probe is built from eugr/spark-vllm-docker
#   fork at pinned ref ab666069935c1f23e8ef56038b4659ac9e8f19f8, then applies mods
#   glm52-sm12x-sparse and glm52-b12x-sparse on top. See the repo's bootstrap.sh
#   for the full build-and-copy workflow.
#
#   1) git clone https://github.com/eugr/spark-vllm-docker
#   2) cd spark-vllm-docker && ./build-and-copy.sh \
#        --vllm-ref ab666069935c1f23e8ef56038b4659ac9e8f19f8 \
#        -t vllm-node-tf5-glm52-b12x:probe --tf5
#   3) Bake mods glm52-sm12x-sparse + glm52-b12x-sparse into the image
#      (see Dockerfile.glm52-consolidated in the repo)
#
# PREREQUISITES:
#   - NCCL 2.30.4 staged on every node at /root/.cache/huggingface/hub/nccl-2.30.4/libnccl.so.2
#     (the image LD_PRELOADs it). Get it with:
#       pip install --no-deps -d /tmp/ncclwheel nvidia-nccl-cu13==2.30.4 && mkdir -p ... && unzip ...
#   - Weights cached on every node:
#       huggingface-cli download QuantTrio/GLM-5.2-Int4-Int8Mix
#       huggingface-cli download CosmicRaisins/GLM-5.2-MTP-INT4-aligned
#
# TUNING NOTES (from community):
#   - gpu-memory-utilization 0.92 is the sweet spot for 350k ctx on 4× GB10 with MTP
#   - MTP k=4 is more reliable on startup than k=5 on GB10
#   - cudagraph FULL requires b12x installed; without it, use PIECEWISE

engine = "vllm"
name = "GLM-5.2 Int4-Int8Mix (vLLM)"
docker_image = "vllm-node-tf5-glm52-b12x:probe"
model = "QuantTrio/GLM-5.2-Int4-Int8Mix"
port = 8210
required_nodes = 4
tensor_parallel_size = 4
pipeline_parallel_size = 1
distributed_backend = "mp"
shm_size = "64gb"

entrypoint = ["vllm", "serve"]

args = [
    # Reasoning & tools
    "--reasoning-parser", "glm45",
    "--tool-call-parser", "glm47",
    "--enable-auto-tool-choice",
    # KV cache & memory
    "--kv-cache-dtype", "fp8_ds_mla",
    "--max-model-len", "262144",
    "--max-num-seqs", "1",
    "--max-num-batched-tokens", "4096",
    "--gpu-memory-utilization", "0.92",
    # Compilation (b12x required for FULL cudagraph)
    "--compilation-config", """
'{"cudagraph_mode":"FULL"}'""",
    # MTP speculative decoding
    "--speculative-config", """
'{"model":"/root/.cache/huggingface/hub/glm52-mtp-int4-aligned","method":"mtp","num_speculative_tokens":4,"attention_backend":"FLASHMLA_SPARSE"}'""",
    # Other
    "--trust-remote-code",
    "--enable-prefix-caching",
    "--enable-flashinfer-autotune",
    "--served-model-name", "glm52",
]

[env]
# Timeout safety net for cold boot (cudagraph + JIT)
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS = "1800"
# NCCL shm_broadcast fix
LD_PRELOAD = "/root/.cache/huggingface/hub/nccl-2.30.4/libnccl.so.2"

# Offline mode — weights must be cached on each node
HF_HOME = "/root/.cache/huggingface"
TRITON_CACHE_DIR = "/root/.cache/huggingface/.tritoncache"
HF_HUB_OFFLINE = "1"
TRANSFORMERS_OFFLINE = "1"

# GLM-5.2 sparse-MLA suite (these guard the Triton fallback paths on sm_121)
VLLM_ALLOW_LONG_MAX_MODEL_LEN = "1"
VLLM_SPARSE_INDEXER_MAX_LOGITS_MB = "256"
GLM52_BIND_HOST_TRITON = "1"
GLM52_MQA_LOGITS_TRITON = "1"
GLM52_PAGED_MQA_TRITON = "1"
GLM52_PAGED_MQA_TOPK_CHUNK_SIZE = "8192"
GLM52_B12X_MLA = "1"
TORCH_CUDA_ARCH_LIST = "12.1a"

# NCCL / RoCE (launcher auto-sets most; recipe env wins last)
NCCL_NET = "IB"
NCCL_IB_DISABLE = "0"
NCCL_IB_GID_INDEX = "3"
NCCL_CROSS_NIC = "1"
NCCL_CUMEM_ENABLE = "0"
NCCL_IGNORE_CPU_AFFINITY = "1"
NCCL_DEBUG = "WARN"

# MoE & compute
OMP_NUM_THREADS = "8"

[mods]
commands = [
    # Copy Triton sparse-MLA kernel files into the vLLM tree inside the container.
    # These are mounted via the mods system at /mods/glm52-sm12x-sparse/ and copied
    # into the vLLM Python package paths before the engine starts.
    "cp /mods/glm52-sm12x-sparse/sparse_mla_kernels.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
    "cp /mods/glm52-sm12x-sparse/sparse_mla_env.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
    "cp /mods/glm52-sm12x-sparse/sm12x_sparse_mla_attn.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
    "cp /mods/glm52-sm12x-sparse/patch_flashmla_ops.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
    "cp /mods/glm52-sm12x-sparse/flashmla_sparse.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/",
    "cp /mods/glm52-sm12x-sparse/sm12x_deep_gemm_fallbacks.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/deepseek_v4_ops/",
    "cp /mods/glm52-sm12x-sparse/sm12x_mqa.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/deepseek_v4_ops/",
    "cp /mods/glm52-sm12x-sparse/b12x_sparse_helpers.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/ops/deepseek_v4_ops/",
    # upstream vLLM #46862: fused indexer Q rope+fp8-quant
    "cp /mods/glm52-sm12x-sparse/sparse_attn_indexer.py /usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/",
    "cp /mods/glm52-sm12x-sparse/deepseek_v2.py /usr/local/lib/python3.12/dist-packages/vllm/model_executor/models/",
]