# syntax=docker/dockerfile:1.7
# GLM-5.3-Flash DFlash2 draft-model support for DGX Spark (SM121/GB10), TP=2.
#
# Stacks on glm53-flash-nvfp4-lab:local (day-0 base + NoPE-MLA + lab
# checkpoint patches) and adds DFlash2 draft-model execution ported from
# local-inference-lab/vllm `dev/jovian-judgement` (see UPSTREAM.md):
#   * qwen3_dflash.py   — refreshed base (fork is a superset of the day-0
#                         file; identical imports, +DFlash2 wiring incl.
#                         native MXFP8 serialized projections)
#   * qwen3_dflash2.py  — the DFlash2DraftModel itself (294 lines)
#   * dflash.py         — v1 spec-decode proposer (+rope is_neox_style
#                         handoff from the target, 1 hunk)
#   * registry.py       — one line: DFlash2DraftModel -> qwen3_dflash2
#
# The day-0 base already ships W4A16_NVFP4, MIXED_PRECISION,
# ModelOptMxFp8LinearMethod, CompressedTensorsW8A8Mxfp8 and the dflash
# proposer — only the DFlash2 model class + registration were missing.

ARG BASE_IMAGE
FROM ${BASE_IMAGE}

COPY qwen3_dflash.py qwen3_dflash2.py dflash.py patch_eagle3.py /tmp/

RUN set -eux; \
    B=/usr/local/lib/python3.12/dist-packages/vllm; \
    cp /tmp/qwen3_dflash.py /tmp/qwen3_dflash2.py "$B/model_executor/models/"; \
    cp /tmp/dflash.py "$B/v1/spec_decode/dflash.py"; \
    grep -q '"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),' \
        "$B/model_executor/models/registry.py"; \
    sed -i '/"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),/a\    "DFlash2DraftModel": ("qwen3_dflash2", "DFlash2Qwen3ForCausalLM"),' \
        "$B/model_executor/models/registry.py"; \
    grep -q '"DFlash2DraftModel"' "$B/model_executor/models/registry.py"; \
    python3 -m py_compile \
        "$B/model_executor/models/qwen3_dflash.py" \
        "$B/model_executor/models/qwen3_dflash2.py" \
        "$B/v1/spec_decode/dflash.py" \
        "$B/model_executor/models/registry.py"; \
    python3 /tmp/patch_eagle3.py; \
    rm -f /tmp/qwen3_dflash.py /tmp/qwen3_dflash2.py /tmp/dflash.py /tmp/patch_eagle3.py