llm: patch 03 - pin the kpool tail block through the unifier
Patch 02 (MLA-multiple max page) cleared the KV page-size unifier: the dflash boot now reaches cudagraph capture. New blocker: the kpool decode kernel asserts tail_kv_cache.shape[2] == pool_size; KpoolTailSpec pins block_size == index_kpool, so the unifier's divisible branch must never grow it. Patch 03 = 02 + a KpoolTailSpec branch that pads page_size_padded instead. Apply-verified: reproduces patched sha 5093e4b7... Runbook updated with the full KV grouping journey.
This commit is contained in:
@@ -0,0 +1,37 @@
|
||||
--- a/vllm/v1/core/kv_cache_utils.py
|
||||
+++ b/vllm/v1/core/kv_cache_utils.py
|
||||
@@ -1395,10 +1395,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
|
||||
@@ -403,6 +403,50 @@ Note also that the KV-cap patch's predecessor arithmetic in
|
||||
"gather-bound" conclusion came from the same throughput arithmetic that this page's
|
||||
profiler window has now falsified.
|
||||
|
||||
### The KV-cache grouping journey: clearing the unifier, then the kpool tail
|
||||
|
||||
With patch 01 v5 the gate passes and the boot moves to KV-cache layout
|
||||
construction — where the next two blockers live, both in
|
||||
`v1/core/kv_cache_utils.py`:
|
||||
|
||||
1. **The page-size unifier refuses GLM's MLA family.** With a draft model the
|
||||
spec dict gains non-MLA drafter attention specs and `HiddenStateCacheSpec`
|
||||
entries, so `_get_kv_cache_groups_glm5_next` (which requires ALL attention
|
||||
specs to be `MLAAttentionSpec`) bails to the generic grouping path — whose
|
||||
`unify_kv_cache_spec_page_size` raises
|
||||
`NotImplementedError: Layer language_model.model.layers.3.self_attn.indexer.k_cache: page size is not divisible… cannot be padded. Padding is only supported for non-MLA attention layers.`
|
||||
MLA layers cannot be *padded* (sparse MLA indexes whole token rows) — but
|
||||
they CAN grow `block_size` when the max page is an exact multiple. The
|
||||
indexer's page is just `tokens_per_state x` the MLA page, so
|
||||
**`glm53-dflash-02-kv-unify-mla-multiple.patch`** bumps the unified max page
|
||||
to a multiple of every `MLAAttentionSpec` page before the loop: all MLA-family
|
||||
layers then take the legal divisible block-growth branch, drafter + mamba
|
||||
layers pad. The 2026-09-30 window proved it: unify passes, memory profiling
|
||||
passes, torch.compile passes, and the boot reaches cudagraph capture — the
|
||||
NotImplementedError is gone.
|
||||
2. **The kpool tail kernel pins its block.** Cudagraph capture then dies in
|
||||
`kpool_ops.kpool_decode_update_and_maybe_write_cache_batched`:
|
||||
`assert tail_kv_cache.shape[2] == pool_size`. The tail cache must view as
|
||||
`[num_blocks, 2, index_kpool, head_dim]` (`KpoolTailSpec(block_size=index_kpool)`,
|
||||
`Glm5NextTailCache` in `glm5next/common/attention.py:159`), but the tail's
|
||||
page (8192 B) *divided* the raised max page, so the unifier's divisible branch
|
||||
grew its `block_size` by the ratio and the view became
|
||||
`[num_blocks, 2, index_kpool x ratio, head_dim]`. Verified not the hybrid
|
||||
promotion: the log carries no "Hybrid KV cache manager is disabled" warning,
|
||||
so `disable_hybrid_kv_cache_manager` is unset and
|
||||
`unify_hybrid_kv_cache_specs` never ran. **`glm53-dflash-03-kv-unify-tailpin.patch`**
|
||||
(= 02 + a `KpoolTailSpec` branch that pads `page_size_padded` instead of ever
|
||||
growing its block) is the fix; supersedes 02 in the repo.
|
||||
|
||||
Both patches mount `v1/core/kv_cache_utils.py` over site-packages via the
|
||||
scratch launcher's `KVUNIFY_PATCH=1` knob (host
|
||||
`/home/dank/dflash-patch2/new/vllm/v1/core/kv_cache_utils.py`, patched sha
|
||||
`5093e4b7…`; apply-verified against the image's pristine file). The window
|
||||
script's decision pipeline was also fixed for a SIGPIPE death (`head -80` on a
|
||||
`podman logs` grep pipeline killed one run before log-save and restore; the log
|
||||
is now captured to a file FIRST, then grepped).
|
||||
|
||||
|
||||
## Measured arms (all 2026-09-30, one boot each)
|
||||
|
||||
| arm | decode | prefill | verdict |
|
||||
|
||||
Reference in New Issue
Block a user