Dockerfile 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. # syntax=docker/dockerfile:1.7
  2. # GLM-5.3-Flash r12 (jovian-judgement) — CLEAN build for DGX Spark (aarch64/sm_121a)
  3. #
  4. # Per operator direction: start from a raw CUDA image (solar-open2-nvfp4
  5. # pattern), NOT the day-0 vllm image. No baked-vllm contamination: torch,
  6. # flashinfer, vLLM and b12x all come from pinned sources in this file.
  7. #
  8. # Stack (rtx6kpro models/glm-5.3-flash.md source contract):
  9. # vLLM = local-inference-lab/vllm @ 54f6e9826c (dev/jovian-judgement base)
  10. # + sequential merges of PR heads 550,552,553,554,555,557
  11. # (the doc's "composition commit" 7d8a09c… is not a public ref)
  12. # b12x = local-inference-lab/b12x @ 139e04048b (master) + PR heads
  13. # 260,267,268,269,270. b12x is pure-Python: the kernels are
  14. # CuTe-DSL / tvm-ffi JIT (nvidia-cutlass-dsl 4.6.2), no CUDA
  15. # compile at install time.
  16. # torch = 2.13.0 cu130 (the fork's requirements/cuda.txt pin)
  17. # flashinfer = 0.6.17 (day-0-era; not a pip dep of the fork since 0.6.14,
  18. # so installed from the flashinfer.ai wheel index)
  19. #
  20. # GB10 sm_121 fix baked in (patch_cuda_sm12x.py): the sm12x auto-priority
  21. # list in vllm/platforms/cuda.py offers only TRITON_MLA (no sparse) and
  22. # FLASHINFER_MLA_SPARSE_SM120 (head 576 only) — neither fits GLM head-512.
  23. # B12X (supports (12,0)/(12,1)) is inserted first.
  24. #
  25. # Build: cd docker-images/glm53-jovian-r12 && bash build.sh (on a Spark node)
  26. ARG CUDA_BASE_IMAGE
  27. ARG CUDA_BASE_DIGEST
  28. ARG UV_VERSION
  29. FROM ${CUDA_BASE_IMAGE}@${CUDA_BASE_DIGEST} AS apt-base
  30. ARG PYPI_INDEX_URL
  31. ENV DEBIAN_FRONTEND=noninteractive \
  32. PIP_BREAK_SYSTEM_PACKAGES=1 \
  33. UV_SYSTEM_PYTHON=1 UV_BREAK_SYSTEM_PACKAGES=1 UV_LINK_MODE=copy
  34. COPY locks/apt-sources.list /tmp/apt-sources.list
  35. COPY locks/apt-packages.txt /tmp/apt-packages.txt
  36. COPY locks/python-bootstrap.txt /tmp/python-bootstrap.txt
  37. RUN rm -f /etc/apt/sources.list.d/*.list /etc/apt/sources.list.d/*.sources \
  38. && cp /tmp/apt-sources.list /etc/apt/sources.list \
  39. && apt-get update \
  40. && grep -vE '^\s*(#|$)' /tmp/apt-packages.txt \
  41. | xargs apt-get install -y --no-install-recommends \
  42. && rm -rf /var/lib/apt/lists/* \
  43. && python3 -m pip install --no-cache-dir --no-deps \
  44. -r /tmp/python-bootstrap.txt --index-url ${PYPI_INDEX_URL}
  45. FROM apt-base AS torch-base
  46. ARG PYPI_INDEX_URL
  47. ARG PYTORCH_INDEX_URL
  48. ARG TORCH_VERSION TORCHVISION_VERSION TORCHAUDIO_VERSION
  49. ARG FLASHINFER_VERSION
  50. RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
  51. uv pip install \
  52. --index-url ${PYPI_INDEX_URL} \
  53. --extra-index-url ${PYTORCH_INDEX_URL} \
  54. --index-strategy unsafe-best-match \
  55. "torch==${TORCH_VERSION}" \
  56. "torchvision==${TORCHVISION_VERSION}" \
  57. "torchaudio==${TORCHAUDIO_VERSION}" \
  58. "flashinfer-python==${FLASHINFER_VERSION}"
  59. # The fork's r12 composition: base branch + sequential merge of the open
  60. # PR heads (ascending). The PR head pseudo-refs are GitHub's stable refs.
  61. FROM torch-base AS deepgemm-builder
  62. ARG DEEPGEMM_REPO DEEPGEMM_REF DEEPGEMM_BRANCH
  63. RUN set -eux; \
  64. git config --global protocol.version 0; \
  65. git clone --branch ${DEEPGEMM_BRANCH} ${DEEPGEMM_REPO} /deps/deepgemm; \
  66. cd /deps/deepgemm; \
  67. git checkout ${DEEPGEMM_REF}; \
  68. git submodule update --init --depth 1 -- third-party/cutlass third-party/fmt
  69. FROM deepgemm-builder AS vllm-builder
  70. ARG VLLM_REPO
  71. ARG VLLM_BASE_REF
  72. ARG VLLM_PR_HEADS
  73. ARG SOURCE_DATE_EPOCH
  74. ARG BUILD_JOBS
  75. ENV TORCH_CUDA_ARCH_LIST=12.1a GIT_PROTOCOL=0 DEEPGEMM_SRC_DIR=/deps/deepgemm
  76. RUN set -eux; \
  77. git config --global user.email build@sm121.lan; \
  78. git config --global user.name sm121-builder; \
  79. git config --global protocol.version 0; \
  80. git clone --recursive ${VLLM_REPO} /workspace/vllm; \
  81. cd /workspace/vllm; \
  82. git checkout ${VLLM_BASE_REF}; \
  83. # PR head commits fetched by SHA (refs/pull/N/head vanishes when a PR
  84. # closes — 550 closed unmerged; raw-SHA fetches stay available).
  85. for pair in ${VLLM_PR_HEADS}; do \
  86. pr="${pair%%:*}"; sha="${pair#*:}"; \
  87. git fetch origin "${sha}:refs/prs/${pr}"; \
  88. git merge --no-ff -m "PR ${pr}" "refs/prs/${pr}"; \
  89. done; \
  90. git rev-parse HEAD
  91. WORKDIR /workspace/vllm
  92. RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
  93. --mount=type=cache,id=ccache,target=/root/.ccache \
  94. uv pip install setuptools-rust setuptools-scm \
  95. && _j="${BUILD_JOBS:-$(nproc)}" \
  96. && export MAX_JOBS="${_j}" CMAKE_BUILD_PARALLEL_LEVEL="${_j}" \
  97. && python3 use_existing_torch.py \
  98. && uv build --no-build-isolation --wheel . --out-dir=/wheels -v
  99. FROM torch-base AS b12x-builder
  100. ARG B12X_REPO
  101. ARG B12X_BASE_REF
  102. ARG B12X_PRS
  103. ARG SOURCE_DATE_EPOCH
  104. RUN set -eux; \
  105. git config --global user.email build@sm121.lan; \
  106. git config --global user.name sm121-builder; \
  107. git config --global protocol.version 0; \
  108. git clone ${B12X_REPO} /workspace/b12x; \
  109. cd /workspace/b12x; \
  110. git checkout ${B12X_BASE_REF}; \
  111. for pr in ${B12X_PRS}; do \
  112. git fetch origin "refs/pull/$pr/head:prs/$pr"; \
  113. git merge --no-ff -m "PR $pr" "prs/$pr"; \
  114. done; \
  115. git rev-parse HEAD
  116. WORKDIR /workspace/b12x
  117. RUN --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
  118. uv build --no-build-isolation --wheel . --out-dir=/wheels -v
  119. FROM torch-base AS runner
  120. ARG PYPI_INDEX_URL
  121. COPY checksums/tiktoken.sha256 /tmp/tiktoken.sha256
  122. COPY patch_cuda_sm12x.py /tmp/patch_cuda_sm12x.py
  123. WORKDIR /workspace
  124. RUN mkdir -p tiktoken_encodings \
  125. && wget -O tiktoken_encodings/o200k_base.tiktoken https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken \
  126. && wget -O tiktoken_encodings/cl100k_base.tiktoken https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken \
  127. && cd tiktoken_encodings && sha256sum -c /tmp/tiktoken.sha256
  128. ENV TIKTOKEN_ENCODINGS_BASE=/workspace/tiktoken_encodings
  129. RUN --mount=type=bind,from=vllm-builder,source=/wheels,target=/vllm-wheels \
  130. --mount=type=bind,from=b12x-builder,source=/wheels,target=/b12x-wheels \
  131. --mount=type=cache,id=uv-cache,target=/root/.cache/uv \
  132. pip install --no-cache-dir --no-build-isolation /vllm-wheels/*.whl \
  133. && pip install --no-cache-dir /b12x-wheels/*.whl \
  134. && python3 /tmp/patch_cuda_sm12x.py \
  135. && find /usr/local/lib/python3.12/dist-packages/vllm -name "__pycache__" \
  136. -type d -prune -exec rm -rf {} + 2>/dev/null || true
  137. # Build gate: the r12 stack must import with the GB10 patches in place.
  138. RUN set -eux; \
  139. python3 -c "import vllm; print('vllm', vllm.__version__)"; \
  140. python3 -c "import b12x; print('b12x ok')"; \
  141. python3 -c "import vllm._flashkda_C; print('flashkda_C ok')"; \
  142. python3 -c "from vllm.config.attention import AttentionBackendEnum as E; \
  143. assert E.B12X.value and 'b12x' in E.B12X.value.lower(); print('B12X enum ok')"; \
  144. grep -q "AttentionBackendEnum.B12X," /usr/local/lib/python3.12/dist-packages/vllm/platforms/cuda.py
  145. ENTRYPOINT []
  146. CMD ["bash"]