Kaynağa Gözat

mimo: vLLM+DFlash experiment PARKED (r39 KV-layout wall); findings spec + recipe banner, patches kept for reopen

Alexander Korolev 3 hafta önce
ebeveyn
işleme
02846405b1

+ 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

+ 24 - 6
recipes/mimo26-flash-mopd-4x-vllm/recipe.toml

@@ -19,12 +19,24 @@
 #   * DFlash via vLLM speculative-config (method dflash, draft from the
 #     checkpoint's dflash/ subfolder, 7 draft tokens) vs sglang's DFlash.
 #
-# Status: first boot under qualification; the lab marks their MiMo DFlash
-# "implemented; runtime qualification unverified", same caveat here.
+# Status: PARKED — dead end on r39 (2026-09-28). Four boot attempts, each one
+# layer deeper: (1) EAGLE3 hook raised on the Omni wrapper — fixed by the
+# image's inner-language-model patch; (2) B12X rejects the draft's non-causal
+# attention — draft moved to TRITON_ATTN; (3) qwen3_dflash hard-imports
+# ModelOptMxFp8LinearMethod that r39's modelopt doesn't ship — patched to be
+# None-tolerant; (4) FATAL: KV-layout validator refuses the MiMo hybrid stack
+# (SWA-128 + GA layers, mixed page sizes 131072/163840) because b12x declares
+# LBNHC only and the check is backend-declared, not env-fixable. r39 was
+# qualified for GLM-5.3, never for MiMo; lil-catalog itself marks their MiMo
+# DFlash "runtime qualification unverified". Reopening recipe + external
+# intel (madeby561 amd64-only rc4, eugr spark-vllm-b12x nightlies, lil-catalog
+# flags): specs/bugs/BUG-2026-09-28-vllm-mimo-dflash-parked.md. Production
+# stays sglang mimo26-flash-mopd-4x.
 
 engine = "vllm"
 name = "MiMo-V2.6-Flash-MOPD (vLLM B12X, 4x Spark, DFlash-7, r39)"
-docker_image = "jovian-judgement:r39-glm53"
+docker_image = "mimo26-spark-vllm:local"
+docker_image_dir = "mimo26-spark-vllm"
 entrypoint = ["vllm", "serve"]
 model = "XiaomiMiMo/MiMo-V2.6-Flash-MOPD"
 parent_model = "XiaomiMiMo/MiMo-V2.6-Flash-MOPD"
