Переглянути джерело

docker-images: mimo26-spark-vllm

Alexander Korolev 1 тиждень тому
батько
коміт
4d470eb805

+ 43 - 0
docker-images/mimo26-spark-vllm/Dockerfile

@@ -0,0 +1,43 @@
+# syntax=docker/dockerfile:1.7
+# mimo26-spark-vllm: thin FROM wrapper over the r39 b12x vLLM image that
+# (1) teaches vLLM's eagle3 aux-hidden-state hook to resolve the interface
+# on the inner language model when the outer model is a multimodal wrapper
+# (MiMoV2OmniForCausalLM), and (2) makes the DFlash MXFP8 guards tolerate a
+# vLLM build without ModelOptMxFp8LinearMethod (r39 predates that class).
+# No kernel changes.
+#
+# Why (2026-09-28): MiMo-V2.6-Flash-MOPD resolves to the Omni architecture
+# because the checkpoint ships visual/audio encoders. For method=dflash the
+# GPU model runner flips use_aux_hidden_state_outputs=True and then
+# set_eagle3_aux_hidden_state_layers() isinstance-checks SupportsEagle3 on
+# the OUTER wrapper — which never implements it — and dies with "Model does
+# not support EAGLE3 interface" right after weight loading. Everything that
+# actually carries aux state (MiMoV2Model + EagleModelMixin) is one level in,
+# and _inner_decoder() already knows how to find it.
+#
+# The patcher is anchor-based and refuses to run against unexpected sources
+# (see patches/eagle3-omni-inner.py), so the build FAILS LOUDLY if the base
+# vLLM ever restructures that function instead of silently no-oping.
+
+ARG BASE_IMAGE
+FROM ${BASE_IMAGE}
+
+COPY patches/eagle3-omni-inner.py /opt/sm121/patches/eagle3-omni-inner.py
+COPY patches/dflash-mxfp8-import-optional.py /opt/sm121/patches/dflash-mxfp8-import-optional.py
+
+RUN set -eux; \
+    python3 /opt/sm121/patches/eagle3-omni-inner.py --apply; \
+    # Fail the build if the patched module no longer imports or the hook
+    # regressed to raising on a wrapper-like model.
+    python3 /opt/sm121/patches/eagle3-omni-inner.py --check; \
+    python3 -c "import vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils as m; \
+assert 'get_language_model' in open(m.__file__).read(), 'patch vanished'"; \
+    # DFlash MXFP8 guards must tolerate a vLLM without ModelOptMxFp8LinearMethod.
+    python3 /opt/sm121/patches/dflash-mxfp8-import-optional.py --apply; \
+    python3 /opt/sm121/patches/dflash-mxfp8-import-optional.py --check; \
+    # vLLM's import-time INFO logs land on stdout — grep the actual path line.
+    SRC="$(python3 -c 'import vllm, os; print(vllm.__path__[0])' 2>/dev/null | grep '^/' | tail -1)"; \
+    python3 -m py_compile \
+      "$SRC/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py" \
+      "$SRC/model_executor/models/qwen3_dflash.py"; \
+    find /usr/local/lib/python3.12/dist-packages/vllm -name '__pycache__' -type d -prune -exec rm -rf {} + 2>/dev/null || true

+ 27 - 0
docker-images/mimo26-spark-vllm/build.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Build the mimo26-spark-vllm wrapper (thin FROM over the r39 b12x image).
+#
+# Uses the classic docker builder (not buildx docker-container): the base
+# image jovian-judgement:r39-glm53 is local-only (never pushed to a
+# registry), and the container-driver builder cannot resolve local FROM
+# references.
+#
+# Generic driver: forwards every KEY=VAL pin in versions.env as a
+# --build-arg to the Dockerfile (no hardcoded variable list).
+set -euo pipefail
+cd "$(dirname "$0")"
+TAG="${TAG:-mimo26-spark-vllm:local}"
+
+ARGS=()
+while IFS= read -r line; do
+    case "$line" in
+        ''|\#*) continue ;;
+    esac
+    key="${line%%=*}"
+    val="${line#*=}"
+    [ -n "$key" ] && ARGS+=(--build-arg "${key}=${val}")
+done < versions.env
+
+echo "=== Building ${TAG} ==="
+docker build "${ARGS[@]}" --progress=plain -t "${TAG}" -f Dockerfile . 2>&1
+echo "=== Done: ${TAG} ==="

