fix

Author: WANDY666

.github/workflows/docker-publish.yml

@@ -118,10 +118,7 @@ jobs:
           # Remove the container
           docker rm "$CONTAINER_ID"
           # Verify wheel file has been copied
-          echo "Exported wheel files:"
-          ls -la ./out
-          WHEEL_COUNT=$(ls -1 ./out/*.whl | wc -l)
-          echo "Total wheel files: $WHEEL_COUNT"
+          ls ./out
 
       - name: Upload .whl to GitHub Release
         if: startsWith(github.ref, 'refs/tags/')
@@ -132,7 +129,7 @@ jobs:
           
           # Get wheel file information
           LIGHTKERNEL_WHEEL=$(ls ./out/lightllm_kernel-*.whl | head -n1 | xargs basename)
-          FLASH_ATTN_WHEEL=$(ls ./out/flash_attn_3-*.whl | head -n1 | xargs basename)
+          FA3_MTP_WHEEL=$(ls ./out/flash_attn_3-*.whl | head -n1 | xargs basename)
           
           gh release create "$TAG_NAME" ./out/*.whl \
             --title "LightKernel Release $TAG_NAME" \
@@ -142,20 +139,20 @@ jobs:
 
           📦 **Packages:**
           - **$LIGHTKERNEL_WHEEL** - Core CUDA kernel library  
-          - **$FLASH_ATTN_WHEEL** - Flash Attention 3.0 (Hopper)
+          - **$FA3_MTP_WHEEL** - Fa3 MTP 3.0 (Hopper)
 
           🚀 **Quick Installation:**
           \`\`\`bash
           # Install both packages
           pip install https://github.com/${{ github.repository }}/releases/download/$TAG_NAME/$LIGHTKERNEL_WHEEL
-          pip install https://github.com/${{ github.repository }}/releases/download/$TAG_NAME/$FLASH_ATTN_WHEEL
+          pip install https://github.com/${{ github.repository }}/releases/download/$TAG_NAME/$FA3_MTP_WHEEL
           \`\`\`
 
           🔧 **Build Environment:**
           - CUDA: 12.6.1 with cuDNN
           - PyTorch: 2.7.1  
           - Python: 3.10
-          - Architecture: CUDA Compute Capabilities 8.0, 8.6, 8.9, 9.0
+          - Architecture: CUDA Compute Capabilities 9.0
 
           🐳 **Docker Image:** \`${{ steps.meta.outputs.tags }}\`
           

Dockerfile

@@ -7,20 +7,15 @@ 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 \
     curl \
     g++ \
     make \
-    git \
-    cmake \
-    ninja-build \
-    build-essential && \
+    git && \
     rm -rf /var/lib/apt/lists/*
 
-# Install Mambaforge
 RUN case ${TARGETPLATFORM} in \
     "linux/arm64")  MAMBA_ARCH=aarch64  ;; \
     *)              MAMBA_ARCH=x86_64   ;; \
@@ -29,24 +24,18 @@ 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
 
-# Install PyTorch with CUDA support
-RUN pip install torch==2.7.1
+WORKDIR /root
 
-# Install build dependencies
-RUN pip install --upgrade pip setuptools wheel build scikit-build-core[pyproject] pybind11 ninja
+RUN pip install torch==2.7.1
 
-# Copy source code to container
 COPY . .
 
 # 🔧 设置 PyTorch 路径，让 CMake 能找到 Torch 配置
@@ -55,18 +44,14 @@ 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"
 
-# Set environment variables for building
-ENV FLASH_ATTENTION_FORCE_BUILD=TRUE \
-    FLASH_ATTENTION_DISABLE_BACKWARD=TRUE \
-    CUDA_HOME=/usr/local/cuda \
+ENV 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)  
@@ -82,17 +67,6 @@ 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=2 NVCC_THREADS=2 FLASH_ATTN_CUDA_ARCHS=90 python setup.py bdist_wheel && \
+    MAX_JOBS=1 NVCC_THREADS=1 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