# MiMo-V2.6-Flash-MOPD (Xiaomi, ~310B total / 12B-active MoE, fp8 attention +
# MXFP4 experts, 1M native context, image+video+audio input) on 2x DGX Spark
# with SGLang + DFlash speculative decoding.
#
# SGLang variant of forum thread 383968 (tonyd2wild's proven vLLM 2x recipe:
# 45.6/155.8 t/s, 1.87M-token fp8 KV pool, 300K ctx, code 70.5 t/s). The
# vLLM recipe needed 4 drop-in patches; the biggest (fused fp8 QKV stored
# pre-sharded for TP4) is already handled natively by SGLang's MiMo-V2
# loader (same approach as vllm#57508 — the provenance of the 8x Pro recipe
# working on this image family), so we run the stock a14 image.
#
# Image: docker-images/mimo26-spark (a14) — same build as the 8x Pro recipe:
# SGLang 06008c17 base + 12 patches (MiMo-V2 day-0, DFlash, RoCEnante, FP8
# tiles) + b12x. The FP8 tile tables are tuned for MiMo's per-rank shapes —
# at TP2 the shapes differ from TP8; stock fallbacks apply if a tile misses.
#
# Model: XiaomiMiMo/MiMo-V2.6-Flash-MOPD (178 GB, same architecture as the
# RL sibling — config.json is byte-identical; MOPD is a post-RL preference
# distillation of the checkpoint). Resolved from the shared NAS hub
# (config hf_cache_dir; both nodes' ~/.cache/huggingface point there). The
# dflash/ subfolder carries the DFlash draft weights — wired as a
# {{dflash}} path ref.
#
# KV math: ~89 GB weights/rank (178 GB / TP2), GMU 0.85 of 128 GB → the pool
# lands ~6 GB under the RL sibling's 362,560 tokens (measured there
# 2026-09-27); expect ~300–340K. Quality/perf of the RL lane
# (recipes/mimo26-flash-rl-2x/reports, 161/184 on v2.7.1.dev7) is the
# comparison baseline until this lane has its own bench.
# Config C (2026-09-28): GMU 0.90 + --context-length 524288 → full pool
# 570,752 tokens, the 512K target met (see the args note for the probe).

engine = "sglang"
name = "MiMo-V2.6-Flash-MOPD (SGLang DFlash, 2x Spark TP2)"
docker_image = "mimo26-spark:a15"
docker_image_dir = "mimo26-spark"
model = "XiaomiMiMo/MiMo-V2.6-Flash-MOPD"
parent_model = "XiaomiMiMo/MiMo-V2.6-Flash-MOPD"
port = 30000
required_nodes = 2
tensor_parallel_size = 2
distributed_backend = "mp"
shm_size = "16gb"

entrypoint = ["python3", "-m", "sglang.launch_server"]

args = [
    "--trust-remote-code",
    "--served-model-name", "mimo-v2.6-flash",
    # Attention / MoE backends (a13/a14 profile, same as the 8x Pro recipe)
    "--attention-backend", "triton",
    "--prefill-attention-backend", "fa4",
    "--mm-attention-backend", "triton_attn",
    "--enable-multimodal",
    "--moe-runner-backend", "marlin",
    "--moe-a2a-backend", "none",
    # Memory / scheduling (128 GB unified per node). Config C (2026-09-28,
    # VERIFIED): 0.88 → 0.90 for the 512K target — pool at 0.88 was 481,600,
    # short of 524,288 by ~43K tokens. Measured at 0.90: full pool 570,752
    # tokens (+ SWA 57,024), context_len 524288 accepted; a 515,001-token
    # completion prefilled + decoded clean (625 s ≈ 825 t/s), >524,288 inputs
    # rejected with the clean context-length error. Host free after the run
    # 7.0/8.1 GB per node, ~8 GB GPU headroom left — no OOM pressure.
    "--mem-fraction-static", "0.90",
    # Model is natively 1M; the measured pool at GMU 0.85 is 340,736 tokens
    # (2x TP2: 83 GB weights/rank leaves ~25 GB/rank; 39/48 SWA layers keep it
    # large in tokens). 320K fits a single long request with headroom; a
    # 524K request is pool-infeasible on 2x (the 8x Pro lane runs 512K).
    "--context-length", "524288",
    "--max-running-requests", "8",
    "--chunked-prefill-size", "8192",
    "--page-size", "64",
    # MiMo hybrid SWA attention
    "--swa-full-tokens-ratio", "0.1",
    # Reasoning / tool parsers
    "--reasoning-parser", "mimo",
    "--tool-call-parser", "mimo",
    # DFlash speculative decoding (draft weights ship inside the model repo)
    "--speculative-algorithm", "DFLASH",
    "--speculative-draft-model-path", "{{dflash}}",
    "--speculative-num-draft-tokens", "8",
    # CUDA graphs
    "--cuda-graph-max-bs-decode", "8",
    "--disable-prefill-cuda-graph",
    # Unified memory: page cache from weight loading competes with the KV pool
    "--weight-loader-drop-cache-after-load",
    # Observability
    "--enable-metrics",
    "--enable-cache-report",
    "--watchdog-timeout", "1800",
]

# DFlash draft weights: a subfolder of the main model repo, resolved per node
# (table header — all bare keys above).
[[extra_mounts]]
repo = "XiaomiMiMo/MiMo-V2.6-Flash-MOPD"
subfolder = "dflash"
container = "dflash"


# JIT/compile cache (verified against the live container, 2026-09-28): SGLang
# keeps everything under /root/.cache/sglang — triton .so (~35 hash dirs),
# inductor (aotautograd/ev/fxgraph), nv/, flashinfer .cache. Bind it to a
# node-local NVMe dir so boots reuse compiled artifacts instead of redoing
# JIT; keys are content/hash based, so mixing images/lanes on one dir is
# safe (and switching 2x<->4x reuses the other's cache). b12x-roce does not
# materialize on TP<8 — nothing else to mount.
[[extra_mounts]]
host = "~/.cache/mimo26-jit/sglang"
container = "/root/.cache/sglang"

[env]
# NOTE: SGLANG8_ROCE_ALLREDUCE is deliberately NOT set — the RoCEnante
# one-shot all-reduce overlay in the a14 image hard-requires the TP8 Gloo
# exchange group ("SG17 RoCEnante requires the TP8 Gloo exchange group"),
# so on TP2 we fall back to stock NCCL over the DAC link.
# a13/a14 tuning switches
MIMO26_FP8_OPROJ = "1"
MIMO26_FP8_DRAFT = "1"
MIMO26_DFLASH_ADAPTIVE = "1"
# NCCL extras from the reference deployment (HCAs/ifname from [nccl."2nodes-dac"])
NCCL_CUMEM_ENABLE = "0"
NCCL_NVLS_ENABLE = "0"
NCCL_PROTO = "^LL128"
NCCL_BUFFSIZE = "1048576"
NCCL_LL128_BUFFSIZE = "262144"
NCCL_MAX_NCHANNELS = "8"
NCCL_DEBUG = "WARN"