diff --git a/Dockerfile b/Dockerfile index 6d2e9c5d..db23aa14 100644 --- a/Dockerfile +++ b/Dockerfile @@ -51,6 +51,11 @@ COPY vllm/patches/layers/csa_attention.py ${VLLM_LAYERS_DIR}/csa_attention.py ARG VLLM_NVFP4_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear/nvfp4 COPY vllm/kernels/linear/nvfp4/cutedsl.py ${VLLM_NVFP4_DIR}/cutedsl.py +# Patch KV cache utils to handle DeepseekV4 SWA page sizes > MLA page sizes +ARG VLLM_CORE_DIR=/usr/local/lib/python3.12/dist-packages/vllm/v1/core +COPY vllm/patches/patch_kv_cache_utils.py /tmp/patch_kv_cache_utils.py +RUN python3 /tmp/patch_kv_cache_utils.py ${VLLM_CORE_DIR}/kv_cache_utils.py && rm /tmp/patch_kv_cache_utils.py + # Register CuTeDSL kernel in vLLM's linear kernel selection ARG VLLM_LINEAR_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/kernels/linear COPY vllm/patches/register_cutedsl_kernel.py /tmp/register_cutedsl_kernel.py diff --git a/vllm/patches/patch_kv_cache_utils.py b/vllm/patches/patch_kv_cache_utils.py new file mode 100644 index 00000000..30ca6fcd --- /dev/null +++ b/vllm/patches/patch_kv_cache_utils.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python3 +"""Patch vLLM kv_cache_utils.py to handle DeepseekV4 SWA page sizes. + +The upstream code asserts that SWA page sizes are <= MLA page sizes, +but DeepseekV4's SWA layers (compress_ratio=1) have page sizes much larger +than compressed MLA layers (C128A, C4A). This patch relaxes the assertion +and handles the case where SWA pages are larger than MLA pages by keeping +them in separate cache groups without padding. +""" +import sys + +def patch(path): + with open(path, 'r') as f: + content = f.read() + + if "CLAWMINE_PATCH_KV_CACHE" in content: + print("Already patched, skipping") + return + + # Replace the assertion + candidate selection with a version that handles + # SWA page sizes being larger than MLA page sizes + old = """ assert max(sm_page_sizes) <= max(all_page_sizes) + + # Unify page size by padding layers' page_size to the nearest larger page_size. + # Compute candidate (nearest larger page_size) for each unique page size. + size_to_candidate: dict[int, int] = {} + for ps in sm_page_sizes: + size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps)""" + + new = """ # CLAWMINE_PATCH_KV_CACHE: relax assertion for DeepseekV4 where + # SWA page sizes can be larger than MLA page sizes. + # When SWA pages exceed all MLA pages, we keep them unpadded. + max_mla_page = max(all_page_sizes) + + # Unify page size by padding layers' page_size to the nearest larger page_size. + # Compute candidate (nearest larger page_size) for each unique page size. + size_to_candidate: dict[int, int] = {} + for ps in sm_page_sizes: + candidates = [x for x in all_page_sizes if x >= ps] + if candidates: + size_to_candidate[ps] = min(candidates) + else: + # No MLA page size large enough — keep SWA page as-is + size_to_candidate[ps] = ps""" + + if old not in content: + print("ERROR: Could not find the code to patch") + sys.exit(1) + + content = content.replace(old, new) + + with open(path, 'w') as f: + f.write(content) + print("Patched kv_cache_utils.py for DeepseekV4 SWA page sizes") + +if __name__ == "__main__": + patch(sys.argv[1])