Fixed aarch64 flash-mla (#12009)

This commit is contained in:
nvjullin
2025-10-24 08:47:04 +08:00
committed by GitHub
parent 6d6e24bcc4
commit 9a71500cfb

View File

@@ -141,16 +141,14 @@ RUN cd /sgl-workspace/DeepEP && \
NVSHMEM_DIR=${NVSHMEM_DIR} TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install . NVSHMEM_DIR=${NVSHMEM_DIR} TORCH_CUDA_ARCH_LIST="${CHOSEN_TORCH_CUDA_ARCH_LIST}" pip install .
# Install flashmla # Install flashmla
RUN if [ "$TARGETARCH" = "amd64" ]; then \ RUN git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla && \
git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla && \
cd flash-mla && \ cd flash-mla && \
git checkout ${FLASHMLA_COMMIT} && \ git checkout ${FLASHMLA_COMMIT} && \
git submodule update --init --recursive && \ git submodule update --init --recursive && \
if [ "$CUDA_VERSION" = "12.6.1" ]; then \ if [ "$CUDA_VERSION" = "12.6.1" ]; then \
export FLASH_MLA_DISABLE_SM100=1; \ export FLASH_MLA_DISABLE_SM100=1; \
fi && \ fi && \
pip install -v . ; \ pip install -v . ;
fi
# Python tools # Python tools
RUN python3 -m pip install --no-cache-dir \ RUN python3 -m pip install --no-cache-dir \