From 8433be2b01ae4d962483927d2c7777841c1bc039 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Wed, 24 Jun 2026 13:23:40 +0000 Subject: [PATCH] build: install flash_mla via uv (no_pypi_wheels group) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit flash_mla (DeepSeek MLA kernels) is declared only in the no_pypi_wheels group, which the extras-based uv sync in docker/Dockerfile.ci.dev does not pull, so it was absent from /opt/venv. Add `--group no_pypi_wheels` to that sync (skipped for the LTS image) so uv builds it from source (no-build-isolation); export FLASH_MLA_DISABLE_SM90=1, NVCC_THREADS=16 and CFLAGS/CXXFLAGS pointing at the CCCL headers (under cccl/ in this base image) so `#include ` resolves. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: oliver könig --- docker/Dockerfile.ci.dev | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile.ci.dev b/docker/Dockerfile.ci.dev index 5405de51142..33be0e6397e 100644 --- a/docker/Dockerfile.ci.dev +++ b/docker/Dockerfile.ci.dev @@ -40,9 +40,18 @@ ENV NVTE_BUILD_NUM_PHILOX_ROUNDS=3 RUN --mount=type=cache,target=/root/.cache/uv \ bash -ex <<"EOF" export NVTE_CUDA_ARCHS="80;90;100" + # flash-mla (no_pypi_wheels group) has no PyPI wheel and is built from source by uv. Point the + # compilers at the CCCL/libcu++ headers (under cccl/ in this base image, not the default CUDA + # include path) and scope the build to the target archs. Skipped for the LTS image. + FLASH_MLA_GROUP="" + if [ "$IMAGE_TYPE" != "lts" ]; then + FLASH_MLA_GROUP="--group no_pypi_wheels" + export FLASH_MLA_DISABLE_SM90=1 NVCC_THREADS=16 \ + CFLAGS="-I/usr/local/cuda/include/cccl" CXXFLAGS="-I/usr/local/cuda/include/cccl" + fi uv venv ${UV_PROJECT_ENVIRONMENT} --system-site-packages uv sync --only-group build - uv sync --extra ${IMAGE_TYPE} --extra mlm --extra ssm --extra te --link-mode copy --locked \ + uv sync --extra ${IMAGE_TYPE} --extra mlm --extra ssm --extra te ${FLASH_MLA_GROUP} --link-mode copy --locked \ --no-install-package torch \ --no-install-package torchvision \ --no-install-package triton \