# DeepSeek-V4-Flash-DSpark — 2× GB10 / DGX Spark, TP=2 over RoCE.
# DSpark speculative decoding on the official DeepSeek-V4-Flash-0731 release
# via the local :fixed overlay on aidendle94's production-3.8 image (see
# docker-images/deepseek-v4-flash-dspark).
#
# Image: sparkrun-vllm-ds4-gb10:fixed
#   production-3.8 clean-overlay base (vLLM 0.21.1, b12x MoE, NCCL 2.30.4,
#   flashinfer 0.6.12, opt-in DSPARK tiny-A8 decode kernel) + patched
#   deepseek_v4 tokenizer (thinking_mode kwarg, anti-injection, strict JSON).
# Model: deepseek-ai/DeepSeek-V4-Flash-0731 — official release.
#   0731 (and the -DSpark preview) both have dspark_block_size = 5. This image's
#   speculator hard-rejects num_speculative_tokens > block_size
#   (v1/worker/gpu/spec_decode/dspark/speculator.py, "must be <= dspark_block_size"),
#   so k=7 (as suggested by the 0731 model card) FAILS to boot. k=5 is the
#   model's native draft block and the correct value here.
#
# Note: this image lineage (vLLM 0.2x with the upstream DSpark speculator) does
# NOT have the shared-expert draft-loader bug reported for the older fork
# (tonyd2wild Patch 4). Its loader already maps mtp.*.shared_experts.w1/w3 into
# gate_up_proj (dspark.py load_weights), so no patch is needed.
#
# Reference: Nvidia Developer Forum — DeepSeek-V4-Flash-0731 recipe
#   https://forums.developer.nvidia.com/t/deepseek-v4-flash-0731-dspark-1m-nvfp4-kv-2x-dgx-spark/378824

engine = "vllm"
name = "DeepSeek V4 Flash DSpark (vLLM)"
docker_image = "sparkrun-vllm-ds4-gb10:fixed"
model = "deepseek-ai/DeepSeek-V4-Flash-0731"
required_nodes = 2
tensor_parallel_size = 2
pipeline_parallel_size = 1
distributed_backend = "mp"
shm_size = "32gb"
docker_image_dir = "deepseek-v4-flash-dspark"

entrypoint = ["vllm", "serve"]

args = [
    "--kv-cache-dtype", "fp8",
    "--block-size", "256",
    "--load-format", "auto",
    "--max-model-len", "393216",
    "--max-num-seqs", "16",
    "--max-num-batched-tokens", "16536",
    "--max-cudagraph-capture-size", "16",
    "--gpu-memory-utilization", "0.9",
    "--served-model-name", "deepseek-v4-flash", "default",
    # Attention: vLLM 0.21 (production-3.8) renamed/removed the SM120 sparse-MLA
    # backend names (FLASHINFER_MLA_SPARSE_DSV4 no longer exists). On GB10 the
    # platform auto-selects SPARSE_MLA_SM120 for MLA; the b12x unified sparse-MLA
    # backend (same fp8_ds_mla envelope) is selected explicitly via
    # VLLM_ATTENTION_BACKEND=B12X_MLA_SPARSE (env below) to pair with the b12x
    # MoE + tiny-decode path.
    "--moe-backend", "b12x",
    # Chunked prefill + prefix caching
    "--enable-chunked-prefill",
    "--enable-prefix-caching",
    "--enable-flashinfer-autotune",
    # Compilation
    "--compilation-config", """
'{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}'""",
    # Async scheduling
    "--async-scheduling",
    "--no-scheduler-reserve-full-isl",
    "--speculative-config", """
'{"method":"dspark","num_speculative_tokens":7,"draft_sample_method":"greedy"}'""",
    # Tokenizer / reasoning
    "--tokenizer-mode", "deepseek_v4",
    "--tool-call-parser", "deepseek_v4",
    "--reasoning-parser", "deepseek_v4",
    "--enable-auto-tool-choice",
    "--generation-config", "vllm",
    # vLLM 0.21 (production-3.8) accepts default_chat_template_kwargs as a
    # single JSON arg (json.loads), not the per-key dotted form used on 3.75.
    "--default-chat-template-kwargs", """
'{"thinking":true,"reasoning_effort":"high","temperature":1,"top_p":0.95}'""",
    # Disable custom all-reduce (not needed on GB10 single-NIC)
    "--disable-custom-all-reduce",
    "--trust-remote-code",
]

# Persistent compile caches (speed up restarts — see guide §2/§7). The HF
# cache is already mounted at /root/.cache/huggingface by the launcher.
[[extra_mounts]]
host = "~/.cache/vllm-ds4-dspark"   # compiled vLLM/attention kernels
container = "/cache"

[[extra_mounts]]
host = "~/.cache/tilelang-ds4"      # DSpark speculative-decode kernels
container = "/root/.tilelang"

[env]
# DSpark speculative decoding
VLLM_DSPARK_REPLICATE_MARKOV_W1 = "1"
# b12x MoE (this image's primary path)
VLLM_USE_B12X_MOE = "1"
# VLLM_USE_B12X_MOE = "0"              # b12x MoE OFF -> cutlass W4A8 (b12x MoE prefill stalls on GB10)

# b12x suite = REQUIRED by this image
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"
B12X_MHC_MAX_TOKENS = "16384"

# vLLM 0.21 (production-3.8): b12x unified sparse-MLA attention for DSV4 on
# GB10 (selects B12xMLASparseBackend; supports fp8_ds_mla KV). No longer
# reachable via --attention-backend (FLASHINFER_MLA_SPARSE_DSV4 removed).
VLLM_ATTENTION_BACKEND = "B12X_MLA_SPARSE"

# DSPARK tiny-A8 decode kernel (production-3.8 /opt/tinydec): routes small-M
# MoE decode (M<=VLLM_B12X_TINY_MAX_M) through the custom w4a8mx Triton
# kernel; stock b12x path otherwise. sitecustomize hook reads these at start.
# Fail-open: any kernel error falls back to stock b12x automatically.
VLLM_B12X_TINY_DECODE = "1"
VLLM_B12X_TINY_MAX_M = "8"
VLLM_B12X_TINY_CHUNK = "8"

# GPU memory management
CUTE_DSL_ARCH = "sm_121a"
PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:True"
OMP_NUM_THREADS = "4"

# Offline mode
HF_HUB_OFFLINE = "1"

# Compile-time arch hints
TORCH_CUDA_ARCH_LIST = "12.1a"

# NCCL overrides (cluster config ib_hca/socket_ifname used for most; recipe wins last)
NCCL_P2P_LEVEL = "SYS"
NCCL_CUMEM_ENABLE = "0"
NCCL_IB_TIMEOUT = "24"
NCCL_IB_QPS_PER_CONNECTION = "8"
NCCL_BUFFERSIZE = "16777216"
NCCL_SOCKET_FAMILY = "inet"

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

# GPU device isolation (single GPU per node)
CUDA_VISIBLE_DEVICES = "0"
CUDA_DEVICE_ORDER = "PCI_BUS_ID"
CUDA_CACHE_MAXSIZE = "4294967296"