# GLM-5.2 NF3-Hybrid — 🧠 CONTEXT LANE (DCP4, 800K max-len)
# 4× DGX Spark (GB10, aarch64/sm_121a), TP=4, DCP=4.
#
# Model: madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid
#   → https://huggingface.co/madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid
#
# Unpruned GLM-5.2 (753B/40B MoE, all 256 experts) in 3-format hybrid quant:
#   327GB — top-64 experts/layer in NVFP4, remaining 192 in NF3 (3-bit),
#   non-expert layers in MXFP8.
#
# EXPERIMENT: Using aidendle94/sparkrun-vllm-ds4-gb10:production-hybrid-1.6
# instead of locally-built vllm-nf3-hybrid:probe. Pre-built on Docker Hub.
# NOTE: production-hybrid-1.6 lacks NF3 hybrid overlay files
# (hybrid_loader.py, nf3_kernel.py, nvfp4_kernel.py, mxfp8_tier.json).
# If it fails, revert to vllm-nf3-hybrid:probe.
#
# DCP4 lane: KV sharded across all 4 nodes → 877K-token pool.
#   Every decode step pays ag_rs allreduce on all 78 layers.
#   MTP k=3 (needs fuse_allreduce_rms) — near-perfect acceptance ~3.0
#   beat k=4/5 by ~8% on long network-taxed steps.
#
# Performance (measured 2026-07-07, 512-tok gens, temp 0, shallow ctx):
#   c1: 21.4–21.6 tok/s  |  c2: 32.9  |  c3: 38.9
#   c4: 45.4  |  c5: 37.4  |  c6: 40.1 aggregate
#   KV pool: 877,056 tokens  |  max-model-len: 800,000
#   Boot ~20 min, MTP k=3 acceptance ~3.0
#
# TUNING NOTES:
#   - gpu-memory-utilization 0.88 (GB10 unified memory; author's 0.96 is
#     dedicated-VRAM territory)
#   - max-num-batched-tokens ≥2048 or NF3 graph capture breaks
#   - max-cudagraph-capture-size must cover seqs×(1+k) = 4×4 = 16
#   - VLLM_DCP_SHARD_DRAFT=1 mandatory for MTP draft KV under DCP
#   - k is lane-specific: CONTEXT lane runs k=3 + fuse_allreduce_rms
#
# PREREQUISITES:
#   - Weights cached on every node:
#       huggingface-cli download madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid
#     (327GB, 334GB on disk, 96 safetensors shards)

engine = "vllm"
name = "GLM-5.2 NF3-Hybrid CONTEXT LANE (vLLM)"
docker_image = "aidendle94/sparkrun-vllm-ds4-gb10:production-hybrid-1.6"
model = "madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid"
required_nodes = 4
tensor_parallel_size = 4
pipeline_parallel_size = 1
distributed_backend = "mp"
# docker_image_dir = "glm52-nf3-hybrid"             # removed for experiment

entrypoint = ["vllm", "serve"]

args = [
    # Quantization & KV cache
    "--quantization", "modelopt_fp4",
    "--kv-cache-dtype", "fp8",
#    "--load-format", "auto",
    "--load-format", "fastsafetensors",
    # Context length (800K max-len, 877K KV pool)
    "--max-model-len", "200000",
    # Batch & memory
    "--max-num-seqs", "4",
    "--max-num-batched-tokens", "4096",
    "--max-cudagraph-capture-size", "32",
    "--gpu-memory-utilization", "0.84",
    # Engine backends
    #"--attention-backend", "B12X_MLA_SPARSE",
    "--moe-backend", "b12x",
    # Scheduling
    "--async-scheduling",
    "--enable-chunked-prefill",
    "--enable-prefix-caching",
    # Compilation (CONTEXT: fuse_allreduce_rms for DCP network tax)
    "--compilation-config", """
'{"pass_config":{"fuse_allreduce_rms":true}}'""",
    # Speculative decode — CONTEXT: MTP k=3 (near-perfect acceptance ~3.0)
    "--speculative-config", """
'{"method":"mtp","num_speculative_tokens":3,"moe_backend":"b12x","draft_sample_method":"probabilistic"}'""",
    # Tool calling & reasoning (verified working from first boot)
    "--tool-call-parser", "glm47",
    "--reasoning-parser", "glm45",
    "--enable-auto-tool-choice",
    # General
    "--trust-remote-code",
    "--enable-flashinfer-autotune",
    "--served-model-name", "glm-5.2",
]

[env]
# Memory stability
#PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:True"
#VLLM_ALLOW_LONG_MAX_MODEL_LEN = "1"

# ── DCP (Decode Context Parallelism) — CONTEXT: DCP4 ──────────────────────
# KV sharded across all 4 nodes — 4× pool, decode pays ag_rs gather
DECODE_CONTEXT_PARALLEL_SIZE = "4"
DCP_COMM_BACKEND = "ag_rs"
DCP_KV_CACHE_INTERLEAVE_SIZE = "1"

# ── b12x suite (required) ─────────────────────────────────────────────────
VLLM_USE_V2_MODEL_RUNNER = "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"
B12X_W4A16_TC_DECODE = "1"
B12X_MOE_FORCE_A16 = "1"
B12X_DENSE_SPLITK_TURBO = "1"

# ── Hybrid loader knobs (v2) ──────────────────────────────────────────────
HYBRID_TIER = "both"
HYBRID_KEPT = "b12x_nf3"
HYBRID_NF3 = "b12x_nf3"
HYBRID_B12X_MAX_TOKENS = "4096"
HYBRID_MXFP8_NATIVE = "1"
# Real path to tier json inside the image (loader hardcodes /opt/venv layout)
HYBRID_MXFP8_TIER_JSON = "/usr/local/lib/python3.12/dist-packages/mxfp8_tier.json"

# ── Architecture (aarch64/sm_121a DGX Spark) ──────────────────────────────
TORCH_CUDA_ARCH_LIST = "12.1a"
CUTE_DSL_ARCH = "sm_121a"
FLASHINFER_CUDA_ARCH_LIST = "12.1a"

# ── Offline mode (weights must be cached on each node) ────────────────────
HF_HOME = "/root/.cache/huggingface"
HF_HUB_OFFLINE = "1"
TRANSFORMERS_OFFLINE = "1"

# ── Timeout safety net (cold boot: cudagraph + JIT) ───────────────────────
VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS = "3600"

# ── NCCL / RoCE (cluster config ib_hca/socket_ifname used for most) ──────
NCCL_CROSS_NIC = "1"
NCCL_CUMEM_ENABLE = "0"
NCCL_IGNORE_CPU_AFFINITY = "1"
NCCL_DEBUG = "WARN"

# ── Compute ────────────────────────────────────────────────────────────────
OMP_NUM_THREADS = "4"