Fix KV cache page size patch: separate groups for large SWA pages
This commit is contained in:
@@ -1,11 +1,17 @@
|
||||
#!/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.
|
||||
DeepseekV4 has three cache types:
|
||||
- C128A (HCA): compress_ratio=128, very small page size
|
||||
- C4A (CSA): compress_ratio=4, medium page size
|
||||
- SWA: compress_ratio=1, large page size
|
||||
|
||||
The upstream code assumes SWA page sizes <= MLA page sizes and pads
|
||||
SWA pages to match MLA. This breaks when SWA pages are LARGER than
|
||||
MLA pages (which is always the case for DeepseekV4).
|
||||
|
||||
Our fix: when SWA pages exceed MLA pages, put them in their own
|
||||
separate cache group without padding.
|
||||
"""
|
||||
import sys
|
||||
|
||||
@@ -17,31 +23,103 @@ def patch(path):
|
||||
print("Already patched, skipping")
|
||||
return
|
||||
|
||||
# Replace the assertion + candidate selection with a version that handles
|
||||
# SWA page sizes being larger than MLA page sizes
|
||||
# The old code: asserts SWA pages <= MLA pages, then pads SWA to MLA
|
||||
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)"""
|
||||
size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps)
|
||||
# Pad and collect layer names per page size.
|
||||
for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
|
||||
current_size = layer_spec.page_size_bytes
|
||||
candidate = size_to_candidate[current_size]
|
||||
if current_size < candidate:
|
||||
object.__setattr__(layer_spec, "page_size_padded", candidate)
|
||||
layers_per_size[candidate].append(layer_name)
|
||||
# NOTE(yifan): for now, inside a UniformKV group, each page_size should
|
||||
# have the same number of layers. This also means we don't need to pad layers
|
||||
# inside a partial-full layer tuple.
|
||||
assert len(set(len(layers) for layers in layers_per_size.values())) == 1
|
||||
num_layers_per_size = len(next(iter(layers_per_size.values())))
|
||||
|
||||
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.
|
||||
# Split layers inside each UniformKV group for aligned #(layers).
|
||||
# See `_get_kv_cache_groups_uniform_page_size` for more details.
|
||||
num_tuple_groups = cdiv(num_layers_per_size, num_layer_tuples)
|
||||
layer_tuples = list(zip(*layers_per_size.values()))
|
||||
for i in range(num_tuple_groups):
|
||||
group_layer_tuples = layer_tuples[i::num_tuple_groups]
|
||||
# Flatten tuples and build dict for from_specs
|
||||
group_layer_names = [
|
||||
name for layer_tuple in group_layer_tuples for name in layer_tuple
|
||||
]
|
||||
group_layer_specs = {
|
||||
name: sm_spec.kv_cache_specs[name] for name in group_layer_names
|
||||
}
|
||||
sub_sm_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
|
||||
assert sub_sm_spec is not None
|
||||
swa_mla_groups.append(
|
||||
KVCacheGroupSpec(
|
||||
layer_names=group_layer_names,
|
||||
kv_cache_spec=sub_sm_spec,
|
||||
)
|
||||
)"""
|
||||
|
||||
# The new code: handle both cases
|
||||
new = """ # CLAWMINE_PATCH_KV_CACHE: Handle DeepseekV4 where SWA page sizes
|
||||
# can be larger than MLA page sizes. Two cases:
|
||||
# 1. All SWA pages <= some MLA page: original padding logic
|
||||
# 2. Some SWA pages > all MLA pages: separate cache group, no padding
|
||||
max_mla_page = max(all_page_sizes)
|
||||
can_pad = max(sm_page_sizes) <= max_mla_page
|
||||
|
||||
# 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 can_pad:
|
||||
# Original logic: pad SWA pages to nearest MLA page
|
||||
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)
|
||||
for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
|
||||
current_size = layer_spec.page_size_bytes
|
||||
candidate = size_to_candidate[current_size]
|
||||
if current_size < candidate:
|
||||
object.__setattr__(layer_spec, "page_size_padded", candidate)
|
||||
layers_per_size[candidate].append(layer_name)
|
||||
assert len(set(len(layers) for layers in layers_per_size.values())) == 1
|
||||
num_layers_per_size = len(next(iter(layers_per_size.values())))
|
||||
num_tuple_groups = cdiv(num_layers_per_size, num_layer_tuples)
|
||||
layer_tuples = list(zip(*layers_per_size.values()))
|
||||
for i in range(num_tuple_groups):
|
||||
group_layer_tuples = layer_tuples[i::num_tuple_groups]
|
||||
group_layer_names = [
|
||||
name for layer_tuple in group_layer_tuples for name in layer_tuple
|
||||
]
|
||||
group_layer_specs = {
|
||||
name: sm_spec.kv_cache_specs[name] for name in group_layer_names
|
||||
}
|
||||
sub_sm_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
|
||||
assert sub_sm_spec is not None
|
||||
swa_mla_groups.append(
|
||||
KVCacheGroupSpec(
|
||||
layer_names=group_layer_names,
|
||||
kv_cache_spec=sub_sm_spec,
|
||||
)
|
||||
)
|
||||
else:
|
||||
# SWA pages are larger than MLA pages.
|
||||
# Put each SWA layer in its own cache group (no padding needed).
|
||||
# This is the DeepseekV4 Blackwell case where compress_ratio=1
|
||||
# layers have much larger pages than compressed layers.
|
||||
for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
|
||||
group_layer_specs = {layer_name: layer_spec}
|
||||
sub_sm_spec = UniformTypeKVCacheSpecs.from_specs(group_layer_specs)
|
||||
if sub_sm_spec is not None:
|
||||
swa_mla_groups.append(
|
||||
KVCacheGroupSpec(
|
||||
layer_names=[layer_name],
|
||||
kv_cache_spec=sub_sm_spec,
|
||||
)
|
||||
)"""
|
||||
|
||||
if old not in content:
|
||||
print("ERROR: Could not find the code to patch")
|
||||
|
||||
Reference in New Issue
Block a user