# DeepSeek-V4-Flash-DSpark — 2× GB10 / DGX Spark, TP=2 over RoCE.
# Uses the DSpark speculative decoding method with aidendle94's production-3.7 image.
#
# Image: aidendle94/sparkrun-vllm-ds4-gb10:production-3.7
#   Latest GB10-native image with DSpark support, b12x MoE, and multi-node fixes.
# Model: deepseek-ai/DeepSeek-V4-Flash-DSpark
#   The DSpark variant tuned for speculative decoding on DGX Spark.
#
# Reference: Nvidia Developer Forum — DeepSeek-V4-Flash-DSpark recipe
#   https://forums.developer.nvidia.com/t/recipe-deepseek-v4-flash-dspark-on-dgx-spark/340994

engine = "vllm"
name = "DeepSeek V4 Flash DSpark (vLLM)"
docker_image = "aidendle94/sparkrun-vllm-ds4-gb10:production-3.75"
model = "deepseek-ai/DeepSeek-V4-Flash-DSpark"
required_nodes = 2
tensor_parallel_size = 2
pipeline_parallel_size = 1
distributed_backend = "mp"
shm_size = "32gb"

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
    "--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",
    "--default-chat-template-kwargs.thinking=true",
    "--default-chat-template-kwargs.reasoning_effort=high",
    "--default-chat-template-kwargs.temperature=0.95",
    "--default-chat-template-kwargs.top_p=0.40",
    # Disable custom all-reduce (not needed on GB10 single-NIC)
    "--disable-custom-all-reduce",
    "--trust-remote-code",
]

[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 = "2"
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"