[Spec decode] --spec-decode-kv-cache-dtype: bf16/fp8 KV for a draft that has its own pool

atom/config.py

@@ -1380,6 +1380,14 @@ class SpeculativeConfig:
     # probability that the first i+1 draft tokens are all accepted), filled in by
     # __post_init__. None => real draft/target rejection sampling.
     synthetic_acceptance_rates: list[float] | None = None
+    # The draft's KV dtype as the user asked for it (--spec-decode-kv-cache-dtype):
+    # None = not asked, else "bf16" / "fp8". Never rewritten; the dtype the draft
+    # will use is `draft_kv_dtype`.
+    kv_cache_dtype: str | None = None
+    # Filled in by `Config.__post_init__` (`resolve_draft_kv_dtype`), once the
+    # target's dtype is settled; every reader (the draft's pool, its attention, the
+    # calibration identity, the compile hash) reads this one.
+    draft_kv_dtype: str | None = field(default=None, init=False)
 
     # model_type → mtp_model_type mapping
     _MTP_TYPE_MAP: ClassVar[dict[str, str]] = {
@@ -1941,6 +1949,23 @@ def qrep_for_step(
     return qrep_enabled and not use_seg_mla and (not is_prefill or prefill_qrep)
 
 
+def resolve_draft_kv_dtype(spec, target_dtype: str, *, own_pool: bool) -> str:
+    """The KV dtype of a speculative draft's pages: the target's, unless the draft has
+    a pool of its own (``own_pool``) and asked for another through
+    ``--spec-decode-kv-cache-dtype``. A draft in the target's pool cannot ask for
+    another dtype."""
+    asked = spec.kv_cache_dtype
+    if not own_pool:
+        if asked is not None and asked != target_dtype:
+            raise ValueError(
+                f"--spec-decode-kv-cache-dtype {asked} differs from "
+                f"--kv-cache-dtype {target_dtype}, but this draft caches in the "
+                "target's KV pool and cannot have its own dtype"
+            )
+        return target_dtype
+    return asked or target_dtype
+
+
 def indexer_cp_unsupported_reason(
     arches,
     tp_size: int,
@@ -2802,6 +2827,23 @@ def __post_init__(self):
             v4_block_size = 256
             if self.kv_cache_block_size != v4_block_size:
                 self.kv_cache_block_size = v4_block_size
+            if self.kv_cache_dtype == "fp8" and self.plugin_config is None:
+                # aiter's V4 prefill / decode implement the fp8 (2buff) KV only on
+                # gfx950 / gfx1250. Elsewhere the KV is bf16, decided here -- before
+                # anything reads `kv_cache_dtype`, the draft's resolved dtype
+                # included -- and not by the builder after the fact. `get_gfx` is the
+                # builder's own probe (rocminfo; it does not initialise CUDA).
+                from aiter.jit.utils.chip_info import get_gfx
+
+                gfx = get_gfx()
+                if gfx not in ("gfx950", "gfx1250"):
+                    logger.warning(
+                        "DeepSeek-V4 --kv_cache_dtype fp8 (2buff) is only supported "
+                        "on gfx950 / gfx1250 (aiter op4/op5); got %r. Falling back to "
+                        "a bf16 KV cache.",
+                        gfx,
+                    )
+                    self.kv_cache_dtype = "bf16"
 
         # GLM-5.3-Flash's indexer caches one pooled key per `index_kpool`
         # tokens, so a KV block of B tokens needs only B // index_kpool index
@@ -2914,6 +2956,26 @@ def __post_init__(self):
                 self.index_cache_dtype = "fp8"
             else:
                 self.index_cache_dtype = self.kv_cache_dtype
+        if self.speculative_config is not None:
+            from atom.spec_decode.draft_kv import draft_owns_kv_pool
+
+            spec = self.speculative_config
+            # The one place the draft's pool and dtype are settled: everything
+            # downstream (the draft builder, the proposer, the calibration identity)
+            # reads `spec.draft_kv_dtype`.
+            spec.draft_kv_dtype = resolve_draft_kv_dtype(
+                spec, self.kv_cache_dtype, own_pool=draft_owns_kv_pool(self)
+            )
+            if spec.draft_kv_dtype != self.kv_cache_dtype and self.kv_transfer_config:
+                # The offload namespace keys the target's KV dtype alone, so a
+                # draft page of another dtype would share a key with the target's.
+                raise ValueError(
+                    f"draft KV dtype {spec.draft_kv_dtype} differs from the "
+                    f"target's {self.kv_cache_dtype}: not supported with KV "
+                    "offload or --kv-transfer-config. Pass "
+                    f"--spec-decode-kv-cache-dtype {self.kv_cache_dtype} to match "
+                    "the target"
+                )
 
         if self.hf_config.model_type == "deepseek_v41_text":
             # CSA2's index plane is paged by the same scorer GLM's is, so a
@@ -3344,7 +3406,9 @@ def compute_hash(self) -> str:
         # steer sampling and stay out.
         spec = self.speculative_config
         factors.append(
-            None if spec is None else (spec.method, spec.num_speculative_tokens)
+            None
+            if spec is None
+            else (spec.method, spec.num_speculative_tokens, spec.draft_kv_dtype)
         )
 
         hash_str = hashlib.md5(

atom/kv_transfer/offload/mp/deployment.py

@@ -167,10 +167,13 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool:
     registering any cache memory.
     """
 
-    if _config_has_own_pool_draft(config):
-        # A DSpark draft whose backend owns a KV pool appends per-rank-sharded
-        # PAGE regions (`draft_kv.py` declares factor 1), so the complete PAGE
-        # object is not replicated even when the target's is.
+    from atom.spec_decode.draft_kv import draft_owns_kv_pool
+
+    if draft_owns_kv_pool(config):
+        # A draft with its own KV pool (an MHA DSpark/EAGLE3 draft,
+        # `draft_owns_kv_pool`) appends per-rank-sharded PAGE regions
+        # (`draft_kv.py` declares factor 1), so the complete PAGE object is not
+        # replicated even when the target's is.
         return False
     hf_config = getattr(config, "hf_config", None)
     # MiniMax-M3 is GQA. Some TP ranks can happen to own the same KV head when
@@ -187,22 +190,6 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool:
     return getattr(hf_config, "kv_lora_rank", None) is not None
 
 
-def _config_has_own_pool_draft(config: Any) -> bool:
-    """Whether a speculative draft caches into a KV pool of its own.
-
-    Mirrors `spec_decode.draft_kv.draft_kv_builder`, which only DSpark calls:
-    the draft's own backend answers through `DRAFT_OWNS_KV_POOL`.
-    """
-
-    speculative = getattr(config, "speculative_config", None)
-    draft_hf = getattr(speculative, "draft_model_hf_config", None)
-    if getattr(speculative, "method", None) != "dspark" or draft_hf is None:
-        return False
-    from atom.utils.selector import attn_family, get_attn_backend
-
-    return bool(get_attn_backend(attn_family(draft_hf)).DRAFT_OWNS_KV_POOL)
-
-
 def _tp_replication_factor(config: Any) -> int:
     """Return the PAGE rank-collapse factor selected before workers start.
 

atom/model_engine/arg_utils.py

@@ -178,6 +178,7 @@ class EngineArgs:
     draft_model: str | None = None
     spec_decode_acceptance_rate: float | None = None
     spec_decode_acceptance_length: float | None = None
+    spec_decode_kv_cache_dtype: str | None = None
     mark_trace: bool = False
     enable_rapidserve: bool = False
     disagg_prefill_max_num_seqs: int | None = None
@@ -470,6 +471,18 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             "DFlash-backbone drafts such as Kimi-K3-DSpark, omitted for "
             "V4-Pro-DSpark which ships inside the target checkpoint).",
         )
+        parser.add_argument(
+            "--spec-decode-kv-cache-dtype",
+            type=str,
+            default=None,
+            choices=["bf16", "fp8"],
+            help="KV dtype of the speculative draft's pages "
+            "(SpeculativeConfig.kv_cache_dtype). A draft with a KV pool of its own "
+            "(an MHA draft: MiniMax-M3.1's DSpark, an MHA EAGLE3 draft) takes "
+            "either; default is --kv_cache_dtype's. A draft in the target's pool "
+            "always has --kv_cache_dtype, and a different value here is an error. "
+            "Needs --method and --num-speculative-tokens > 0.",
+        )
         parser.add_argument(
             "--spec-decode-acceptance-length",
             type=float,
@@ -928,6 +941,7 @@ def _get_engine_kwargs(self) -> dict:
                 method=method,
                 model=draft_model or self.model,
                 num_speculative_tokens=num_spec_tokens,
+                kv_cache_dtype=kwargs.pop("spec_decode_kv_cache_dtype"),
                 synthetic_acceptance_rate=synthetic_acceptance_rate,
                 synthetic_acceptance_length=synthetic_acceptance_length,
             )
@@ -937,6 +951,11 @@ def _get_engine_kwargs(self) -> dict:
             kwargs.pop("draft_model")
             kwargs.pop("spec_decode_acceptance_rate")
             kwargs.pop("spec_decode_acceptance_length")
+            if kwargs.pop("spec_decode_kv_cache_dtype") is not None:
+                raise ValueError(
+                    "--spec-decode-kv-cache-dtype needs --method and "
+                    "--num-speculative-tokens > 0."
+                )
             kwargs["speculative_config"] = None
 
         # --enable-tbo [prefill|all] → enable_tbo + enable_tbo_decode

atom/model_ops/attentions/aiter_attention.py

@@ -147,10 +147,6 @@ def _mtp_prepare_decode_metadata_kernel(
 
 
 class AiterBackend(AttentionBackend):
-    # An MHA draft keeps its own K and V; there is no target tensor its rows
-    # could be.
-    DRAFT_OWNS_KV_POOL: ClassVar[bool] = True
-
     @staticmethod
     def get_name() -> str:
         return "ATOM_ATTENTION"

atom/model_ops/attentions/backends.py

@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
 from contextlib import nullcontext
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 if TYPE_CHECKING:
     from atom.distributed.pcp_utils import PcpBalGroup
@@ -63,12 +63,6 @@
 class AttentionBackend(ABC):
     """Abstract class for attention backends."""
 
-    #: Whether a *draft* of this flavor caches into a pool of its own. False
-    #: means its rows are the target's -- an MLA draft's latent is the target's
-    #: latent -- and `make_kv_pool` below is never asked. A property of the
-    #: flavor, so a class answer and not a call.
-    DRAFT_OWNS_KV_POOL: ClassVar[bool] = False
-
     @staticmethod
     @abstractmethod
     def get_name() -> str:
@@ -94,11 +88,11 @@ def make_kv_pool(
     ):
         """A pool of `layers` rows at this config's geometry.
 
-        Asked of a *draft* model's backend, and only where `DRAFT_OWNS_KV_POOL`
-        — the target's layers are its builder's own business. Both counts are
-        told rather than read: `layers` only the walk knows, and the block only
-        the target builder, whose it is and not this backend's for the reason
-        in `DraftKvBuilder.kv_pool`.
+        Asked of a *draft* model's backend, and only for a family in
+        `selector.DRAFT_POOL_FAMILIES` — the target's layers are its builder's
+        own business. Both counts are told rather than read: `layers` only the
+        walk knows, and the block only the target builder, whose it is and not
+        this backend's for the reason in `DraftKvBuilder.kv_pool`.
         """
         raise NotImplementedError
 

atom/model_ops/attentions/deepseek_v4_attn.py

@@ -531,25 +531,8 @@ def __init__(self, model_runner):
         # rope pool is always bf16.
         self._kv_fp8 = model_runner.kv_cache_dtype == "fp8"
         # aiter prefill (op4) / decode (op5) implement the fp8 (2buff) path only
-        # on gfx950 / gfx1250. On any other arch, transparently fall back to a
-        # bf16 KV cache instead of hard-failing. Flipping self._kv_fp8 here (before
-        # the *_dtype attrs are read) keeps the whole V4 path consistent: pool
-        # sizing (sub_pool_specs / swa_block_bytes_per_layer), quant_mode,
-        # and module.kv_fp8 (build_kv_cache_tensor) all key off self._kv_fp8 /
-        # these dtype attrs. Sync model_runner.kv_cache_dtype (and the shared
-        # config) so any generic reader / log line agrees.
-        if self._kv_fp8 and get_gfx() not in ("gfx950", "gfx1250"):
-            logger.warning(
-                "DeepSeek-V4 --kv_cache_dtype fp8 (2buff) is only supported on "
-                "gfx950 / gfx1250 (aiter op4/op5); got %r. Falling back to a "
-                "bf16 KV cache.",
-                get_gfx(),
-            )
-            self._kv_fp8 = False
-            model_runner.kv_cache_dtype = "bf16"
-            cfg = getattr(model_runner, "config", None)
-            if cfg is not None and getattr(cfg, "kv_cache_dtype", None) == "fp8":
-                cfg.kv_cache_dtype = "bf16"
+        # on gfx950 / gfx1250; elsewhere `Config` has already made the KV bf16
+        # (`Config.__post_init__`), so this reads the settled dtype.
         if self._kv_fp8:
             self._swa_dtype = dtypes.fp8
             self._classical_dtype = dtypes.fp8

atom/spec_decode/calibration.py

@@ -13,7 +13,9 @@
 def calibration_identity(config, device):
     model = Path(config.model)
     files = ("config.json", "model.safetensors.index.json")
-    return {
+    spec = config.speculative_config
+    draft_dtype = spec.draft_kv_dtype if spec is not None else config.kv_cache_dtype
+    identity = {
         "model_files": {
             name: hashlib.sha256((model / name).read_bytes()).hexdigest()
             for name in files
@@ -27,6 +29,11 @@ def calibration_identity(config, device):
         "torch_version": torch.__version__,
         "hip_version": torch.version.hip,
     }
+    # Only when it is not the target's, so profiles from before it was a choice
+    # still match.
+    if draft_dtype != config.kv_cache_dtype:
+        identity["draft_kv_cache_dtype"] = draft_dtype
+    return identity
 
 
 def validate_calibration(profile, identity, *, width, max_batch):

atom/spec_decode/draft_kv.py

@@ -12,30 +12,66 @@
 DRAFT_KV_ROWS = "draft_kv"
 
 
+# The methods whose proposer calls `draft_kv_builder`; the others (MTP, ngram) run
+# the target's own layers or have no draft model.
+DRAFT_POOL_METHODS = ("eagle3", "dspark")
+
+
+def draft_owns_kv_pool(config) -> bool:
+    """Whether `config`'s speculative draft caches into a KV pool of its own.
+
+    What `draft_kv_builder` will answer for the same draft, decided from the config
+    alone so `Config` and the offload deployment checks can ask it before any runner
+    exists. EAGLE3 and DSpark (`DRAFT_POOL_METHODS`) build one where the draft's family
+    owns it, and only the native runner builds it: under a plugin the framework owns
+    the KV cache, so the answer there is no. Read off the
+    draft's attention family (`DRAFT_POOL_FAMILIES`), not by resolving its backend,
+    which would import aiter into the API process that builds `Config`.
+    """
+    from atom.plugin import is_plugin_mode
+
+    # First: under a plugin the offload adapter hands over vLLM's own speculative
+    # config unchanged, which has no `draft_model_hf_config` to read.
+    if is_plugin_mode():
+        return False
+    spec = config.speculative_config
+    if (
+        spec is None
+        or spec.method not in DRAFT_POOL_METHODS
+        or spec.draft_model_hf_config is None
+    ):
+        return False
+    from atom.utils.selector import DRAFT_POOL_FAMILIES, attn_family
+
+    return attn_family(spec.draft_model_hf_config) in DRAFT_POOL_FAMILIES
+
+
 def draft_kv_builder(model_runner, draft_hf):
     """The KV builder a draft needs of its own, or None if it needs none.
 
     A draft is not an attention flavor; it is a model that has one. So the
     flavor comes from the draft's own config through the same selector the
-    target uses, and that backend says whether a draft of its flavor caches
-    into a pool of its own (`DRAFT_OWNS_KV_POOL`). Where it does not, the
-    runner never sees a draft builder at all.
+    target uses, and the family table (`selector.DRAFT_POOL_FAMILIES`) says
+    whether a draft of its flavor caches into a pool of its own. Where it does
+    not, the runner never sees a draft builder at all.
 
     Which is why speculative decoding imports no attention anywhere: adding a
-    flavor is a flag and a `make_kv_pool` on that backend.
+    flavor is a table entry and a `make_kv_pool` on that backend.
     """
     from aiter import dtypes
 
-    from atom.utils.selector import attn_family, get_attn_backend
+    from atom.utils.selector import DRAFT_POOL_FAMILIES, attn_family, get_attn_backend
 
-    backend = get_attn_backend(attn_family(draft_hf))
-    if not backend.DRAFT_OWNS_KV_POOL:
+    family = attn_family(draft_hf)
+    if family not in DRAFT_POOL_FAMILIES:
         return None
+    backend = get_attn_backend(family)
     make_pool = partial(
         backend.make_kv_pool,
         draft_hf,
         world_size=model_runner.world_size,
-        kv_dtype=dtypes.d_dtypes[model_runner.config.kv_cache_dtype],
+        # settled once, by `Config`
+        kv_dtype=dtypes.d_dtypes[model_runner.config.speculative_config.draft_kv_dtype],
     )
     # A way to build it, not the pool: how many rows it has is the walk's to
     # say, and this runs from the proposer's `__init__` -- `build_drafter`
@@ -175,7 +211,7 @@ def build_kv_cache_tensor(self, module):
         idx = self.pool_rows[DRAFT_KV_ROWS][module]
         k_cache, v_cache = self.kv_pool.kv_views(idx)
         module.max_model_len = runner.config.max_model_len
-        if runner.config.kv_cache_dtype == "fp8":
+        if self.kv_pool.scale_fields:  # the layout's, as the target binder asks
             module.k_scale, module.v_scale = self.kv_pool.scale_views(idx)
         module.k_cache = k_cache
         module.v_cache = v_cache