# GLM-5.2 QuantTrio Int4-Int8Mix — 4× DGX Spark / GB10 (sm_121), TP=4.
#
# EXPERIMENT: Using aidendle94/sparkrun-vllm-ds4-gb10:production-hybrid-1.6
# instead of locally-built sparkrun-vllm-ds4-gb10:glm52.
# Pre-built on Docker Hub with CUDA 13.2.1, custom NCCL, b12x suite.
#
# Model: QuantTrio/GLM-5.2-Int4-Int8Mix (mixed Int4/Int8, 10% expert-pruned)
#   → https://huggingface.co/QuantTrio/GLM-5.2-Int4-Int8Mix
#   Includes in-checkpoint MTP draft model.
#
# PREREQUISITES:
#   - Weights cached on every node:
#       huggingface-cli download QuantTrio/GLM-5.2-Int4-Int8Mix
#
# REFERENCE:
#   https://forums.developer.nvidia.com/t/glm-5-2-int4-int8-on-8x-gb10-1200-t-s-prefill-33-54-t-s-avg-decode-generic-coding-structured/376831
#   https://github.com/ciprianveg/gb10-glm-5.2

engine = "vllm"
name = "GLM-5.2 Int4-Int8Mix (vLLM, DS4 image)"
docker_image = "aidendle94/sparkrun-vllm-ds4-gb10:production-hybrid-1.6"
model = "QuantTrio/GLM-5.2-Int4-Int8Mix"
parent_model = "zai-org/GLM-5.2"
required_nodes = 4
tensor_parallel_size = 4
pipeline_parallel_size = 1
distributed_backend = "mp"
shm_size = "64gb"
# docker_image_dir = "glm52"             # removed for experiment

entrypoint = ["vllm", "serve"]

args = [
    # KV cache & memory
    "--kv-cache-dtype", "fp8",
    "--max-model-len", "150000",
    "--max-num-seqs", "1",
    "--max-num-batched-tokens", "4096",
    "--gpu-memory-utilization", "0.91",
    # Engine backends
    "--attention-backend", "B12X_MLA_SPARSE",
    "--moe-backend", "b12x",
    # Compilation
    "--compilation-config", """
'{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}'""",
    # Chunked prefill + prefix caching
    "--enable-chunked-prefill",
    "--enable-prefix-caching",
    "--enable-flashinfer-autotune",
    # MTP speculative decoding (in-checkpoint draft model)
#    "--speculative-config", """
#'{"method":"mtp","num_speculative_tokens":4,"moe_backend":"b12x","draft_sample_method":"probabilistic"}'""",
    # Tokenizer / reasoning (GLM-5.2 family)
    "--tool-call-parser", "glm47",
    "--reasoning-parser", "glm45",
    "--enable-auto-tool-choice",
    "--trust-remote-code",
    "--served-model-name", "glm52",
]

[env]
# b12x suite (image requires this)
VLLM_USE_B12X_MOE = "1"
VLLM_USE_B12X_MHC = "1"
VLLM_USE_B12X_WO_PROJECTION = "1"
VLLM_USE_B12X_FP8_GEMM = "1"
VLLM_USE_B12X_SPARSE_INDEXER = "1"
VLLM_USE_V2_MODEL_RUNNER = "1"
VLLM_USE_FLASHINFER_SAMPLER = "1"
VLLM_USE_AOT_COMPILE = "1"
VLLM_USE_MEGA_AOT_ARTIFACT = "0"
VLLM_USE_BREAKABLE_CUDAGRAPH = "0"
VLLM_MEMORY_PROFILE_INCLUDE_ATTN = "1"
B12X_MLA_SM120_UNIFIED = "1"
B12X_DENSE_SPLITK_TURBO = "1"
B12X_W4A16_TC_DECODE = "1"
B12X_MHC_MAX_TOKENS = "16384"

# GLM-5.2 support flags
VLLM_ALLOW_LONG_MAX_MODEL_LEN = "1"
VLLM_SPARSE_INDEXER_MAX_LOGITS_MB = "256"

# Timeout safety net for cold boot (cudagraph + JIT)
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS = "1800"

# Offline mode — weights must be cached on each node
HF_HUB_OFFLINE = "1"

# Architecture
TORCH_CUDA_ARCH_LIST = "12.1a"
CUTE_DSL_ARCH = "sm_121a"
OMP_NUM_THREADS = "4"

# GPU memory management
PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:True"

# NCCL overrides
NCCL_CROSS_NIC = "1"
NCCL_CUMEM_ENABLE = "0"
NCCL_IGNORE_CPU_AFFINITY = "1"
NCCL_DEBUG = "WARN"

# PCIe allreduce not needed on GB10 single-NIC
VLLM_ENABLE_PCIE_ALLREDUCE = "0"
VLLM_PCIE_ALLREDUCE_BACKEND = "cpp"

# GPU device isolation
CUDA_VISIBLE_DEVICES = "0"
CUDA_DEVICE_ORDER = "PCI_BUS_ID"
CUDA_CACHE_MAXSIZE = "4294967296"