58 lines
2.2 KiB
Python
58 lines
2.2 KiB
Python
#!/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])
|