ARG BASE_IMAGE=vllm/vllm-openai@sha256:a551e05307cd2e0092139d84db32af9c97e67d2eeeff072d21e429131d8c23f0
FROM ${BASE_IMAGE}

# GLM-5.3-Flash (glm5_next, NoPE MLA: qk_rope_head_dim=0) on H3C MegaCube SM121.
# The day-0 image's only capability-12 sparse-MLA backend is
# FLASHINFER_MLA_SPARSE_SM120, whose packed fp8_ds_mla cache layout hard-requires
# DeepSeek's pe_dim=64 and dies in warmup ("pe_dim must be 64 for fp8_ds_mla").
# The SM90 NoPE backend (FlashInfer BatchMLAPagedAttentionWrapper, plain bf16
# cache) supports kpe=0, and its FA2 path passed a direct probe on SM121 with
# GLM's real shape (32 heads/rank, ckv 512, kpe 0, page_size 1). Extend it to
# capability 12, select FA2 off-Hopper, and scope the FlashInfer>=0.6.18
# feature gate to the fp8-cache path it actually protects.
RUN python3 - <<'PY'
from pathlib import Path

base = Path("/usr/local/lib/python3.12/dist-packages/vllm")

p = base / "platforms/cuda.py"
s = p.read_text()
old = """        elif device_capability.major == 12:
            return [
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
            ]"""
new = """        elif device_capability.major == 12:
            return [
                AttentionBackendEnum.TRITON_MLA,
                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM90,
                AttentionBackendEnum.FLASHINFER_MLA_SPARSE_SM120,
            ]"""
if s.count(old) != 1:
    raise SystemExit("unexpected cuda.py capability-12 MLA candidate list; refusing to patch")
p.write_text(s.replace(old, new))

p = base / "v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py"
s = p.read_text()

old = "    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:\n        return capability.major == 9\n"
new = "    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:\n        return capability.major in (9, 12)\n"
if s.count(old) != 1:
    raise SystemExit("unexpected sm90 capability gate; refusing to patch")
s = s.replace(old, new)

old = '            backend="fa3",\n'
new = '            backend=("fa3" if torch.cuda.get_device_capability()[0] == 9 else "fa2"),\n'
if s.count(old) != 1:
    raise SystemExit("unexpected sm90 wrapper backend literal; refusing to patch")
s = s.replace(old, new)

old = """        if not has_flashinfer_sm90_nope_mla():
            return (
                "FLASHINFER_MLA_SPARSE_SM90 requires FlashInfer with SM90 "
                "MLA support (ckv_scale_arr in "
                "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)"
            )"""
new = """        if kv_cache_dtype in ("fp8", "fp8_e4m3") and not has_flashinfer_sm90_nope_mla():
            return (
                "FLASHINFER_MLA_SPARSE_SM90 fp8 KV requires FlashInfer with "
                "SM90 MLA support (ckv_scale_arr in "
                "BatchMLAPagedAttentionWrapper.run, FlashInfer >= 0.6.18)"
            )"""
if s.count(old) != 1:
    raise SystemExit("unexpected sm90 flashinfer version gate; refusing to patch")
s = s.replace(old, new)

p.write_text(s)
print("sm121 NoPE MLA patches applied")
PY


# PDL (Programmatic Dependent Launch) is enabled for any capability >= 9,
# including SM121, in the Triton kernels that carry KDA recurrent state
# (causal_conv1d, gather_initial_states, scatter_states). On GB10/sm_121a
# (Triton 3.7.1 + CUDA 13.0, unvalidated combination) this produces a
# timing-dependent read-before-store race: boots NaN or don't depending on
# kernel launch timing (full-device-sync debug hooks reliably mask it).
# Gate PDL to the arches it was validated on.
RUN python3 - <<'PY'
from pathlib import Path

p = Path("/usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py")
s = p.read_text()
old = """    @classmethod
    def is_arch_support_pdl(cls) -> bool:
        try:
            device = torch.cuda.current_device()
            major, _ = torch.cuda.get_device_capability(device)
        except Exception:
            return False
        return major >= 9
"""
new = """    @classmethod
    def is_arch_support_pdl(cls) -> bool:
        try:
            device = torch.cuda.current_device()
            major, _ = torch.cuda.get_device_capability(device)
        except Exception:
            return False
        # PDL lowering is unvalidated on SM12x (GB10) and races on KDA
        # state kernels there; keep it to Hopper/Blackwell-datacenter.
        return major in (9, 10)
"""
if s.count(old) != 1:
    raise SystemExit("unexpected is_arch_support_pdl source; refusing to patch")
