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

# --privileged makes --shm-size redundant; kept for documentation
# shm_size = "64gb"

entrypoint = ["vllm", "serve"]

args = [
    "--kv-cache-dtype", "fp8",
    "--block-size", "256",
    "--max-model-len", "524288",
    "--max-num-seqs", "6",
    "--max-num-batched-tokens", "8192",
    "--max-cudagraph-capture-size", "192",
    "--attention-backend", "B12X_MLA_SPARSE",
    "--moe-backend", "flashinfer_cutlass",
    "--linear-backend", "b12x",
    "--enable-chunked-prefill",
    "--gpu-memory-utilization", "0.75",
    "--compilation-config", """
'{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}'""",
    "--enable-prefix-caching",
    "--load-format", "safetensors",
    "--speculative-config", """
'{"method":"mtp","num_speculative_tokens":2}'""",
    "--tokenizer-mode", "deepseek_v4",
    "--tool-call-parser", "deepseek_v4",
    "--enable-auto-tool-choice",
    "--reasoning-parser", "deepseek_v4",
    "--reasoning-config", """
'{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}'""",
    "--default-chat-template-kwargs", """
'{"thinking":true,"preserve_thinking":true,"reasoning_effort":"high"}'""",
    "--override-generation-config", """
'{"temperature": 1.0, "top_p": 1.0}'""",
    "--enable-flashinfer-autotune",
    "--trust-remote-code",
    "--served-model-name", "deepseek-v4-flash", "default",
]

[env]
# Ray memory management — prevents OOM during cudagraph compilation
RAY_memory_monitor_refresh_ms = "0"
RAY_num_prestart_python_workers = "0"
RAY_object_store_memory = "1073741824"

# vLLM features for DeepSeek V4 on Blackwell (GB10 / sm_120)
VLLM_ALLOW_LONG_MAX_MODEL_LEN = "1"
# --- backend selection (the working config) ---
USES_B12X = "True"
VLLM_USE_B12X_MOE = "0"              # b12x MoE OFF -> cutlass W4A8 (b12x MoE prefill stalls on GB10)
MOE_BACKEND = "flashinfer_cutlass"   # W4A8 cutlass MoE
ATTN_BACKEND = "B12X_MLA_SPARSE"     # the only sparse-MLA that runs on GB10/sm_121a here
ASYNC_SCHED = "0"                    # async-scheduling crashes 2-node TP — KEEP OFF
# --- b12x suite = REQUIRED by this image (disabling any crashes engine init) ---
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 = "1"
VLLM_USE_BREAKABLE_CUDAGRAPH = "0"   # breakable mode regresses prefill on GB10 — keep off
VLLM_MEMORY_PROFILE_INCLUDE_ATTN = "1"
VLLM_PREFIX_CACHE_RETENTION_INTERVAL = "4096"
VLLM_CACHE_ROOT = "/cache/jit/vllm-cache"   # persist AOT/compile artifacts
B12X_MLA_SM120_UNIFIED = "1"
B12X_DENSE_SPLITK_TURBO = "1"
B12X_W4A16_TC_DECODE = "1"
B12X_MHC_MAX_TOKENS = "16384"
VLLM_SPARSE_INDEXER_MAX_LOGITS_MB = "256"

# Compile-time arch hints for PyTorch and FlashInfer on GB10
TORCH_CUDA_ARCH_LIST = "12.1a"
FLASHINFER_CUDA_ARCH_LIST = "12.1a"
CUTE_DSL_ARCH = "sm_121a"
OMP_NUM_THREADS = "4"

#CUDA_HOME = "/opt/env/targets/sbsa-linux"
FLASHINFER_DISABLE_VERSION_CHECK = "1"
TILELANG_CLEANUP_TEMP_FILES = "1"
DG_JIT_USE_NVRTC = "0"
DG_JIT_PRINT_COMPILER_COMMAND = "1"
#VLLM_NCCL_SO_PATH = "/opt/env/lib/python3.12/site-packages/nvidia/nccl/lib/libnccl.so.2"
HF_HUB_OFFLINE = "1"
TRANSFORMERS_OFFLINE = "1"