[Spec Decode] Add support for num draft layers > 1 for Eagle

tests/v1/core/test_kv_cache_utils.py

@@ -3,6 +3,7 @@
 import hashlib
 import importlib
 from collections.abc import Callable
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
@@ -2731,3 +2732,127 @@ def test_resolve_block_hashes_rejects_mismatched_view():
     mismatched = BlockHashListWithBlockSize(raw, 2, 8)
     with pytest.raises(AssertionError):
         resolve_block_hashes(mismatched, 2, 4)
+
+
+# ---------------------------------------------------------------------------
+# Eagle draft KV cache grouping (multi_draft / fold_draft / separate_draft).
+# These tests exercise _get_kv_cache_groups_uniform_page_size directly with
+# synthetic specs so they never load real model weights.
+# ---------------------------------------------------------------------------
+
+
+def _build_eagle_specs(
+    num_full: int,
+    num_sw: int,
+    num_draft_layers: int,
+    draft_matches: bool = True,
+) -> dict[str, KVCacheSpec]:
+    """Build a synthetic kv_cache_spec with target full + sliding-window layers
+    and ``num_draft_layers`` eagle draft layers.
+
+    When ``draft_matches`` is True the draft layers share the target
+    sliding-window spec (so they can be folded); otherwise they use a unique
+    spec (different sliding window) that no target layer matches.
+    """
+    spec: dict[str, KVCacheSpec] = {}
+    for i in range(num_full):
+        spec[f"target.full.{i}"] = new_kv_cache_spec()
+    for i in range(num_sw):
+        spec[f"target.sw.{i}"] = new_sliding_window_spec(sliding_window=1)
+    draft_window = 1 if draft_matches else 999
+    for i in range(num_draft_layers):
+        spec[f"eagle_draft.sw.{i}"] = new_sliding_window_spec(
+            sliding_window=draft_window
+        )
+    return spec
+
+
+def _fake_vllm_config(pp_size: int) -> Any:
+    # The grouping routine only reads parallel_config.pipeline_parallel_size,
+    # so a lightweight stand-in avoids building a full (GPU-bound) VllmConfig.
+    return SimpleNamespace(
+        parallel_config=SimpleNamespace(pipeline_parallel_size=pp_size)
+    )
+
+
+def _draft_names(spec: dict[str, KVCacheSpec]) -> set[str]:
+    return {ln for ln in spec if "eagle" in ln}
+
+
+@pytest.mark.parametrize(
+    "num_draft_layers,pp_size,draft_matches,expected_mode",
+    [
+        # D <= 1: default grouping, regardless of PP.
+        (1, 1, True, "default"),
+        (1, 2, True, "default"),
+        # D > 1, P == 1: fold into a same-spec target group.
+        (3, 1, True, "fold"),
+        # D > 1, P > 1: standalone (separate) group.
+        (3, 2, True, "separate"),
+        # D > 1, P == 1 but no matching target spec: fall back to separate.
+        (3, 1, False, "separate"),
+    ],
+)
+def test_eagle_draft_kv_cache_grouping(
+    num_draft_layers, pp_size, draft_matches, expected_mode
+):
+    spec = _build_eagle_specs(
+        num_full=4,
+        num_sw=13,
+        num_draft_layers=num_draft_layers,
+        draft_matches=draft_matches,
+    )
+    groups = kv_cache_utils._get_kv_cache_groups_uniform_page_size(
+        _fake_vllm_config(pp_size), spec
+    )
+
+    # Every layer must appear in exactly one group. This is the key regression
+    # guard: the previous logic could process draft layers twice (folded AND
+    # inserted as a separate group), duplicating them here.
+    all_layers = [ln for g in groups for ln in g.layer_names]
+    assert sorted(all_layers) == sorted(spec.keys())
+
+    draft_names = _draft_names(spec)
+    draft_groups = [g for g in groups if any("eagle" in ln for ln in g.layer_names)]
+    # All draft layers must live in exactly one group.
+    assert len(draft_groups) == 1
+    draft_group = set(draft_groups[0].layer_names)
+    assert draft_names <= draft_group
+
+    if expected_mode == "separate":
+        # Standalone group containing exactly the draft layers.
+        assert draft_group == draft_names
+    else:
+        # default / fold: the draft shares its group with target layers.
+        assert draft_group > draft_names
+        if expected_mode == "fold":
+            # The folded group holds exactly group_size (== 4) layers:
+            # (group_size - num_draft_layers) target layers + all draft layers.
+            assert len(draft_group) == 4
+
+
+def test_eagle_default_group_size_uses_target_only():
+    # Target: 3 full + 3 sliding -> target-only group_size == 3. If the single
+    # default draft layer were counted when computing group_size, the
+    # sliding type would have 4 layers and the 1.5x heuristic would bump
+    # group_size to 4. Verify the draft does NOT perturb group_size.
+    spec = _build_eagle_specs(num_full=3, num_sw=3, num_draft_layers=1)
+    groups = kv_cache_utils._get_kv_cache_groups_uniform_page_size(
+        _fake_vllm_config(1), spec
+    )
+    all_layers = [ln for g in groups for ln in g.layer_names]
+    assert sorted(all_layers) == sorted(spec.keys())
+    # No group exceeds the target-only group_size of 3 real layers.
+    assert max(len(g.layer_names) for g in groups) == 3
+
+
+def test_eagle_no_draft_grouping_unaffected():
+    # Sanity check: with no eagle draft layers the grouping is the plain
+    # hybrid split (one full group + one sliding group here).
+    spec = _build_eagle_specs(num_full=4, num_sw=4, num_draft_layers=0)
+    groups = kv_cache_utils._get_kv_cache_groups_uniform_page_size(
+        _fake_vllm_config(1), spec
+    )
+    all_layers = [ln for g in groups for ln in g.layer_names]
+    assert sorted(all_layers) == sorted(spec.keys())
+    assert not any("eagle" in ln for ln in all_layers)

