Browse Source

feat(deepseek-v4): add eugr b12x recipe variant with baked draft-loader mod

New recipe variant based on NVIDIA forum comment #28: eugr's prebuilt
eugr/spark-vllm-b12x image + the b12x-branch deepseek-v4-flash-0731.yaml.

- recipes/deepseek-v4-flash-eugr/recipe.toml: faithful translation of
  eugr's recipe (B12X_MLA_SPARSE attention + linear-backend b12x,
  --load-format instanttensor, dspark k=5 probabilistic, reasoning-config
  pinned to deepseek_v4 parser, eugr env set verbatim). No [mods] block.
- docker-images/deepseek-v4-flash-eugr/: FROM overlay that bakes the
  instanttensor-hybrid-draft-loader mod into the image at build time
  (vendored patcher runs check->apply->check under set -eux; refuses
  incompatible vLLM, so a future image bump fails the build loudly).
  No versions.env — prebuilt base, like deepseek-v4-flash-dspark.

Build verified on 2nodes-dac/dgx1.lan (5m07s): image stamped with checksum
label, patcher confirmed applied to vLLM model_loader/__init__.py, dry-run
serve sees the image as fresh (no STALE rebuild).

Ref: https://forums.developer.nvidia.com/t/376220/28
Alexander Korolev 2 months ago
parent
commit
a8e4932fc4

+ 43 - 0
docker-images/deepseek-v4-flash-eugr/Dockerfile

@@ -0,0 +1,43 @@
+# syntax=docker/dockerfile:1.7
+# DeepSeek-V4-Flash-0731 — eugr/spark-vllm-b12x overlay with the
+# instanttensor-hybrid-draft-loader mod baked in.
+#
+# Base: eugr/spark-vllm-b12x — the prebuilt b12x-branch image published in
+# NVIDIA forum comment #28 (https://forums.developer.nvidia.com/t/376220/28).
+# Ubuntu 24.04 + system Python 3.12; vLLM is installed into the system
+# site-packages (the mod's default /usr/local/lib/python3.12/dist-packages,
+# but the path is resolved dynamically below so this is robust to a moved
+# install).
+#
+# eugr's recipe applies `mods/instanttensor-hybrid-draft-loader` at
+# run-recipe.sh time (it patches vLLM's model_loader get_model() to keep
+# InstantTensor for the target while letting an embedded MTP/dspark draft
+# use lazy safetensors). sm121 has no runtime mod-upload step for Docker
+# recipes, so this overlay bakes the patch in at build time instead: the
+# running container already has the hybrid loader, and the recipe carries
+# no [mods] block.
+#
+# Build safety: the patcher refuses to patch an incompatible module (it
+# anchors on vLLM's exact get_model signature and self-compiles the result).
+# If a future eugr image bumps vLLM and that signature changes, `docker
+# build` fails here instead of silently shipping an unpatched image.
+#
+# Build: cd docker-images/deepseek-v4-flash-eugr && bash build.sh
+# (the sm121 launcher builds on demand and stamps sm121.checksum)
+
+FROM eugr/spark-vllm-b12x
+
+# Vendor the mod's patcher (vendored from eugr/spark-vllm-docker b12x branch,
+# mods/instanttensor-hybrid-draft-loader/patch_model_loader.py) so the build is
+# self-contained and the file is folded into the sm121 build checksum.
+COPY patches/patch_model_loader.py /tmp/patch_model_loader.py
+
+# Resolve the model_loader module path from Python itself (do not hardcode the
+# dist-packages prefix), then check -> apply -> check, matching the mod's run.sh.
+RUN set -eux; \
+    TARGET="$(python3 -c 'import vllm, os; print(os.path.join(vllm.__path__[0], "model_executor/model_loader/__init__.py"))')"; \
+    python3 /tmp/patch_model_loader.py --check "$TARGET"; \
+    python3 /tmp/patch_model_loader.py "$TARGET"; \
+    python3 /tmp/patch_model_loader.py --check "$TARGET"; \
+    find "$(dirname "$TARGET")" -name '__pycache__' -type d -prune -exec rm -rf {} + 2>/dev/null || true; \
+    rm -f /tmp/patch_model_loader.py

+ 10 - 0
docker-images/deepseek-v4-flash-eugr/build.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+# Build for the deepseek-v4-flash-eugr overlay (eugr/spark-vllm-b12x + baked
+# instanttensor-hybrid-draft-loader mod). No versions.env — this is a FROM
+# overlay on a prebuilt public image, like deepseek-v4-flash-dspark.
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-eugr/spark-vllm-b12x:sm121}"
+echo "=== Building ${TAG} ==="
+docker build --progress=plain -t "${TAG}" . 2>&1
+echo "=== Done: ${TAG} ==="

