biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

more fixes

Browse Source
This commit is contained in:
biondizzle
2026-05-14 19:35:39 +00:00
parent 3be051e140
commit 1ceff541b0
3 changed files with 18 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
Dockerfile
View File

@@ -15,6 +15,8 @@ ENV TORCH_CUDA_ARCH_LIST="10.0"
# Clone latest CUTLASS (has NVFP4 block-scaled MMA support)
RUN git clone --depth 1 https://github.com/NVIDIA/cutlass.git /root/cutlass
ARG CACHE_BUSTER=${TIMESTAMP}
# Copy and install the NVFP4 mega_moe kernel (from this repo)
COPY src/ /root/nvfp4-megamoe-kernel/src/
COPY pyproject.toml /root/nvfp4-megamoe-kernel/pyproject.toml
@@ -32,18 +34,16 @@ RUN pip install tilelang
ENV PYTHONPATH="/root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm:/root/nvfp4-megamoe-kernel:${PYTHONPATH}"
# Copy vLLM patches
COPY vllm/patches/deepseek_v4.py /tmp/patches/deepseek_v4.py
COPY vllm/patches/staging_kernel.py /tmp/patches/staging_kernel.py
COPY vllm/patches/deepseek_v4_attention.py /tmp/patches/deepseek_v4_attention.py
# Patch vLLM — overwrite model files and register architecture
ARG VLLM_MODELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models
ARG VLLM_LAYERS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers
# Apply patches
RUN VLLM_MODELS_DIR=$(python3 -c "import vllm.model_executor.models; import os; print(os.path.dirname(vllm.model_executor.models.__file__))") && \
VLLM_LAYERS_DIR=$(python3 -c "import vllm.model_executor.layers; import os; print(os.path.dirname(vllm.model_executor.layers.__file__))") && \
cp /tmp/patches/deepseek_v4.py "$VLLM_MODELS_DIR/deepseek_v4.py" && \
cp /tmp/patches/staging_kernel.py "$VLLM_MODELS_DIR/staging_kernel.py" && \
cp /tmp/patches/deepseek_v4_attention.py "$VLLM_LAYERS_DIR/deepseek_v4_attention.py" && \
rm -rf /tmp/patches
COPY vllm/patches/deepseek_v4.py ${VLLM_MODELS_DIR}/deepseek_v4.py
COPY vllm/patches/staging_kernel.py ${VLLM_MODELS_DIR}/staging_kernel.py
COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attention.py
RUN sed -i 's/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),\n "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),/' \
${VLLM_MODELS_DIR}/registry.py
# Verify
RUN python3 -c "import torch; import cutlass_nvfp4_gemm._C; print('CUTLASS NVFP4 OK')" && \