recipes/docker-images/mimo26-spark-vllm/patches/dflash-mxfp8-import-optional.py

#!/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())