Dockerfile.vllm-gemma4

# lobes vLLM-Gemma4 — custom image that SERVES the Gemma 4 12B gemma4_unified gear.
#
# Gemma 4 12B's arch (`gemma4_unified`) is EARLY-FUSION multimodal with
# HETEROGENEOUS per-layer head sizes: 40 sliding-attention layers at head_dim=256
# and 8 full-attention layers at global_head_dim=512. Serving it needs vLLM's
# NATIVE `Gemma4UnifiedForConditionalGeneration` class, which gives the two
# attention types different KV block sizes and auto-forces TRITON_ATTN.
#
# That native class only exists in vLLM NIGHTLY (>= 0.23.1rc1). Released vLLM
# <= 0.22.1 — including the NGC 26.06 base this file used before — has no native
# unified class, so it falls back to the generic Transformers modeling backend,
# which builds every layer with a SINGLE head_size (256) and crashes the 8
# full-attention layers' o_proj (marlin_gemm: a.size(1)=4096 != size_k=8192).
# No attention-backend flag fixes that (proven live #71: TRITON engaged via the
# --attention-backend CLI flag and it crashed identically) — the fix is the
# native class, i.e. nightly vLLM. See docs/gemma-4-12b-nvfp4.md.
#
# Base: the OFFICIAL vLLM nightly image (multi-arch arm64+amd64), pinned by
# DIGEST for reproducibility. This digest resolves to vLLM 0.23.1rc1.dev672 +
# transformers 5.12.1 (registers gemma4_unified) on Blackwell-capable torch, and
# was live-validated on the DGX Spark GB10 (sm_121) on 2026-07-01 (#71): the gear
# serves and answers text + image + audio requests. Bump the digest DELIBERATELY
# (and re-validate) rather than floating the moving :nightly tag.
FROM vllm/vllm-openai@sha256:7c5a10e9a8b3c8642f4d0463a41215176c0dd834b4f0967287c7e3e517cf1be9

# ---------------------------------------------------------------------------
# Audio input needs the vllm[audio] extra. Gemma 4's audio path decodes with
# soundfile/librosa and RESAMPLES with `av` (PyAV) — the stock nightly image
# ships none of these, so audio requests 500 with "Please install vllm[audio]"
# (verified #71: text+image worked without them; audio needed `av` specifically).
# Install with uv (ships in the nightly image), not bare pip. --system writes
# into the image's site-packages where vLLM's import paths resolve.
#
# Versions are PINNED to the set live-validated with this digest on the DGX Spark
# GB10 (2026-07-01, #71). Bump them DELIBERATELY (and re-validate) rather than
# floating — same reproducibility contract as the digest-pinned base above.
# ---------------------------------------------------------------------------
RUN uv pip install --system --no-cache-dir \
    librosa==0.11.0 soundfile==0.14.0 av==17.1.0 soxr==1.1.0

# ---------------------------------------------------------------------------
# Build-stage verification (network-free): assert the base ships the NATIVE
# gemma4_unified class (the whole point of the nightly bump) and that vLLM
# imports. Single logical line — each statement ends with `;` and physical lines
# are continued with a trailing backslash (a multi-line python3 -c WITHOUT
# continuations makes Docker parse each body line as its own instruction; see
# tests/test_gemma4_dockerfile.py for the lint that guards this).
# ---------------------------------------------------------------------------
RUN python3 -c "import vllm; \
    archs = vllm.ModelRegistry.get_supported_archs(); \
    unified = [a for a in archs if 'Gemma4Unified' in a]; \
    assert unified, 'Gemma4UnifiedForConditionalGeneration not in vllm.ModelRegistry — base image predates the native gemma4_unified class (needs vLLM >= 0.23.1rc1)'; \
    import transformers.models.auto.configuration_auto as cfg; \
    assert 'gemma4_unified' in cfg.CONFIG_MAPPING, 'transformers does not register gemma4_unified'; \
    print('OK — native gemma4_unified in ModelRegistry:', unified, '; import vllm OK')"