--- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1210,17 +1210,38 @@ for name, spec in kv_cache_spec.items() if isinstance(spec, KpoolTailSpec) } + hidden_specs = { + name: spec + for name, spec in kv_cache_spec.items() + if isinstance(spec, HiddenStateCacheSpec) + } attn_specs = { name: spec for name, spec in kv_cache_spec.items() - if not isinstance(spec, (MambaSpec, KpoolTailSpec)) + if not isinstance( + spec, (MambaSpec, KpoolTailSpec, HiddenStateCacheSpec) + ) } - if not mamba_specs or not all( - type(spec) is MLAAttentionSpec for spec in attn_specs.values() - ): + if not mamba_specs or not attn_specs: + return None + # FORK-LOCAL (patches/glm53-dflash-04-kv-groups-draft): with a dflash + # drafter, attn_specs also carries the drafter's attention layers. The + # GLM MLA layers keep their own group (no page unification -- per-group + # pages are the glm5 design); the drafter layers must form one uniform + # group of their own, and the hidden-state layers get per-layer groups + # like the generic tail. + mla_specs = { + name: spec + for name, spec in attn_specs.items() + if type(spec) is MLAAttentionSpec + } + draft_specs = { + name: spec for name, spec in attn_specs.items() if name not in mla_specs + } + if not mla_specs: + return None + if draft_specs and UniformTypeKVCacheSpecs.from_specs(draft_specs) is None: return None - - mla_specs = cast(dict[str, MLAAttentionSpec], attn_specs) idx_pages = { spec.page_size_bytes for spec in mla_specs.values() if spec.tokens_per_state > 1 } @@ -1233,7 +1254,7 @@ mla_pages = {mla_specs[name].page_size_bytes for name in mla_names} assert len(mla_pages) == 1 mla_page = mla_pages.pop() - uniform_spec = UniformTypeKVCacheSpecs.from_specs(attn_specs) + uniform_spec = UniformTypeKVCacheSpecs.from_specs(mla_specs) assert uniform_spec is not None tail_group: KVCacheGroupSpec | None = None @@ -1270,10 +1291,20 @@ for index, name in enumerate(mamba_specs): mamba_grouped_names[index % num_groups].append(name) + draft_group: KVCacheGroupSpec | None = None + if draft_specs: + draft_uniform = UniformTypeKVCacheSpecs.from_specs(draft_specs) + assert draft_uniform is not None + draft_group = KVCacheGroupSpec(list(draft_specs), draft_uniform) + hidden_groups = [ + KVCacheGroupSpec([name], spec) for name, spec in hidden_specs.items() + ] return ( - [KVCacheGroupSpec(list(attn_specs), uniform_spec)] + [KVCacheGroupSpec(list(mla_specs), uniform_spec)] + ([tail_group] if tail_group is not None else []) + + ([draft_group] if draft_group is not None else []) + create_kv_cache_group_specs(padded_specs, mamba_grouped_names) + + hidden_groups ) @@ -1289,10 +1320,17 @@ int, list[str], int, + KVCacheGroupSpec | None, + list[KVCacheGroupSpec], ] | None ): - """Recognize the GLM-5.3-Flash grouping after optional PP projection.""" + """Recognize the GLM-5.3-Flash grouping after optional PP projection. + + FORK-LOCAL (patches/glm53-dflash-04-kv-groups-draft): also recognizes the + drafter's uniform group and the per-layer hidden-state groups, returned + as ``draft_group`` / ``hidden_groups`` so the pool builder and the memory + estimator can size their tensors beside the aliased GLM pool.""" uniform_groups = [ group for group in kv_cache_groups @@ -1301,17 +1339,28 @@ mamba_groups = [ group for group in kv_cache_groups if isinstance(group.kv_cache_spec, MambaSpec) ] + hidden_groups = [ + group + for group in kv_cache_groups + if isinstance(group.kv_cache_spec, HiddenStateCacheSpec) + ] attn_group: KVCacheGroupSpec | None = None tail_group: KVCacheGroupSpec | None = None + draft_group: KVCacheGroupSpec | None = None for group in uniform_groups: inner = cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).kv_cache_specs if all(type(spec) is MLAAttentionSpec for spec in inner.values()): attn_group = group elif all(isinstance(spec, KpoolTailSpec) for spec in inner.values()): tail_group = group + else: + draft_group = group if attn_group is None or not mamba_groups: return None - if len(uniform_groups) + len(mamba_groups) != len(kv_cache_groups): + if ( + len(uniform_groups) + len(mamba_groups) + len(hidden_groups) + != len(kv_cache_groups) + ): return None attn_uniform = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) @@ -1362,6 +1411,8 @@ idx_page, tail_names, tail_page, + draft_group, + hidden_groups, ) @@ -1395,10 +1446,34 @@ return kv_cache_spec max_page_size = max(page_sizes) + # GLM-5.3-Flash + a draft model: the sparse indexer k_cache page is + # tokens_per_state x the MLA page (same per-token bytes), and MLA layers + # cannot be PADDED to an arbitrary page size (sparse MLA indexes the cache + # in whole token rows) -- but they CAN grow their block size when the max + # page is an exact multiple of the MLA page. Bump the max page to a + # multiple of every MLA-family page so all of them take the divisible + # block-growth branch below instead of the padding refusal, while the + # draft and mamba layers simply pad to the (slightly larger) page. + mla_pages = { + spec.page_size_bytes + for spec in kv_cache_spec.values() + if isinstance(spec, MLAAttentionSpec) + } + if mla_pages: + mla_common = math.lcm(*mla_pages) + max_page_size = -(-max_page_size // mla_common) * mla_common new_kv_cache_spec = {} for layer_name, layer_spec in kv_cache_spec.items(): if layer_spec.page_size_bytes == max_page_size: new_kv_cache_spec[layer_name] = layer_spec + elif isinstance(layer_spec, KpoolTailSpec): + # The kpool tail kernel pins the tail block to the kpool size + # (tail_kv_cache.view(...) must be [blocks, 2, index_kpool, + # head_dim]); growing block_size breaks that contract. Pad the + # physical page instead -- the tail is its own group and the pad + # wastes bytes only. + new_spec = replace(layer_spec, page_size_padded=max_page_size) + new_kv_cache_spec[layer_name] = new_spec elif isinstance(layer_spec, MambaSpec): # MambaSpec's page size is determined by its state shapes and does # not scale with block_size, so pad the page instead. This is the @@ -1593,8 +1668,31 @@ ) -> int: """Return the largest cache group's bytes per block.""" if (glm5_layout := _glm5_next_tensor_layout(kv_cache_groups)) is not None: - _, _, mla_names, idx_names, mla_page, idx_page, _, _ = glm5_layout - return len(mla_names) * mla_page + len(idx_names) * idx_page + ( + _, + _, + mla_names, + idx_names, + mla_page, + idx_page, + _, + _, + draft_group, + hidden_groups, + ) = glm5_layout + bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page + if draft_group is not None: + d_inner = cast( + UniformTypeKVCacheSpecs, draft_group.kv_cache_spec + ).kv_cache_specs + bytes_per_block += sum( + spec.page_size_bytes for spec in d_inner.values() + ) + for group in hidden_groups: + bytes_per_block += ( + group.kv_cache_spec.page_size_bytes * len(group.layer_names) + ) + return bytes_per_block bytes_per_block = max( sum( @@ -1701,12 +1799,30 @@ idx_page, tail_names, _, + draft_group, + hidden_groups, ) = glm5_layout bytes_per_block = len(mla_names) * mla_page + len(idx_names) * idx_page + extra_bytes_per_block = 0 + if draft_group is not None: + d_inner = cast( + UniformTypeKVCacheSpecs, draft_group.kv_cache_spec + ).kv_cache_specs + extra_bytes_per_block += sum( + spec.page_size_bytes for spec in d_inner.values() + ) + for group in hidden_groups: + extra_bytes_per_block += ( + group.kv_cache_spec.page_size_bytes * len(group.layer_names) + ) num_blocks = may_override_num_blocks( - vllm_config, available_memory // bytes_per_block + vllm_config, + available_memory // (bytes_per_block + extra_bytes_per_block), ) - size = bytes_per_block * num_blocks + # One backing allocation shared by every tensor (allocate_kv_cache + # asserts all tensor sizes equal): the aliased GLM pool plus the + # draft/hidden regions appended after it. + size = (bytes_per_block + extra_bytes_per_block) * num_blocks attn_specs = cast( UniformTypeKVCacheSpecs, attn_group.kv_cache_spec ).kv_cache_specs @@ -1745,6 +1861,37 @@ ).kv_cache_specs add_tensor(tail_name, tail_specs[tail_name], offset) + # Draft + hidden-state layers get their own tensors appended after + # the aliased pool (no aliasing: independent block tables). + offset_cursor = bytes_per_block * num_blocks + if draft_group is not None: + d_inner = cast( + UniformTypeKVCacheSpecs, draft_group.kv_cache_spec + ).kv_cache_specs + d_page = next(iter(d_inner.values())).page_size_bytes + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=list(draft_group.layer_names), + layer_stride=d_page * num_blocks, + block_stride=d_page, + offset=offset_cursor, + ) + ) + offset_cursor += d_page * num_blocks * len(draft_group.layer_names) + for group in hidden_groups: + h_page = group.kv_cache_spec.page_size_bytes + kv_cache_tensors.append( + KVCacheTensor( + size=size, + layers=list(group.layer_names), + layer_stride=h_page * num_blocks, + block_stride=h_page, + offset=offset_cursor, + ) + ) + offset_cursor += h_page * num_blocks * len(group.layer_names) + return KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=kv_cache_tensors, @@ -2440,6 +2587,8 @@ idx_page, tail_names, _, + draft_group, + hidden_groups, ) = glm5_layout uniform_spec = cast(UniformTypeKVCacheSpecs, attn_group.kv_cache_spec) total_blocks = uniform_spec.max_memory_usage_pages(vllm_config) @@ -2452,7 +2601,30 @@ ) if tail_names: total_blocks += 1 - return total_blocks * (len(mla_names) * mla_page + len(idx_names) * idx_page) + pool_bytes = total_blocks * ( + len(mla_names) * mla_page + len(idx_names) * idx_page + ) + extra_bytes = 0 + if draft_group is not None: + d_inner = cast( + UniformTypeKVCacheSpecs, draft_group.kv_cache_spec + ).kv_cache_specs + d_spec = next(iter(d_inner.values())) + d_blocks = cdiv( + d_spec.max_memory_usage_bytes(vllm_config), d_spec.page_size_bytes + ) + extra_bytes += d_blocks * d_spec.page_size_bytes * len( + draft_group.layer_names + ) + for group in hidden_groups: + h_spec = group.kv_cache_spec + h_blocks = cdiv( + h_spec.max_memory_usage_bytes(vllm_config), h_spec.page_size_bytes + ) + extra_bytes += ( + h_blocks * h_spec.page_size_bytes * len(group.layer_names) + ) + return pool_bytes + extra_bytes bytes_per_block = _pool_bytes_per_block(kv_cache_groups) total_blocks = 0