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...
This commit is contained in:
Frank committed 2026-09-30 19:19:59 -07:00
1 parent 2bc2f8c833
commit d56db43aa3
2 files changed
+340

No files matched your search

@@ -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"):