dank/serverAppConfigs
1
0
Fork 0
Code Issues 6 Pull Requests Actions Packages Projects Releases Wiki Activity

llm: patch 04 - glm5 KV pool carries draft + hidden-state groups

Browse Source 
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...
This commit is contained in:
Frank
2026-09-30 19:19:59 -07:00
parent 2bc2f8c833
commit d56db43aa3
2 changed files with 340 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/patches/glm53-dflash-04-kv-groups-draft.patch
+317
View File
@@ -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
scripts/patches/glm53-dflash-04b-kv-registry-dflash.patch
+23
View File
@@ -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"):