serverAppConfigs/scripts/patches/glm53-dflash-04-kv-groups-draft.patch

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