p.write_text(s.replace(old, new))
print("PDL gated off on SM12x")
PY


# Indexer hardening (second NaN class): top-k destination was torch.empty and
# the top-k kernels only guarantee the first min(k, valid) entries; rows with
# fewer valid pools carried uninitialized pool ids -> bogus token indices ->
# MLA gathers uninitialized KV -> NaN lottery. Init to -1 + clamp expansion.
COPY patch_v7.py /tmp/patch_v7.py
RUN python3 /tmp/patch_v7.py && rm /tmp/patch_v7.py


# fp8 KV cache for the FA2 NoPE sparse-MLA path on SM12x (see patch script).
COPY patch_v8_fp8.py /tmp/patch_v8_fp8.py
RUN python3 /tmp/patch_v8_fp8.py && rm /tmp/patch_v8_fp8.py

RUN python3 - <<'PY'
from pathlib import Path
p=Path('/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers/sparse_attn_indexer_kpool.py')
s=p.read_text()
old='        if current_platform.is_cuda() and select_k in (512, 1024, 2048):'
new='        if current_platform.is_cuda() and select_k in (512, 1024, 2048) and torch.cuda.get_device_properties(0).multi_processor_count >= 78:'
assert s.count(old)==1
p.write_text(s.replace(old,new))
PY
RUN pip uninstall -y flashinfer-jit-cache
RUN python3 -c "import importlib.metadata as m; print({p:m.version(p) for p in ['vllm','transformers','flashinfer-python','flashinfer-cubin']})"

ENV CPATH=""
RUN test -f /usr/local/lib/python3.12/dist-packages/nvidia/cu13/include/nvrtc.h && \
    ln -s /usr/local/lib/python3.12/dist-packages/nvidia/cu13/include/nvrtc.h /usr/local/cuda/include/nvrtc.h && \
    test -f /usr/local/cuda/lib64/libnvrtc.so.13 && \
    (test -e /usr/local/cuda/lib64/libnvrtc.so || ln -s libnvrtc.so.13 /usr/local/cuda/lib64/libnvrtc.so)

COPY patch_glm_aux_capture.py /tmp/patch_glm_aux_capture.py
COPY patch_dflash_kv_nightly.py /tmp/patch_dflash_kv_nightly.py
COPY patch_prefix_cache_draft_group.py /tmp/patch_prefix_cache_draft_group.py
RUN python3 /tmp/patch_glm_aux_capture.py && python3 /tmp/patch_dflash_kv_nightly.py && python3 /tmp/patch_prefix_cache_draft_group.py
RUN python3 -c 'from vllm.model_executor.models.registry import ModelRegistry; assert "DFlash2DraftModel" in ModelRegistry.get_supported_archs(); print("native DFlash2 registry verified")'

COPY patch_mla_kv_dtype.py /tmp/patch_mla_kv_dtype.py
RUN python3 /tmp/patch_mla_kv_dtype.py

COPY patch_mtp_bf16.py /opt/glm53-patches/patch_mtp_bf16.py
RUN python3 /opt/glm53-patches/patch_mtp_bf16.py

COPY patch_local_autotune.py /tmp/patch_local_autotune.py
RUN python3 /tmp/patch_local_autotune.py

COPY patch_warmup_cache.py /opt/glm53-patches/patch_warmup_cache.py
RUN python3 /opt/glm53-patches/patch_warmup_cache.py

# Fresh customer machines compile the large MoE extension here, without weights.
# Runtime retains these immutable build products; no historical host cache is needed.
COPY precompile.py /opt/megacube-glm/precompile.py
RUN FLASHINFER_CUDA_ARCH_LIST=12.1a TORCH_CUDA_ARCH_LIST=12.1a MAX_JOBS=2 \
    python3 /opt/megacube-glm/precompile.py
LABEL megacube.model=GLM-5.3-Flash-NVFP4 megacube.runtime.revision=glm53-dflash3-tp2-256k-cn-v4 megacube.glm.precompiled=1
