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:
Frank
2026-09-30 18:38:53 -07:00
parent a1f14a1c0f
commit 2bc2f8c833
2 changed files with 81 additions and 0 deletions
@@ -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