@@ -55,10 +67,16 @@ args = [
     "--load-format", "instanttensor",
     "--kernel-config", """
 '{"enable_cutedsl_warmup":false,"enable_flashinfer_autotune":false}'""",
-    # DFlash draft from the checkpoint's dflash/ subfolder (7 tokens, B12X
-    # draft attention) — the vLLM mirror of the sglang lane's DFLASH flags.
+    # DFlash draft from the checkpoint's dflash/ subfolder (7 tokens). Draft
+    # attention must be TRITON_ATTN, NOT B12X: the block-diffusion draft runs
+    # NON-CAUSAL attention and b12x (r39) rejects that ("non-causal attention
+    # not supported") — verified in boot log. The draft is two layers over an
+    # 8-token block (speed irrelevant); the target keeps B12X everywhere.
     "--speculative-config", """
-'{"method":"dflash","model":"{{dflash}}","num_speculative_tokens":7,"attention_backend":"B12X"}'""",
+'{"method":"dflash","model":"{{dflash}}","num_speculative_tokens":7,"attention_backend":"TRITON_ATTN"}'""",
+    # vLLM resolves this checkpoint to MiMoV2OmniForCausalLM (the multimodal
+    # wrapper); the eagle3-hook fix lives in docker-images/mimo26-spark-vllm
+    # (patched image), which is what this lane runs.
     "--language-model-only",
     "--reasoning-parser", "mimo",
     "--tool-call-parser", "mimo",

+ 58 - 0
specs/bugs/BUG-2026-09-28-vllm-mimo-dflash-parked.md

@@ -0,0 +1,58 @@
+# MiMo-V2.6-Flash-MOPD on vLLM+r39 — the 2026-09-28 experiment (PARKED)
+
+Lane: `recipes/mimo26-flash-mopd-4x-vllm` + image `docker-images/mimo26-spark-vllm`
+(thin FROM wrapper on `jovian-judgement:r39-glm53`, two anchor patches).
+
+**Verdict: dead end for now.** Every layer of the vLLM DFlash path had to be
+forced onto a fork that was qualified for GLM-5.3, not MiMo, and each boot
+revealed one more incompatibility. None were ours to fix cleanly:
+
+1. **EAGLE3 hook vs Omni wrapper** — MiMo-V2.6 resolves to
+   `MiMoV2OmniForCausalLM` (mm wrapper); `set_eagle3_aux_hidden_state_layers`
+   isinstance-checks the OUTER class and raises. Fixable (and fixed) with a
+   language-model fallback patch; the protocol's own method unwraps the same
+   way. Patched, works:
+   `"Eagle3 interface resolved on MiMoV2FlashForCausalLM inside MiMoV2OmniForCausalLM"`,
+   aux layers (1,12,24,36,48) taken from the draft config.
+2. **B12X rejects non-causal attention** — DFlash's block-diffusion draft runs
+   non-causal; b12x in r39 declares it unsupported. Workaround: draft on
+   `TRITON_ATTN` (target keeps B12X). Works, draft loads (2.7 GB, mask
+   embedding loaded).
+3. **`ModelOptMxFp8LinearMethod` ImportError** — qwen3_dflash.py hard-imports
+   a class r39's modelopt doesn't ship (fork predates it). Patched to a
+   None-tolerant getattr (bf16 draft unaffected).
+4. **KV-layout wall (fatal):** during cudagraph-memory profiling,
+   `validate_kv_cache_layout` refuses: *"KV cache layout LBNHC cannot express
+   this model's mixed page sizes ([131072, 163840]); set
+   VLLM_KV_CACHE_LAYOUT=BLHNC"*. The hint is env-only, but BLHNC is a
+   backend-declared-layout question — b12x in r39 declares LBNHC only. This
+   is the architecture mismatch the whole experiment kept routing around:
+   **the MiMo hybrid stack (SWA-128 + 9 GA layers, different page sizes per
+   group) needs a vLLM/b12x generation qualified for it; r39 is not.**
+
+## External intel (2026-09-28)
+
+* `local-inference-lab/lil-catalog` MiMo DFlash entry: "implemented; runtime
+  qualification unverified" — they never ran it either.
+* madeby561/vllm `mimo-v26-flash-b12x-20260923-rc4`: MiMo RL + b12x + FP8 KV +
+  DFlash-7, with adaptive verification — but **linux/amd64 only** (RTX 6000
+  Pro). Not runnable here; its compose is a good flag corpus
+  (`enable_adaptive_verification`, `kv_cache_dtype: fp8` for target+draft,
+  `--max-num-scheduled-tokens 2048`, `prefill-compute-share`, `BLHNC`-class
+  layouts would all apply if a Spark-native equivalent appears).
+* `eugr/spark-vllm-b12x` publishes arm64 nightlies daily (latest 2026-09-28)
+  — first candidate when someone asks whether MiMo landed there.
+
+## How to reopen
+
+1. Rebase the wrapper onto a MiMo-qualified vLLM (eugr nightly that carries
+   MiMo + b12x, or a new lab rev with upstream's MiMo/b12x reconciliation —
+   the upstream PR #46104 lineage) and retry the same recipe. The two
+   patches (eagle3-inner + mxfp8-optional) stay valid until that vLLM ships
+   the equivalents; then delete them.
+2. If the new stack boots, run the same long-context slope sweep used for
+   sglang (`~/.cache/sm121/bench/slope-28x`: 77.6 → 23.5 t/s d0→256K) and the
+   serial quality bench (anchor: sglang 168/184, 91/100).
+3. sglang 4x (`mimo26-flash-mopd-4x`) is the production lane in the meantime;
+   the FROZEN_KV_MTP experiment is the next sglang-side lever for the
+   long-context cliff.