# 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.75 image (see
# docker-images/deepseek-v4-flash-dspark).
#
# Image: sparkrun-vllm-ds4-gb10:fixed
#   production-3.75 base (DSpark, b12x MoE, patched NCCL 2.30.4, multi-node
#   fixes) + instanttensor for faster model loading.
# 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", "524288",
    "--max-num-seqs", "16",
    "--max-num-batched-tokens", "16536",
    "--max-cudagraph-capture-size", "16",
    "--gpu-memory-utilization", "0.84",
    "--served-model-name", "deepseek-v4-flash", "default",
    # Engine backends
    "--attention-backend", "FLASHINFER_MLA_SPARSE_DSV4",
    "--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",
    # DSpark speculative decoding — k=5 is DeepSeek-V4-Flash's native
    # dspark_block_size (0731 == preview). Do NOT raise to k=7 (0731 model-card
    # rec): this image hard-rejects k > block_size at boot. draft_sample_method
    # is accepted by the speculative-config parser (no-op on the DSpark path).
    "--speculative-config", """
'{"method":"dspark","num_speculative_tokens":5,"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",
    "--default-chat-template-kwargs.thinking=true",
    "--default-chat-template-kwargs.reasoning_effort=high",
    "--default-chat-template-kwargs.temperature=1",
    "--default-chat-template-kwargs.top_p=1",
    # 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"
VLLM_MEMORY_PROFILE_INCLUDE_ATTN = "1"
VLLM_PREFIX_CACHE_RETENTION_INTERVAL = "4096"
B12X_MLA_SM120_UNIFIED = "1"
B12X_DENSE_SPLITK_TURBO = "1"
B12X_W4A16_TC_DECODE = "1"
B12X_MHC_MAX_TOKENS = "16384"

# 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"