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

Also replace layers/mhc.py CustomOp dispatch

Browse Source 
The original layers/mhc.py forward_cuda calls
torch.ops.vllm.mhc_pre_tilelang which triggers TileLang JIT.
Replace with our torch implementations in forward_cuda.
This is what the CustomOp dispatch routes through.
This commit is contained in:
biondizzle
2026-05-19 05:31:05 +00:00
parent 5e6d459145
commit e404e18efb
2 changed files with 106 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
Dockerfile
View File

@@ -40,8 +40,11 @@ COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attent
COPY vllm/patches/layers/deepseek_compressor.py ${VLLM_LAYERS_DIR}/deepseek_compressor.py
# Replace MHC TileLang kernels with pure PyTorch (avoids TileLang JIT on Blackwell)
# Install our torch op implementations and patch the MHC kernels __init__
# 1. Patch layers/mhc.py — CustomOp dispatch uses torch impls instead of tilelang
# 2. Install our torch op registrations (mhc_torch_ops.py)
# 3. Patch kernels/mhc/__init__.py to not import tilelang
ARG VLLM_MHC_KERNELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/mhc
COPY vllm/patches/layers/mhc.py ${VLLM_LAYERS_DIR}/mhc.py
COPY vllm/patches/kernels/mhc_torch_ops.py ${VLLM_MHC_KERNELS_DIR}/mhc_torch_ops.py
RUN echo 'from .torch import *' > ${VLLM_MHC_KERNELS_DIR}/__init__.py && \
echo 'from .mhc_torch_ops import *' >> ${VLLM_MHC_KERNELS_DIR}/__init__.py