+ 130 - 0
docker-images/mimo26-spark-vllm/patches/dflash-mxfp8-import-optional.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+"""Make vLLM's DFlash MXFP8 guards tolerant of vLLM builds that do not ship
+ModelOptMxFp8LinearMethod (the r39 fork predates that class).
+
+Root cause (2026-09-28, mimo26 vLLM lane): the DFlash draft loader reaches
+the class via two `from vllm.model_executor.layers.quantization.modelopt
+import (ModelOptMxFp8LinearMethod,)` statements (qwen3_dflash.py,
+_build_context_kv_buffers / process_weights_after_loading). r39's modelopt
+module ships ModelOptMxFp8Config/ModelOptMxFp8FusedMoE but no *LinearMethod*,
+so the import raises ImportError inside _build_fused_kv_buffers right after
+the bf16 draft weights load — the guards themselves only matter for MXFP8
+drafts, ours is bf16.
+
+Both sites become getattr-on-the-module with an explicit None check, so a
+real MXFP8 checkpoint on a vLLM that DOES ship the class still takes the
+packed-scale path unchanged, and a bf16 draft just skips it. Anchor-based,
+refuses on drift.
+"""
+
+import importlib.util
+import sys
+
+import vllm  # noqa: F401  (resolve the install from the running vllm)
+
+MODULE = "vllm.model_executor.models.qwen3_dflash"
+
+OLD1 = """        from vllm.model_executor.layers.quantization.modelopt import (
+            ModelOptMxFp8LinearMethod,
+        )
+
+        quant_methods = [a.qkv_proj.quant_method for a in layers_attn]
+        uses_mxfp8 = [
+            isinstance(method, ModelOptMxFp8LinearMethod) for method in quant_methods
+        ]
+"""
+
+NEW1 = """        import importlib
+
+        ModelOptMxFp8LinearMethod = getattr(
+            importlib.import_module(
+                "vllm.model_executor.layers.quantization.modelopt"
+            ),
+            "ModelOptMxFp8LinearMethod",
+            None,
+        )
+
+        quant_methods = [a.qkv_proj.quant_method for a in layers_attn]
+        uses_mxfp8 = [
+            ModelOptMxFp8LinearMethod is not None
+            and isinstance(method, ModelOptMxFp8LinearMethod)
+            for method in quant_methods
+        ]
+"""
+
+OLD2 = """        from vllm.model_executor.layers.quantization.modelopt import (
+            ModelOptMxFp8LinearMethod,
+        )
+
+        quant_method = self.layers[0].self_attn.qkv_proj.quant_method
+        if not isinstance(quant_method, ModelOptMxFp8LinearMethod):
+            return
+"""
+
+NEW2 = """        import importlib
+
+        ModelOptMxFp8LinearMethod = getattr(
+            importlib.import_module(
+                "vllm.model_executor.layers.quantization.modelopt"
+            ),
+            "ModelOptMxFp8LinearMethod",
+            None,
+        )
+
+        quant_method = self.layers[0].self_attn.qkv_proj.quant_method
+        if (
+            ModelOptMxFp8LinearMethod is None
+            or not isinstance(quant_method, ModelOptMxFp8LinearMethod)
+        ):
+            return
+"""
+
+TARGETS = [(OLD1, NEW1), (OLD2, NEW2)]
+
+
+def _consistent(src: str, old: str, new: str) -> bool:
+    co, cn = src.count(old), src.count(new)
+    if cn >= 1:
+        return cn == 1 and co == 0
+    return co == 1
+
+
+def main() -> int:
+    check = "--check" in sys.argv
+    spec = importlib.util.find_spec(MODULE)
+    if not spec or not spec.origin:
+        print(f"{MODULE} not importable", file=sys.stderr)
+        return 1
+    path = spec.origin
+    with open(path, encoding="utf-8") as f:
+        src = f.read()
+
+    if check:
+        if all(_consistent(src, old, new) for old, new in TARGETS):
+            state = "patched" if src.count(NEW1) else "pristine"
+            print(f"CHECK OK ({state}): {path}")
+            return 0
+        print(f"CHECK FAIL: {path} is neither pristine nor patched", file=sys.stderr)
+        return 1
+
+    if all(src.count(new) == 1 for _, new in TARGETS):
+        print(f"ALREADY PATCHED: {path}")
+        return 0
+    if src.count(OLD1) != 1 or src.count(OLD2) != 1:
+        print(
+            f"refusing: {path} anchors matched {src.count(OLD1)}/{src.count(OLD2)}x;"
+            " re-sync patch",
+            file=sys.stderr,
+        )
+        return 1
+
+    for old, new in TARGETS:
+        src = src.replace(old, new, 1)
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(src)
+    print(f"APPLIED: {path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

+ 90 - 0
docker-images/mimo26-spark-vllm/patches/eagle3-omni-inner.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Patch vLLM's set_eagle3_aux_hidden_state_layers to resolve the eagle3
+interface on the inner language-model decoder when the outer model is a
+wrapper (e.g. MiMoV2OmniForCausalLM, whose text stack is
+MiMoV2ForCausalLM/MiMoV2Model and does implement SupportsEagle3).
+
+Root cause (2026-09-28, mimo26 vLLM lane): MiMo-V2.6-Flash-MOPD resolves to
+MiMoV2OmniForCausalLM (multimodal wrapper). For method=dflash the GPU model
+runner marks use_aux_hidden_state_outputs=True and calls
+set_eagle3_aux_hidden_state_layers(target, spec), which isinstance-checks
+SupportsEagle3 against the OUTER class and raises
+"Model does not support EAGLE3 interface" before ever consulting
+_inner_decoder(). Everything downstream (aux slots, _maybe_add_hidden_state)
+lives on the inner MiMoV2Model, and with PP=1 the relay helpers are no-ops,
+so an inner-resolved check is exactly the correct behaviour here.
+
+Idempotent, anchor-based: `--check` exits 0 only when the patch is present,
+2 when the file is still patchable (and the build must apply it), 1 when the
+anchor is absent — it refuses silently-wrong sources. `--apply` writes once
+(no backup file; the container layer is disposable and the image rebuilds
+from this dir).
+"""
+
+import sys
+
+MODULE = "vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils"
+
+ANCHOR = (
+    "    if not supports_eagle3(model):\n"
+    '        raise RuntimeError("Model does not support EAGLE3 interface")\n'
+)
+
+PATCHED = (
+    "    if not supports_eagle3(model):\n"
+    "        # Multimodal wrappers (e.g. MiMoV2OmniForCausalLM) keep the\n"
+    "        # eagle3 carrier one level in — the language model whose method\n"
+    "        # set_aux_hidden_state_layers() already knows how to reach the\n"
+    "        # EagleModelMixin holder. Resolve it before failing (the\n"
+    "        # _inner_decoder() helper does NOT fit: it returns the bare\n"
+    "        # decoder, which lacks the protocol's method).\n"
+    "        _lang = None\n"
+    "        if hasattr(model, \"get_language_model\"):\n"
+    "            _lang = model.get_language_model()\n"
+    "        elif hasattr(model, \"language_model\"):\n"
+    "            _lang = model.language_model\n"
+    "        if _lang is None or not supports_eagle3(_lang):\n"
+    '            raise RuntimeError("Model does not support EAGLE3 interface")\n'
+    "        logger.info(\n"
+    '            "Eagle3 interface resolved on %s inside %s",\n'
+    "            type(_lang).__name__,\n"
+    "            type(model).__name__,\n"
+    "        )\n"
+    "        model = _lang\n"
+)
+
+
+def target_path() -> str:
+    import importlib
+
+    mod = importlib.import_module(MODULE)
+    path = getattr(mod, "__file__", None)
+    if not path:
+        raise SystemExit(f"{MODULE} has no __file__ (frozen install?)")
+    return path
+
+
+def main() -> int:
+    check = "--check" in sys.argv
+    path = target_path()
+    with open(path, encoding="utf-8") as f:
+        src = f.read()
+
+    if PATCHED in src:
+        print(f"ALREADY PATCHED: {path}")
+        return 0
+    if ANCHOR not in src:
+        print(f"ANCHOR NOT FOUND (unexpected source): {path}", file=sys.stderr)
+        return 1
+    if check:
+        print(f"NOT PATCHED YET: {path}", file=sys.stderr)
+        return 2
+
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(src.replace(ANCHOR, PATCHED, 1))
+    print(f"APPLIED: {path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

+ 5 - 0
docker-images/mimo26-spark-vllm/versions.env

@@ -0,0 +1,5 @@
+# versions.env - build input for the mimo26-spark-vllm wrapper.
+# Thin FROM wrapper: the only inputs are the base image (jovian-judgement
+# :r39-glm53, built by docker-images/jovian-judgement-r39 / present on the
+# Spark nodes) and the eagle3-omni patcher. Convention: KEY=VAL.
+BASE_IMAGE=jovian-judgement:r39-glm53