# GLM-5.2 NF3-Hybrid — 🏎️ FAST LANE (DCP1, 200K max-len)
# 4× DGX Spark (GB10, aarch64/sm_121a), TP=4, DCP=1.
#
# 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.
#
# DCP1 lane: each sequence's KV cache lives entirely on one node → zero
#   cross-node attention traffic per decode step. ~1/4 the KV pool of DCP4.
#   Best for interactive chat, agents, throughput-sensitive workloads.
#
# Performance (measured 2026-07-07, 512-tok gens, temp 0, MTP acceptance 3.3–4.0):
#   c1: 24–29 tok/s (content-dependent, peaks to 29.1 @ accept 4.02)
#   c2: 36.6  |  c3: 45.7  |  c4: 54.2  |  c5: 59.2  |  c6: 67.4
#   KV pool: 219,264 tokens  |  max-model-len: 200,000
#
# TUNING NOTES:
#   - gpu-memory-utilization 0.88 (GB10 unified memory; author's 0.96 is
#     dedicated-VRAM territory; 0.88 = safe for 200K + 12e9 KV allocation)
#   - kv-cache-memory-bytes 12e9 per node caps KV pool at 219K tokens
#   - max-num-batched-tokens ≥2048 or NF3 graph capture breaks
#   - max-cudagraph-capture-size must cover seqs×(1+k) = 6×5 = 30
#   - FAST lane uses MTP k=4, NO fuse_allreduce_rms (no benefit without
#     DCP network tax)
#   - k is lane-specific: FAST runs k=4, no fuse
#
# 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 FAST LANE (vLLM)"
docker_image = "aidendle94/sparkrun-vllm-ds4-gb10:production-hybrid-1.6"
model = "madeby561/GLM-5.2-MXFP8-NVFP4-NF3-Hybrid"
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-nf3-hybrid"             # removed for experiment

entrypoint = ["vllm", "serve"]

args = [
    # Quantization & KV cache
    "--quantization", "modelopt_fp4",
    "--kv-cache-dtype", "fp8",
    "--load-format", "fastsafetensors",
    # Context length (200K max-len, 219K KV pool)
    "--max-model-len", "200000",
    # Batch & memory — FAST lane: 6 concurrent, 12e9 KV bytes
    "--max-num-seqs", "6",
    "--max-num-batched-tokens", "4096",
    "--max-cudagraph-capture-size", "32",
    "--gpu-memory-utilization", "0.88",
    "--kv-cache-memory-bytes", "12000000000",
    # Engine backends
    "--attention-backend", "B12X_MLA_SPARSE",
    "--moe-backend", "b12x",
    # Scheduling
    "--async-scheduling",
    "--enable-chunked-prefill",
    "--enable-prefix-caching",
    # Speculative decode — FAST: MTP k=4 (accept 3.3–4.0, best aggregate)
    "--speculative-config", """
'{"method":"mtp","num_speculative_tokens":4,"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]
# ── DCP (Decode Context Parallelism) — FAST: DCP1 ────────────────────────
# KV lives locally — zero cross-node attention traffic, ~1/4 pool of DCP4
DECODE_CONTEXT_PARALLEL_SIZE = "1"
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"
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"