vllm/model_executor/models/cohere_eagle.py

@@ -146,9 +146,13 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         target_layer_num = vllm_config.model_config.get_num_layers(
             vllm_config.parallel_config
         )
+
+        # Name the eagle draft model so it can be
+        # distinguished from the target model
+        # in the KV cache groups.
         self.model = CohereEagleModel(
             vllm_config=vllm_config,
-            prefix=maybe_prefix(prefix, "model"),
+            prefix=maybe_prefix(prefix, "eagle"),
             start_layer_id=target_layer_num,
         )
 

vllm/model_executor/models/commandr.py

@@ -137,7 +137,10 @@ def __init__(
         self.hidden_size = config.hidden_size
         self.total_num_heads = config.num_attention_heads
         self.num_heads = self.total_num_heads // tp_size
-        self.head_dim = self.hidden_size // self.total_num_heads
+        if hasattr(config, "head_dim"):
+            self.head_dim = config.head_dim
+        else:
+            self.head_dim = self.hidden_size // self.total_num_heads
         self.total_num_kv_heads = config.num_key_value_heads
         if self.total_num_kv_heads >= tp_size:
             # Number of KV heads is greater than TP size, so we partition

vllm/v1/core/kv_cache_utils.py

@@ -1135,7 +1135,20 @@
     return not kv_cache_spec
 
 
+def _warn_padding_layers(num_layers: int, group_size: int) -> None:
+    """Warn about wasted KV cache memory when a layer type does not divide
+    evenly into ``group_size`` and must be padded to fill the last group."""
+    num_padding_layers = group_size - num_layers % group_size
+    if num_padding_layers != group_size:
+        logger.warning(
+            "Add %d padding layers, may waste at most %.2f%% KV cache memory",  # noqa
+            num_padding_layers,
+            num_padding_layers / num_layers * 100,
+        )
+
+
 def _get_kv_cache_groups_uniform_page_size(
+    vllm_config: VllmConfig,
     kv_cache_spec: dict[str, KVCacheSpec],
 ) -> list[KVCacheGroupSpec]:
     """
@@ -1200,12 +1213,40 @@
     Returns:
         The generated KVCacheGroupSpecs
     """
-    # Group all layers by kv_cache_spec.
+
+    # --- Eagle draft layer placement --------------------------------------
+    # Eagle speculative-decoding draft layers must all end up in a single KV
+    # cache group (the proposer shares one AttentionMetadata across them). How
+    # we place them depends on the number of draft layers (D) and the
+    # pipeline-parallel size (P):
+    #   * D <= 1           -> "default": let the lone draft layer flow through
+    #     the normal per-spec grouping (legacy behavior). A single layer never
+    #     scatters and can reuse an existing group.
+    #   * D > 1 and P == 1 -> "fold": pack the draft layers into one existing
+    #     same-spec target group so they stay together without adding an extra
+    #     block-churn stream (an extra stream collapses prefix-cache hit rate).
+    #   * D > 1 and P > 1  -> "separate": give the draft its own group. Folding
+    #     breaks the layers[i::n] PP stage-balancing below, and >1 draft layers
+    #     would otherwise scatter across the strided groups.
+    draft_layer_group = [
+        layer_name for layer_name in kv_cache_spec if "eagle" in layer_name
+    ]
+    num_draft = len(draft_layer_group)
+    pp_size = vllm_config.parallel_config.pipeline_parallel_size
+    multi_draft = num_draft > 1
+    fold_draft = multi_draft and pp_size == 1
+    separate_draft = multi_draft and pp_size > 1
+
+    # Group the target (non-draft) layers by kv_cache_spec. Draft layers are
+    # always excluded here so group_size is derived purely from the target
+    # model; a default (D <= 1) draft layer is re-added after group_size is
+    # computed so it cannot perturb it.
     # E.g., 2 full attention layers and 3 sliding window attention layers,
     # -> (full.0, full.1), (sw.0, sw.1, sw.2).
     same_type_layers: dict[KVCacheSpec, list[str]] = defaultdict(list)
     for layer_name, layer_spec in kv_cache_spec.items():
-        same_type_layers[layer_spec].append(layer_name)
+        if "eagle" not in layer_name:
+            same_type_layers[layer_spec].append(layer_name)
 
     # Split each group into smaller groups, to make the number of layers in each
     # group identical. Add padding to the last group of each type if necessary.
@@ -1232,14 +1273,46 @@
         # extra layers to one attention type.
         group_size = max_num_layers
     grouped_layers = []
-    for layers in same_type_layers.values():
-        num_padding_layers = group_size - len(layers) % group_size
-        if num_padding_layers != group_size:
-            logger.warning(
-                "Add %d padding layers, may waste at most %.2f%% KV cache memory",  # noqa
-                num_padding_layers,
-                num_padding_layers / len(layers) * 100,
-            )
+
+    draft_spec = kv_cache_spec[draft_layer_group[0]] if draft_layer_group else None
+
+    if multi_draft:
+        # Fold and separate both place all draft layers into one group, so they
+        # must fit within a single group (there are only group_size layer slots).
+        assert num_draft <= group_size, (
+            f"Number of eagle draft layers ({num_draft}) must be <= group_size "
+            f"({group_size}) so all draft layers fit in one group."
+        )
+
+    if fold_draft and draft_spec not in same_type_layers:
+        # Folding needs a target group with an identical spec, otherwise
+        # create_kv_cache_group_specs' merge() would fail. Fall back to a
+        # separate draft group.
+        logger.warning(
+            "No target layer shares the draft KV spec; "
+            "falling back to a separate draft group."
+        )
+        fold_draft = False
+        separate_draft = True
+
+    if num_draft == 1:
+        # Default: re-add the single draft layer to its matching target spec
+        # (or, if none matches, it forms its own single-layer group) so it is
+        # split into a group below alongside the target layers. Done after
+        # group_size is computed so the draft cannot perturb it.
+        same_type_layers[draft_spec].append(draft_layer_group[0])
+
+    for spec, layers in same_type_layers.items():
+        if fold_draft and spec == draft_spec:
+            # Fold: build one group of exactly group_size layers, made of
+            # (group_size - num_draft) target layers plus all draft layers; the
+            # remaining target layers of this spec fall through to the split.
+            _warn_padding_layers(len(layers) + num_draft, group_size)
+            reserved = max(group_size - num_draft, 0)
+            grouped_layers.append(layers[:reserved] + draft_layer_group)
+            layers = layers[reserved:]
+        else:
+            _warn_padding_layers(len(layers), group_size)
         num_groups = cdiv(len(layers), group_size)
         # In PP case, say if we have
         # - stage 0: full.0, sw.0, sw.1
@@ -1254,6 +1327,18 @@
         # instead of layers[i * group_size: (i + 1) * group_size]
         for i in range(num_groups):
             grouped_layers.append(layers[i::num_groups])
+
+    if separate_draft:
+        # Insert the standalone draft group adjacent to the same-type target
+        # groups so group ids of each attention type stay contiguous (required
+        # by HybridKVCacheCoordinator's non-interleaving check).
+        draft_spec_type = type(kv_cache_spec[draft_layer_group[0]])
+        insert_idx = 0
+        for idx, group in enumerate(grouped_layers):
+            if isinstance(kv_cache_spec[group[0]], draft_spec_type):
+                insert_idx = idx + 1
+        grouped_layers.insert(insert_idx, draft_layer_group)
+
     return create_kv_cache_group_specs(kv_cache_spec, grouped_layers)
 
 
@@ -1781,7 +1866,7 @@
     # the page size of the layers. For cases cannot be unified, this function
     # will raise an error.
     filtered_spec = unify_kv_cache_spec_page_size(filtered_spec)
-    groups = _get_kv_cache_groups_uniform_page_size(filtered_spec)
+    groups = _get_kv_cache_groups_uniform_page_size(vllm_config, filtered_spec)
 
     # Add hidden-state layers back with page aligned to the common page.
     if hidden_specs: