From b04bff7e8ba751671e1e528d852d12c1974b4f53 Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 16 May 2026 03:50:07 +0000 Subject: [PATCH] feat: clean Dockerfile, docker-compose, import fixes for CuTeDSL build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dockerfile: - Removed: C++ CUTLASS extension build, TileLang install, CUTLASS clone - Added: nvidia-cutlass-dsl==4.5.0 install, cutedsl/ copy - Copy nvfp4_cutedsl.py to vllm models dir - Verify step checks cutlass import docker-compose.yml: - Removed stale env vars (MEGA_MOE_DEBUG, MEGA_MOE_STATIC, etc.) deepseek_v4.py: - Fix import: vllm.nvfp4_cutedsl → vllm.model_executor.models.nvfp4_cutedsl README.md: - Updated results: 0% weight loss confirmed (bit-identical view-cast) - 1.1% cosine loss is entirely from activation quantization --- Dockerfile | 29 ++++++++++++----------------- README.md | 13 ++++++++----- docker-compose.yml | 8 -------- vllm/patches/deepseek_v4.py | 2 +- 4 files changed, 21 insertions(+), 31 deletions(-) diff --git a/Dockerfile b/Dockerfile index 8155124b..bcc7b5d3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -# DeepSeek V4 NVFP4 vLLM + CUTLASS NVFP4 Mega MoE Kernel +# DeepSeek V4 NVFP4 vLLM + CuTeDSL NVFP4 MoE Kernel FROM vllm/vllm-openai:nightly-x86_64 # Remove broken nixl_ep (built against CUDA 12, image is CUDA 13) @@ -7,32 +7,25 @@ RUN pip uninstall -y nixl-ep; rm -rf /usr/local/lib/python3.12/dist-packages/nix RUN apt-get update && apt-get install -y git screen cmake libcusolver-dev-13-0 libcusparse-dev-13-0 libcublas-dev-13-0 libcurand-dev-13-0 libcufft-dev-13-0 libnvjitlink-dev-13-0 && rm -rf /var/lib/apt/lists/* # Remove the broken symlink if it exists -RUN rm -f /usr/local/cuda/lib64/libcudart.so.12 +RUN rm -f /usr/local/cuda/lib64/libcudrt.so.12 ENV CUDA_HOME=/usr/local/cuda 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 +# Install CuTeDSL (NVFP4 block-scaled GEMM kernel framework) +RUN pip install nvidia-cutlass-dsl==4.5.0 nvidia-cutlass-dsl-libs-base==4.5.0 ARG CACHE_BUSTER=${TIMESTAMP} -# Copy and install the NVFP4 mega_moe kernel (from this repo) +# Copy the NVFP4 mega_moe Python kernel (no C++ build needed) COPY src/ /root/nvfp4-megamoe-kernel/src/ COPY pyproject.toml /root/nvfp4-megamoe-kernel/pyproject.toml RUN cd /root/nvfp4-megamoe-kernel && pip install -e . -# Build the CUTLASS NVFP4 block-scaled GEMM extension -RUN cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm && \ - mkdir -p cutlass_nvfp4_gemm && \ - CUTLASS_INCLUDE_DIR=/root/cutlass/include \ - TORCH_CUDA_ARCH_LIST=10.0 \ - python3 setup.py build_ext --inplace +# Copy the CuTeDSL kernel and bridge layer +COPY cutedsl/ /root/nvfp4-megamoe-kernel/cutedsl/ -# Install TileLang (for potential future use) -RUN pip install tilelang - -ENV PYTHONPATH="/root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm:/root/nvfp4-megamoe-kernel:${PYTHONPATH}" +ENV PYTHONPATH="/root/nvfp4-megamoe-kernel:${PYTHONPATH}" # Patch vLLM — overwrite model files and register architecture ARG VLLM_MODELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models @@ -40,11 +33,13 @@ ARG VLLM_LAYERS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/ COPY vllm/patches/deepseek_v4.py ${VLLM_MODELS_DIR}/deepseek_v4.py COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attention.py +COPY vllm/nvfp4_cutedsl.py ${VLLM_MODELS_DIR}/nvfp4_cutedsl.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')" && \ +RUN python3 -c "import torch; print(f'PyTorch {torch.__version__} OK')" && \ python3 -c "import vllm; print('vLLM OK')" && \ - python3 -c "import nvfp4_megamoe_kernel; print('NVFP4 kernel OK')" + python3 -c "import nvfp4_megamoe_kernel; print('NVFP4 kernel OK')" && \ + python3 -c "import cutlass; print('CuTeDSL OK')" diff --git a/README.md b/README.md index da2e9922..c8edece3 100644 --- a/README.md +++ b/README.md @@ -36,8 +36,11 @@ NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLIR → PT **Results with real DeepSeek-V4 layer 0 weights:** - L1 GEMM alone: cosine 0.995 - Full MoE pipeline (L1→SiLU→L2→scatter): cosine 0.989 +- Weight loading: **0% loss** — direct uint8→float4_e2m1fn_x2 view-cast, bit-identical to checkpoint +- Activation quantization: ~1.1% cosine loss (dynamic BF16→NVFP4 — inherent to the format, unavoidable) +- GEMM kernel: 0% loss (CuTeDSL is correct) -The 0.989 is expected quantization loss — we dequantize the checkpoint NVFP4→BF16 then re-quantize BF16→NVFP4 (double quantization). Future optimization: load checkpoint FP4 bytes directly into `float4_e2m1fn_x2` tensors. +The 0.989 cosine is entirely from activation quantization. The weights are bit-identical to the checkpoint — no BF16 round-trip, no precision loss. ### Key Lessons @@ -112,14 +115,14 @@ python3 layertest.py ### Phase 2: vLLM Integration (IN PROGRESS) - Wire `cutedsl/moe_pipeline.py` into the vLLM DeepSeek-V4 model -- Replace `nvfp4_mega_moe_full()` call with `run_nvfp4_moe()` -- Handle weight loading: checkpoint uint8 → float4_e2m1fn_x2 directly (no BF16 round-trip) +- Replace `nvfp4_mega_moe_full()` call with `CuTeDSLMoERunner.run()` +- Weight loading: checkpoint uint8 → float4_e2m1fn_x2 view-cast (bit-preserving, no BF16 round-trip) +- Block scales (float8_e4m3fn) and global scales (float32) pass through directly from checkpoint +- L1 dual global scale handling: normalize to max(gate_gs, up_gs), fold ratio into block scales - Remove C++ CUTLASS extension build from Dockerfile - Add CuTeDSL dependency to the Docker build ### Phase 3: Optimization -- Load checkpoint FP4 bytes directly into float4_e2m1fn_x2 (skip BF16 round-trip) -- Use checkpoint block scales directly (skip dequant→requant) - Explore larger tile sizes for better occupancy - Profile end-to-end inference on full model diff --git a/docker-compose.yml b/docker-compose.yml index a9c2c547..b195d0f3 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -9,14 +9,6 @@ services: - OMP_NUM_THREADS=128 - CUDA_LAUNCH_BLOCKING=0 - TORCH_SHOW_CPP_STACKTRACES=0 - - MEGA_MOE_DEBUG=0 - - MEGA_MOE_STATIC=0 - - NVFP4_DEBUG=0 - - NVFP4_DEBUG_SYNC=0 - - SKIP_ATTENTION=0 - - MEGA_MOE_USE_CUTLASS=1 - - DG_JIT_DEBUG=0 - - DEEP_GEMM_JIT_DEBUG=0 command: - /model - --trust-remote-code diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py index 9836456b..54cb5017 100644 --- a/vllm/patches/deepseek_v4.py +++ b/vllm/patches/deepseek_v4.py @@ -430,7 +430,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): # (one multiply + float8 round-trip on the *ratio only* — much better # than dequantizing the entire weight matrix through BF16). - from vllm.nvfp4_cutedsl import CuTeDSLMoERunner + from vllm.model_executor.models.nvfp4_cutedsl import CuTeDSLMoERunner l1_fp4, l1_sf, l1_gs = [], [], [] l2_fp4, l2_sf, l2_gs = [], [], []