Patch KV cache utils: handle DeepseekV4 SWA page sizes > MLA page sizes

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

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])