fix

Author: WANDY666

Dockerfile

@@ -7,6 +7,7 @@ ARG TARGETPLATFORM
 ENV PATH=/opt/conda/bin:$PATH \
     CONDA_PREFIX=/opt/conda
 
+# Install system dependencies
 RUN chmod 777 -R /tmp && apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
     ca-certificates \
     libssl-dev \
@@ -19,6 +20,7 @@ RUN chmod 777 -R /tmp && apt-get update && DEBIAN_FRONTEND=noninteractive apt-ge
     build-essential && \
     rm -rf /var/lib/apt/lists/*
 
+# Install Mambaforge
 RUN case ${TARGETPLATFORM} in \
     "linux/arm64")  MAMBA_ARCH=aarch64  ;; \
     *)              MAMBA_ARCH=x86_64   ;; \
@@ -27,18 +29,24 @@ RUN case ${TARGETPLATFORM} in \
     bash ~/mambaforge.sh -b -p /opt/conda && \
     rm ~/mambaforge.sh
 
+# Install Python
 RUN case ${TARGETPLATFORM} in \
     "linux/arm64")  exit 1 ;; \
     *)              /opt/conda/bin/conda update -y conda &&  \
     /opt/conda/bin/conda install -y "python=${PYTHON_VERSION}" ;; \
     esac && \
     /opt/conda/bin/conda clean -ya
 
+# Set working directory
+WORKDIR /workspace
 
-WORKDIR /root
+# Install PyTorch with CUDA support
+RUN pip install torch==2.7.1
 
-RUN pip install torch==2.7.1 numpy
+# Install build dependencies
+RUN pip install --upgrade pip setuptools wheel build scikit-build-core[pyproject] pybind11 ninja
 
+# Copy source code to container
 COPY . .
 
 # 🔧 设置 PyTorch 路径，让 CMake 能找到 Torch 配置
@@ -47,14 +55,18 @@ RUN python -c "import torch; print(f'PyTorch installed at: {torch.__path__[0]}')
     TORCH_PATH=$(python -c "import torch; print(torch.utils.cmake_prefix_path)") && \
     echo "Torch CMAKE path: $TORCH_PATH"
 
-ENV CUDA_HOME=/usr/local/cuda \
+# Set environment variables for building
+ENV FLASH_ATTENTION_FORCE_BUILD=TRUE \
+    FLASH_ATTENTION_DISABLE_BACKWARD=TRUE \
+    CUDA_HOME=/usr/local/cuda \
     CUDA_ROOT=/usr/local/cuda
 
 # 🎯 关键修复：设置 CMAKE_PREFIX_PATH 让 CMake 找到 PyTorch
 RUN TORCH_CMAKE_PATH=$(python -c "import torch; print(torch.utils.cmake_prefix_path)") && \
     echo "export CMAKE_PREFIX_PATH=$TORCH_CMAKE_PATH:\$CMAKE_PREFIX_PATH" >> ~/.bashrc && \
     echo "CMAKE_PREFIX_PATH=$TORCH_CMAKE_PATH" >> /etc/environment
 
+# Create output directory
 RUN mkdir -p /out
 
 # Build lightllm-kernel package (main project)  
@@ -70,6 +82,17 @@ RUN echo "🔧 Building lightllm-kernel package..." && \
 # Build flash_attn_3 package (hopper)
 RUN echo "🔧 Building flash_attn_3 package..." && \
     cd flash-attention/hopper && \
-    MAX_JOBS=1 NVCC_THREADS=1 FLASH_ATTN_CUDA_ARCHS=90 FLASH_ATTENTION_DISABLE_SM80=TRUE python setup.py bdist_wheel && \
+    MAX_JOBS=2 NVCC_THREADS=2 FLASH_ATTN_CUDA_ARCHS=90 FLASH_ATTENTION_DISABLE_SM80=TRUE python setup.py bdist_wheel && \
     cp dist/*.whl /out/ && \
     echo "✅ flash_attn_3 build completed"
+
+# Verify all wheels are built
+RUN echo "📦 Final wheel packages:" && \
+    ls -la /out/ && \
+    WHEEL_COUNT=$(ls -1 /out/*.whl | wc -l) && \
+    echo "Total wheels built: $WHEEL_COUNT" && \
+    if [ "$WHEEL_COUNT" -ne 2 ]; then \
+        echo "❌ Error: Expected 2 wheels, found $WHEEL_COUNT" && exit 1; \
+    else \
+        echo "✅ Successfully built all wheel packages"; \
+    fi