diff --git a/vllm/patches/patch_kv_cache_utils.py b/vllm/patches/patch_kv_cache_utils.py index 30ca6fcd..bd7faab6 100644 --- a/vllm/patches/patch_kv_cache_utils.py +++ b/vllm/patches/patch_kv_cache_utils.py @@ -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")