--- 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