From 1619cf6c35e1803e761205eb3813cfe8def6f159 Mon Sep 17 00:00:00 2001 From: ganyi Date: Thu, 12 Nov 2026 02:05:23 -0600 Subject: [PATCH 1/3] Add --spec-decode-kv-cache-dtype: the KV dtype of a speculative draft's own pool A speculative draft that caches into a KV pool of its own (an MHA draft: MiniMax-M3.1's DSpark, an MHA EAGLE3 draft) can now have its pages at bf16 or fp8 independently of the target's --kv_cache_dtype. A draft in the target's pool (MTP, an MLA draft) always has the target's dtype, and a different request is an error. - SpeculativeConfig.kv_cache_dtype is what was asked (never rewritten); SpeculativeConfig.draft_kv_dtype is what Config settled, the one value the draft builder, the draft's attention, the DSpark/EAGLE proposers, the calibration identity and the compile hash read. - Whether a draft owns a pool is decided once in Config from the attention family (selector.DRAFT_POOL_FAMILIES), not by resolving a backend, so the API process does not import aiter; DRAFT_OWNS_KV_POOL is gone and a test pins the table to the backends that define make_kv_pool. EAGLE3 is covered as well as DSpark. - DeepSeek-V4's fp8 -> bf16 fallback off gfx950/gfx1250 moves from the attention builder into Config, before the draft's dtype is resolved, so a V4 draft in the target's pool is not left asking for fp8 against a bf16 pool. - A draft dtype different from the target's is refused with --kv-transfer-config: the offload namespace keys the target's dtype only. - Docs: SpeculativeConfig.kv_cache_dtype, the flag row, the --method row. Co-Authored-By: Claude --- atom/config.py | 66 +++++- atom/kv_transfer/offload/mp/deployment.py | 20 +- atom/model_engine/arg_utils.py | 19 ++ atom/model_ops/attentions/aiter_attention.py | 4 - atom/model_ops/attentions/backends.py | 18 +- atom/model_ops/attentions/deepseek_v4_attn.py | 21 +- atom/spec_decode/calibration.py | 9 +- atom/spec_decode/draft_kv.py | 50 ++++- atom/spec_decode/dspark_proposer.py | 23 +- atom/spec_decode/eagle_proposer.py | 4 + atom/utils/selector.py | 9 + docs/configuration_guide.md | 4 +- tests/test_arg_utils_spec.py | 31 +++ tests/test_compile_hash_spec.py | 1 + tests/test_draft_kv_builder.py | 142 ++++++++++++- tests/test_draft_kv_dtype.py | 197 ++++++++++++++++++ tests/test_dspark.py | 47 +++++ tests/test_dspark_profile.py | 34 ++- tests/test_kv_pool_contract.py | 69 +++--- tests/test_lmcache_mp.py | 5 +- tests/test_lmcache_mp_native_worker.py | 1 + tests/test_v4_indexer_defaults.py | 44 +++- 22 files changed, 705 insertions(+), 113 deletions(-) create mode 100644 tests/test_draft_kv_dtype.py diff --git a/atom/config.py b/atom/config.py index e4241bc694..2abd760582 100644 --- a/atom/config.py +++ b/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( diff --git a/atom/kv_transfer/offload/mp/deployment.py b/atom/kv_transfer/offload/mp/deployment.py index 8f1482d425..c96763e352 100644 --- a/atom/kv_transfer/offload/mp/deployment.py +++ b/atom/kv_transfer/offload/mp/deployment.py @@ -167,7 +167,9 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool: registering any cache memory. """ - if _config_has_own_pool_draft(config): + from atom.spec_decode.draft_kv import draft_owns_kv_pool + + if draft_owns_kv_pool(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. @@ -187,22 +189,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. diff --git a/atom/model_engine/arg_utils.py b/atom/model_engine/arg_utils.py index 2e9d91adce..4b0cb190e3 100644 --- a/atom/model_engine/arg_utils.py +++ b/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 diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index 11d140b683..ae7a17d8cf 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/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" diff --git a/atom/model_ops/attentions/backends.py b/atom/model_ops/attentions/backends.py index 5cc46ffa4d..cec01d77aa 100644 --- a/atom/model_ops/attentions/backends.py +++ b/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 diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index db8e81b852..5150016038 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/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 diff --git a/atom/spec_decode/calibration.py b/atom/spec_decode/calibration.py index d67fb132de..e6af3e8c93 100644 --- a/atom/spec_decode/calibration.py +++ b/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): diff --git a/atom/spec_decode/draft_kv.py b/atom/spec_decode/draft_kv.py index d8d1c5ad7c..33e81f82ce 100644 --- a/atom/spec_decode/draft_kv.py +++ b/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 + into a pool of its own (`selector.DRAFT_POOL_FAMILIES`). 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 diff --git a/atom/spec_decode/dspark_proposer.py b/atom/spec_decode/dspark_proposer.py index a0f361743f..eeca3249ab 100644 --- a/atom/spec_decode/dspark_proposer.py +++ b/atom/spec_decode/dspark_proposer.py @@ -426,7 +426,7 @@ def _build_draft_model(self, model_class) -> nn.Module: # Standalone draft: build from the DRAFT's own hf_config, exactly as # EagleProposer does for eagle3. Shallow-copy rather than deepcopy -- # atom_config can hold non-picklable cuda.Stream objects, and only - # hf_config / compilation_config are mutated here. + # hf_config / kv_cache_dtype / compilation_config are mutated here. import copy from atom.config import CompilationLevel @@ -434,6 +434,8 @@ def _build_draft_model(self, model_class) -> nn.Module: draft_hf = self.speculative_config.draft_model_hf_config draft_atom_config = copy.copy(self.config) draft_atom_config.hf_config = draft_hf + # Before the model is built: its attention reads this at construction. + draft_atom_config.kv_cache_dtype = self.speculative_config.draft_kv_dtype draft_atom_config.compilation_config = copy.copy(self.config.compilation_config) draft_atom_config.compilation_config.level = CompilationLevel.NO_COMPILATION model = model_class( @@ -517,7 +519,8 @@ def _resolve_dtype_q(self, forward_context) -> "tuple[torch.dtype, bool]": # d_dtypes maps the "auto" cache dtype to None, and torch.empty would # silently read that as float32 rather than the model dtype. - from_config = dtypes.d_dtypes.get(self.config.kv_cache_dtype) or self.dtype + kv_cache_dtype = self.speculative_config.draft_kv_dtype + from_config = dtypes.d_dtypes.get(kv_cache_dtype) or self.dtype bound = self._bound_draft_k_cache(forward_context) if bound is None: @@ -528,15 +531,17 @@ def _resolve_dtype_q(self, forward_context) -> "tuple[torch.dtype, bool]": if bound.dtype != from_config: layer_num = self.model.first_attention_layer_num logger.warning( - "DSpark draft layer_%d is bound to a %s KV cache, but " - "--kv_cache_dtype=%s implies %s. Using the bound tensor's dtype " - "for q_out so the fused write agrees -- but the two should not " - "be able to differ, since the draft binds into the pool the " - "engine allocated from that same flag. Check that layer_%d " - "resolved to a draft row and not to a target layer.", + "DSpark draft layer_%d is bound to a %s KV cache, but the draft's " + "KV dtype (speculative_config.draft_kv_dtype: the target's " + "--kv-cache-dtype, or --spec-decode-kv-cache-dtype for a draft with " + "a pool of its own, resolved once by Config) is %s, implying %s. " + "Using the bound tensor's dtype for q_out so the fused write agrees " + "-- but the pool the draft binds into is allocated at that same " + "resolved dtype, so the two should not be able to differ. Check that " + "layer_%d resolved to a draft row and not to a target layer.", layer_num, bound.dtype, - self.config.kv_cache_dtype, + kv_cache_dtype, from_config, layer_num, ) diff --git a/atom/spec_decode/eagle_proposer.py b/atom/spec_decode/eagle_proposer.py index 5c5d07e37b..22012ceefe 100644 --- a/atom/spec_decode/eagle_proposer.py +++ b/atom/spec_decode/eagle_proposer.py @@ -245,6 +245,10 @@ def _build_draft_model(self, model_class) -> nn.Module: self.config.compilation_config ) draft_atom_config.compilation_config.level = CompilationLevel.NO_COMPILATION + # The draft's attention reads its cache dtype from this config, and + # its pool is built at `speculative_config.draft_kv_dtype`: they + # must be the one value (DSpark does the same). + draft_atom_config.kv_cache_dtype = self.speculative_config.draft_kv_dtype # Draft attention layer_num must continue from the target model's # layer count so it maps to the correct kv_cache_data entry. model = model_class( diff --git a/atom/utils/selector.py b/atom/utils/selector.py index e309c2e35e..9c3f407a36 100644 --- a/atom/utils/selector.py +++ b/atom/utils/selector.py @@ -67,6 +67,15 @@ def has_recurrent_state(self) -> bool: return self in (Family.KIMI_MLA, Family.QSA_GDN, Family.GDN, Family.CSA2) +#: The families whose *draft* caches into a KV pool of its own: an MHA draft keeps +#: its own K and V, there is no target tensor its rows could be. Every other family's +#: draft binds into the target's pool (an MLA / CSA draft's latent is the target's). +#: A table on the family, not a flag on the backend class, so `Config` (which runs in +#: the API process, where importing a backend would import aiter) and the draft +#: builder read the same answer; tests/test_kv_pool_contract.py pins it to the +#: backends that define `make_kv_pool`. +DRAFT_POOL_FAMILIES = frozenset({Family.MHA}) + # What makes a hybrid one is which of its layers are linear, and no field says # that -- so these two are named, and everything else is read off the shape. _KIMI_MLA_TYPES = ("kimi_linear", "glm5_next_text") diff --git a/docs/configuration_guide.md b/docs/configuration_guide.md index 6c6a34e6e3..05e2b2006b 100644 --- a/docs/configuration_guide.md +++ b/docs/configuration_guide.md @@ -292,6 +292,7 @@ method with `num_speculative_tokens=1` is supported. | `model` | `Optional[str]` | `None` | Draft model name or path (typically the same as the target model for MTP) | | `num_speculative_tokens` | `Optional[int]` | `None` | Number of speculative tokens per iteration; **must be `1`** | | `draft_model_hf_config` | `Optional[PretrainedConfig]` | `None` | HuggingFace config for the draft model; auto-loaded from `model` when `None` | +| `kv_cache_dtype` | `Optional[str]` | `None` | KV dtype of the draft's pages as requested (`"bf16"` or `"fp8"`), set by `--spec-decode-kv-cache-dtype`; never rewritten. `Config` stores the dtype the draft will use in `draft_kv_dtype`. A draft with a KV pool of its own (one whose attention is MHA: MiniMax-M3.1's standalone DSpark, an MHA EAGLE3 draft) takes `"bf16"` or `"fp8"`; omitted, it has the target's dtype. A draft in the target's pool (MTP, an MLA draft) always has the target's dtype, and a different request is an error. A draft dtype that differs from the target's is refused with `--kv-transfer-config` (KV offload included), whose namespace keys the target's dtype only. | ### Table-driven MTP config @@ -386,8 +387,9 @@ all flags via `add_cli_args()` and converts them into a `Config` via | `--moe-backend` | | `str` | `"standard"` | MoE implementation; choices: `standard`, `mega`. `mega` runs FlyDSL MegaMoE and requires `--enable-expert-parallel`; with `-sp N` it instead runs the Ulysses SP MXFP4 MoE through AITER `MegaMoeTP` on gfx950, on every step whose native MoE quantizes activations to FP4 (prefill, and decode from 256 global rows), and needs MORI built with `BUILD_CCO_SDMA=ON`. Startup fails when a layer or setting is outside what it serves. Each launch covers at most `min(ceil(max_num_batched_tokens / N), 8192)` local rows | | `--torch-profiler-dir` | | `str` | `None` | Directory for torch profiler traces | | `--enable-dp-attention` | | flag | `False` | Enable DP attention | -| `--method` | | `str` | `None` | Speculative method; choices: `mtp` | +| `--method` | | `str` | `None` | Speculative method; choices: `mtp`, `eagle3`, `dspark` | | `--num-speculative-tokens` | | `int` | `1` | Number of speculative tokens per iteration | +| `--spec-decode-kv-cache-dtype` | | `str` | `None` | The draft's KV dtype, `SpeculativeConfig.kv_cache_dtype` (native ATOM has no `--spec-decode-config` JSON; this flag is its setting); choices: `bf16`, `fp8`. Needs `--method` and `--num-speculative-tokens` > 0. See the `SpeculativeConfig` table for what each value means per draft | | `--max-num-batched-tokens` | | `int` | `16384` | Maximum number of tokens to batch in the async engine | | `--state-checkpoint-interval-tokens` | | `int` | `8192` | Tokens between rungs of the per-request state checkpoint ladder; must be a multiple of the prefix-cache hash block size. `0` turns checkpointing off entirely; **`-1` turns off only the ladder**, leaving the prompt-end anchor and the demand rung as the placements. Prompts shorter than one interval publish no rung. A rung also quantizes prefill chunk boundaries, since a checkpoint is only valid where a forward ends exactly on one — which is most of what `-1` buys back | | `--state-checkpoint-demand` / `--no-state-checkpoint-demand` | | `bool` | on | Let a prefix hit that was refused for want of a checkpoint place a rung of its own. Off leaves the prompt-end anchor as the only placement. `ATOM_STATE_CHECKPOINT_DEMAND=0` overrides this without touching a launch script | diff --git a/tests/test_arg_utils_spec.py b/tests/test_arg_utils_spec.py index 175932983b..4d56903094 100644 --- a/tests/test_arg_utils_spec.py +++ b/tests/test_arg_utils_spec.py @@ -228,11 +228,42 @@ def test_method_mtp_valid_tokens_builds_speculative_config(self): method="mtp", model=args.model, num_speculative_tokens=3, + kv_cache_dtype=None, synthetic_acceptance_rate=None, synthetic_acceptance_length=None, ) assert kwargs["speculative_config"] is fake_spec_config + def test_spec_decode_kv_cache_dtype_reaches_speculative_config(self): + """--spec-decode-kv-cache-dtype is SpeculativeConfig.kv_cache_dtype; the + native CLI has no --spec-decode-config JSON to carry it.""" + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + argv = ["--method", "mtp", "--spec-decode-kv-cache-dtype", "bf16"] + with patch("atom.model_engine.arg_utils.SpeculativeConfig") as mock_cls: + EngineArgs.from_cli_args(parser.parse_args(argv))._get_engine_kwargs() + assert mock_cls.call_args.kwargs["kv_cache_dtype"] == "bf16" + + @pytest.mark.parametrize( + "argv", [[], ["--method", "mtp", "--num-speculative-tokens", "0"]] + ) + def test_spec_decode_kv_cache_dtype_without_speculation_is_refused(self, argv): + """Both ways speculation is off: no --method, or one given with no tokens.""" + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = EngineArgs.from_cli_args( + parser.parse_args(["--spec-decode-kv-cache-dtype", "fp8", *argv]) + ) + with pytest.raises( + ValueError, match="needs --method and --num-speculative-tokens > 0" + ): + args._get_engine_kwargs() + + def test_spec_decode_kv_cache_dtype_has_no_fp4(self, capsys): + """No draft pool has an FP4 layout, so the flag does not offer one.""" + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + with pytest.raises(SystemExit): + parser.parse_args(["--spec-decode-kv-cache-dtype", "fp4"]) + assert "invalid choice" in capsys.readouterr().err + class TestEngineArgsIndexCacheDtype: """Regression tests for index-cache dtype CLI and deferred defaults.""" diff --git a/tests/test_compile_hash_spec.py b/tests/test_compile_hash_spec.py index a33f264706..2aa7caec88 100644 --- a/tests/test_compile_hash_spec.py +++ b/tests/test_compile_hash_spec.py @@ -32,6 +32,7 @@ def _hash(spec=None): def _spec(method="dspark", k=3, **kw): + kw.setdefault("draft_kv_dtype", "bf16") # a field of every SpeculativeConfig return SimpleNamespace(method=method, num_speculative_tokens=k, **kw) diff --git a/tests/test_draft_kv_builder.py b/tests/test_draft_kv_builder.py index 31b44b8ec4..309e7bab56 100644 --- a/tests/test_draft_kv_builder.py +++ b/tests/test_draft_kv_builder.py @@ -23,7 +23,8 @@ from torch import nn from atom.model_ops.attentions.mha_kv_pool import MhaKvPool -from atom.spec_decode.draft_kv import DRAFT_KV_ROWS, DraftKvBuilder +from atom.spec_decode.draft_kv import DRAFT_KV_ROWS, DraftKvBuilder, draft_owns_kv_pool +from atom.utils.selector import Family class _PagedAttention(nn.Module): @@ -224,3 +225,142 @@ def test_the_count_tracks_the_model_it_walked(draft_layers): assert builder.kv_pool is not None assert asked == [(draft_layers, TARGET_BLOCK)] + + +class _FakePool: + def __init__(self, scale_fields): + self.scale_fields = scale_fields + + def kv_views(self, idx): + return torch.zeros(1), torch.zeros(1) + + def scale_views(self, idx): + return torch.ones(1), torch.ones(1) + + +@pytest.mark.parametrize("scale_fields", [[object()], []]) +def test_scales_are_bound_when_the_pools_layout_has_them(scale_fields): + """The layout decides, as the target binder asks it: not a dtype string.""" + runner = _runner(draft_layers=1) + runner.config = SimpleNamespace(max_model_len=64) + builder = DraftKvBuilder(runner, None) + builder._kv_pool = _FakePool(scale_fields) + module = next(iter(builder.pool_rows[DRAFT_KV_ROWS])) + module.layer_num = 0 + + tensor = builder.build_kv_cache_tensor(module) + + assert (tensor.k_scale is not None) is bool(scale_fields) + + +def _spec_config(method="dspark", draft_hf="draft-hf"): + return SimpleNamespace( + speculative_config=SimpleNamespace( + method=method, draft_model_hf_config=draft_hf + ) + ) + + +@pytest.mark.parametrize( + "config, owns, want", + [ + (_spec_config(), True, True), + (_spec_config(), False, False), # the draft's backend shares the target's pool + (_spec_config("eagle3"), True, True), # a draft model of its own, MHA: a pool + ( + _spec_config("eagle3"), + False, + False, + ), # ... an MLA one binds into the target's + (_spec_config("mtp"), True, False), # MTP runs the target's layers + (_spec_config(draft_hf=None), True, False), + (SimpleNamespace(speculative_config=None), True, False), + ], +) +def test_a_native_draft_owns_a_pool_when_its_backend_says_so( + monkeypatch, config, owns, want +): + from atom.utils import selector + + monkeypatch.setattr( + selector, "attn_family", lambda _hf: Family.MHA if owns else Family.MLA + ) + assert draft_owns_kv_pool(config) is want + + +def test_a_plugin_draft_owns_no_pool_and_the_selector_is_not_asked(monkeypatch): + """The native runner builds a draft's pool; under vLLM / SGLang the selector has + no backend for every draft flavor (CSA2 raises), so it is not consulted.""" + from atom.plugin import prepare + from atom.utils import selector + + def refuse(*_a): + raise AssertionError("the selector was asked") + + monkeypatch.setattr(prepare, "_CURRENT_FRAMEWORK", "vllm") + monkeypatch.setattr(selector, "attn_family", refuse) + assert draft_owns_kv_pool(_spec_config()) is False + + +def test_a_plugin_config_without_native_draft_fields_answers_no(monkeypatch): + """The vLLM offload adapter passes vLLM's SpeculativeConfig through unchanged: it + has `method` but no `draft_model_hf_config`, and the answer must not read it.""" + from atom.plugin import prepare + + monkeypatch.setattr(prepare, "_CURRENT_FRAMEWORK", "vllm") + vllm_spec = SimpleNamespace(method="eagle3", draft_model_config=object()) + assert draft_owns_kv_pool(SimpleNamespace(speculative_config=vllm_spec)) is False + + +@pytest.mark.parametrize( + "target, asked, pool_dtype", + [ + ("fp8", "bf16", "bf16"), # the inverse overrides mismatch too + ("bf16", "fp8", "fp8"), + ("fp8", None, "fp8"), + ], +) +def test_an_eagle3_draft_attention_is_built_at_the_dtype_of_its_pool( + target, asked, pool_dtype +): + """The draft's attention reads its cache dtype off the config it is built + with, and its pool is built at `speculative_config.draft_kv_dtype` (what + `draft_kv_builder` hands `make_kv_pool`). They are one value; the target's + config keeps its own.""" + # eagle_proposer pulls aiter in through atom.distributed.pcp_utils + pytest.importorskip("aiter", reason="the proposer imports aiter at module load") + from atom.config import CompilationConfig, resolve_draft_kv_dtype + from atom.spec_decode.eagle_proposer import EagleProposer + + spec = SimpleNamespace( + method="eagle3", + draft_model_hf_config=SimpleNamespace(), + kv_cache_dtype=asked, + ) + spec.draft_kv_dtype = resolve_draft_kv_dtype(spec, target, own_pool=True) + assert spec.draft_kv_dtype == pool_dtype + + p = EagleProposer.__new__(EagleProposer) + p.speculative_config = spec + p.runner = SimpleNamespace() + p.config = SimpleNamespace( + kv_cache_dtype=target, + hf_config=SimpleNamespace(num_hidden_layers=3), + compilation_config=CompilationConfig(), + ) + seen = {} + + def model_class(cfg, layer_offset): + seen["dtype"] = cfg.kv_cache_dtype + return object() + + import atom.spec_decode.eagle_proposer as mod + + orig, mod.draft_kv_builder = mod.draft_kv_builder, lambda runner, hf: None + try: + p._build_draft_model(model_class) + finally: + mod.draft_kv_builder = orig + + assert seen == {"dtype": pool_dtype} + assert p.config.kv_cache_dtype == target diff --git a/tests/test_draft_kv_dtype.py b/tests/test_draft_kv_dtype.py new file mode 100644 index 0000000000..ba90073c7b --- /dev/null +++ b/tests/test_draft_kv_dtype.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +"""The KV dtype of a speculative draft: ``--spec-decode-kv-cache-dtype`` (what was +asked, ``SpeculativeConfig.kv_cache_dtype``) and what ``Config`` settles it to +(``draft_kv_dtype``, the one every reader uses). GPU-free: a real ``Config`` over a +stubbed HF config.""" + +import types + +import pytest +from transformers import PretrainedConfig + +from atom import config as config_module +from atom.config import ( + CompilationConfig, + Config, + SpeculativeConfig, + resolve_draft_kv_dtype, +) +from tests.test_v4_indexer_defaults import _GfxProbe, _install_chip_info_stub + +OFFLOAD = {"kv_connector": "lmcache_offload", "kv_role": "offload"} + + +def _config(monkeypatch, *, architectures, **overrides): + hf_config = PretrainedConfig( + architectures=architectures, max_position_embeddings=8192 + ) + monkeypatch.setattr( + config_module, "get_hf_config", lambda _m, trust_remote_code=False: hf_config + ) + monkeypatch.setattr(config_module, "get_generation_config", lambda _m: None) + monkeypatch.setattr( + config_module, + "QuantizationConfig", + lambda *_a, **_k: types.SimpleNamespace( + exclude_layers=[], compute_hash=lambda: "" + ), + ) + monkeypatch.setattr(config_module, "is_plugin_mode", lambda: False) + _install_chip_info_stub(monkeypatch, _GfxProbe("gfx950")) + return Config( + model="test-model", + torch_profiler_dir=None, + compilation_config=CompilationConfig(level=0, use_cudagraph=False), + **{"kv_cache_dtype": "bf16", **overrides}, + ) + + +def _llama(monkeypatch, **overrides): + return _config(monkeypatch, architectures=["LlamaForCausalLM"], **overrides) + + +@pytest.fixture(autouse=True) +def no_synthetic_acceptance_import(monkeypatch): + """`SpeculativeConfig` pulls in triton (the rejection sampler) to resolve a + forced-acceptance knob these tests never set.""" + monkeypatch.setattr(SpeculativeConfig, "_resolve_synthetic_acceptance", lambda s: 0) + + +def _draft(method="dspark", asked=None): + """A speculative config as `Config` holds it: what `resolve_draft_kv_dtype` reads.""" + return types.SimpleNamespace(method=method, kv_cache_dtype=asked) + + +@pytest.mark.parametrize("target", ["bf16", "fp8"]) +@pytest.mark.parametrize("asked", [None, "bf16", "fp8"]) +def test_an_own_pool_draft_keeps_the_targets_dtype_unless_asked(target, asked): + assert resolve_draft_kv_dtype(_draft(asked=asked), target, own_pool=True) == ( + asked or target + ) + + +def test_a_shared_pool_draft_has_the_targets_dtype(): + assert resolve_draft_kv_dtype(_draft("mtp"), "fp8", own_pool=False) == "fp8" + assert resolve_draft_kv_dtype(_draft("mtp", "bf16"), "bf16", own_pool=False) == ( + "bf16" + ) + with pytest.raises(ValueError, match="caches in the target's KV pool"): + resolve_draft_kv_dtype(_draft("mtp", "bf16"), "fp8", own_pool=False) + + +def _dspark_spec(**kw): + """A standalone MHA DSpark draft (MiniMax-M3.1's), built the way the CLI builds it.""" + hf = PretrainedConfig( + architectures=["DSparkMiniMaxDraftModel"], + model_type="minimax_m3", + dspark_target_layer_ids=[1, 2], + dspark_noise_token_id=0, + ) + return SpeculativeConfig( + method="dspark", + model="draft", + num_speculative_tokens=3, + draft_model_hf_config=hf, + **kw, + ) + + +def _eagle3_spec(**kw): + """An EAGLE3 draft that is an MHA model (llama): it has a pool of its own.""" + hf = PretrainedConfig(architectures=["LlamaForCausalLM"], model_type="llama") + return SpeculativeConfig( + method="eagle3", + model="draft", + num_speculative_tokens=3, + draft_model_hf_config=hf, + **kw, + ) + + +@pytest.mark.parametrize("target", ["bf16", "fp8"]) +@pytest.mark.parametrize("asked", [None, "bf16", "fp8"]) +def test_config_settles_the_draft_dtype_and_keeps_what_was_asked( + monkeypatch, target, asked +): + spec = _dspark_spec(kv_cache_dtype=asked) + config = _llama(monkeypatch, kv_cache_dtype=target, speculative_config=spec) + + assert spec.kv_cache_dtype == asked + assert spec.draft_kv_dtype == (asked or target) + assert config.kv_cache_dtype == target + + +def test_the_drafts_dtype_keys_the_compile_cache(monkeypatch): + hashes = { + asked: _llama( + monkeypatch, + speculative_config=_dspark_spec(kv_cache_dtype=asked), + ).compute_hash() + for asked in ("bf16", "fp8") + } + assert hashes["bf16"] != hashes["fp8"] + + +@pytest.mark.parametrize("asked", [None, "bf16", "fp8"]) +def test_an_mha_eagle3_draft_has_a_pool_of_its_own_and_takes_its_own_dtype( + monkeypatch, asked +): + spec = _eagle3_spec(kv_cache_dtype=asked) + _llama(monkeypatch, kv_cache_dtype="fp8", speculative_config=spec) + assert spec.draft_kv_dtype == (asked or "fp8") + + +def test_an_mtp_draft_shares_the_targets_pool(monkeypatch): + spec = _eagle3_spec(kv_cache_dtype="bf16") + spec.method = "mtp" + with pytest.raises(ValueError, match="caches in the target's KV pool"): + _llama(monkeypatch, kv_cache_dtype="fp8", speculative_config=spec) + + +def test_a_plugin_config_with_a_v41_dspark_draft_never_asks_for_a_backend( + monkeypatch, +): + """V4.1 DSpark is a supported plugin deployment and vLLM / SGLang have no CSA2 + backend: whether the draft owns a pool is a native runner's question, so `Config` + must not route it through a backend lookup (which raises for CSA2), nor import + one (aiter) into the API process.""" + from atom.plugin import prepare + from atom.utils import selector + + def refuse(*_a): + raise AssertionError("a backend was resolved") + + hf = PretrainedConfig( + architectures=["DeepseekV41DSparkModel"], model_type="deepseek_v41_dspark" + ) + monkeypatch.setattr(selector, "get_attn_backend", refuse) + for framework in ("vllm", "atom"): + monkeypatch.setattr(prepare, "_CURRENT_FRAMEWORK", framework) + spec = SpeculativeConfig( + method="dspark", + model="draft", + num_speculative_tokens=3, + draft_model_hf_config=hf, + ) + _llama(monkeypatch, speculative_config=spec) + assert spec.draft_kv_dtype == "bf16" + + +def test_a_draft_dtype_beside_the_targets_cannot_share_the_offload_namespace( + monkeypatch, +): + # the offload namespace keys the target's KV dtype alone + with pytest.raises( + ValueError, + match="draft KV dtype fp8 differs.*--spec-decode-kv-cache-dtype bf16", + ): + _llama( + monkeypatch, + kv_transfer_config=OFFLOAD, + speculative_config=_dspark_spec(kv_cache_dtype="fp8"), + ) + # the remedy it names works + spec = _dspark_spec(kv_cache_dtype="bf16") + _llama(monkeypatch, kv_transfer_config=OFFLOAD, speculative_config=spec) + assert spec.draft_kv_dtype == "bf16" diff --git a/tests/test_dspark.py b/tests/test_dspark.py index 2ece86b5ca..c1a9cf4ac8 100644 --- a/tests/test_dspark.py +++ b/tests/test_dspark.py @@ -2658,3 +2658,50 @@ def test_minimax_m31_dspark_refuses_tp_below_its_kv_heads(monkeypatch, tp): config = types.SimpleNamespace(num_attention_heads=64, num_key_value_heads=4) with pytest.raises(ValueError, match="needs TP >= 4"): m3_dspark.M3DSparkAttention(None, config, None, layer_num=0) + + +def test_standalone_draft_is_built_with_its_own_kv_dtype_target_untouched(monkeypatch): + """The draft model reads its KV dtype off the config it is constructed + with, so the shallow copy carries the draft's (fp8 under a bf16 target) and + the target's config keeps its own.""" + from atom.config import CompilationConfig + from atom.spec_decode import draft_kv + from atom.spec_decode.dspark_proposer import DSparkProposer + + seen = {} + monkeypatch.setattr(draft_kv, "draft_kv_builder", lambda runner, hf: None) + + p = DSparkProposer.__new__(DSparkProposer) + spec = types.SimpleNamespace( + draft_model_hf_config=types.SimpleNamespace(), + draft_kv_dtype="fp8", # as `Config` resolved it, under a bf16 target + use_dspark_with_draft=lambda: True, + ) + p.speculative_config = spec + p.runner = types.SimpleNamespace() + p.config = types.SimpleNamespace( + kv_cache_dtype="bf16", + hf_config=types.SimpleNamespace(num_hidden_layers=3), + compilation_config=CompilationConfig(), + ) + + def model_class(cfg, layer_offset): + seen["dtype"] = cfg.kv_cache_dtype + return object() + + p._build_draft_model(model_class) + + assert seen == {"dtype": "fp8"} + assert p.config.kv_cache_dtype == "bf16" + + +def test_dtype_q_follows_the_drafts_dtype_not_the_targets(): + from atom.spec_decode.dspark_proposer import DSparkProposer + + p = DSparkProposer.__new__(DSparkProposer) + p.config = types.SimpleNamespace(kv_cache_dtype="fp8") + p.dtype = torch.bfloat16 + p.speculative_config = types.SimpleNamespace(draft_kv_dtype="bf16") + p._bound_draft_k_cache = lambda fc: None + + assert p._resolve_dtype_q(None) == (torch.bfloat16, False) diff --git a/tests/test_dspark_profile.py b/tests/test_dspark_profile.py index 3ea9853404..4ea040f5dc 100644 --- a/tests/test_dspark_profile.py +++ b/tests/test_dspark_profile.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: MIT from copy import deepcopy +from types import SimpleNamespace import pytest +import torch -from atom.spec_decode.calibration import validate_calibration +from atom.spec_decode.calibration import calibration_identity, validate_calibration def profile(): @@ -42,3 +44,33 @@ def test_incompatible_or_invalid_profile_is_rejected(field, value): data[field] = value with pytest.raises(ValueError): validate_calibration(data, identity, width=5, max_batch=4) + + +def identity_of(tmp_path, monkeypatch, **kw): + for name in ("config.json", "model.safetensors.index.json"): + (tmp_path / name).write_text("{}") + monkeypatch.setattr(torch.cuda, "get_device_name", lambda device: "gpu") + config = SimpleNamespace( + model=str(tmp_path), + hf_config=SimpleNamespace(model_type="m"), + tensor_parallel_size=4, + kv_cache_dtype="fp8", + index_cache_dtype="fp8", + enforce_eager=False, + # as `Config` resolves it: the target's, unless the draft owns a pool + speculative_config=SimpleNamespace(draft_kv_dtype="fp8"), + ) + for key, value in kw.items(): + setattr(config, key, value) + return calibration_identity(config, 0) + + +def test_a_draft_dtype_beside_the_targets_changes_the_identity(tmp_path, monkeypatch): + def with_draft(dtype): + spec = SimpleNamespace(draft_kv_dtype=dtype) + return identity_of(tmp_path, monkeypatch, speculative_config=spec) + + # the target's own dtype leaves it as it was, so older profiles still match + assert with_draft("fp8") == identity_of(tmp_path, monkeypatch) + assert "draft_kv_cache_dtype" not in with_draft("fp8") + assert with_draft("bf16") != with_draft("fp8") diff --git a/tests/test_kv_pool_contract.py b/tests/test_kv_pool_contract.py index 213b8d00fd..6bcda9dc9e 100644 --- a/tests/test_kv_pool_contract.py +++ b/tests/test_kv_pool_contract.py @@ -166,14 +166,14 @@ def test_neither_has_a_default(path, cls, fn, arg): # ── Who builds a draft a pool of its own ─────────────────────────────────── # -# Two members that have to agree: the flag `draft_kv_builder` reads, and the -# factory it calls once the flag says yes. Getting the flag wrong is the silent -# half -- a backend that builds pools but forgets to say so leaves its drafts -# sharing the target's rows, which is a real layout for MLA and a wrong one -# here. The other half is loud: the base factory raises. +# Two members that have to agree: the families `DRAFT_POOL_FAMILIES` says own a +# draft pool, and the backends that define the factory `draft_kv_builder` calls +# for them. A family listed without a factory is the loud half (the base factory +# raises); a factory no family reaches is the silent one -- a backend that builds +# pools but is not in the table leaves its drafts sharing the target's rows, a real +# layout for MLA and a wrong one here. POOL_FACTORY = "make_kv_pool" -DRAFT_FLAG = "DRAFT_OWNS_KV_POOL" # The declaration, which defines the factory in order to refuse it. FACTORY_DECLARATION = "AttentionBackend" @@ -191,42 +191,43 @@ def _classes_defining_the_factory() -> list[tuple[str, ast.ClassDef]]: ] -def _sets_the_flag_true(cls: ast.ClassDef) -> bool: - """Whether this class body assigns `DRAFT_OWNS_KV_POOL = True`, annotated - or not. Its own body: an inheriting backend gets the answer with the - factory, and the two cannot part.""" - for stmt in cls.body: - targets = ( - [stmt.target] - if isinstance(stmt, ast.AnnAssign) - else getattr(stmt, "targets", []) - ) - if any(isinstance(t, ast.Name) and t.id == DRAFT_FLAG for t in targets): - return stmt.value is not None and getattr(stmt.value, "value", None) is True - return False +def _backend_class_name(family) -> str: + from atom.utils.selector import get_attn_backend_cls + + return get_attn_backend_cls(family, False, False).rsplit(".", 1)[-1] -def test_the_factory_is_declared_where_the_flag_is(): - """Both on `AttentionBackend`, so a backend inherits a coherent pair - (no pool, and a factory that says so) rather than half of one.""" +def test_the_factory_is_declared_on_the_backend_base(): + """So a backend outside the table inherits a factory that refuses, not none.""" declarations = _classes_defining_the_factory() assert FACTORY_DECLARATION in [cls.name for _, cls in declarations] -@pytest.mark.parametrize( - "cls", - [ - pytest.param(cls, id=f"{stem}.{cls.name}") - for stem, cls in _classes_defining_the_factory() +def test_every_family_that_owns_a_draft_pool_has_a_backend_that_builds_it(): + from atom.utils.selector import DRAFT_POOL_FAMILIES + + builders = { + cls.name + for _, cls in _classes_defining_the_factory() if cls.name != FACTORY_DECLARATION - ], -) -def test_a_backend_that_can_build_a_draft_pool_says_so(cls): - """The silent direction. Without the flag `draft_kv_builder` returns None - and the draft binds into the target's rows -- no error, no log, and a - draft reading K and V that are not its own.""" - assert _sets_the_flag_true(cls) + } + assert {_backend_class_name(f) for f in DRAFT_POOL_FAMILIES} <= builders + + +def test_a_backend_that_can_build_a_draft_pool_is_in_the_table(): + """The silent direction: out of the table, `draft_kv_builder` returns None and + the draft binds into the target's rows -- no error, no log.""" + from atom.utils.selector import DRAFT_POOL_FAMILIES, Family + + builders = { + cls.name + for _, cls in _classes_defining_the_factory() + if cls.name != FACTORY_DECLARATION + } + listed = {_backend_class_name(f) for f in DRAFT_POOL_FAMILIES} + # `get_attn_backend_cls` is asked per family with no plugin and no unified attn + assert builders == listed, (builders, listed, list(Family)) def test_a_geometry_row_space_names_itself_readably(): diff --git a/tests/test_lmcache_mp.py b/tests/test_lmcache_mp.py index 62d50d0ff3..223ec66bb8 100644 --- a/tests/test_lmcache_mp.py +++ b/tests/test_lmcache_mp.py @@ -287,11 +287,10 @@ def test_auto_rank_collapse_is_off_for_a_draft_with_its_own_pool(monkeypatch): from atom.utils import selector owns_pool = {"value": True} - monkeypatch.setattr(selector, "attn_family", lambda _hf: "draft-family") monkeypatch.setattr( selector, - "get_attn_backend", - lambda _family: SimpleNamespace(DRAFT_OWNS_KV_POOL=owns_pool["value"]), + "attn_family", + lambda _hf: selector.Family.MHA if owns_pool["value"] else selector.Family.MLA, ) config = _config(tp=8, kv_lora_rank=512) assert deployment._tp_replication_factor(config) == 8 diff --git a/tests/test_lmcache_mp_native_worker.py b/tests/test_lmcache_mp_native_worker.py index 9539e33730..772b015b38 100644 --- a/tests/test_lmcache_mp_native_worker.py +++ b/tests/test_lmcache_mp_native_worker.py @@ -94,6 +94,7 @@ def config(**extra): kv_cache_block_size=4, tensor_parallel_size=2, hf_config=SimpleNamespace(model_type="reusable_test_model", kv_lora_rank=512), + speculative_config=None, kv_transfer_config={ "kv_connector": "lmcache_mp", "kv_role": "offload", diff --git a/tests/test_v4_indexer_defaults.py b/tests/test_v4_indexer_defaults.py index 0ba6d1497f..fe6a846a77 100644 --- a/tests/test_v4_indexer_defaults.py +++ b/tests/test_v4_indexer_defaults.py @@ -12,7 +12,7 @@ from transformers import PretrainedConfig import atom.config as config_module -from atom.config import CompilationConfig, Config +from atom.config import CompilationConfig, Config, SpeculativeConfig _V4_KERNELS_INIT = ( Path(__file__).resolve().parents[1] / "atom/model_ops/v4_kernels/__init__.py" @@ -117,17 +117,55 @@ def _make_config( probe = _GfxProbe(gfx) _install_chip_info_stub(monkeypatch, probe) + # a V4 asking for fp8 KV is probed for its arch (the fp8 -> bf16 fallback), which + # the tests of the indexer default's own probing leave out of their counts + kv = "bf16" if architectures is None or "Llama" not in architectures[0] else "fp8" config = Config( model="test-model", - kv_cache_dtype="fp8", torch_profiler_dir=None, compilation_config=CompilationConfig(level=0, use_cudagraph=False), - **overrides, + **{"kv_cache_dtype": kv, **overrides}, ) config.quant_config = None return config, probe +@pytest.mark.parametrize( + "gfx, kv", [("gfx950", "fp8"), ("gfx1250", "fp8"), ("gfx942", "bf16")] +) +def test_v4_fp8_kv_is_settled_by_config_with_the_arch_aiter_serves_it_on( + monkeypatch, gfx, kv +): + """aiter's V4 fp8 (2buff) KV runs on gfx950 / gfx1250 only; elsewhere `Config` + makes the KV bf16 itself, so a speculative draft resolved after it (it shares the + target's pool) is not left asking for fp8 against a bf16 pool.""" + monkeypatch.setattr(SpeculativeConfig, "_resolve_synthetic_acceptance", lambda s: 0) + spec = SpeculativeConfig( + method="dspark", + model="draft", + num_speculative_tokens=3, + draft_model_hf_config=PretrainedConfig(architectures=["DeepseekV4DSparkModel"]), + ) + config, _ = _make_config( + monkeypatch, gfx=gfx, kv_cache_dtype="fp8", speculative_config=spec + ) + + assert config.kv_cache_dtype == kv + assert spec.draft_kv_dtype == kv + + +def test_v4_fp8_fallback_is_not_applied_under_a_plugin(monkeypatch): + config, probe = _make_config( + monkeypatch, + gfx="gfx942", + kv_cache_dtype="fp8", + plugin_config=SimpleNamespace(is_vllm=False, vllm_config=None), + ) + + assert config.kv_cache_dtype == "fp8" + assert probe.calls == 0 + + @pytest.mark.parametrize("gfx", ["gfx950", "gfx1250", "future_gfx"]) def test_native_v4_defaults_indexer_to_fp4(monkeypatch, gfx): config, probe = _make_config(monkeypatch, gfx=gfx) From 9d7ad4c442ebf687f6a363795fa68d915135ab3a Mon Sep 17 00:00:00 2001 From: ganyi Date: Thu, 12 Nov 2026 04:53:59 -0600 Subject: [PATCH 2/3] Pin unified attention off in the pool-contract tests; reword three stale comments The default-backend checks in test_kv_pool_contract.py named AiterBackend, which ATOM_USE_UNIFIED_ATTN=1 replaces with TritonMHABackend (same factory, inherited): an autouse fixture now sets it to 0. The comments in deployment.py, eagle_proposer.py and draft_kv.py said a backend flag / two attributes where the family table and three attributes are meant. No behaviour change. Co-Authored-By: Claude --- atom/kv_transfer/offload/mp/deployment.py | 7 ++++--- atom/spec_decode/draft_kv.py | 6 +++--- atom/spec_decode/eagle_proposer.py | 4 ++-- tests/test_kv_pool_contract.py | 8 +++++++- 4 files changed, 16 insertions(+), 9 deletions(-) diff --git a/atom/kv_transfer/offload/mp/deployment.py b/atom/kv_transfer/offload/mp/deployment.py index c96763e352..be63f95b2a 100644 --- a/atom/kv_transfer/offload/mp/deployment.py +++ b/atom/kv_transfer/offload/mp/deployment.py @@ -170,9 +170,10 @@ def _config_has_fully_replicated_tp_pages(config: Any) -> bool: from atom.spec_decode.draft_kv import draft_owns_kv_pool if draft_owns_kv_pool(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. + # 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 diff --git a/atom/spec_decode/draft_kv.py b/atom/spec_decode/draft_kv.py index 33e81f82ce..d81520ceb6 100644 --- a/atom/spec_decode/draft_kv.py +++ b/atom/spec_decode/draft_kv.py @@ -51,9 +51,9 @@ def draft_kv_builder(model_runner, draft_hf): 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 (`selector.DRAFT_POOL_FAMILIES`). 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 table entry and a `make_kv_pool` on that backend. diff --git a/atom/spec_decode/eagle_proposer.py b/atom/spec_decode/eagle_proposer.py index 22012ceefe..0b686b091a 100644 --- a/atom/spec_decode/eagle_proposer.py +++ b/atom/spec_decode/eagle_proposer.py @@ -237,8 +237,8 @@ def _build_draft_model(self, model_class) -> nn.Module: # Shallow-copy instead of deepcopy: with MLA targets (K2.6), the # atom_config holds non-picklable cuda.Stream objects under # downstream fields that deepcopy can't traverse. We only mutate - # hf_config and compilation_config.level on the draft, so - # isolating just those two attrs is sufficient. + # hf_config, compilation_config.level and kv_cache_dtype on the + # draft, so isolating just those attrs is sufficient. draft_atom_config = copy.copy(self.config) draft_atom_config.hf_config = draft_model_hf_config draft_atom_config.compilation_config = copy.copy( diff --git a/tests/test_kv_pool_contract.py b/tests/test_kv_pool_contract.py index 6bcda9dc9e..f48db81066 100644 --- a/tests/test_kv_pool_contract.py +++ b/tests/test_kv_pool_contract.py @@ -191,6 +191,13 @@ def _classes_defining_the_factory() -> list[tuple[str, ast.ClassDef]]: ] +@pytest.fixture(autouse=True) +def _default_backends(monkeypatch): + # With ATOM_USE_UNIFIED_ATTN=1 the MHA family selects TritonMHABackend, + # which inherits the factory; these checks name the default backends. + monkeypatch.setenv("ATOM_USE_UNIFIED_ATTN", "0") + + def _backend_class_name(family) -> str: from atom.utils.selector import get_attn_backend_cls @@ -226,7 +233,6 @@ def test_a_backend_that_can_build_a_draft_pool_is_in_the_table(): if cls.name != FACTORY_DECLARATION } listed = {_backend_class_name(f) for f in DRAFT_POOL_FAMILIES} - # `get_attn_backend_cls` is asked per family with no plugin and no unified attn assert builders == listed, (builders, listed, list(Family)) From a3802ba4993cf747d20dead3be2991e6222cf2f4 Mon Sep 17 00:00:00 2001 From: ganyi Date: Thu, 12 Nov 2026 06:23:49 -0600 Subject: [PATCH 3/3] Refuse an unsupported speculative kv_cache_dtype in the resolver argparse restricts the flag to bf16/fp8, a Python caller of SpeculativeConfig could pass anything and would fail later at the dtype lookup. Co-Authored-By: Claude --- atom/config.py | 5 +++++ tests/test_draft_kv_dtype.py | 6 ++++++ 2 files changed, 11 insertions(+) diff --git a/atom/config.py b/atom/config.py index 2abd760582..9a02d7a0f2 100644 --- a/atom/config.py +++ b/atom/config.py @@ -1955,6 +1955,11 @@ def resolve_draft_kv_dtype(spec, target_dtype: str, *, own_pool: bool) -> str: ``--spec-decode-kv-cache-dtype``. A draft in the target's pool cannot ask for another dtype.""" asked = spec.kv_cache_dtype + # argparse restricts the flag; a Python caller can pass anything + if asked not in (None, "bf16", "fp8"): + raise ValueError( + f"speculative kv_cache_dtype must be 'bf16' or 'fp8', got {asked!r}" + ) if not own_pool: if asked is not None and asked != target_dtype: raise ValueError( diff --git a/tests/test_draft_kv_dtype.py b/tests/test_draft_kv_dtype.py index ba90073c7b..dd540e6a24 100644 --- a/tests/test_draft_kv_dtype.py +++ b/tests/test_draft_kv_dtype.py @@ -71,6 +71,12 @@ def test_an_own_pool_draft_keeps_the_targets_dtype_unless_asked(target, asked): ) +@pytest.mark.parametrize("own_pool", [True, False]) +def test_an_unsupported_draft_dtype_is_refused_whoever_asks(own_pool): + with pytest.raises(ValueError, match="must be 'bf16' or 'fp8'"): + resolve_draft_kv_dtype(_draft(asked="int4"), "bf16", own_pool=own_pool) + + def test_a_shared_pool_draft_has_the_targets_dtype(): assert resolve_draft_kv_dtype(_draft("mtp"), "fp8", own_pool=False) == "fp8" assert resolve_draft_kv_dtype(_draft("mtp", "bf16"), "bf16", own_pool=False) == (