Patch compressor cache for Blackwell (no FlashMLA alignment) - fixes 91 missing layers

Dockerfile

@@ -64,6 +64,8 @@ COPY vllm/patches/patch_swa_cache.py /tmp/patch_swa_cache.py
 RUN python3 /tmp/patch_swa_cache.py ${VLLM_SPARSE_SWA_DIR}/sparse_swa.py && rm /tmp/patch_swa_cache.py
 COPY vllm/patches/patch_indexer_cache.py /tmp/patch_indexer_cache.py
 RUN python3 /tmp/patch_indexer_cache.py ${VLLM_LAYERS_DIR2}/deepseek_v4_attention.py && rm /tmp/patch_indexer_cache.py
+COPY vllm/patches/patch_compressor_cache.py /tmp/patch_compressor_cache.py
+RUN python3 /tmp/patch_compressor_cache.py ${VLLM_LAYERS_DIR2}/deepseek_compressor.py && rm /tmp/patch_compressor_cache.py
 
 # Debug: print layer name mismatch
 ARG VLLM_WORKER_DIR=/usr/local/lib/python3.12/dist-packages/vllm/v1/worker

vllm/patches/patch_compressor_cache.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+"""Patch DeepseekCompressor cache for Blackwell: remove FlashMLA alignment."""
+import sys
+
+def patch(path):
+    with open(path, 'r') as f:
+        content = f.read()
+
+    if "CLAWMINE_PATCH_COMPRESSOR" in content:
+        print("Already patched, skipping")
+        return
+
+    old = """        return SlidingWindowMLASpec(  # only has one vector instead of K + V
+            block_size=self.block_size,
+            num_kv_heads=1,
+            head_size=self.state_dim,
+            dtype=self.dtype,
+            sliding_window=self.sliding_window,
+            alignment=576,  # NOTE: FlashMLA requires 576B alignment
+        )"""
+
+    new = """        # CLAWMINE_PATCH_COMPRESSOR: No FlashMLA alignment on Blackwell
+        from vllm.platforms import current_platform
+        _is_blackwell = (
+            current_platform.get_device_capability() is not None
+            and current_platform.get_device_capability().major >= 10
+        )
+        return SlidingWindowMLASpec(  # only has one vector instead of K + V
+            block_size=self.block_size,
+            num_kv_heads=1,
+            head_size=self.state_dim,
+            dtype=self.dtype,
+            sliding_window=self.sliding_window,
+            alignment=None if _is_blackwell else 576,
+        )"""
+
+    if old not in content:
+        print("ERROR: Could not find the code to patch in " + path)
+        sys.exit(1)
+
+    content = content.replace(old, new)
+
+    with open(path, 'w') as f:
+        f.write(content)
+    print("Patched DeepseekCompressor for Blackwell")
+
+if __name__ == "__main__":
+    patch(sys.argv[1])