Dockerfile 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. # DeepSeek-V4.1-Flash — DGX Spark (arm64, sm_121a), clone-at-pin.
  2. #
  3. # One Dockerfile for tonyd2wild's overlay1/3/5 chain (MIT; see versions.env
  4. # header + patches/README note). No fallbacks: every gate fails loud.
  5. #
  6. # base pinned vllm/vllm-openai nightly (arm64 digest)
  7. # vLLM vllm-project/vllm @ dsv41-feat commit+tree (day-0 DSV4.1 support)
  8. # ext _C_stable_libtorch rebuilt for sm_121a (stable-only cmake build)
  9. # FI FlashInfer 0.7.0rc1 from pinned source (0.6.18 lacks the SM120
  10. # sparse-MLA decode kernel for page_block_size=32)
  11. # warm sparse-MLA + mxfp8 JIT kernels prebuilt in the image runtime env
  12. # patches the 7 boot-10 file replacements, baked over dist-packages,
  13. # md5-gated against the published manifest
  14. #
  15. # Model: deepseek-ai/DeepSeek-V4.1-Flash (510 GB; Engram tables stay on disk
  16. # via the baked engram.py patch — DSV41_ENGRAM_DISK=1 at serve time).
  17. ARG VLLM_BASE_IMAGE
  18. FROM ${VLLM_BASE_IMAGE}
  19. ARG VLLM_REPO
  20. ARG VLLM_REF
  21. ARG VLLM_REF_TREE
  22. ARG VLLM_PRETEND_VERSION
  23. ARG CUTLASS_TAG
  24. ARG CUTLASS_URL
  25. ARG FI_SHA
  26. ARG FI_CUTLASS_SHA
  27. ARG FI_CCCL_SHA
  28. ARG FI_SPDLOG_SHA
  29. ENV DEBIAN_FRONTEND=noninteractive \
  30. PIP_DISABLE_PIP_VERSION_CHECK=1 \
  31. MAX_JOBS=4
  32. # ---- vLLM python tree @ pin (clone-at-pin: commit + tree hash gates) ----
  33. RUN set -eux; \
  34. apt-get update -qq && apt-get install -y --no-install-recommends git ca-certificates >/dev/null && rm -rf /var/lib/apt/lists/*; \
  35. git config --global protocol.version 0; \
  36. git init -q /src/vllm; \
  37. git -C /src/vllm fetch --depth 1 ${VLLM_REPO} ${VLLM_REF}; \
  38. git -C /src/vllm checkout -q ${VLLM_REF}; \
  39. test "$(git -C /src/vllm rev-parse HEAD)" = "${VLLM_REF}"; \
  40. test "$(git -C /src/vllm rev-parse 'HEAD^{tree}')" = "${VLLM_REF_TREE}"; \
  41. echo "vLLM pinned at ${VLLM_REF} / tree ${VLLM_REF_TREE}"
  42. # ---- _C_stable_libtorch rebuilt for sm_121a (their build_stable_ext.sh) ----
  43. # stable-only build: external_projects includes commented out, cutlass fetched
  44. # at the pinned tag, Release build of the single C extension target.
  45. RUN set -eux; \
  46. apt-get update -qq && apt-get install -y --no-install-recommends cmake ninja-build g++ >/dev/null && rm -rf /var/lib/apt/lists/*; \
  47. cd /src/vllm; \
  48. mkdir -p build/_deps/cutlass-src; \
  49. curl -sL -m 600 ${CUTLASS_URL} | tar xz -C build/_deps/cutlass-src --strip-components=1; \
  50. ls build/_deps/cutlass-src/include/cutlass/cutlass.h >/dev/null; \
  51. cp CMakeLists.txt CMakeLists.txt.orig; \
  52. sed -i -E "s|^(\s*)include\(cmake/external_projects/|\1# STABLE-ONLY: include(cmake/external_projects/|" CMakeLists.txt; \
  53. PYPATH=$(python3 -c "import sys;print(':'.join(p for p in sys.path if p))"); \
  54. TORCH_PREFIX=$(python3 -c "import torch;print(torch.utils.cmake_prefix_path)"); \
  55. NVRTC=$(ls /usr/local/cuda/lib64/libnvrtc.so /usr/local/cuda/lib64/libnvrtc.so.* /usr/local/lib/python3.12/dist-packages/nvidia/*/lib/libnvrtc.so* /usr/lib/aarch64-linux-gnu/libnvrtc.so* 2>/dev/null | head -1); \
  56. test -n "${NVRTC}"; \
  57. cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DVLLM_TARGET_DEVICE=cuda \
  58. -DVLLM_PYTHON_EXECUTABLE=$(which python3) -DVLLM_PYTHON_PATH="$PYPATH" \
  59. -DFETCHCONTENT_BASE_DIR=build/_deps -DFETCHCONTENT_SOURCE_DIR_CUTLASS=build/_deps/cutlass-src \
  60. -DCMAKE_PREFIX_PATH="$TORCH_PREFIX" -DNVCC_THREADS=2 -DCUDA_nvrtc_LIBRARY="$NVRTC" -DTORCH_CUDA_ARCH_LIST=12.1a; \
  61. cmake --build build --target _C_stable_libtorch -j 20; \
  62. find build vllm -maxdepth 2 -name "_C_stable_libtorch*.so" | tee /tmp/stable_ext.so; \
  63. # place the built extension in the package dir (vLLM imports vllm._C_stable_libtorch)
  64. cp -v $(find build vllm -maxdepth 2 -name "_C_stable_libtorch*.so" | head -1) vllm/
  65. # ---- merge the branch python tree (+ the sm121 .so) over the base install ----
  66. RUN set -eux; \
  67. SITE=/usr/local/lib/python3.12/dist-packages/vllm; \
  68. cp -a /src/vllm/vllm/. ${SITE}/; \
  69. find ${SITE} -name "__pycache__" -type d -prune -exec rm -rf {} +; \
  70. python3 -c "import vllm; print('vllm', vllm.__version__)"; \
  71. # existence check only: dlopen of the CUDA extension needs libcuda.so.1
  72. # (GPU driver), which the build container does not have — the real import
  73. # happens at serve time on the GPU node
  74. ls ${SITE}/_C_stable_libtorch*.so
  75. # ---- FlashInfer 0.7.0rc1 from pinned source (0.6.18 has no SM120 sparse-MLA) ----
  76. RUN set -eux; \
  77. pip uninstall -y -q flashinfer-jit-cache flashinfer-cubin flashinfer-python || true; \
  78. rm -rf /root/.cache/flashinfer; \
  79. mkdir -p /opt/fi-src; \
  80. curl -sL -m 900 "https://codeload.github.com/flashinfer-ai/flashinfer/tar.gz/${FI_SHA}" | tar xz -C /opt/fi-src --strip-components=1; \
  81. mkdir -p /opt/fi-src/3rdparty/cutlass /opt/fi-src/3rdparty/cccl /opt/fi-src/3rdparty/spdlog; \
  82. curl -sL -m 900 "https://codeload.github.com/NVIDIA/cutlass/tar.gz/${FI_CUTLASS_SHA}" | tar xz -C /opt/fi-src/3rdparty/cutlass --strip-components=1; \
  83. curl -sL -m 900 "https://codeload.github.com/NVIDIA/cccl/tar.gz/${FI_CCCL_SHA}" | tar xz -C /opt/fi-src/3rdparty/cccl --strip-components=1; \
  84. curl -sL -m 900 "https://codeload.github.com/gabime/spdlog/tar.gz/${FI_SPDLOG_SHA}" | tar xz -C /opt/fi-src/3rdparty/spdlog --strip-components=1; \
  85. ls /opt/fi-src/3rdparty/cutlass/include/cutlass/cutlass.h /opt/fi-src/3rdparty/cccl/README.md /opt/fi-src/3rdparty/spdlog/include/spdlog/spdlog.h; \
  86. cd /opt/fi-src && BUILD_NVEP=0 FLASHINFER_BUILD_NO_PIP=1 pip install --no-deps --no-build-isolation -q .; \
  87. rm -rf /opt/fi-src/build; \
  88. python3 -c "import flashinfer; assert '0.7.0rc1' in flashinfer.__version__, flashinfer.__version__; print('flashinfer', flashinfer.__version__)"
  89. ENV VLLM_HAS_FLASHINFER_CUBIN=1
  90. # ---- prebuild the SM120 sparse-MLA + mxfp8 GEMM JIT kernels ----
  91. # Exact runtime env (no verbose/debug nvcc flags — those change the cache key).
  92. # The compiled modules persist in /root/.cache/flashinfer inside the image.
  93. COPY patches/prewarm5.py /tmp/prewarm5.py
  94. RUN set -eux; \
  95. FLASHINFER_CUDA_ARCH_LIST=12.1a TORCH_CUDA_ARCH_LIST=12.1a \
  96. FLASHINFER_DISABLE_VERSION_CHECK=1 VLLM_HAS_FLASHINFER_CUBIN=1 \
  97. FLASHINFER_NVCC_THREADS=1 MAX_JOBS=4 \
  98. python3 /tmp/prewarm5.py | tee /tmp/prewarm.log; \
  99. grep -q "SPARSE-BUILT+LOADED" /tmp/prewarm.log || { echo "sparse MLA prewarm FAILED"; exit 1; }
  100. # ---- bake the 7 boot-10 patch files (md5-gated against their manifest) ----
  101. COPY patches/ /tmp/patches/
  102. RUN set -eux; \
  103. python3 - <<'PY'
  104. import hashlib, os, shutil, sys
  105. SITE = "/usr/local/lib/python3.12/dist-packages/vllm"
  106. # md5 prefixes from tonyd2wild patch/README.md (boot-10 manifest)
  107. MD5 = {
  108. "engram.py": "c0329107", "model_state.py": "0a14bee6",
  109. "weight_utils.py": "7e1027f1", "attention.py": "da9ef196",
  110. "flashinfer_sparse.py": "af0f8447", "sparse_swa.py": "cc419353",
  111. "sparse_attn_indexer.py": "a9b73756",
  112. }
  113. bad = []
  114. for line in open("/tmp/patches/mounts.txt"):
  115. line = line.strip()
  116. if not line or line.startswith("#"):
  117. continue
  118. f, rel = line.split()
  119. src = f"/tmp/patches/{f}"
  120. m = hashlib.md5(open(src, "rb").read()).hexdigest()[:8]
  121. if m != MD5.get(f):
  122. bad.append(f"{f}: {m} != {MD5.get(f)}")
  123. continue
  124. dst = os.path.join(SITE, rel)
  125. os.makedirs(os.path.dirname(dst), exist_ok=True)
  126. shutil.copyfile(src, dst)
  127. print(f"baked {f} -> {rel} (md5 {m})")
  128. sys.exit("\n".join(bad) if bad else None)
  129. PY
  130. # ---- final gates ----
  131. RUN set -eux; \
  132. python3 -c "import vllm; print('vllm', vllm.__version__)"; \
  133. ls /usr/local/lib/python3.12/dist-packages/vllm/_C_stable_libtorch*.so; \
  134. python3 -c "import flashinfer; print('flashinfer', flashinfer.__version__)"; \
  135. grep -q "DSV41_ENGRAM_DISK" /usr/local/lib/python3.12/dist-packages/vllm/models/deepseek_v4_1/common/engram.py; \
  136. echo "=== image gates passed ==="