+ 194 - 0
docker-images/deepseek-v4-flash-eugr/patches/patch_model_loader.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""Add an opt-in hybrid loader policy for speculative draft models."""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import sys
+from pathlib import Path
+
+MARKER = "# spark-vllm mod: instanttensor-hybrid-draft-loader v1"
+
+GET_MODEL_ANCHOR = """
+def get_model(
+    *,
+    vllm_config: VllmConfig,
+    model_config: ModelConfig | None = None,
+    prefix: str = "",
+    load_config: LoadConfig | None = None,
+) -> nn.Module:
+    loader = get_model_loader(load_config or vllm_config.load_config)
+    if model_config is None:
+        model_config = vllm_config.model_config
+    return loader.load_model(
+        vllm_config=vllm_config, model_config=model_config, prefix=prefix
+    )
+"""
+
+GET_MODEL_REPLACEMENT = f"""
+{MARKER}
+def _instanttensor_draft_load_config(
+    vllm_config: VllmConfig,
+    model_config: ModelConfig,
+    load_config: LoadConfig | None,
+) -> LoadConfig:
+    \"\"\"Resolve the loader for one model without changing the target loader.\"\"\"
+    import os
+
+    from vllm.config import replace
+
+    mode = os.environ.get("INSTANTTENSOR_DRAFT_LOADER", "auto").strip().lower()
+    allowed_modes = ("auto", "safetensors", "instanttensor")
+    if mode not in allowed_modes:
+        raise ValueError(
+            "INSTANTTENSOR_DRAFT_LOADER must be one of "
+            f"{{', '.join(allowed_modes)}}; got {{mode!r}}"
+        )
+
+    effective = load_config or vllm_config.load_config
+    load_format = getattr(effective.load_format, "value", effective.load_format)
+    if mode == "instanttensor" or str(load_format).lower() != "instanttensor":
+        return effective
+
+    speculative_config = getattr(vllm_config, "speculative_config", None)
+    draft_model_config = getattr(speculative_config, "draft_model_config", None)
+    if draft_model_config is None or model_config is not draft_model_config:
+        return effective
+
+    if mode == "auto":
+        target_model_config = (
+            getattr(speculative_config, "target_model_config", None)
+            or vllm_config.model_config
+        )
+        draft_source = (
+            getattr(draft_model_config, "model", None),
+            getattr(draft_model_config, "revision", None),
+        )
+        target_source = (
+            getattr(target_model_config, "model", None),
+            getattr(target_model_config, "revision", None),
+        )
+        if draft_source != target_source:
+            return effective
+
+    logger.info_once(
+        "Hybrid draft loading: using lazy safetensors for speculative draft "
+        "weights while preserving InstantTensor for the target model "
+        "(INSTANTTENSOR_DRAFT_LOADER=%s).",
+        mode,
+    )
+    return replace(
+        effective,
+        load_format="safetensors",
+        safetensors_load_strategy="lazy",
+    )
+
+
+def get_model(
+    *,
+    vllm_config: VllmConfig,
+    model_config: ModelConfig | None = None,
+    prefix: str = "",
+    load_config: LoadConfig | None = None,
+) -> nn.Module:
+    if model_config is None:
+        model_config = vllm_config.model_config
+    resolved_load_config = _instanttensor_draft_load_config(
+        vllm_config, model_config, load_config
+    )
+    loader = get_model_loader(resolved_load_config)
+    return loader.load_model(
+        vllm_config=vllm_config, model_config=model_config, prefix=prefix
+    )
+"""
+
+
+def validate_shape(text: str, *, patched: bool) -> None:
+    tree = ast.parse(text)
+    functions = {
+        node.name
+        for node in tree.body
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
+    }
+    required = {"get_model", "get_model_loader"}
+    if patched:
+        required.add("_instanttensor_draft_load_config")
+    missing = sorted(required - functions)
+    if missing:
+        raise ValueError(
+            "unexpected model-loader module; missing " + ", ".join(missing)
+        )
+
+
+def patched_text(text: str) -> str:
+    validate_shape(text, patched=MARKER in text)
+    if MARKER in text:
+        if text.count(MARKER) != 1:
+            raise ValueError("hybrid-loader marker occurs more than once")
+        if "resolved_load_config = _instanttensor_draft_load_config(" not in text:
+            raise ValueError("mod marker exists but get_model does not use the helper")
+        compile(text, "<patched model_loader/__init__.py>", "exec")
+        return text
+
+    count = text.count(GET_MODEL_ANCHOR)
+    if count != 1:
+        raise ValueError(
+            f"expected exactly one supported get_model function; found {count}"
+        )
+    patched = text.replace(GET_MODEL_ANCHOR, GET_MODEL_REPLACEMENT, 1)
+    validate_shape(patched, patched=True)
+    compile(patched, "<patched model_loader/__init__.py>", "exec")
+    return patched
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("target", type=Path)
+    parser.add_argument(
+        "--check", action="store_true", help="validate compatibility without writing"
+    )
+    args = parser.parse_args()
+
+    if not args.target.is_file():
+        print(
+            "[instanttensor-hybrid-draft-loader ERROR] "
+            f"target not found: {args.target}",
+            file=sys.stderr,
+        )
+        return 1
+
+    original = args.target.read_text()
+    try:
+        patched = patched_text(original)
+    except (SyntaxError, ValueError) as exc:
+        print(
+            "[instanttensor-hybrid-draft-loader ERROR] "
+            f"refusing to patch {args.target}: {exc}",
+            file=sys.stderr,
+        )
+        return 1
+
+    if args.check:
+        state = "already patched" if patched == original else "compatible"
+        print(f"[instanttensor-hybrid-draft-loader] {args.target} is {state}.")
+        return 0
+
+    if patched == original:
+        print(
+            "[instanttensor-hybrid-draft-loader] "
+            "Model loader is already patched; skipping."
+        )
+        return 0
+
+    temporary = args.target.with_suffix(
+        args.target.suffix + ".instanttensor-draft-mod.tmp"
+    )
+    temporary.write_text(patched)
+    temporary.replace(args.target)
+    print(f"[instanttensor-hybrid-draft-loader] Patched {args.target}.")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

+ 137 - 0
recipes/deepseek-v4-flash-eugr/recipe.toml

@@ -0,0 +1,137 @@
+# DeepSeek-V4-Flash-0731 DSpark — 2× GB10 / DGX Spark, TP=2 over RoCE.
+#
+# This variant uses eugr's prebuilt b12x image (eugr/spark-vllm-b12x) and
+# mirrors eugr's recipes/deepseek-v4-flash-0731.yaml on the b12x branch — the
+# stack published in NVIDIA forum comment #28:
+#   https://forums.developer.nvidia.com/t/instructions-for-running-deepseek-v4-flash-with-dspark-using-eugrs-repo/376220/28
+# ("While I'm working on integrating b12x into the build/test pipeline and
+# merging into main, I pushed prebuilt container as eugr/spark-vllm-b12x …
+# this version supports 3x Sparks with TP=3! To use, just add --tp 3.")
+#
+# Image: eugr/spark-vllm-b12x:sm121 — the prebuilt base with the
+# instanttensor-hybrid-draft-loader mod baked in at build time (see
+# docker-images/deepseek-v4-flash-eugr). eugr's recipe applies that mod via
+# run-recipe.sh --apply-mod; sm121 has no runtime mod-upload step for Docker
+# recipes, so the overlay patches vLLM's model_loader get_model() at image
+# build time instead. The recipe therefore carries no [mods] block.
+#
+# The mod keeps InstantTensor for the primary model while letting the embedded
+# dspark/MTP draft use lazy safetensors (avoids a second InstantTensor
+# GPU-streaming pass over the checkpoint). It activates only when
+# --load-format instanttensor is set (which this recipe does).
+#
+# Model: deepseek-ai/DeepSeek-V4-Flash-0731 — official release; native
+# dspark_block_size = 5, so num_speculative_tokens = 5.
+#
+# Differences vs recipes/deepseek-v4-flash (aidendle94 production-3.75 image):
+#   - attention backend B12X_MLA_SPARSE (eugr) vs FLASHINFER_MLA_SPARSE_DSV4
+#   - --linear-backend b12x + --load-format instanttensor (the mod's trigger)
+#   - --reasoning-config pinned to deepseek_v4 parser w/ empty start/end strs
+#   - draft_sample_method probabilistic + draft attention_backend B12X_MLA_SPARSE
+#   - max_num_seqs=8 / batched=8192 / cudagraph=64 (eugr defaults)
+#   - VLLM_USE_MEGA_AOT_ARTIFACT="-1" (eugr) vs "0"
+#
+# 3× Sparks (TP=3): bump tensor_parallel_size + required_nodes to 3 (the eugr
+# image supports it per comment #28). k=5 is unchanged (model-native).
+
+engine = "vllm"
+name = "DeepSeek V4 Flash DSpark (eugr b12x)"
+docker_image = "eugr/spark-vllm-b12x:sm121"
+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-eugr"
+
+entrypoint = ["vllm", "serve"]
+
+args = [
+    "--kv-cache-dtype", "fp8",
+    "--block-size", "256",
+    "--load-format", "instanttensor",
+    "--max-model-len", "auto",
+    "--max-num-seqs", "8",
+    "--max-num-batched-tokens", "8192",
+    "--max-cudagraph-capture-size", "64",
+    "--gpu-memory-utilization", "0.85",
+    "--served-model-name", "deepseek-v4-flash",
+    # Engine backends — b12x across attention / MoE / linear
+    "--attention-backend", "B12X_MLA_SPARSE",
+    "--moe-backend", "b12x",
+    "--linear-backend", "b12x",
+    # Prefix caching
+    "--enable-prefix-caching",
+    # Compilation
+    "--compilation-config", """
+'{"cudagraph_mode":"FULL_AND_PIECEWISE","custom_ops":["all"]}'""",
+    # DSpark speculative decoding — k=5 (DeepSeek-V4-Flash native
+    # dspark_block_size). draft_sample_method per eugr recipe; draft attention
+    # backend B12X_MLA_SPARSE. The baked mod makes the embedded draft load via
+    # lazy safetensors while the target stays on InstantTensor.
+    "--speculative-config", """
+'{"method":"dspark","num_speculative_tokens":5,"draft_sample_method":"probabilistic","attention_backend":"B12X_MLA_SPARSE"}'""",
+    # Tokenizer / reasoning — reasoning-config pins the parser + empty
+    # reasoning markers (deepseek_v4 tokenizer mode emits them itself).
+    "--tokenizer-mode", "deepseek_v4",
+    "--tool-call-parser", "deepseek_v4",
+    "--reasoning-parser", "deepseek_v4",
+    "--reasoning-config", """
+'{"reasoning_parser":"deepseek_v4","reasoning_start_str":"","reasoning_end_str":""}'""",
+    "--enable-auto-tool-choice",
+    "--generation-config", "vllm",
+    "--default-chat-template-kwargs.thinking=true",
+    "--default-chat-template-kwargs.reasoning_effort=high",
+    "--disable-custom-all-reduce",
+    "--trust-remote-code",
+]
+
+# Persistent vLLM compile cache (speeds up restarts). Mounted at vLLM's
+# default cache root for root (/root/.cache/vllm); the HF cache is already
+# mounted at /root/.cache/huggingface by the launcher.
+[[extra_mounts]]
+host = "~/.cache/vllm-ds4-eugr"
+container = "/root/.cache/vllm"
+
+[env]
+# --- b12x + compile (from eugr's recipe env, verbatim) ---
+CUTE_DSL_ARCH = "sm_121a"
+VLLM_USE_AOT_COMPILE = "1"
+VLLM_USE_BREAKABLE_CUDAGRAPH = "0"
+VLLM_USE_MEGA_AOT_ARTIFACT = "-1"
+VLLM_MEMORY_PROFILE_INCLUDE_ATTN = "1"
+VLLM_USE_FLASHINFER_SAMPLER = "1"
+VLLM_USE_B12X_WO_PROJECTION = "1"
+VLLM_USE_B12X_MHC = "1"
+VLLM_USE_B12X_FP8_GEMM = "1"
+VLLM_USE_B12X_MOE = "1"
+VLLM_USE_B12X_SPARSE_INDEXER = "1"
+VLLM_USE_V2_MODEL_RUNNER = "1"
+B12X_MLA_SM120_UNIFIED = "1"
+B12X_MOE_FORCE_A8 = "1"
+
+# --- arch / memory / GPU isolation (operational on the sm121 cluster) ---
+TORCH_CUDA_ARCH_LIST = "12.1a"
+PYTORCH_CUDA_ALLOC_CONF = "expandable_segments:True"
+OMP_NUM_THREADS = "4"
+CUDA_VISIBLE_DEVICES = "0"
+CUDA_DEVICE_ORDER = "PCI_BUS_ID"
+CUDA_CACHE_MAXSIZE = "4294967296"
+
+# --- offline ---
+HF_HUB_OFFLINE = "1"
+
+# --- NCCL overrides (GB10 single-NIC tuning; image-agnostic — mirrors
+#     recipes/deepseek-v4-flash. The launcher still injects IB/HCA/socket
+#     vars from the cluster config; these are recipe-level extras.) ---
+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"