From 39310c357dca211c364b9af3d0c1b9a8b57c4357 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Tue, 19 May 2026 09:52:23 +0000
Subject: [PATCH] Patch compressor cache for Blackwell (no FlashMLA alignment)
 - fixes 91 missing layers

---
 Dockerfile                             |  2 ++
 vllm/patches/patch_compressor_cache.py | 48 ++++++++++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 vllm/patches/patch_compressor_cache.py

diff --git a/Dockerfile b/Dockerfile
index 7b838e44..c6f00da2 100644
--- a/Dockerfile
+++ b/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
diff --git a/vllm/patches/patch_compressor_cache.py b/vllm/patches/patch_compressor_cache.py
new file mode 100644
index 00000000..d278ae59
--- /dev/null
+++ b/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])