From d56db43aa3710e62d1c3a85372f3d93f5c39bf21 Mon Sep 17 00:00:00 2001 From: Frank Date: Wed, 30 Sep 2026 19:19:59 -0700 Subject: [PATCH] llm: patch 04 - glm5 KV pool carries draft + hidden-state groups The dflash boot cleared the unifier and the kpool tail (02/03) and died on KV capacity: on the generic path the mamba states are NOT aliased into the MLA pages (the glm5 tensor layout's whole point), so the capacity estimator demanded ~30.6 GiB fixed against 0.92 GiB available (the plan's KV reserve registry only knows the target's 13,600 B/token). Patch 04 keeps the aliased glm5 pool and appends the drafter's uniform group and per-layer hidden-state groups beside it: grouping partitions hidden + draft specs, the layout returns a 10-tuple, the builder sizes one shared backing allocation (pool + extras), and the memory estimator adds the extras. Patch 04b extends the GLM registry kv_bytes_per_token by the drafter (4,096 B/token) and hidden-state (5 x hidden x 2 B/token) costs so the expert plan reserves enough. Local smoke tests (CPU torch + real spec classes) validate grouping shape, layout recognition, builder block math and memory accounting against synthetic GLM+draft+hidden spec sets. Apply-verified: reproduces shas 5521e50b... / 1a623d5c... --- .../glm53-dflash-04-kv-groups-draft.patch | 317 ++++++++++++++++++ .../glm53-dflash-04b-kv-registry-dflash.patch | 23 ++ 2 files changed, 340 insertions(+) create mode 100644 scripts/patches/glm53-dflash-04-kv-groups-draft.patch create mode 100644 scripts/patches/glm53-dflash-04b-kv-registry-dflash.patch diff --git a/scripts/patches/glm53-dflash-04-kv-groups-draft.patch b/scripts/patches/glm53-dflash-04-kv-groups-draft.patch new file mode 100644 index 0000000..9db8786 --- /dev/null +++ b/scripts/patches/glm53-dflash-04-kv-groups-draft.patch @@ -0,0 +1,317 @@ +--- 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 diff --git a/scripts/patches/glm53-dflash-04b-kv-registry-dflash.patch b/scripts/patches/glm53-dflash-04b-kv-registry-dflash.patch new file mode 100644 index 0000000..4229db2 --- /dev/null +++ b/scripts/patches/glm53-dflash-04b-kv-registry-dflash.patch @@ -0,0 +1,23 @@ +--- a/vllm/model_executor/offloader/expert_activation_registry.py ++++ b/vllm/model_executor/offloader/expert_activation_registry.py +@@ -316,7 +316,19 @@ + + + def _glm53_kv_bytes_per_token(cfg: "VllmConfig") -> int: +- return 13600 ++ per_tok = 13600 ++ # FORK-LOCAL (patches/glm53-dflash-04-kv-groups-draft): with a dflash/dspark ++ # drafter, the KV spec set gains the drafter's attention layers and the aux ++ # hidden-state cache (one bf16 hidden vector per target aux layer per ++ # token). Both are full-length per-token costs the aliased glm5 pool ++ # cannot absorb, so the reserve must include them. The drafter is a 5-layer ++ # qwen3 (8 kv heads x 128 x K+V x bf16 = 4096 B/token); the aux count is ++ # the dflash target_layer_ids count (5). ++ sc = getattr(cfg, "speculative_config", None) ++ if sc is not None and getattr(sc, "use_aux_hidden_state_outputs", False): ++ hidden = int(cfg.model_config.hf_config.hidden_size) ++ per_tok += 5 * hidden * 2 + 4096 ++ return per_tok + + + for _arch in ("Glm5NextForConditionalGeneration", "Glm5NextForCausalLM"):