diff --git a/.github/benchmark/models_accuracy.json b/.github/benchmark/models_accuracy.json index 43a57cc45f..1b4e356b32 100644 --- a/.github/benchmark/models_accuracy.json +++ b/.github/benchmark/models_accuracy.json @@ -247,6 +247,19 @@ "accuracy_baseline_model": "moonshotai/Kimi-K3 + Inferact/Kimi-K3-DSpark", "_baseline_note": "Kimi-K3 DSpark block spec-decode (7 tokens) on the Kimi-K3 target with the Inferact/Kimi-K3-DSpark draft. Spec-decode is lossless w.r.t. the target, so baseline reuses the Kimi-K3 FP8 base (0.95); threshold 0.94 matches the target. mtp_accept_threshold intentionally omitted until the first CI run reports the DSpark acceptance rate -- add it once measured." }, + { + "model_name": "Kimi-K3 DSpark PP2", + "model_path": "moonshotai/Kimi-K3", + "extraArgs": "--kv_cache_dtype fp8 -tp 4 --pipeline-parallel-size 2 --enforce-eager --trust-remote-code --gpu-memory-utilization 0.9 --block-size 128 --enable_prefix_caching --method dspark --draft-model Inferact/Kimi-K3-DSpark --num-speculative-tokens 7 --max-num-seqs 64 --online_quant_config '{\"global_quant_config\": \"ptpc_fp8\", \"exclude_layer\": [\"lm_head\", \"model.embed_tokens\", \"*self_attn.[qkv]_conv1d*\", \"*block_sparse_moe.experts*\", \"*block_sparse_moe.routed_expert_*\", \"*vision_tower*\", \"*mm_projector*\", \"context_proj\", \"layers\"]}'", + "env_vars": "", + "runner": "linux-atom-do-mi350x-8", + "test_level": "pr", + "accuracy_threshold": 0.94, + "accuracy_timeout_minutes": 60, + "accuracy_baseline": 0.95, + "accuracy_baseline_model": "moonshotai/Kimi-K3 + Inferact/Kimi-K3-DSpark", + "_baseline_note": "The only pipeline-parallel accuracy case: Kimi-K3 pp2 x tp4 with DSpark (7 tokens) and prefix caching on, which together exercise the PP-only paths -- the undeferred verify window, the cross-stage aux relay, the checkpoint hold and a KDA state fork landing on a verify step. Eager because PP decode has no CUDA graphs yet, hence the longer timeout. ptpc_fp8 online quant on the target only, as the Kimi-K3 TP8 case runs it (that case has no drafter); \"context_proj\" and \"layers\" (the draft's unprefixed module names, which no target module matches) keep the DSpark draft in BF16 as the Kimi-K3 DSpark TP8 case runs it. Spec decode is lossless, so baseline and threshold match Kimi-K3." + }, { "model_name": "Kimi-K3.1-Pro", "model_path": "moonshotai/Kimi-K3.1-Pro", diff --git a/atom/config.py b/atom/config.py index 127bb9539e..239dc79c7e 100644 --- a/atom/config.py +++ b/atom/config.py @@ -2624,6 +2624,7 @@ def __post_init__(self): # builds the model. self.offload_kv_encoding = snapshot_kv_encoding(self.hf_config) self._resolve_dynamic_chunking() + self._check_pp_support() if self.speculative_config is not None: num_spec = self.speculative_config.num_speculative_tokens @@ -2835,6 +2836,52 @@ def __post_init__(self): self.state_replay_tokens = state_replay_tokens(self) + def _check_pp_support(self) -> None: + """Refuse what PP does not wire yet: TBO, P/D consumers, and speculation + an architecture does not declare in ``pp_speculative_methods`` (a P/D + producer does not speculate); anything else fails mid-serving or drafts + from zero buffers. + """ + if self.pipeline_parallel_size <= 1: + return + # UBatchWrapper.forward takes no intermediate_tensors. + if self.enable_tbo: + raise ValueError( + "TBO (enable_tbo) is not supported with pipeline_parallel_size > 1." + ) + from atom.kv_transfer.disaggregation.pd_producer import ( + disaggregated_kv_role, + ) + + # A consumer's first decode leaves T0 plus mtp_k placeholders, not the + # mtp_k + 1 the PP verify window assumes, and the PP loop has no + # remote-KV path at all. + role = disaggregated_kv_role(self) + if role == "kv_consumer": + raise ValueError( + "A P/D kv_consumer is not supported with pipeline_parallel_size > 1." + ) + # A producer never speculates; it loads the drafter only to write the + # draft layer's KV. + if self.speculative_config is None or role == "kv_producer": + return + from atom.models.registry import pp_speculative_methods + + archs = getattr(self.hf_config, "architectures", None) or () + spec = self.speculative_config + method = "dspark" if spec.use_dspark() else spec.method + if not any(method in pp_speculative_methods.get(a, ()) for a in archs): + raise ValueError( + "pipeline_parallel_size > 1 with speculative decoding supports " + "only the methods registry.pp_speculative_methods declares, got " + f"architectures={list(archs)}, method={method!r}." + ) + if self.dspark.confidence_schedule: + raise ValueError( + "The DSpark confidence scheduler (confidence_schedule) is not " + "supported with pipeline_parallel_size > 1." + ) + def _resolve_dynamic_chunking(self) -> None: """Validate dynamic chunking, or turn it off where it cannot size a chunk. diff --git a/atom/distributed/pp_comm.py b/atom/distributed/pp_comm.py index 1e4f719686..642e662cb2 100644 --- a/atom/distributed/pp_comm.py +++ b/atom/distributed/pp_comm.py @@ -22,7 +22,18 @@ # Keys carried between pipeline stages. `sparse_kv_indices` is optional: present # only when a PP boundary splits a DSA IndexShare group (see model_runner). -_PP_PROXY_KEYS = ("hidden_states", "residual", "sparse_kv_indices") +# `block_residual` is Kimi-K3's AttnRes block window [T, k, H]. `PP_AUX_KEY` +# relays a drafter's target aux hidden states (`PPAuxRelay`, set/popped by +# model_runner). Keys missing here are silently dropped by the send. +PP_AUX_KEY = "pp_aux" + +_PP_PROXY_KEYS = ( + "hidden_states", + "residual", + "block_residual", + "sparse_kv_indices", + PP_AUX_KEY, +) def pp_send_allgather_group(): @@ -125,15 +136,21 @@ def _async_send_object( def async_send_intermediate_tensors( - it: IntermediateTensors, + it: IntermediateTensors, owned: tuple[str, ...] = () ) -> list[P2PWork]: - """Non-blocking send of hidden_states/residual to the next PP stage.""" + """Non-blocking send of hidden_states/residual to the next PP stage. + + ``owned`` keys name tensors the caller hands over and never touches + again: they are sent as they are. The rest are cloned, since the model + may overwrite them (a graph's output buffer, say) before the send lands. + """ pp = get_pp_group() dst_local = (pp.rank_in_group + 1) % pp.world_size dst_global = pp.ranks[dst_local] tensors = {k: it.tensors[k] for k in _PP_PROXY_KEYS if k in it.tensors} metadata_list, tensor_list = _split_tensor_dict(tensors) + owned_ids = {id(tensors[k]) for k in owned if k in tensors} works = _async_send_object(metadata_list, dst_global, pp.cpu_group) @@ -144,9 +161,10 @@ def async_send_intermediate_tensors( for tensor in tensor_list: if tensor.numel() == 0: continue + is_owned = id(tensor) in owned_ids if ag_group is not None and tensor.numel() % ag_size == 0: tensor = tensor.reshape(ag_size, -1)[ag_rank] - buf = tensor.clone() + buf = tensor if is_owned else tensor.clone() if buf.device.type == "cpu": w = torch.distributed.isend(buf, dst=dst_global, group=pp.cpu_group) else: diff --git a/atom/model_engine/model_runner.py b/atom/model_engine/model_runner.py index 3ee2f6fdb7..21cdaa20f8 100644 --- a/atom/model_engine/model_runner.py +++ b/atom/model_engine/model_runner.py @@ -38,6 +38,7 @@ pcp_round_robin_split, ) from atom.distributed.pp_comm import ( + PP_AUX_KEY, async_send_intermediate_tensors, commit_pp_send_work, recv_intermediate_tensors, @@ -54,7 +55,11 @@ from atom.model_engine.output_transfer import D2HQueue from atom.model_engine.page_unit_checkpoint import PagedStateCheckpointSpec from atom.model_engine.run_labels import build_run_label -from atom.model_engine.scheduler import ScheduledBatch, ScheduledBatchOutput +from atom.model_engine.scheduler import ( + ScheduledBatch, + ScheduledBatchOutput, + defers_output, +) from atom.model_engine.sequence import ( Sequence, SequenceStatus, @@ -100,6 +105,7 @@ from atom.sampling_params import RequiredChannel from atom.spec_decode.drafter import Drafter from atom.spec_decode.factory import build_drafter +from atom.spec_decode.pp_aux_relay import PPAuxRelay from atom.spec_decode.status import SpecStatusQueue from atom.utils import ( CpuGpuBuffer, @@ -200,17 +206,7 @@ def __init__( # The scheduler never speculates for a P/D producer, though the runner # keeps its drafter loaded: every decode it runs is q=1. self.is_pd_producer = _kv_config_has_producer(kv_cfg) - # P/D hands off prompt-end state plus the first sampled token. Deferred - # output makes the producer decode once more to surface that token, so - # the consumer applies T0 to a state that already has it. Only a - # recurrent state notices -- a paged write lands at T0's own position, - # so repeating it is idempotent. - hands_off_recurrent_state = ( - self.is_pd_producer and runner.attn_family.has_recurrent_state - ) - self.is_deferred_out = ( - not self.is_pipeline_parallel and not hands_off_recurrent_state - ) + self.is_deferred_out = defers_output(runner.config, runner.attn_family) self.runner = runner device = runner.device @@ -236,6 +232,8 @@ def __init__( self.batch_info = CpuGpuBuffer( 3, dtype=torch.int32, device=device, publication_group="input_ids" ) + # On every PP stage: each runs the verify windows (draft-id staging, + # state rollback), though only the last one owns the drafter. self.use_spec = use_spec self.num_spec_tokens = num_spec_tokens # The drafter's `draft_tokens_per_seq`, set once it is loaded. @@ -345,6 +343,20 @@ def clean(self): self.num_rejected: np.ndarray | None = None self.num_bonus: np.ndarray | None = None + def verify_ctx_lens(self, context_lens: np.ndarray) -> np.ndarray: + """`batch.context_lens` moved back to the start of this verify window. + + Deferred output: the last step's `num_rejected` slots are still in + `context_lens`. Undeferred (PP): postprocess already left one run of + `mtp_k + 1` placeholders, so the shift is 1. Before any sampler output + there is nothing to shift. Only a verify step has a window, so every + caller sits behind `max_seqlen_q > 1`. + """ + if self.num_rejected is None: + return context_lens + shift = self.num_rejected if self.is_deferred_out else 1 + return context_lens - np.asarray(shift, dtype=context_lens.dtype) + @staticmethod def _batch_process_token_ids(token_ids: list) -> list[tuple[int, ...]]: """Batch process token_ids: vectorized -1 truncation using numpy.""" @@ -563,25 +575,40 @@ def _stage_input_ids( token_ids = scheduled_tokens[ total_tokens_prefill : total_tokens_prefill + total_tokens_decode ] - # No spec path supports pipeline parallel yet; it wants the deferred - # branch's per-request staging. Gate on PP itself rather than on - # `use_spec`: a P/D producer also reaches this branch, and the flag - # is still set there (the runner keeps the drafter loaded to write - # the draft's context KV at prefill) even though the scheduler - # turned speculation off for it. - if self.use_spec and self.is_pipeline_parallel: - raise NotImplementedError("pipeline parallel + speculative decode") - # PP verifies nothing locally and a P/D producer never speculates - # (its DP dummy included), so no row here carries drafts. One that - # did would need the deferred branch's per-request staging: under - # DP a block drafter's `max_seqlen_q` is the group's maximum, not - # this batch's row width. - assert batch.num_spec_step == 0, ( - "non-deferred decode got a speculative batch " - f"(num_spec_step={batch.num_spec_step})" - ) - self.input_ids.np[:total_tokens_decode] = token_ids + # A P/D producer never speculates (its DP dummy included), so no + # row here carries drafts. One that did would need the deferred + # branch's per-request staging: under DP a block drafter's + # `max_seqlen_q` is the group's maximum, not this batch's row + # width. Under PP only a producer is held to it; a PP consumer + # verifies its own drafts below. + if not self.is_pipeline_parallel or self.is_pd_producer: + assert batch.num_spec_step == 0, ( + "non-deferred decode got a speculative batch " + f"(num_spec_step={batch.num_spec_step})" + ) + if not self.is_pipeline_parallel: + return ( + self._publish_input_ids(total_tokens_decode, publication_group), + None, + ) + if self.use_spec and batch.num_spec_step > 0: + # PP: the host already holds each row's anchor, so fill only + # the draft columns. + bs, lens, cu_np = self.runner.attn_metadata_builder.decode_spans(batch) + spec = batch.scheduled_spec_decode_tokens[:bs] + cols = np.arange(spec.shape[1]) + # Row i's drafts land right after its anchor at cu_np[i]. + drafted = cols < (lens[:, None] - 1) + self.input_ids.np[(cu_np[:bs, None] + 1 + cols)[drafted]] = spec[ + drafted + ] + # Previous step's acceptance counts, read off the batch so every + # stage (not just the sampling one) rolls its linear-attention + # state back correctly; unset, GDN assumes 1 accepted token. + if self.use_spec: + self.num_rejected = batch.num_rejected + self.num_bonus = batch.num_bonus return ( self._publish_input_ids(total_tokens_decode, publication_group), None, @@ -811,7 +838,9 @@ def __init__(self, rank: int, config: Config): default_dtype = self.config.torch_dtype torch.set_default_dtype(default_dtype) torch.set_default_device(self.device) - use_spec = bool(self.config.speculative_config) and get_pp_group().is_last_rank + # Every PP stage runs the verify window, drafter or not: its width + # sizes the KDA conv state and MLA work metadata. + use_spec = bool(self.config.speculative_config) self.num_spec_tokens = ( self.config.speculative_config.num_speculative_tokens if use_spec else 0 ) @@ -927,6 +956,7 @@ def __init__(self, rank: int, config: Config): logger.info("TBO enabled: model wrapped with UBatchWrapper") if getattr(self, "drafter", None) is not None: self.drafter.arm_aux_capture(self.model) + self.pp_aux_relay = PPAuxRelay.build(self.config, self.model) installer = mono_decode_installers.get(hf_config.architectures[0]) if installer is not None: self.model = resolve_obj_by_qualname(installer)( @@ -3191,7 +3221,7 @@ def run_model( if pp_enabled: self._setup_pp_shared_indexer() - intermediate_tensors = None + intermediate_tensors = recv_aux = None if pp_enabled and not pp_group.is_first_rank: intermediate_tensors = recv_intermediate_tensors() # GLM-5.2 IndexShare: load prior rank's top-k for leading @@ -3202,6 +3232,11 @@ def run_model( if recv_sparse is not None and self._pp_recv_needs_sparse: tgt = self.attn_metadata_builder._sparse_kv_indices_gpu tgt[: recv_sparse.numel()].copy_(recv_sparse) + # DSpark aux hidden states from target layers on earlier + # stages. Popped for the same reason as the top-k above. + recv_aux = intermediate_tensors.tensors.pop(PP_AUX_KEY, None) + if self.pp_aux_relay is not None and pp_group.is_last_rank: + self.pp_aux_relay.absorb(recv_aux, self.drafter) # Time only the chunk computation between PP receive and send. chunk_sample = self._dynamic_chunking.start_sample(batch) @@ -3228,10 +3263,18 @@ def run_model( model_output.tensors["sparse_kv_indices"] = ( self.attn_metadata_builder._sparse_kv_indices_gpu[:n] ) + if self.pp_aux_relay is not None: + aux = self.pp_aux_relay.pack( + recv_aux, + model_output.tensors["hidden_states"].shape[0], + ) + if aux is not None: + model_output.tensors[PP_AUX_KEY] = aux if self._pp_pending_send: commit_pp_send_work(self._pp_pending_send) + # The relay hands its payload over: no clone needed. self._pp_pending_send = async_send_intermediate_tensors( - model_output + model_output, owned=(PP_AUX_KEY,) ) hidden_states = None logits = None @@ -3566,8 +3609,16 @@ def postprocess( prev_rejected_num = np.zeros(0, dtype=np.int32) prev_bonus_num = np.zeros(0, dtype=np.int32) else: - prev_rejected_num = np.zeros(batch.total_seqs_num, dtype=np.int32) - prev_bonus_num = np.zeros(batch.total_seqs_num, dtype=np.int32) + # Undeferred: report THIS step's counts; the next batch carries them + # back for linear-attention state rollback. Zeros if no drafts scored. + if spec_decode_metadata is not None: + # From the host ids, no D2H: accepted drafts + 1, then -1s. + emitted = np.fromiter(map(len, token_ids_out), np.int32) + prev_bonus_num = emitted - 1 + prev_rejected_num = sampled_tokens.shape[-1] - emitted + else: + prev_rejected_num = np.zeros(batch.total_seqs_num, dtype=np.int32) + prev_bonus_num = np.zeros(batch.total_seqs_num, dtype=np.int32) # PP stages and P/D producers still run the drafter. if hasattr(self, "drafter"): # Mid-prompt sequences get their anchor corrected inside diff --git a/atom/model_engine/pp_engine_core.py b/atom/model_engine/pp_engine_core.py index 6e436bd94e..1b870e2647 100644 --- a/atom/model_engine/pp_engine_core.py +++ b/atom/model_engine/pp_engine_core.py @@ -167,8 +167,7 @@ def _pp_head_step(self): if not needs_output and not scheduled_batch.pp_completion_ack: self._in_flight.popleft() self.scheduler.release_pp_inflight(scheduled_batch) - if self._defer_prefix_hash: - self._pending_prefix_hash.append((scheduled_batch, seqs)) + self._queue_or_flush_prefill_hashes(scheduled_batch, seqs) continue completion = self.pp_transport.recv_completion(timeout_ms=poll_ms) if completion is None: @@ -185,8 +184,7 @@ def _pp_head_step(self): self._in_flight.popleft() self.scheduler.release_pp_inflight(scheduled_batch) if not needs_output: - if self._defer_prefix_hash: - self._pending_prefix_hash.append((scheduled_batch, seqs)) + self._queue_or_flush_prefill_hashes(scheduled_batch, seqs) continue assert fwd_out is not None, "final PP chunk completed without output" @@ -210,6 +208,18 @@ def _pp_head_step(self): if finished_seqs: self.output_queue.put_nowait(finished_seqs) + def _queue_or_flush_prefill_hashes(self, batch, seqs): + """Queue a no-output chunk's hashes; flush the whole queue if it carries + a held seq's chunk.""" + if not self._defer_prefix_hash: + return + self._pending_prefix_hash.append((batch, seqs)) + # A seq waiting on this chunk's state checkpoint is out of the prefill + # scan, so no later step would flush for it. Filing a checkpoint copies + # nothing; it only needs FIFO order, so flush now. + if self.scheduler.carries_held_checkpoint(batch): + self._flush_pending_prefix_hashes() + def _flush_pending_prefix_hashes(self): while self._pending_prefix_hash: batch, seqs = self._pending_prefix_hash.popleft() diff --git a/atom/model_engine/scheduler.py b/atom/model_engine/scheduler.py index e99d70a4d7..965abf268d 100644 --- a/atom/model_engine/scheduler.py +++ b/atom/model_engine/scheduler.py @@ -67,7 +67,8 @@ from atom.observability import tracing from atom.sampling_params import RequiredChannel from atom.structured_output import StructuredRequest -from atom.utils import envs +from atom.utils import envs, get_hf_text_config +from atom.utils.selector import attn_family if TYPE_CHECKING: from atom.model_engine.prefill_delayer import PrefillDelayer @@ -77,6 +78,28 @@ _DYNAMIC_CHUNKING_SUPPLY_WINDOW = 16 +def defers_output(config, family=None) -> bool: + """Whether each step's sampled ids reach the host a step late, read on + device by the next forward. The scheduler and the runner's + `tokenIDProcessor` both take it from here: the two modes leave different + trailing placeholder counts, so they must agree. + + Not under PP, where the ids cross stages on the host. Not for a P/D + producer with recurrent state either: it hands off prompt-end state plus + the first sampled token, and deferring makes it decode once more to surface + that token, so the consumer applies T0 to a state that already has it. A + paged write lands at T0's own position, so only a recurrent state notices. + `family` is the config's `attn_family`, resolved here when not given. + """ + if getattr(config, "pipeline_parallel_size", 1) > 1: + return False + if not kv_config_has_producer(getattr(config, "kv_transfer_config", {}) or {}): + return True + if family is None: + family = attn_family(get_hf_text_config(config.hf_config)) + return not family.has_recurrent_state + + class ChunkKind(enum.Enum): """Which prefill chunk a sizing call is for, as dynamic chunking sees it. @@ -208,6 +231,7 @@ def __init__( is_final_chunk: list[bool] | None = None, next_token_ids: list[int] | None = None, state_maintenance_ops: StateMaintenanceOps | None = None, + is_deferred_out: bool = True, ): if scheduled_spec_decode_tokens is None: scheduled_spec_decode_tokens = {} @@ -378,8 +402,12 @@ def __init__( for i, (seq, num) in enumerate(zip(seqs.values(), num_scheduled_tokens)): if seq.type == SequenceType.PREFILL: offset = self.num_cached_tokens[i] - else: + elif is_deferred_out or num_spec_step == 0: offset = seq.num_tokens - num_rejected[i] - num + else: + # Undeferred: the tail is the one placeholder run postprocess + # recorded (rejects already refilled); the anchor is before it. + offset = seq.num_tokens - seq.num_placeholder_tokens - 1 staged.extend(seq.token_ids[offset : offset + num]) # Checked here because nothing downstream will: the array below wraps # whatever length was staged, and a sequence too short to fill its @@ -766,20 +794,18 @@ def __init__( self.drafter_needs_next_token = has_spec and not ( config.speculative_config.use_dspark() ) - # True when this engine both drafts and verifies; False under PP - # (which only drafts for handoff to the decode node). pp_size = getattr(config, "pipeline_parallel_size", 1) - self.spec_decode_local = self.use_spec and pp_size == 1 # Positions past `len(seq)` each forward's draft round writes KV to. # Allocated before the forward, with the step's own blocks: see # `draft_kv_lookahead` for the arithmetic and what an unallocated one # costs. self.draft_lookahead = draft_kv_lookahead(self.mtp_k) self.block_manager.draft_lookahead = self.draft_lookahead - if self.use_spec and not self.spec_decode_local: + self.is_deferred_out = defers_output(config) + if self.use_spec and pp_size > 1: logger.info( - "Speculative decoding: drafting only (pipeline_parallel_size=%d). " - "Drafts are produced for handoff; this engine verifies none.", + "Speculative decoding: verifying locally under pipeline " + "parallelism (pipeline_parallel_size=%d).", pp_size, ) # `engine_stats` (spec / cache / throughput sections) is constructed @@ -845,6 +871,13 @@ def __init__( # them until the head releases the id after postprocess, so a seq is # never decoded against a token not yet appended. self._pp_inflight_token_block: set[int] = set() + # Seqs whose last chunk ended on a state-checkpoint rung not yet filed, + # by that chunk's end; prefill skips them. Under PP the next chunk would + # otherwise overwrite the state slot before it is filed under the rung's + # hash. Keyed on advance_on_schedule, which only PP sets. + # shortcut: holds prefill until the rung's chunk completes; #2709 files + # the checkpoint at schedule time and drops the hold for most requests. + self._checkpoint_holds: dict[int, int] = {} from atom.utils.forward_context import get_kvconnector @@ -1038,6 +1071,9 @@ def _probe_local_prefill_work(self) -> tuple[bool, int]: for seq in self.running: if not seq.is_partial_prefill: continue + # Phase 1 skips a held seq, so its chunk is not pending work. + if seq.id in self._checkpoint_holds: + continue resume = seq.num_cached_tokens start = self.block_manager.replay_start(seq, resume, hit=False) chunk = self._partial_prefill_chunk( @@ -1495,7 +1531,7 @@ def _next_decode_completion_steps(self) -> int | None: """Decode steps until the first running decode reaches max_tokens if every step accepts all its draft tokens; None without decode. EOS can come sooner, rejected drafts make it later. Only an engine that - verifies drafts (`spec_decode_local`) gains more than a token a step.""" + verifies drafts (`use_spec`) gains more than a token a step.""" fewest = None for seq in self.running: if seq.is_partial_prefill: @@ -1505,7 +1541,7 @@ def _next_decode_completion_steps(self) -> int | None: fewest = left if fewest is None: return None - tokens_per_step = 1 + (self.mtp_k if self.spec_decode_local else 0) + tokens_per_step = 1 + (self.mtp_k if self.use_spec else 0) return -(-fewest // tokens_per_step) def _oldest_waiting_prefill_age_ms(self) -> float: @@ -2157,6 +2193,8 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: break if not seq.is_partial_prefill: continue + if seq.id in self._checkpoint_holds: + continue resume = seq.num_cached_tokens start = self.block_manager.replay_start(seq, resume, hit=False) chunk = self._partial_prefill_chunk( @@ -2796,7 +2834,7 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: num_seqs_decode = 0 num_decode_tokens = 0 # anchor + drafts if verifying locally, anchor alone otherwise. - spec_width = self.mtp_k if self.spec_decode_local else 0 + spec_width = self.mtp_k if self.use_spec else 0 tokens_per_decode_seq = spec_width + 1 num_new_tokens = spec_width + 1 remote_kv_blocks: set[int] = set() @@ -2833,7 +2871,7 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: if blocked_by_pinned_save: break if not preempted_current: - if self.spec_decode_local and seq.spec_token_ids.size > 0: + if self.use_spec and seq.spec_token_ids.size > 0: scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids num_seqs_decode += 1 num_decode_tokens += num_new_tokens @@ -2915,7 +2953,8 @@ def _schedule(self) -> tuple[ScheduledBatch, dict[int, Sequence]] | None: total_seqs_num_prefill=num_seqs_prefill, total_seqs_num_decode=num_seqs_decode, connector_meta_output=connector_meta_output, - num_spec_step=self.mtp_k if self.spec_decode_local else 0, + num_spec_step=self.mtp_k if self.use_spec else 0, + is_deferred_out=self.is_deferred_out, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, remote_kv_block_ids=sorted(remote_kv_blocks) if remote_kv_blocks else [], remote_kv_seq_blocks=remote_kv_seq_blocks, @@ -3596,6 +3635,7 @@ def _finalize_prefill_chunk(self, seq: Sequence, start: int, chunk: int) -> int: # chunk's start on the block grid instead, or every snapshot # position in it is off the grid too. chunk = bm.align_midstep_chunk(seq, start, chunk) + on_rung = bool(target) # `cancel_state_fork` runs only when the first two hold, and returning # False is what leaves the fork in place — hence holding the chunk open. if ( @@ -3604,6 +3644,12 @@ def _finalize_prefill_chunk(self, seq: Sequence, start: int, chunk: int) -> int: and not bm.cancel_state_fork(seq) ): chunk = bm.state.min_fork_tokens + on_rung = False + # With several forwards in flight, hold the next chunk until this + # checkpoint is filed. PP refuses PAGE checkpoints, so a rung here is + # strictly inside the prompt and there always is a next chunk. + if on_rung and self.advance_on_schedule and bm.enable_prefix_caching: + self._checkpoint_holds[seq.id] = start + chunk return chunk def _checkpoint_room(self, seq: Sequence, finished: bool) -> int: @@ -3616,13 +3662,10 @@ def _checkpoint_room(self, seq: Sequence, finished: bool) -> int: batch to issue a copy on either; - the seq is still on its prompt, where the prefill call site has already decided using the prompt's own remainder; - - speculative decode, for a class that checkpoints by forking. Two - reasons, and either is enough. The spec path's state index has no - read-side counterpart (see - `GDNAttentionMetadataBuilder.prepare_state_indices`), so a fork must - never reach it. And a spec step commits `1 + accepted_drafts` tokens, - which is what a fork's successor actually gets — the rest is rolled - back and re-forwarded — so no promise made here can be kept. + - speculative decode, for a class that checkpoints by forking. A spec + step commits `1 + accepted_drafts` tokens, which is what a fork's + successor actually gets — the rest is rolled back and re-forwarded — + so no promise made here can be kept. Checkpointing during *prefill* stays safe on the same models: `min_fork_tokens` guarantees prompt is left over, and prompt always forwards down the non-spec path. @@ -3975,8 +4018,8 @@ def preempt(self, seq: Sequence) -> bool: # coherent answer that ends before it answers anything. # # `num_placeholder_tokens` is the only width that describes what is - # actually there: postprocess appends `mtp_k + is_deferred_out - - # num_rejected` slots and records exactly that count. Re-deriving it + # actually there: postprocess records exactly what it appended, which + # differs between output modes. Re-deriving it # here as `mtp_k + num_rejected` agrees only when deferred output is # off AND nothing was rejected -- on a TP-only MTP engine it leaves one # `eos_token_id` behind (the case above), and with `num_rejected = r` @@ -4081,6 +4124,10 @@ def register_prefill_hashes( self, batch: ScheduledBatch, seqs: Iterable[Sequence] | None = None ) -> None: """Hash blocks for middle chunked-prefill chunks that skip postprocess.""" + # Before any bail-out or hashing that can raise, so no seq of this + # batch (preempted, or after a failing one) is held forever. + for req_id in self._checkpoint_holds_owned_by(batch): + del self._checkpoint_holds[req_id] if not self.block_manager.enable_prefix_caching: return if batch.is_final_chunk is None: @@ -4104,6 +4151,27 @@ def register_prefill_hashes( replay_until=_scheduled_replay_until(batch, i), ) + def carries_held_checkpoint(self, batch: ScheduledBatch) -> bool: + """Whether `batch` carries the chunk a held seq waits on. + + If so, its prefix hashes must be filed now rather than deferred: a held + seq is out of the prefill scan until then and would stall. + """ + return bool(self._checkpoint_holds_owned_by(batch)) + + def _checkpoint_holds_owned_by(self, batch: ScheduledBatch) -> list[int]: + """Held seqs whose held chunk is in `batch`; an earlier chunk of the + same seq, still in flight, must not release the hold.""" + if not self._checkpoint_holds: + return [] + return [ + req_id + for i, req_id in enumerate(batch.req_ids) + if req_id in self._checkpoint_holds + and int(batch.num_cached_tokens[i]) + int(batch.num_scheduled_tokens[i]) + >= self._checkpoint_holds[req_id] + ] + def postprocess( self, seqs: list[Sequence], @@ -4152,6 +4220,9 @@ def postprocess( # See _batch_seq_lookup: seq may have left running. seq_by_id = self._batch_seq_lookup(seqs) final = batch.is_final_chunk + # Before the loop; see register_prefill_hashes. + for req_id in self._checkpoint_holds_owned_by(batch): + del self._checkpoint_holds[req_id] for i, req_id in enumerate(batch.req_ids): seq = seq_by_id.get(req_id) if seq is None or final is None or not seq.block_table: @@ -4216,13 +4287,14 @@ def postprocess( stream_outputs = [] num_new_generation_tokens = 0 - need_placeholder = is_deferred_out or self.spec_decode_local + need_placeholder = is_deferred_out or self.use_spec # Drafts occupy trailing slots only on an engine that verifies them; a - # drafting-only engine's tokens are all real. - num_placeholder_width = self.mtp_k if self.spec_decode_local else 0 - num_placeholder = self.mtp_k - if is_deferred_out: - num_placeholder += 1 + # drafting-only engine's tokens are all real. Undeferred, `num_rejected` + # alone names the tail. + num_placeholder_width = self.mtp_k if (self.use_spec and is_deferred_out) else 0 + # Up to `mtp_k` accepted drafts plus one bonus (or, deferred without + # spec, the one-step lag slot). + num_placeholder = self.mtp_k + 1 for seq in self.running: error = getattr(fwd_output, "structured_errors", {}).get(seq.id) @@ -4354,9 +4426,7 @@ def postprocess( continue # In-place overwrite only when placeholders already exist. - if is_deferred_out or ( - self.spec_decode_local and seq.num_placeholder_tokens > 0 - ): + if is_deferred_out or (self.use_spec and seq.num_placeholder_tokens > 0): # int() casts strip the np.int32 wrapper coming out of # fwd_output's np.ndarray indexing. Without these, the values # propagate into seq.num_rejected / seq.num_bonus_tokens, then @@ -4364,7 +4434,13 @@ def postprocess( # contaminating downstream logs and arithmetic with np.int32. num_rejected = int(fwd_output.num_rejected[idx]) num_bonus = int(fwd_output.num_bonus[idx]) - offset = 0 if (num_new_token + num_rejected) == 1 else self.mtp_k + # Deferred: two runs outstanding, this step's starts `mtp_k` + # further back. Undeferred: only one. + offset = ( + self.mtp_k + if (is_deferred_out and (num_new_token + num_rejected) != 1) + else 0 + ) # Align stats with vLLM: only count steps that actually ran # speculation (drafts proposed and validated). Skip the # prefill-only step where no draft tokens were scored against @@ -4611,11 +4687,22 @@ def postprocess( if need_placeholder: # placeholder for the each decode step for seq in seqs: - if seq.status == SequenceStatus.RUNNING and not seq.is_partial_prefill: + # Frozen chunk state under schedule-time advancement, as above. + mid_prefill = ( + seq.id in pp_middle_chunk_ids + if self.advance_on_schedule + else seq.is_partial_prefill + ) + if seq.status == SequenceStatus.RUNNING and not mid_prefill: # Clamped because `preempt()` strips exactly this many: # `range()` on a negative width appends nothing, so storing # one unclamped would claim placeholders that are not there. - num = max(0, num_placeholder - seq.num_rejected) + if is_deferred_out: + num = max(0, num_placeholder - seq.num_rejected) + else: + # Refill the run to full width; deriving it from rejects + # overgrows it by `mtp_k` on a step that scored no drafts. + num = max(0, num_placeholder - seq.num_placeholder_tokens) for _ in range(num): seq.append_token(self.eos_token_id) # Added to what the overwrite above left unconsumed, not @@ -5191,7 +5278,6 @@ def __init__(self, config: Config, disagg_cu_shm_name: str = ""): self.running: deque[Sequence] = deque() # spec decode not used on prefill side self.use_spec = False - self.spec_decode_local = False self.mtp_k = 0 # Only the throughput section applies to the prefill side: it does not # speculate, and it has no BlockManager to source prefix-cache hits. @@ -5578,7 +5664,7 @@ def _schedule(self): self.preempt(seq) break else: - if self.spec_decode_local and seq.spec_token_ids.size > 0: + if self.use_spec and seq.spec_token_ids.size > 0: scheduled_spec_decode_tokens[seq.id] = seq.spec_token_ids num_new_tokens = self.mtp_k + 1 self.block_manager.may_append( @@ -5616,7 +5702,8 @@ def _schedule(self): total_tokens_num_decode=total_tokens_num_decode, total_seqs_num=len(scheduled_seqs), total_seqs_num_decode=len(scheduled_seqs), - num_spec_step=self.mtp_k if self.spec_decode_local else 0, + num_spec_step=self.mtp_k if self.use_spec else 0, + is_deferred_out=self.is_deferred_out, scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, cu_stream_fraction=self.cu_fraction, state_maintenance_ops=self.block_manager.take_state_maintenance_ops(), diff --git a/atom/model_ops/attention_gdn.py b/atom/model_ops/attention_gdn.py index 8f090699f6..47066619b3 100644 --- a/atom/model_ops/attention_gdn.py +++ b/atom/model_ops/attention_gdn.py @@ -168,6 +168,8 @@ def _spec_conv1d_update( ): """Update speculative convolution state using the native cache layout.""" spec_state_indices_tensor = gdn_metadata.spec_state_indices_tensor + spec_state_indices_in = gdn_metadata.spec_state_indices_in_tensor + nsd = gdn_metadata.num_spec_decodes return causal_conv1d_update( mixed_qkv_spec, conv_state, @@ -176,9 +178,12 @@ def _spec_conv1d_update( self.num_v_heads * self.head_v_dim // self.tp_size, self.conv1d.bias, self.activation, - conv_state_indices=spec_state_indices_tensor[:, 0][ - : gdn_metadata.num_spec_decodes - ], + conv_state_indices=spec_state_indices_tensor[:nsd, 0], + conv_state_indices_in=( + None + if spec_state_indices_in is None + else spec_state_indices_in[:nsd, 0] + ), num_accepted_tokens=num_accepted_tokens, query_start_loc=gdn_metadata.spec_query_start_loc, # The verify window, which sizes the conv rollback window @@ -428,6 +433,7 @@ def forward( is_kda=False, route=gdn_metadata.replayssm_route, o=(core_attn_out[: query_spec.shape[1]] if spec_only else None), + fork_slots=gdn_metadata.replayssm_fork_slots, ) last_recurrent_state = None else: @@ -616,6 +622,7 @@ def forward( use_qk_l2norm_in_kernel=True, is_kda=False, route=gdn_metadata.replayssm_route, + fork_slots=gdn_metadata.replayssm_fork_slots, ) last_recurrent_state = None elif gdn_metadata.num_decodes > 0: diff --git a/atom/model_ops/attention_mla.py b/atom/model_ops/attention_mla.py index bedfd3f11f..9f63654e24 100644 --- a/atom/model_ops/attention_mla.py +++ b/atom/model_ops/attention_mla.py @@ -602,7 +602,11 @@ def __init__( self.dtype = dtype self.min_query_heads = kwargs.get("min_query_heads", _MLA_MIN_HEADS) - self.padded_num_heads = max(num_heads, self.min_query_heads) + # aiter MLA only serves multiples of 16 heads; round up so a per-rank + # count like 24 (K3 at tp4) does not abort in get_mla_metadata. + self.padded_num_heads = mla_kernel_num_heads( + max(num_heads, self.min_query_heads) + ) # Heads past num_heads are dead lanes: the MLA kernels compute them and # `_restore_query_heads` throws them away. They are zeros rather than # repeats of the real heads because zeros are what a producer can write diff --git a/atom/model_ops/attentions/aiter_attention.py b/atom/model_ops/attentions/aiter_attention.py index b077c1ea49..c1d0858bff 100644 --- a/atom/model_ops/attentions/aiter_attention.py +++ b/atom/model_ops/attentions/aiter_attention.py @@ -1202,9 +1202,9 @@ def prepare_decode( block_tables = batch.block_tables if max_seqlen_q > 1: - num_rejected = self.model_runner.tokenID_processor.num_rejected - if num_rejected is not None: - context_lens -= num_rejected + context_lens = self.model_runner.tokenID_processor.verify_ctx_lens( + context_lens + ) positions = decode_positions(context_lens, max_seqlen_q) max_seqlen_k = np.max(context_lens) diff --git a/atom/model_ops/attentions/aiter_mla.py b/atom/model_ops/attentions/aiter_mla.py index 131cb019b0..cb52b354e4 100644 --- a/atom/model_ops/attentions/aiter_mla.py +++ b/atom/model_ops/attentions/aiter_mla.py @@ -50,6 +50,7 @@ MLAAttention, mla_dcp_kernel_num_heads, mla_dcp_sparse_prefill_num_heads, + mla_kernel_num_heads, ) from atom.model_ops.attentions.sparse_mla_head_groups import ( HEADS as HEAD_GROUP_HEADS, @@ -465,7 +466,11 @@ def __init__(self, model_runner): config = model_runner.config hf_config = config.hf_config # `self.num_attention_heads` set by CommonAttentionBuilder.__init__. - self.padded_num_attention_heads = max(self.num_attention_heads, _MLA_MIN_HEADS) + # Must equal `MLAAttention.padded_num_heads`: that pads the query, this + # sizes the kernel's work descriptors for it. + self.padded_num_attention_heads = mla_kernel_num_heads( + max(self.num_attention_heads, _MLA_MIN_HEADS) + ) self.is_sparse = model_runner.has_mla_indexer self.index_topk = hf_config.index_topk if self.is_sparse else -1 # GLM-5.3's pooled indexer selects `index_topk // index_kpool` POOLS -- @@ -2873,12 +2878,12 @@ def prepare_decode( context_lens = np.asarray(batch.context_lens, dtype=np.int32) block_tables = batch.block_tables if not batch.is_dummy_run and max_seqlen_q > 1: - # Get num_rejected (already mapped to current batch order in prepare_input_ids) - num_rejected = self.model_runner.tokenID_processor.num_rejected - if num_rejected is not None: - context_lens -= num_rejected - num_blocks = cdiv(context_lens, self.model_runner.block_size) - block_tables = [bt[:n] for bt, n in zip(block_tables, num_blocks)] + # Already in current batch order (prepare_input_ids). + context_lens = self.model_runner.tokenID_processor.verify_ctx_lens( + context_lens + ) + num_blocks = cdiv(context_lens, self.model_runner.block_size) + block_tables = [bt[:n] for bt, n in zip(block_tables, num_blocks)] positions = decode_positions(context_lens, max_seqlen_q) # Before the slots, not after: `slot_mapping` reads this packed table. @@ -2968,7 +2973,7 @@ def prepare_decode( ) for i, (ctx_len, is_first) in enumerate( zip( - batch.context_lens, + context_lens, batch.is_first_decode_without_local_prefill, ) ) diff --git a/atom/model_ops/attentions/deepseek_v41/backend.py b/atom/model_ops/attentions/deepseek_v41/backend.py index 093e88864f..76d7459764 100644 --- a/atom/model_ops/attentions/deepseek_v41/backend.py +++ b/atom/model_ops/attentions/deepseek_v41/backend.py @@ -1180,10 +1180,9 @@ def prepare_decode(self, batch, running_bs, running_tokens, max_seqlen_q): if self.geometry.speculative_tokens and not batch.is_dummy_run: # The scheduler reserves a full draft span, including placeholders # from the previous step. Ragged verification takes its head. - starts = np.asarray(batch.context_lens) - (batch.num_spec_step + 1) - rejected = self.model_runner.tokenID_processor.num_rejected - if rejected is not None: - starts = starts - rejected + starts = self.model_runner.tokenID_processor.verify_ctx_lens( + np.asarray(batch.context_lens) - (batch.num_spec_step + 1) + ) return self._prepare( batch, running_bs, diff --git a/atom/model_ops/attentions/deepseek_v4_attn.py b/atom/model_ops/attentions/deepseek_v4_attn.py index fd5d1139f6..5aa6cb9333 100644 --- a/atom/model_ops/attentions/deepseek_v4_attn.py +++ b/atom/model_ops/attentions/deepseek_v4_attn.py @@ -2487,22 +2487,18 @@ def prepare_decode( # on the batch (= num_spec_step+1 for plain MTP, or the DSpark q-bucket # when shrunk). positions/attn use this so the (running_bs, q) graph is # selected. See `ForwardMode.max_seqlen_q`. - # MTP: roll back ctx by `num_rejected` so this fwd's positions overwrite - # last fwd's rejected-draft slots (matches aiter_mla.py:701 / - # aiter_attention.py:542). `batch.context_lens` = `seq.num_tokens` - # which the scheduler advances by `mtp_k - num_rejected` placeholders - # per fwd (scheduler.py:789); without this rollback, MTP-k positions - # would skip ahead by `num_rejected` and the rejected slots would - # never be overwritten with the corrected K/V. `num_rejected` is None - # on dummy runs and on the first fwd before any sampler output. + # MTP: `seq.num_tokens` runs past this fwd's window by + # `verify_ctx_lens()` (`num_rejected` deferred, 1 undeferred/PP); + # roll ctx back so positions overwrite last fwd's rejected-draft slots. + # None on dummy runs and before the first sampler output. # The rolled-back ctx is also what anchors `positions` (at `ctx - # full_q`, below), and every compress count is now `visible_*(pos)`, so # a rejected slot's KV falls out of range on its own — `block_tables` # needs no truncation here. if not batch.is_dummy_run and max_seqlen_q > 1: - num_rejected = self.model_runner.tokenID_processor.num_rejected - if num_rejected is not None: - context_lens_np = context_lens_np - num_rejected.astype(np.int32) + context_lens_np = self.model_runner.tokenID_processor.verify_ctx_lens( + context_lens_np + ) # DSpark q-shrink: anchor the forwarded q tokens to the draft span HEAD # (ctx-full_q), not the tail, so they stay in [ctx-full_q .. ctx-1] (never # OOB); dropped tail slots are re-drafted next step (lossless). No-op when diff --git a/atom/model_ops/attentions/gdn_attn.py b/atom/model_ops/attentions/gdn_attn.py index 94b3e384e2..d438910ab4 100644 --- a/atom/model_ops/attentions/gdn_attn.py +++ b/atom/model_ops/attentions/gdn_attn.py @@ -115,15 +115,16 @@ class GDNAttentionMetadata: ) spec_state_indices_tensor: torch.Tensor | None = None # shape: [batch, num_spec] - # Optional read-side mapping when the framework allocates a different - # physical block for the first speculative token at a block boundary. + # Optional read-side mapping: a state fork's source in column 0 (or, in the + # plugins, a different physical block for the first speculative token at a + # block boundary). Same shape and strides as `spec_state_indices_tensor`. spec_state_indices_in_tensor: torch.Tensor | None = None non_spec_state_indices_tensor: torch.Tensor | None = ( None # shape: [batch - num_spec_decodes,] ) # Slots the incoming state is READ from, when a state fork makes that differ - # from `non_spec_state_indices_tensor`. Same tensor otherwise; None on the - # spec path, which never carries a fork. + # from `non_spec_state_indices_tensor`. Same tensor otherwise; on the spec + # path, set only when the step carries a fork. non_spec_state_indices_in_tensor: torch.Tensor | None = None # Prefill-path state bookkeeping built once per step instead of once per # GDN layer: `~has_initial_state`, and the write / read slots as int64 so @@ -131,6 +132,9 @@ class GDNAttentionMetadata: no_initial_state: torch.Tensor | None = None non_spec_state_indices_i64: torch.Tensor | None = None non_spec_state_indices_in_i64: torch.Tensor | None = None + # True if any sequence reads its state from a slot other than the one it + # writes. Host-side so backends can skip fork handling without a device sync. + has_state_fork: bool = False spec_sequence_masks: torch.Tensor | None = None # shape: [batch,] spec_token_indx: torch.Tensor | None = None non_spec_token_indx: torch.Tensor | None = None @@ -165,6 +169,8 @@ class GDNAttentionMetadata: replayssm_cache_len: int = 0 replayssm_route: str = "auto" replayssm_max_query_len: int = 1 + # A state fork's (write slot, source slot) rows, int64 [2, F]; None unforked. + replayssm_fork_slots: torch.Tensor | None = None # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -172,6 +178,14 @@ class GDNAttentionMetadata: token_chunk_offset_ptr: torch.Tensor | None = None +def fork_read_table(write: torch.Tensor, read: torch.Tensor) -> torch.Tensor: + """The spec read table: a forked row (its read slot differs from its + committed one) reads the source in EVERY column, since the kernel resumes + from column `num_accepted - 1`; other rows read what they write.""" + forked = (read != write[:, 0]).unsqueeze(1) + return torch.where(forked, read.unsqueeze(1), write) + + class GDNStateMixin(PoolRowsMixin): def __init__(self, model_runner, **kwargs): super().__init__(model_runner=model_runner, **kwargs) @@ -222,9 +236,9 @@ def _init_gdn_state( lin.get("short_conv_kernel_size", 4), ) - self.num_spec = 0 - if hasattr(model_runner, "drafter"): - self.num_spec = model_runner.drafter.mtp_k + # Not the drafter's: under PP only the last stage has one, but every + # stage runs the verify window over its own recurrent state. + self.num_spec = model_runner.num_spec_tokens self.use_spec_decode = self.num_spec > 0 # --- ReplaySSM ------------------------------------------------------ @@ -1196,10 +1210,16 @@ def prepare_state_indices(self, batch: ScheduledBatch, with_spec: bool = False): non_spec_state_indices_in = self.non_spec_state_indices_in_tensor.np spec_state_indices = self.spec_state_indices_tensor.np fork_srcs = getattr(batch, "state_fork_srcs", None) or () - assert not (with_spec and any(s >= 0 for s in fork_srcs)), ( - "state fork on the spec-decode path: spec_state_indices_tensor has " - "no read-side counterpart (BlockManager only forks onto prefill)" - ) + self._has_state_fork = any(s >= 0 for s in fork_srcs) + assert not ( + with_spec and self._has_state_fork and not self.spec_honours_state_fork() + ), "state fork on the spec-decode path, which this model does not read" + # (write slot, source slot) of the forked rows, for ReplaySSM. + self._fork_slots = [ + (slots[0], src) + for slots, src in zip(batch.state_slots, fork_srcs) + if src >= 0 + ] num_state_rows = len(batch.state_slots) non_spec_state_indices[num_state_rows:] = PAD_SLOT_ID non_spec_state_indices_in[num_state_rows:] = PAD_SLOT_ID @@ -1219,6 +1239,10 @@ def prepare_state_indices(self, batch: ScheduledBatch, with_spec: bool = False): non_spec_state_indices_in[idx] = src if src >= 0 else committed else: spec_state_indices[idx, : len(slots)] = slots + # Same fork contract as the non-spec branch above; the + # builder turns it into `spec_state_indices_in_tensor`. + src = fork_srcs[idx] if idx < len(fork_srcs) else -1 + non_spec_state_indices_in[idx] = src if src >= 0 else committed if self.replayssm: # ReplaySSM holds one slot per request, and that same index # addresses the checkpoint pool and the record buffers @@ -1227,6 +1251,11 @@ def prepare_state_indices(self, batch: ScheduledBatch, with_spec: bool = False): # fills only the 2-D one. non_spec_state_indices[idx] = committed + def spec_honours_state_fork(self) -> bool: + """Whether the model's spec path reads `spec_state_indices_in_tensor`, + so a state fork may land on a verify step (pp>1). Opt in per backend.""" + return False + def prepare_num_accepted_tokens(self, batch: ScheduledBatch): """Publish `num_bonus + 1` per request, 1 everywhere else, in one copy. @@ -1279,7 +1308,7 @@ def prepare_gdn_metadata( self.prepare_state_indices(batch, with_spec=False) spec_token_indx = None non_spec_token_indx = None - spec_state_indices_tensor = None + spec_state_indices_tensor = spec_state_indices_in_tensor = None non_spec_state_indices_tensor = ( self.non_spec_state_indices_tensor.copy_to_gpu(num_reqs) ) @@ -1318,7 +1347,16 @@ def prepare_gdn_metadata( # 1-D tensor too; ship that to the device for `slot_idx`. self.non_spec_state_indices_tensor.copy_to_gpu(num_reqs) non_spec_state_indices_tensor = None - non_spec_state_indices_in_tensor = None + # Only when forked: every other spec step ships None, as before. + # A fork runs eagerly (pp>1), so a fresh read-side table is fine. + non_spec_state_indices_in_tensor = spec_state_indices_in_tensor = None + if self._has_state_fork: + non_spec_state_indices_in_tensor = ( + self.non_spec_state_indices_in_tensor.copy_to_gpu(num_reqs) + ) + spec_state_indices_in_tensor = fork_read_table( + spec_state_indices_tensor, non_spec_state_indices_in_tensor + ) spec_query_start_loc = query_start_loc non_spec_query_start_loc = None num_accepted_tokens = self.num_accepted_tokens[:num_reqs] @@ -1371,11 +1409,13 @@ def prepare_gdn_metadata( spec_query_start_loc=spec_query_start_loc, non_spec_query_start_loc=non_spec_query_start_loc, spec_state_indices_tensor=spec_state_indices_tensor, + spec_state_indices_in_tensor=spec_state_indices_in_tensor, non_spec_state_indices_tensor=non_spec_state_indices_tensor, non_spec_state_indices_in_tensor=non_spec_state_indices_in_tensor, no_initial_state=no_initial_state, non_spec_state_indices_i64=non_spec_state_indices_i64, non_spec_state_indices_in_i64=non_spec_state_indices_in_i64, + has_state_fork=self._has_state_fork, spec_sequence_masks=spec_sequence_masks, spec_token_indx=spec_token_indx, non_spec_token_indx=non_spec_token_indx, @@ -1425,6 +1465,16 @@ def _attach_replayssm( write_pos.index_fill_(0, live, -1) return + if self._fork_slots: + # A fork's write slot starts from the source checkpoint alone (the + # kernel wrapper seeds it), so park its cursor like a prefill's: + # otherwise it folds the slot's previous tenant's records. + # Pinned: a pageable H2D syncs the stream. + fork = torch.tensor(list(zip(*self._fork_slots)), pin_memory=True) + fork = fork.to(write_pos.device, non_blocking=True) + write_pos.index_fill_(0, fork[0], -1) + md.replayssm_fork_slots = fork + # Decode: apply the PREVIOUS step's accepted counts. For non-spec # decode `num_accepted_tokens` stays at its initialised value of 1, # which is exactly one committed token per step. diff --git a/atom/model_ops/attentions/kimi_mla_gdn_attn.py b/atom/model_ops/attentions/kimi_mla_gdn_attn.py index 940c2a4e78..e5b8639b05 100644 --- a/atom/model_ops/attentions/kimi_mla_gdn_attn.py +++ b/atom/model_ops/attentions/kimi_mla_gdn_attn.py @@ -50,6 +50,10 @@ class _KimiMLAGDNCommon(PageUnitGeometryMixin, GDNStateMixin): def __init__(self, model_runner): super().__init__(model_runner=model_runner) + def spec_honours_state_fork(self) -> bool: + # K3's KDA forward reads the fork table; glm5_next's does not. + return self.model_runner.hf_text_config.model_type == "kimi_linear" + def _module_kinds(self, module) -> tuple: """Two row spaces: the MLA pool's rows, and the KDA state slots. diff --git a/atom/model_ops/fla_ops/fused_sigmoid_gating.py b/atom/model_ops/fla_ops/fused_sigmoid_gating.py index 526c9dd0dd..b7b274d77e 100644 --- a/atom/model_ops/fla_ops/fused_sigmoid_gating.py +++ b/atom/model_ops/fla_ops/fused_sigmoid_gating.py @@ -39,6 +39,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( ht, cu_seqlens, ssm_state_indices, + ssm_state_indices_in, num_accepted_tokens, scale, N: tl.int64, # num of sequences @@ -112,10 +113,10 @@ def fused_sigmoid_gating_delta_rule_update_kernel( i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 else: i_t = 0 - # Load state index and check for PAD_SLOT_ID (-1) - state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( - tl.int64 - ) + # Load the read-side state index and check for PAD_SLOT_ID (-1) + state_idx = tl.load( + ssm_state_indices_in + i_n * stride_indices_seq + i_t + ).to(tl.int64) # Skip if state index is invalid (PAD_SLOT_ID = -1) if state_idx < 0: return @@ -212,6 +213,7 @@ def fused_sigmoid_gating_delta_rule_update( num_accepted_tokens: torch.Tensor | None = None, use_qk_l2norm_in_kernel: bool = False, is_kda: bool = False, + ssm_state_indices_in: torch.Tensor | None = None, lower_bound: float | None = None, ): """ @@ -269,6 +271,17 @@ def fused_sigmoid_gating_delta_rule_update( stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1 else: stride_indices_seq, stride_indices_tok = ssm_state_indices.stride() + # The initial state is read at `ssm_state_indices_in` (a state fork's + # source) and written at `ssm_state_indices`; one set of strides serves + # both, so the read side must be laid out exactly like the write side. + if ssm_state_indices_in is None: + ssm_state_indices_in = ssm_state_indices + else: + assert ( + ssm_state_indices is not None + and ssm_state_indices_in.shape == ssm_state_indices.shape + and ssm_state_indices_in.stride() == ssm_state_indices.stride() + ), "ssm_state_indices_in must match ssm_state_indices in shape and stride" grid = (NK, NV, N * HV) fused_sigmoid_gating_delta_rule_update_kernel[grid]( @@ -287,6 +300,7 @@ def fused_sigmoid_gating_delta_rule_update( ht=final_state, cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, + ssm_state_indices_in=ssm_state_indices_in, num_accepted_tokens=num_accepted_tokens, scale=scale, N=N, diff --git a/atom/model_ops/fla_ops/replayssm.py b/atom/model_ops/fla_ops/replayssm.py index 8c4b42419d..a8958c20e2 100644 --- a/atom/model_ops/fla_ops/replayssm.py +++ b/atom/model_ops/fla_ops/replayssm.py @@ -438,6 +438,18 @@ def _replayssm_commit_kernel( tl.store(write_pos + slot, base + tl.load(num_accepted + i_n)) +def seed_forked_slots(ckpt: torch.Tensor, fork_slots: torch.Tensor | None) -> None: + """Copy a state fork's source checkpoint onto its write slot. + + The kernels read and write `slot_idx` only, and a fork's write slot must + hold the base its records fold onto. The builder parks that slot's cursor + at -1, so nothing is folded over it. `fork_slots` is the builder's int64 + `[2, F]` (write, source) pairs, so this indexes without a host sync. + """ + if fork_slots is not None: + ckpt[fork_slots[0]] = ckpt[fork_slots[1]] + + def replayssm_commit( write_pos: torch.Tensor, slot_idx: torch.Tensor, @@ -1369,6 +1381,7 @@ def replayssm_gated_delta_rule( is_kda: bool = False, route: str = "auto", o: torch.Tensor | None = None, + fork_slots: torch.Tensor | None = None, ) -> torch.Tensor: """ReplaySSM forward for one linear-attention layer. @@ -1390,6 +1403,8 @@ def replayssm_gated_delta_rule( o: optional contiguous ``[T_tot, HV, V]`` destination (e.g. a slice of the layer's output buffer); the kernel writes into it directly instead of a fresh tensor, saving the caller a device copy. + fork_slots: a state fork's (write, source) slot pairs, see + :func:`seed_forked_slots`. Returns: ``o`` of shape ``[1, T_tot, HV, V]`` -- same leading-batch convention @@ -1397,6 +1412,7 @@ def replayssm_gated_delta_rule( replacement at the call sites in ``attention_gdn.py``. """ assert q.shape[0] == 1, "varlen layout expected (B == 1)" + seed_forked_slots(ckpt, fork_slots) _, T_tot, H, K = q.shape HV, V = v.shape[2], v.shape[3] N = cu_seqlens.numel() - 1 @@ -1778,6 +1794,7 @@ def replayssm_sigmoid_gating_delta_rule( scale: float | None = None, use_qk_l2norm_in_kernel: bool = False, lower_bound: float | None = None, + fork_slots: torch.Tensor | None = None, ) -> torch.Tensor: """ReplaySSM for the KDA (Kimi-K3) linear-attention layer. @@ -1786,9 +1803,11 @@ def replayssm_sigmoid_gating_delta_rule( the pool holds one checkpoint per request instead of one state per speculative token. See the module docstring for the algebra. - ``o`` may be passed to write the output in place (K3 does this). + ``o`` may be passed to write the output in place (K3 does this); + ``fork_slots`` as in :func:`replayssm_gated_delta_rule`. """ assert q.shape[0] == 1, "varlen layout expected (B == 1)" + seed_forked_slots(ckpt, fork_slots) _, T_tot, H, K = q.shape HV, V = v.shape[2], v.shape[3] N = cu_seqlens.numel() - 1 diff --git a/atom/models/deepseek_v4/mono/session.py b/atom/models/deepseek_v4/mono/session.py index c4d1a193d4..42a2d4b52b 100644 --- a/atom/models/deepseek_v4/mono/session.py +++ b/atom/models/deepseek_v4/mono/session.py @@ -1280,7 +1280,7 @@ def bind_aux(): from atom.mono.runtime.aux_taps import bind_aux_taps from atom.spec_decode.dspark_proposer import DSparkProposer - spec = drafter._aux_capture_spec(self.model) + spec = drafter._aux_capture_spec() if ( spec is None or spec.capture != "output" diff --git a/atom/models/kimi_k3/model.py b/atom/models/kimi_k3/model.py index c0d8e9daa8..1947caa90b 100644 --- a/atom/models/kimi_k3/model.py +++ b/atom/models/kimi_k3/model.py @@ -1365,11 +1365,12 @@ def _forward_impl( # for exactly one forward: the prefix-cache hit that forks off a # checkpoint (BlockManager._attach_state_group reads `state_fork_src` # and writes a freshly popped group). Reading the write slot there - # resumes from whatever the recycled group still held. Only prefill can - # carry a fork -- prepare_state_indices asserts it, and min_fork_tokens - # keeps the chunk long enough -- so the decode branches below stay on - # `state_indices`. Falling back to the write slot leaves every non-fork - # forward bit-identical. Mirrors attention_gdn.py. + # resumes from whatever the recycled group still held. Under pp=1 only + # prefill carries a fork, since the checkpoint is taken between prompt + # chunks; under pp>1 it lands after the prompt, so the decode and + # spec-decode branches below read it too. + # Falling back to the write slot leaves every non-fork forward + # bit-identical. Mirrors attention_gdn.py. state_indices_in = kda_metadata.non_spec_state_indices_in_tensor if state_indices_in is None: state_indices_in = state_indices @@ -1427,6 +1428,7 @@ def _forward_impl( # Slice the per-token cache-slot indices once (used for both the # conv update and the fused recurrence below). decode_state_indices = state_indices[:num_actual_tokens] + decode_state_indices_in = state_indices_in[:num_actual_tokens] q, k, v = causal_conv1d_update( mixed_qkv, conv_state, @@ -1436,6 +1438,7 @@ def _forward_impl( None, self.activation, conv_state_indices=decode_state_indices, + conv_state_indices_in=decode_state_indices_in, validate_data=False, ) q = rearrange(q, "t (h d) -> 1 t h d", d=self.head_dim) @@ -1473,6 +1476,7 @@ def _forward_impl( o=out, use_qk_l2norm_in_kernel=True, lower_bound=self._kda_gate_lower_bound, + fork_slots=kda_metadata.replayssm_fork_slots, ) else: fused_sigmoid_gating_delta_rule_update( @@ -1488,6 +1492,7 @@ def _forward_impl( inplace_final_state=True, cu_seqlens=query_start_loc[: kda_metadata.num_decodes + 1], ssm_state_indices=decode_state_indices, + ssm_state_indices_in=decode_state_indices_in, use_qk_l2norm_in_kernel=True, is_kda=True, lower_bound=self._kda_gate_lower_bound, @@ -1497,6 +1502,10 @@ def _forward_impl( spec_state_indices = kda_metadata.spec_state_indices_tensor spec_query_start_loc = kda_metadata.spec_query_start_loc num_accepted_tokens = kda_metadata.num_accepted_tokens + # Under pp>1 the first verify after the prompt can fork off a + # checkpoint: the builder then ships the source as the read side. + spec_state_indices_in = kda_metadata.spec_state_indices_in_tensor + nsd = kda_metadata.num_spec_decodes q, k, v = causal_conv1d_update( mixed_qkv, conv_state, @@ -1507,9 +1516,12 @@ def _forward_impl( self.activation, # First reserved slot per seq holds the resume state; the kernel # walks forward via num_accepted_tokens + query_start_loc. - conv_state_indices=spec_state_indices[:, 0][ - : kda_metadata.num_spec_decodes - ], + conv_state_indices=spec_state_indices[:nsd, 0], + conv_state_indices_in=( + None + if spec_state_indices_in is None + else spec_state_indices_in[:nsd, 0] + ), num_accepted_tokens=num_accepted_tokens, query_start_loc=spec_query_start_loc, # Verify window: sizes the conv rollback window and hence the @@ -1527,7 +1539,6 @@ def _forward_impl( k = rearrange(k, "t (h d) -> 1 t h d", d=self.head_dim) v = rearrange(v, "t (h d) -> 1 t h d", d=self.head_dim) if getattr(kda_metadata, "replayssm", False): - nsd = kda_metadata.num_spec_decodes replayssm_sigmoid_gating_delta_rule( q, k, @@ -1547,6 +1558,7 @@ def _forward_impl( o=out, use_qk_l2norm_in_kernel=True, lower_bound=self._kda_gate_lower_bound, + fork_slots=kda_metadata.replayssm_fork_slots, ) else: fused_sigmoid_gating_delta_rule_update( @@ -1567,6 +1579,7 @@ def _forward_impl( # num_accepted_tokens the kernel reads the resume state from # slot[num_accepted-1] and writes a snapshot after each token. ssm_state_indices=spec_state_indices, + ssm_state_indices_in=spec_state_indices_in, num_accepted_tokens=num_accepted_tokens, use_qk_l2norm_in_kernel=True, is_kda=True, @@ -1795,7 +1808,12 @@ def __init__(self, atom_config: Config, prefix: str = ""): self.config = config self.vocab_size = config.vocab_size - if get_pp_group().is_first_rank: + # The drafter runs on the last stage and shares this table, and not + # every draft checkpoint (Kimi-K3.1-DSpark) ships its own. + pp = get_pp_group() + if pp.is_first_rank or ( + pp.is_last_rank and atom_config.speculative_config is not None + ): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size ) diff --git a/atom/models/kimi_k3/mono/plan.py b/atom/models/kimi_k3/mono/plan.py index fa25314712..1f66b7bcc3 100644 --- a/atom/models/kimi_k3/mono/plan.py +++ b/atom/models/kimi_k3/mono/plan.py @@ -51,6 +51,11 @@ def step_shape(metadata, rows: int) -> tuple[int | None, str]: return None, "no KDA metadata" if metadata.num_prefills: return None, f"{metadata.num_prefills} prefills in the batch" + if metadata.has_state_fork: + # A forked step reads one slot and writes another. The production + # kernels take that read-side index; K1 reads `st_idx` only, so it + # would run on the unseeded write slot. + return None, "state fork (K1 has no read-side slot index)" if getattr(metadata, "spec_sequence_masks", None) is None: return _plain_decode_shape(metadata, rows) if metadata.num_decodes: @@ -108,8 +113,7 @@ def state_slots(metadata, ones: torch.Tensor) -> tuple[torch.Tensor, torch.Tenso A plain decode is that with one column and one accepted token: it reads and writes the request's one slot in place, as the production decode does. The production decode reads ``non_spec_state_indices_in_tensor`` instead where - a state fork makes the two differ, but forks happen only on prefill - (``GDNStateMixin.prepare_state_indices``), and a served step has none. + a state fork makes the two differ, but ``step_shape`` refuses a forked step. ``ones`` is int32, at least as long as the step has requests. """ if getattr(metadata, "spec_sequence_masks", None) is not None: diff --git a/atom/models/registry.py b/atom/models/registry.py index 602debfeaf..60c5eee92d 100644 --- a/atom/models/registry.py +++ b/atom/models/registry.py @@ -65,3 +65,8 @@ "atom.models.minimax_m3.mono.dispatch.install_mono_decode" ), } +# Speculative methods each architecture serves under pipeline_parallel_size > 1 +# (its drafts verified per stage, its drafter's aux taps relayed across stages). +pp_speculative_methods = { + "KimiK3ForConditionalGeneration": ("dspark",), +} diff --git a/atom/plugin/vllm/attention/metadata.py b/atom/plugin/vllm/attention/metadata.py index 05d768bc2a..6a3bc85c8e 100644 --- a/atom/plugin/vllm/attention/metadata.py +++ b/atom/plugin/vllm/attention/metadata.py @@ -17,7 +17,11 @@ from atom.config import get_current_atom_config from atom.distributed.dcp_utils import dcp_persistent_supported -from atom.model_ops.attention_mla import _MLA_MIN_HEADS, mla_dcp_kernel_num_heads +from atom.model_ops.attention_mla import ( + _MLA_MIN_HEADS, + mla_dcp_kernel_num_heads, + mla_kernel_num_heads, +) from atom.plugin.vllm.attention.layer_mla import ( disabled_mla_persistent_metadata, mla_fold_kv_metadata_triton, @@ -1340,8 +1344,10 @@ def __init__( self.num_attention_heads, ) - self.padded_num_attention_heads = max( - self.num_attention_heads, self.min_query_heads + # Must equal `MLAAttention.padded_num_heads`, which the plugin's MLA + # layer inherits: that pads the query, this plans the work for it. + self.padded_num_attention_heads = mla_kernel_num_heads( + max(self.num_attention_heads, self.min_query_heads) ) # Whether DCP decode can run in persistent mode on this GPU (gfx950 has # the lse persistent kernel, gfx942 does not — see dcp_utils). Cached diff --git a/atom/spec_decode/drafter.py b/atom/spec_decode/drafter.py index 4aeb781121..b36e8323d0 100644 --- a/atom/spec_decode/drafter.py +++ b/atom/spec_decode/drafter.py @@ -17,6 +17,7 @@ pcp_pad_len, ) from atom.model_loader.loader import load_model +from atom.models.utils import PPMissingLayer from atom.spec_decode.draft_graph import DraftGraph from atom.utils import CpuGpuBuffer, resolve_obj_by_qualname from atom.utils.forward_context import ( @@ -113,6 +114,16 @@ def _pick(module): return embed +def tap_is_resident(layers: nn.Module, lid: int) -> bool: + """Whether aux tap ``lid`` runs on this PP stage: a layer owned by another + stage is a ``PPMissingLayer``. The embedding runs only on the first stage, + though the last may hold a copy for its drafter. Out-of-range ids are not + resident.""" + if lid == -1: + return get_pp_group().is_first_rank + return 0 <= lid < len(layers) and not isinstance(layers[lid], PPMissingLayer) + + def _identity_extract(output: Any, _module: nn.Module) -> torch.Tensor: """Aux extractor for the ``-1`` (embedding) tap: the embedding output is already ``[N, hidden]``, so take it verbatim (unwrapping a (tensor, scale) @@ -120,6 +131,71 @@ def _identity_extract(output: Any, _module: nn.Module) -> torch.Tensor: return output[0] if isinstance(output, tuple) else output +def _aux_hook(extract, write): + """A target-layer forward hook handing ``write(tensor, row_offset)`` the + tap's aux tensor, on target forwards only.""" + + def _hook(module, _inputs, output): + ctx = get_forward_context().context + # Target forwards only. The `-1` tap sits on the embedding, which the + # draft model shares — its forward_spec would otherwise clobber these + # rows with noise-token embeddings. + if ctx.is_draft: + return + tensor = extract(output, module) + if tensor is None: + return + # The offset is 0 except under TBO, where this hook fires once per + # micro-batch on a disjoint token slice — writing at row 0 + # unconditionally would make the ubatches overwrite each other. + write(tensor, ctx.ubatch_token_offset) + + return _hook + + +def install_aux_taps( + target_model: nn.Module, spec: AuxCaptureSpec, layer_ids, make_hook +): + """Hook each of ``layer_ids`` (a subset of ``spec.layer_ids``) on the target + with ``make_hook(lid, extract)``. + + Reference convention (deepspec validate_target_layer_ids): ``spec``'s ids + must be strictly increasing and each in {-1} U [0, n_layers). Each id + resolves to a (module, extract) tap: -1 taps the embedding output + (hidden_states[0], taken verbatim); k taps decoder layer k through + ``spec.extract``. + """ + layers = _resolve_decoder_layers(target_model) + n_layers = len(layers) + prev = None + for lid in spec.layer_ids: + if not (lid == -1 or 0 <= lid < n_layers): + raise ValueError( + f"aux capture layer id {lid} out of range {{-1}} U [0,{n_layers})." + ) + if prev is not None and lid <= prev: + raise ValueError( + "aux capture layer ids must be strictly increasing, got " + f"{spec.layer_ids}." + ) + prev = lid + for lid in layer_ids: + module, extract = ( + (_resolve_embedding(target_model), _identity_extract) + if lid == -1 + else (layers[lid], spec.extract) + ) + hook = make_hook(lid, extract) + if lid != -1 and spec.capture == "input": + + def pre_hook(module, inputs, hook=hook): + hook(module, (), inputs) + + module.register_forward_pre_hook(pre_hook) + else: + module.register_forward_hook(hook) + + # Draft-model architecture registry: maps a draft checkpoint's architecture # string to the ATOM wrapper class. Covers all drafter flavors (serial MTP, # EAGLE3, DSpark), so it lives on the shared base rather than a per-flavor file. @@ -158,6 +234,9 @@ class Drafter(abc.ABC): ModelRunner used to scatter through its code. """ + # The leading taps' rows relayed from earlier PP stages, as one block. + _relayed_aux: torch.Tensor | None = None + def __init__(self, atom_config: Config, device: torch.device, runner): self.config = atom_config self.speculative_config = self.config.speculative_config @@ -320,12 +399,18 @@ def _build_draft_model(self, model_class) -> nn.Module: the target config; ``EagleProposer`` overrides for the eagle3 arch.""" return model_class(self.config) - def _aux_capture_spec(self, target_model: nn.Module) -> AuxCaptureSpec | None: - """Declare which target layers to tap for aux hidden states and how to - extract them. Default: no capture. Hook-based drafters (DSpark) override - this; the base ``arm_aux_capture`` turns the spec into forward hooks.""" + @classmethod + def aux_capture_spec(cls, config: Config) -> AuxCaptureSpec | None: + """The aux capture this drafter class makes, from config alone (None: + no capture). What a drafterless PP stage taps and relays on.""" return None + def _aux_capture_spec(self) -> AuxCaptureSpec | None: + """Declare which target layers to tap for aux hidden states and how to + extract them. Default: the config-only ``aux_capture_spec``. The base + ``arm_aux_capture`` turns the spec into forward hooks.""" + return self.aux_capture_spec(self.config) + @abc.abstractmethod def propose( self, @@ -375,51 +460,30 @@ def arm_aux_capture(self, target_model: nn.Module) -> None: this drafter's ``_aux_capture_spec``, writing each tapped layer's aux tensor into a preallocated buffer in-place (cudagraph-safe). Called once after ``load_model``. No-op when the drafter declares no spec.""" - spec = self._aux_capture_spec(target_model) + spec = self._aux_capture_spec() if spec is None: return + layers = _resolve_decoder_layers(target_model) + # A tap on an earlier PP stage reaches `take_relayed_aux` instead, so it + # gets no rows and no hook. install_aux_taps still rejects bad ids. + own = tuple(lid for lid in spec.layer_ids if tap_is_resident(layers, lid)) self._aux_buffers = [ torch.zeros( - self.max_num_tokens, + self.max_num_tokens if lid in own else 0, spec.hidden_size, device=self.device, dtype=self.dtype, ) - for _ in spec.layer_ids + for lid in spec.layer_ids ] - layers = _resolve_decoder_layers(target_model) - n_layers = len(layers) - # Reference convention (deepspec validate_target_layer_ids): ids must be - # strictly increasing and each in {-1} U [0, n_layers). Each id resolves to - # a (module, extract) tap funneled through the SAME hook: -1 taps the - # embedding output (hidden_states[0], taken verbatim); k taps decoder - # layer k's output through the drafter's mHC extract. - prev = None - for buf_idx, lid in enumerate(spec.layer_ids): - if not (lid == -1 or 0 <= lid < n_layers): - raise ValueError( - f"aux capture layer id {lid} out of range {{-1}} U [0,{n_layers})." - ) - if prev is not None and lid <= prev: - raise ValueError( - "aux capture layer ids must be strictly increasing, got " - f"{spec.layer_ids}." - ) - prev = lid - module, extract = ( - (_resolve_embedding(target_model), _identity_extract) - if lid == -1 - else (layers[lid], spec.extract) - ) - hook = self._make_aux_hook(buf_idx, extract) - if lid != -1 and spec.capture == "input": - - def pre_hook(module, inputs, hook=hook): - hook(module, (), inputs) - - module.register_forward_pre_hook(pre_hook) - else: - module.register_forward_hook(hook) + install_aux_taps( + target_model, + spec, + own, + lambda lid, extract: self._make_aux_hook( + spec.layer_ids.index(lid), extract + ), + ) self._captures_aux = True # A target that runs some layers on a subset of rows (V4.1 bounded # replay) moves those layers' captures to their forward rows itself. @@ -437,24 +501,11 @@ def pre_hook(module, inputs, hook=hook): def _make_aux_hook(self, buf_idx: int, extract): buffer = self._aux_buffers[buf_idx] - def _hook(module, _inputs, output): - ctx = get_forward_context().context - # Target forwards only. The `-1` tap sits on the embedding, which the - # draft model shares — its forward_spec would otherwise clobber these - # rows with noise-token embeddings. - if ctx.is_draft: - return - tensor = extract(output, module) - if tensor is None: - return - # In-place write into the fixed buffer (cudagraph-safe). Offset is 0 - # except under TBO, where this hook fires once per micro-batch on a - # disjoint token slice — writing at row 0 unconditionally would make - # the ubatches overwrite each other. - off = ctx.ubatch_token_offset + def write(tensor, off): + # In-place write into the fixed buffer (cudagraph-safe). buffer[off : off + tensor.shape[0]].copy_(tensor) - return _hook + return _aux_hook(extract, write) def aux_for(self, target_hidden_states: torch.Tensor) -> list[torch.Tensor] | None: """The target aux hidden states captured during the last forward, @@ -469,16 +520,23 @@ def aux_for(self, target_hidden_states: torch.Tensor) -> list[torch.Tensor] | No """ if not self._captures_aux: return None + blocks = self._aux_buffers + if self._relayed_aux is not None: + relayed = self._relayed_aux.shape[-1] // blocks[0].shape[-1] + blocks = [self._relayed_aux, *blocks[relayed:]] n = target_hidden_states.shape[0] pcp_ws = get_pcp_world_size() ctx = getattr(get_forward_context(), "context", None) if pcp_ws > 1 and n > 0 and getattr(ctx, "target_pcp_split", False): n_local = pcp_pad_len(n, pcp_ws) // pcp_ws - return [ - pcp_allgather_rerange(buf[:n_local], pcp_ws)[:n] - for buf in self._aux_buffers - ] - return [buf[:n] for buf in self._aux_buffers] + return [pcp_allgather_rerange(buf[:n_local], pcp_ws)[:n] for buf in blocks] + return [buf[:n] for buf in blocks] + + def take_relayed_aux(self, block: torch.Tensor) -> None: + """This forward's leading taps, relayed from earlier PP stages as one + ``[N, k * hidden_size]`` block in ``layer_ids`` order; ``aux_for`` + serves it in place of their capture buffers, uncopied.""" + self._relayed_aux = block def anchors_to_gpu(self, anchors: list[int]) -> torch.Tensor: """Scheduler-supplied anchors as an int32 GPU tensor, without blocking. diff --git a/atom/spec_decode/dspark_proposer.py b/atom/spec_decode/dspark_proposer.py index eb7b2ed314..e039fa6ee7 100644 --- a/atom/spec_decode/dspark_proposer.py +++ b/atom/spec_decode/dspark_proposer.py @@ -509,25 +509,42 @@ def verify_scheduler(self): return self._verify_scheduler # ---- aux-hidden-state ownership (declarative; base owns the hook machinery) ---- - def _aux_capture_spec(self, target_model: nn.Module) -> AuxCaptureSpec: - """Resolve the draft's feature contract; the base owns capture buffers.""" - draft_cfg = self.speculative_config.draft_model_hf_config + @classmethod + def aux_capture_spec(cls, config) -> AuxCaptureSpec | None: + """Default DSpark aux-capture spec, from config alone (None if no DSpark).""" + spec_cfg = getattr(config, "speculative_config", None) + if spec_cfg is None or not spec_cfg.use_dspark(): + return None + draft_cfg = spec_cfg.draft_model_hf_config layer_ids = tuple( - int(i) for i in getattr(draft_cfg, "dspark_target_layer_ids", ()) + int(i) for i in (getattr(draft_cfg, "dspark_target_layer_ids", None) or ()) ) if not layer_ids: raise ValueError( "DSpark requires dspark_target_layer_ids on the draft config." ) - own = getattr(self.model, "target_aux_capture_spec", None) - if own is not None: - return own(layer_ids, self.config.hf_config.hidden_size) return AuxCaptureSpec( layer_ids=layer_ids, - hidden_size=self.config.hf_config.hidden_size, - extract=self._extract_layer_hidden, + hidden_size=config.hf_config.hidden_size, + extract=cls._extract_layer_hidden, ) + def _aux_capture_spec(self) -> AuxCaptureSpec: + """Resolve the draft's feature contract; the base owns capture buffers.""" + spec = self.aux_capture_spec(self.config) + own = getattr(self.model, "target_aux_capture_spec", None) + if own is None: + return spec + if int(getattr(self.config, "pipeline_parallel_size", 1) or 1) > 1: + # Drafterless PP stages can only rebuild the default spec; relaying + # it to a draft with its own contract would be silently wrong. + raise NotImplementedError( + f"{type(self.model).__name__} declares its own " + "target_aux_capture_spec, which pipeline parallelism cannot " + "reconstruct on a drafterless stage." + ) + return own(spec.layer_ids, spec.hidden_size) + @staticmethod def _extract_layer_hidden(output, block: nn.Module): """Reconstruct a target layer's post-layer hidden state ``[N, dim]``. @@ -614,7 +631,11 @@ def compute_draft_kv( return forward_context = get_forward_context() scheduled_bs = forward_context.context.scheduled_bs - main_hidden_all = torch.cat(aux_hidden_states, dim=-1) + # One block when every tap was relayed from earlier PP stages. + if len(aux_hidden_states) == 1: + main_hidden_all = aux_hidden_states[0] + else: + main_hidden_all = torch.cat(aux_hidden_states, dim=-1) with record_function( f"draft_kv[bs={scheduled_bs} tok={main_hidden_all.shape[0]}]" ): diff --git a/atom/spec_decode/factory.py b/atom/spec_decode/factory.py index a05f6f2529..8e9f9e62f5 100644 --- a/atom/spec_decode/factory.py +++ b/atom/spec_decode/factory.py @@ -11,6 +11,11 @@ def build_drafter(config: Config, device, runner) -> Drafter: downstream call-site depends on the ``Drafter`` contract, not the concrete type. """ + return drafter_class(config)(config, device, runner) + + +def drafter_class(config: Config) -> type[Drafter]: + """The drafter class this config speculates with, on every PP stage.""" if config.speculative_config.use_dspark(): - return DSparkProposer(config, device, runner) - return EagleProposer(config, device, runner) + return DSparkProposer + return EagleProposer diff --git a/atom/spec_decode/pp_aux_relay.py b/atom/spec_decode/pp_aux_relay.py new file mode 100644 index 0000000000..5a67381946 --- /dev/null +++ b/atom/spec_decode/pp_aux_relay.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Relay a block drafter's target aux hidden states across PP stages. + +The drafter lives only on the last PP stage, so hooks on taps owned by earlier +stages land on ``PPMissingLayer`` and never fire, leaving zero aux buffers and +garbage drafts. Earlier stages capture their own taps and forward them in the +intermediate-tensor payload, concatenated along the hidden axis in the +spec's ``layer_ids`` order, which ``install_aux_taps`` requires strictly +increasing (-1, the embedding, first). Stages own contiguous ascending layer +ranges, so every stage derives the same columns from ``layer_ids`` alone. +""" + +import logging + +import torch +from aiter.dist.parallel_state import get_pp_group + +from atom.models.utils import PPMissingLayer +from atom.spec_decode.drafter import ( + AuxCaptureSpec, + _aux_hook, + _resolve_decoder_layers, + install_aux_taps, + tap_is_resident, +) +from atom.spec_decode.factory import drafter_class + +logger = logging.getLogger("atom") + + +class PPAuxRelay: + """Per-stage aux relay, built on every stage when speculating with pp > 1. + + Sends ``incoming_ids + own_ids`` columns, in ascending tap order. + """ + + def __init__( + self, + spec: AuxCaptureSpec, + target_model: torch.nn.Module, + ) -> None: + pp = get_pp_group() + layers = _resolve_decoder_layers(target_model) + stage_start = next( + i for i, layer in enumerate(layers) if not isinstance(layer, PPMissingLayer) + ) + taps = spec.layer_ids + + self.spec = spec + self.own_ids = tuple(lid for lid in taps if tap_is_resident(layers, lid)) + # Every tap below this stage's first layer was captured upstream, which + # is exactly what the previous stage sent, even if it owned none itself. + self.incoming_ids = tuple( + lid + for lid in taps + if not tap_is_resident(layers, lid) and lid < stage_start + ) + self.outgoing_ids = () if pp.is_last_rank else self.incoming_ids + self.own_ids + + # The send payload, allocated per forward at its row count (a resident + # max_num_batched_tokens buffer is ~0.5 GB per rank) by the lowest own + # tap, which fires first; every own tap captures into its columns and + # pack() copies the received taps into the leading ones. The last + # stage's hooks write into the drafter's own buffers instead. + # The write reads no bookkeeping: inside the compiled forward Dynamo + # freezes it to the warmup trace. + # shortcut: a recorded graph would bake the capture-time buffer, so + # this holds only while PP stages run their target eagerly. + self._send: torch.Tensor | None = None + self._col = {lid: col for col, lid in enumerate(self.outgoing_ids)} + + @classmethod + def build(cls, config, target_model) -> "PPAuxRelay | None": + """This stage's armed relay, or None when no tap crosses a PP boundary. + + Every stage takes the drafter class's config-only spec, so all agree on + the columns; a draft with its own contract refuses PP when it arms. + """ + if get_pp_group().world_size <= 1 or config.speculative_config is None: + return None + spec = drafter_class(config).aux_capture_spec(config) + if spec is None: + return None + relay = cls(spec, target_model) + if not relay.incoming_ids and not relay.outgoing_ids: + return None + relay.arm(target_model) + return relay + + def arm(self, target_model: torch.nn.Module) -> None: + """Install capture hooks for this stage's taps (no-op on the last stage, + whose taps the drafter's own ``arm_aux_capture`` hooks).""" + if get_pp_group().is_last_rank: + return + install_aux_taps( + target_model, + self.spec, + self.own_ids, + lambda lid, extract: _aux_hook(extract, self._capture_into(lid)), + ) + logger.info( + "PPAuxRelay: capturing target layers %s, relaying %s forward", + list(self.own_ids), + list(self.outgoing_ids), + ) + + def _capture_into(self, lid): + h, col = self.spec.hidden_size, self._col[lid] + width = len(self.outgoing_ids) * h + + opens = lid == self.own_ids[0] + + def write(tensor, _off): + # TBO is refused under PP, so a forward is never split into ubatches. + if opens: + self._send = torch.empty( + tensor.shape[0], width, device=tensor.device, dtype=tensor.dtype + ) + self._send[:, col * h : (col + 1) * h].copy_(tensor) + + return write + + def pack(self, received: torch.Tensor | None, num_tokens: int): + """Payload to send on (received block, then own captures), or None.""" + if not self.outgoing_ids: + return None + if self.incoming_ids: + self._check_received(received) + if not self.own_ids: + return received[:num_tokens] + # Hand the step's buffer over: the send path owns it from here. + send, self._send = self._send, None + assert ( + send is not None and send.shape[0] == num_tokens + ), f"PPAuxRelay: no capture of {num_tokens} rows to send" + if self.incoming_ids: + send[:, : received.shape[-1]].copy_(received[:num_tokens]) + return send + + def _check_received(self, received: torch.Tensor | None) -> None: + if received is None: + raise RuntimeError( + f"PPAuxRelay expected aux rows for target layers " + f"{list(self.incoming_ids)} from the previous stage, got none." + ) + expected = len(self.incoming_ids) * self.spec.hidden_size + if received.shape[-1] != expected: + raise RuntimeError( + f"PPAuxRelay expected {expected} aux columns for target layers " + f"{list(self.incoming_ids)}, got {received.shape[-1]}." + ) + + def absorb(self, received: torch.Tensor | None, drafter) -> None: + """Hand the relayed block to the drafter, which reads it in place.""" + if not self.incoming_ids: + return + self._check_received(received) + drafter.take_relayed_aux(received) diff --git a/docs/distributed_guide.md b/docs/distributed_guide.md index 830e9bc51d..f437d5c1a3 100644 --- a/docs/distributed_guide.md +++ b/docs/distributed_guide.md @@ -504,6 +504,9 @@ This means each GPU runs an independent attention computation (no TP AllReduce f | `atom/model_ops/linear.py` | `ColumnParallelLinear`, `RowParallelLinear`, `QKVParallelLinear`, `MergedColumnParallelLinear` | | `atom/model_ops/moe.py` | `FusedMoE`, `FusedMoEParallelConfig` (EP configuration and expert distribution) | | `atom/model_ops/fused_moe/mori_prepare_finalize.py` | `MoriPrepareAndFinalize` (MORI dispatch/combine for EP) | +| `atom/model_engine/pp_engine_core.py` | Pipeline-parallel stage loops (head / downstream) | +| `atom/distributed/pp_comm.py` | What crosses a PP stage boundary: `_PP_PROXY_KEYS` (`hidden_states`, `residual`, `block_residual`, ...) plus `PP_AUX_KEY` (`pp_aux`) | +| `atom/spec_decode/pp_aux_relay.py` | `PPAuxRelay`: speculative decoding under PP. The drafter lives on the last stage; earlier stages capture the target aux hidden states they own and send them on under `pp_aux` | --- diff --git a/docs/model_support_guide.md b/docs/model_support_guide.md index f1d91fa4cc..e4f7d19512 100644 --- a/docs/model_support_guide.md +++ b/docs/model_support_guide.md @@ -157,6 +157,7 @@ ATOM resolves the HuggingFace `architectures` field from a model's `config.json` - **MoE:** MXFP4 latent MoE (`KimiSparseMoeBlock`) with SiTU activation and optional dual-stream shared/routed overlap (`ATOM_K3_SHARED_EXPERT_OVERLAP`). - **Vision:** `atom/models/kimi_k3_vl.py` implements MoonViT3d — patch embed with a bilinearly resampled learnable position grid, 27 blocks with packed `wqkv` and complex 2D RoPE over non-causal varlen attention, then an `sd2_tpool` 2x2 merge and a `patchmergerv2` projector into the 7168-wide text space. Replicated per TP rank (~0.9 GB bf16), built only on the first pipeline rank. - **Image tokens:** The HF processor emits one `<|media_pad|>` per image and expands it inside the model. ATOM instead expands it during input processing (`build_kimi_k3_inputs` in `atom/models/kimi_k3/image_processing.py`, using the shared helper in `atom/multimodal/processing.py`) so the scheduler, KV blocks and positions see the true prompt length. It also returns each run as an `embedding_spans` entry, so image prefills chunk like text through `atom/model_engine/multimodal_runtime.py`. +- **Pipeline parallel:** pp2 x tp4 runs with or without the DSpark drafter (`--enforce-eager`: PP decode has no CUDA graphs yet). The drafter sits on the last stage; `PPAuxRelay` (`atom/spec_decode/pp_aux_relay.py`) brings it the aux hidden states captured on earlier stages. ### Kimi-K3.1 (`KimiK31ForConditionalGeneration`) diff --git a/recipes/Kimi-K3.md b/recipes/Kimi-K3.md index 5b3def1b97..a9d17376fe 100644 --- a/recipes/Kimi-K3.md +++ b/recipes/Kimi-K3.md @@ -180,6 +180,30 @@ change only `ATOM_MONO_ENABLE` between `0` and `1`. Pin `--spec-decode-acceptance-length` on both arms, or the two runs do different amounts of work and the comparison means nothing. +### Pipeline parallel (PP2 × TP4) with DSpark + +```bash +python -m atom.entrypoints.openai_server \ + --model Kimi-K3 \ + --kv_cache_dtype fp8 -tp 4 --pipeline-parallel-size 2 \ + --enforce-eager \ + --trust-remote-code \ + --gpu-memory-utilization 0.9 \ + --block-size 128 \ + --enable_prefix_caching \ + --method dspark \ + --draft-model "${DRAFT_MODEL_PATH}" \ + --num-speculative-tokens 7 \ + --max-num-seqs 64 \ + --online_quant_config '{"global_quant_config": "ptpc_fp8", "exclude_layer": ["lm_head", "model.embed_tokens", "*self_attn.[qkv]_conv1d*", "*block_sparse_moe.experts*", "*block_sparse_moe.routed_expert_*", "*vision_tower*", "*mm_projector*", "context_proj", "layers"]}' +``` + +`--enforce-eager` is required: PP decode has no CUDA graphs yet. The drafter +lives on the last stage, and earlier stages relay the target hidden states it +reads (`atom/spec_decode/pp_aux_relay.py`). `"context_proj"` and `"layers"` are +the draft's own unprefixed module names; they keep the draft out of online +quant without touching the target. This is the CI case `Kimi-K3 DSpark PP2`. + --- ## Accuracy test diff --git a/tests/attentions/deepseek_v41/test_metadata.py b/tests/attentions/deepseek_v41/test_metadata.py index 2162e3d2a0..825756ac4c 100644 --- a/tests/attentions/deepseek_v41/test_metadata.py +++ b/tests/attentions/deepseek_v41/test_metadata.py @@ -524,7 +524,9 @@ def test_parent_preparation_preserves_phase(device, phase, position, length): builder.cache = PagedAttentionCache(geo, 4, 1, device, max_tokens=8) builder.model_runner = SimpleNamespace( forward_vars=metadata_buffers(1, 8, 4, device, geo), - tokenID_processor=SimpleNamespace(num_rejected=None), + tokenID_processor=SimpleNamespace( + num_rejected=None, verify_ctx_lens=lambda ctx: ctx + ), ) batch = SimpleNamespace( is_dummy_run=False, diff --git a/tests/models/deepseek_v41/test_dspark_integration.py b/tests/models/deepseek_v41/test_dspark_integration.py index 6643de1e39..a0d053f917 100644 --- a/tests/models/deepseek_v41/test_dspark_integration.py +++ b/tests/models/deepseek_v41/test_dspark_integration.py @@ -201,9 +201,13 @@ def capture_fixture(monkeypatch): proposer = DSparkProposer.__new__(DSparkProposer) proposer.model = draft proposer.speculative_config = SimpleNamespace( - draft_model_hf_config=SimpleNamespace(dspark_target_layer_ids=(0,)) + draft_model_hf_config=SimpleNamespace(dspark_target_layer_ids=(0,)), + use_dspark=lambda: True, + ) + proposer.config = SimpleNamespace( + hf_config=SimpleNamespace(hidden_size=8), + speculative_config=proposer.speculative_config, ) - proposer.config = SimpleNamespace(hf_config=SimpleNamespace(hidden_size=8)) proposer.max_num_tokens, proposer.device, proposer.dtype = 9, "cpu", torch.float32 target = nn.Module() target.layers = nn.ModuleList([ShiftBlock()]) @@ -236,7 +240,7 @@ def test_input_capture_uses_stream_mean_and_respects_offset_and_draft(monkeypatc def test_output_capture_remains_default(monkeypatch): proposer, target, _ = capture_fixture(monkeypatch) - proposer._aux_capture_spec = lambda _: AuxCaptureSpec( + proposer._aux_capture_spec = lambda: AuxCaptureSpec( (0,), 8, lambda output, block: output.residual.mean(-2).squeeze(0) ) proposer.arm_aux_capture(target) @@ -269,7 +273,11 @@ def test_decode_positions_use_accepted_prefix_and_full_reservation(): builder.cache = PagedAttentionCache(builder.geometry, 20, 5, "cpu") builder.block_size, builder.device = 16, "cpu" builder.model_runner = SimpleNamespace( - tokenID_processor=SimpleNamespace(num_rejected=np.array([0, 4])), + # Deferred output (pp == 1): the verify window's shift is num_rejected. + tokenID_processor=SimpleNamespace( + num_rejected=np.array([0, 4]), + verify_ctx_lens=lambda ctx: ctx - np.array([0, 4]), + ), forward_vars=metadata_buffers(2, 4, 10, geometry=builder.geometry), ) batch = SimpleNamespace( diff --git a/tests/models/test_kimi_k3_mono_max_rows.py b/tests/models/test_kimi_k3_mono_max_rows.py index 79907d3a8a..272328a7c8 100644 --- a/tests/models/test_kimi_k3_mono_max_rows.py +++ b/tests/models/test_kimi_k3_mono_max_rows.py @@ -31,6 +31,7 @@ def _metadata(requests: int, qlen: int): spec_state_indices_tensor=torch.zeros(requests, qlen, dtype=torch.int32), num_accepted_tokens=torch.ones(requests, dtype=torch.int32), replayssm=False, + has_state_fork=False, ) diff --git a/tests/models/test_kimi_k3_mono_plan.py b/tests/models/test_kimi_k3_mono_plan.py index a175b0a56f..d8a6d96958 100644 --- a/tests/models/test_kimi_k3_mono_plan.py +++ b/tests/models/test_kimi_k3_mono_plan.py @@ -88,6 +88,7 @@ def _metadata(**over): "spec_state_indices_tensor": torch.zeros(2, 4, dtype=torch.int32), "num_accepted_tokens": torch.ones(2, dtype=torch.int32), "replayssm": False, + "has_state_fork": False, } base.update(over) return SimpleNamespace(**base) @@ -109,6 +110,7 @@ def test_step_rows_accepts_a_uniform_spec_verify(): {"num_accepted_tokens": None}, {"num_actual_tokens": 7}, # disagrees with the rows {"spec_state_indices_tensor": torch.zeros(2, 4, dtype=torch.int64)}, + {"has_state_fork": True}, # K1 has no read-side slot index ], ) def test_step_rows_refuses_every_other_shape(over): @@ -126,6 +128,7 @@ def _plain_metadata(requests=3, **over): "num_spec_decodes": 0, "num_actual_tokens": requests, "non_spec_state_indices_tensor": torch.arange(requests, dtype=torch.int32), + "has_state_fork": False, } base.update(over) return SimpleNamespace(**base) @@ -148,6 +151,7 @@ def test_step_rows_accepts_a_plain_decode_as_one_token_a_request(): {"non_spec_state_indices_tensor": torch.arange(3, dtype=torch.int64)}, # a strided slot table: K1 reads it at a stride of one {"non_spec_state_indices_tensor": torch.arange(6, dtype=torch.int32)[::2]}, + {"has_state_fork": True}, ], ) def test_step_rows_refuses_every_other_plain_shape(over): diff --git a/tests/test_aiter_flydsl_gdn.py b/tests/test_aiter_flydsl_gdn.py index c187b749fa..53992c42bd 100644 --- a/tests/test_aiter_flydsl_gdn.py +++ b/tests/test_aiter_flydsl_gdn.py @@ -437,6 +437,8 @@ def test_cache_policy_forward_dispatch(monkeypatch, mode): spec_state_indices_tensor=None, non_spec_state_indices_tensor=idx, non_spec_state_indices_in_tensor=idx, + has_state_fork=False, + replayssm_fork_slots=None, num_actual_tokens=3, num_accepted_tokens=None, num_prefills=0, diff --git a/tests/test_gdn_state_fork.py b/tests/test_gdn_state_fork.py new file mode 100644 index 0000000000..e7db14cead --- /dev/null +++ b/tests/test_gdn_state_fork.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""A state fork reaching a decode forward, on both the plain and spec paths. + +Under pp=1 the checkpoint is taken between prompt chunks, so the forward that +consumes the fork is the next chunk. Under pp>1 it lands after the prompt, so +the fork falls on the first decode -- and with DSpark on, that decode is a +spec-verify step. `prepare_state_indices` is the host half of both: it has to +publish the slot the state is READ from, whichever path is running. + +`prepare_state_indices` reads a handful of attributes off the builder, so it +is driven here unbound over a stand-in rather than through a real one. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +pytest.importorskip( + "atom.model_ops.attentions.gdn_attn", + reason="requires AITER to import the attention backends", + exc_type=ImportError, +) + +import torch + +from atom.model_ops.attentions.gdn_attn import ( + PAD_SLOT_ID, + GDNStateMixin, + fork_read_table, +) +from atom.model_ops.fla_ops.replayssm import seed_forked_slots + +WIDTH = 4 # mtp_k + 1 reserved slots per request + + +def _builder(rows=8): + def buf(shape): + return SimpleNamespace(np=np.full(shape, -7, dtype=np.int32)) + + return SimpleNamespace( + _check_state_indices_writable=lambda: None, + non_spec_state_indices_tensor=buf(rows), + non_spec_state_indices_in_tensor=buf(rows), + spec_state_indices_tensor=buf((rows, WIDTH)), + replayssm=False, + spec_honours_state_fork=lambda: True, + ) + + +def _batch(fork_srcs, slots): + return SimpleNamespace(state_fork_srcs=list(fork_srcs), state_slots=list(slots)) + + +def _slots(n): + return [list(range(i * WIDTH, (i + 1) * WIDTH)) for i in range(n)] + + +def test_decode_fork_publishes_the_source_as_the_read_slot(): + b = _builder() + GDNStateMixin.prepare_state_indices(b, _batch([-1, 30], _slots(2)), with_spec=False) + assert b._has_state_fork + np.testing.assert_array_equal(b.non_spec_state_indices_tensor.np[:2], [0, 4]) + # Row 0 has no fork and reads what it writes; row 1 reads the checkpoint. + np.testing.assert_array_equal(b.non_spec_state_indices_in_tensor.np[:2], [0, 30]) + + +def test_spec_fork_publishes_the_source_too(): + b = _builder() + GDNStateMixin.prepare_state_indices(b, _batch([-1, 30], _slots(2)), with_spec=True) + assert b._has_state_fork + # The spec kernels resume from column 0, which is the write slot... + np.testing.assert_array_equal(b.spec_state_indices_tensor.np[:2, 0], [0, 4]) + # ...and the builder reads the source published here in its place. + np.testing.assert_array_equal(b.non_spec_state_indices_in_tensor.np[:2], [0, 30]) + # ReplaySSM parks the forked write slot's cursor and seeds it from 30. + assert b._fork_slots == [(4, 30)] + + +def test_spec_fork_reads_the_source_whatever_was_accepted(): + # The kernel resumes from column `num_accepted - 1`, so a bonus on a forked + # verify must still land on the source, not a recycled write slot. + b = _builder() + GDNStateMixin.prepare_state_indices(b, _batch([-1, 30], _slots(2)), with_spec=True) + table = fork_read_table( + torch.from_numpy(b.spec_state_indices_tensor.np[:2]), + torch.from_numpy(b.non_spec_state_indices_in_tensor.np[:2]), + ) + assert table.tolist() == [[0, 1, 2, 3], [30, 30, 30, 30]] + + +def test_spec_fork_is_refused_unless_the_backend_opts_in(): + b = _builder() + b.spec_honours_state_fork = lambda: False + with pytest.raises(AssertionError): + GDNStateMixin.prepare_state_indices( + b, _batch([-1, 30], _slots(2)), with_spec=True + ) + GDNStateMixin.prepare_state_indices(b, _batch([-1, 30], _slots(2)), with_spec=False) + + +def test_no_fork_reads_the_slot_it_writes_on_both_paths(): + for with_spec in (False, True): + b = _builder() + GDNStateMixin.prepare_state_indices( + b, _batch([-1, -1], _slots(2)), with_spec=with_spec + ) + assert not b._has_state_fork + np.testing.assert_array_equal(b.non_spec_state_indices_in_tensor.np[:2], [0, 4]) + + +def test_rows_past_the_batch_are_padded_on_every_table(): + b = _builder(rows=8) + GDNStateMixin.prepare_state_indices(b, _batch([-1], _slots(1)), with_spec=True) + assert (b.non_spec_state_indices_tensor.np[1:] == PAD_SLOT_ID).all() + assert (b.non_spec_state_indices_in_tensor.np[1:] == PAD_SLOT_ID).all() + assert (b.spec_state_indices_tensor.np[1:] == PAD_SLOT_ID).all() + + +def test_replayssm_seeds_only_the_forked_write_slots(): + ckpt = torch.arange(6.0).view(6, 1) + seed_forked_slots(ckpt, torch.tensor([[4], [2]])) + torch.testing.assert_close(ckpt.view(-1), torch.tensor([0.0, 1, 2, 3, 2, 5])) diff --git a/tests/test_h2d_attention_publication.py b/tests/test_h2d_attention_publication.py index a91ce6cce0..705bcaffd5 100644 --- a/tests/test_h2d_attention_publication.py +++ b/tests/test_h2d_attention_publication.py @@ -85,7 +85,9 @@ def make_runner( eos_token_id=2, ), ) - runner.tokenID_processor = SimpleNamespace(num_rejected=None) + runner.tokenID_processor = SimpleNamespace( + num_rejected=None, verify_ctx_lens=lambda ctx: ctx + ) runner.forward_vars = { name: CpuGpuBuffer(64, dtype=torch.int32, device=runner.device) for name in ("input_ids", "decode_src", "batch_info") diff --git a/tests/test_h2d_runner_publication.py b/tests/test_h2d_runner_publication.py index bb8b8ce3fd..ff2c9aa399 100644 --- a/tests/test_h2d_runner_publication.py +++ b/tests/test_h2d_runner_publication.py @@ -38,6 +38,7 @@ def runner_with_buffers( hf_config=SimpleNamespace(hidden_size=8), torch_dtype=torch.float32, ) + runner.attn_family = SimpleNamespace(has_recurrent_state=False) runner.model = SimpleNamespace() runner.use_mrope = True runner.enforce_eager = True diff --git a/tests/test_lmcache_mp_builder_views.py b/tests/test_lmcache_mp_builder_views.py index 63af70a564..1fc7da22e7 100644 --- a/tests/test_lmcache_mp_builder_views.py +++ b/tests/test_lmcache_mp_builder_views.py @@ -153,6 +153,7 @@ def noop(*args, **kwargs): MLAAttention=type("MLAAttention", (), {}), mla_dcp_kernel_num_heads=noop, mla_dcp_sparse_prefill_num_heads=noop, + mla_kernel_num_heads=lambda heads: max(16, -(-heads // 16) * 16), ), "atom.model_ops.glm5_next.geometry": _module( "atom.model_ops.glm5_next.geometry", diff --git a/tests/test_mooncake_store_offload.py b/tests/test_mooncake_store_offload.py index 5e5612169a..26d6b9361a 100644 --- a/tests/test_mooncake_store_offload.py +++ b/tests/test_mooncake_store_offload.py @@ -2915,6 +2915,7 @@ def call_func(name, *args, **_kwargs): for name, owner in ( ("_dispatch_connector_only_batch", PPEngineCoreProc), ("_dispatch_idle_offload_work", PPEngineCoreProc), + ("_queue_or_flush_prefill_hashes", PPEngineCoreProc), ("_advance_idle_kv_transfer", EngineCore), ): setattr(head, name, getattr(owner, name).__get__(head)) diff --git a/tests/test_pcp_ops.py b/tests/test_pcp_ops.py index b3f6b02bad..58dccd92ee 100644 --- a/tests/test_pcp_ops.py +++ b/tests/test_pcp_ops.py @@ -268,7 +268,9 @@ def fake_gather(x, pcp_size): monkeypatch.setattr(dr, "get_pcp_world_size", lambda: ws) monkeypatch.setattr(dr, "pcp_allgather_rerange", fake_gather) monkeypatch.setattr(dr, "get_forward_context", lambda: ctx) - drafter = SimpleNamespace(_captures_aux=True, _aux_buffers=[torch.ones(16, 2)]) + drafter = SimpleNamespace( + _captures_aux=True, _aux_buffers=[torch.ones(16, 2)], _relayed_aux=None + ) hidden = torch.zeros(n, 2) (aux,) = dr.Drafter.aux_for(drafter, hidden) diff --git a/tests/test_pd_chunked_transfer.py b/tests/test_pd_chunked_transfer.py index e211b698ef..7180c5ee66 100644 --- a/tests/test_pd_chunked_transfer.py +++ b/tests/test_pd_chunked_transfer.py @@ -524,7 +524,7 @@ def test_real_scheduler_pins_from_first_chunk_through_final_handoff( publish(state, int(batch.context_lens[0])) result = ScheduledBatchOutput( [seq.id], - [(42, 43, 44, 45) if sched.spec_decode_local else (42,)], + [(42, 43, 44, 45) if sched.use_spec else (42,)], np.array([0]), np.array([0]), np.array([[43, 44, 45]]) if draft_tokens else None, diff --git a/tests/test_pd_pp.py b/tests/test_pd_pp.py index a74ad0a74c..fdca7f1a21 100644 --- a/tests/test_pd_pp.py +++ b/tests/test_pd_pp.py @@ -1533,6 +1533,9 @@ def _fake_head(batch): head._dispatch_connector_only_batch = ( PPEngineCoreProc._dispatch_connector_only_batch.__get__(head) ) + head._queue_or_flush_prefill_hashes = ( + PPEngineCoreProc._queue_or_flush_prefill_hashes.__get__(head) + ) PPEngineCoreProc._pp_head_step(head) return head @@ -1720,6 +1723,9 @@ def _head_for_batch(batch, *, dynamic_chunking_active=False): head.scheduler.take_rejected.return_value = None head.scheduler.dynamic_chunking_active = dynamic_chunking_active head._dispatch_connector_only_batch = MagicMock() + head._queue_or_flush_prefill_hashes = ( + PPEngineCoreProc._queue_or_flush_prefill_hashes.__get__(head) + ) return PPEngineCoreProc, head diff --git a/tests/test_pp.py b/tests/test_pp.py index 18579df465..e219e0458c 100644 --- a/tests/test_pp.py +++ b/tests/test_pp.py @@ -364,6 +364,24 @@ def test_async_send_full_tensor_when_not_divisible(monkeypatch): assert sent[2:] == [35] +def test_async_send_hands_owned_tensors_over_uncloned(monkeypatch): + _fake_async_pp_group(monkeypatch, rank_in_group=0, world_size=2) + monkeypatch.setattr(pp_comm.envs, "ATOM_PP_SEND_ALLGATHER", False) + sent = [] + monkeypatch.setattr( + pp_comm.torch.distributed, + "isend", + lambda buf, dst, group: sent.append(buf) or MagicMock(), + ) + hidden, aux = torch.randn(4, 8), torch.randn(4, 16) + pp_comm.async_send_intermediate_tensors( + IntermediateTensors({"hidden_states": hidden, pp_comm.PP_AUX_KEY: aux}), + owned=(pp_comm.PP_AUX_KEY,), + ) + assert sent[2] is not hidden and torch.equal(sent[2], hidden) + assert sent[3] is aux + + # --------------------------------------------------------------------------- # Inter-stage transport (PPStageTransport ZMQ) # --------------------------------------------------------------------------- @@ -1143,3 +1161,153 @@ def test_dynamic_pp_prefill_registers_prefix_hashes(): assert seq.num_cached_tokens == prompt hit = sched.block_manager.can_allocate(_make_seq(prompt, block_size=bs)) assert hit == prompt // bs - 1 + + +# --- _checkpoint_holds: skipped while held, released when its batch files -- + + +def _held_after_first_chunk(prompt=400, bs=16): + """A PP seq with one middle chunk in flight, held as a rung would hold it.""" + # Prefix caching on: the hold is only ever taken with it (`_partial_prefill_chunk`). + sched = Scheduler( + _pp_config( + max_num_batched_tokens=128, + kv_cache_block_size=bs, + enable_prefix_caching=True, + ) + ) + seq = _make_seq(prompt, block_size=bs) + sched.add(seq) + batch, seqs = sched.schedule() + assert seq.id in batch.req_ids and seq.is_partial_prefill + sched._checkpoint_holds[seq.id] = _chunk_end(batch, seq.id) + return sched, seq, batch, seqs + + +def _chunk_end(batch, req_id): + i = list(batch.req_ids).index(req_id) + return int(batch.num_cached_tokens[i]) + int(batch.num_scheduled_tokens[i]) + + +def _prefill_ids(sched): + res = sched.schedule() + if res is None: + return set() + batch, _ = res + return set(batch.req_ids[: batch.total_seqs_num_prefill]) + + +def test_held_seq_is_skipped_by_both_prefill_scans(): + sched, seq, _, _ = _held_after_first_chunk() + pending, _ = sched._local_prefill_pending_work() + assert not pending, "the probe counted a held chunk as pending work" + assert seq.id not in _prefill_ids(sched) + assert seq.id in sched._checkpoint_holds + + +def test_hold_is_released_when_its_batch_files_its_hashes(): + sched, seq, batch, seqs = _held_after_first_chunk() + assert sched.carries_held_checkpoint(batch) + sched.register_prefill_hashes(batch, seqs.values()) + assert seq.id not in sched._checkpoint_holds + assert not sched.carries_held_checkpoint(batch) + assert seq.id in _prefill_ids(sched) + + +def test_hold_is_released_only_by_the_chunk_that_holds_it(): + # An earlier chunk of the same seq, in flight when the held one was + # scheduled, files first: the held chunk's checkpoint is not filed yet. + sched, seq, first, first_seqs = _held_after_first_chunk() + del sched._checkpoint_holds[seq.id] + held, held_seqs = sched.schedule() + assert seq.id in held.req_ids and seq.is_partial_prefill + sched._checkpoint_holds[seq.id] = _chunk_end(held, seq.id) + assert not sched.carries_held_checkpoint(first) + sched.register_prefill_hashes(first, first_seqs.values()) + assert seq.id in sched._checkpoint_holds + assert seq.id not in _prefill_ids(sched) + assert sched.carries_held_checkpoint(held) + sched.register_prefill_hashes(held, held_seqs.values()) + assert seq.id not in sched._checkpoint_holds + + +@pytest.mark.parametrize("acked", [False, True]) +def test_head_flushes_a_held_chunk_whether_or_not_it_is_acked(acked): + # Dynamic chunking acks middle chunks: that branch must flush a hold too, + # or the held seq never re-enters the prefill scan and stalls. + import queue + + from aiter_stub import stubbed_aiter + + with stubbed_aiter(): + from atom.model_engine.pp_engine_core import PPEngineCoreProc + + sched, seq, batch, seqs = _held_after_first_chunk() + batch.pp_completion_ack = acked + sched.mark_pp_inflight(batch) + core = PPEngineCoreProc.__new__(PPEngineCoreProc) + core.scheduler = sched + core.pp_size = 1 # the in-flight chunk fills the window: schedule nothing + core._in_flight = deque([(batch, seqs, False)]) + core._pending_prefix_hash = deque() + core._defer_prefix_hash = True + core.output_queue = queue.Queue() + core._poll_kv_transfer_progress = lambda: None + core._poll_dynamic_chunking_calibration = lambda: None + core.runner_mgr = types.SimpleNamespace(call_func=lambda *a, **k: None) + completions = deque([(batch.req_ids, None)] if acked else []) + core.pp_transport = types.SimpleNamespace( + recv_completion=lambda **k: completions.popleft() if completions else None + ) + core._pp_head_step() + assert not core._in_flight and not core._pending_prefix_hash + assert seq.id not in sched._checkpoint_holds + assert seq.id in _prefill_ids(sched) + + +@pytest.mark.parametrize("filer", ["register_prefill_hashes", "postprocess"]) +def test_hold_is_released_even_if_the_seq_left_running(filer): + # Preempted, aborted or retired while its chunk was in flight: the batch + # still files, and the hold must go with it or the id is held forever. + import numpy as np + + sched, seq, batch, _ = _held_after_first_chunk() + sched.running = deque(s for s in sched.running if s.id != seq.id) + if filer == "register_prefill_hashes": + sched.register_prefill_hashes(batch, []) + else: + n = len(batch.req_ids) + out = ScheduledBatchOutput( + req_ids=list(batch.req_ids), + token_ids=[(0,)] * n, + num_rejected=np.zeros(n, dtype=np.int32), + num_bonus=np.zeros(n, dtype=np.int32), + draft_token_ids=None, + is_deferred_out=False, + ) + sched.postprocess([], out, batch=batch) + assert seq.id not in sched._checkpoint_holds + + +def test_hold_is_released_even_if_hashing_raises(): + # The head loop logs a failed filing and carries on; the seqs after the + # one that raised must not stay out of the prefill scan forever. + sched = Scheduler( + _pp_config( + max_num_batched_tokens=512, + kv_cache_block_size=16, + enable_prefix_caching=True, + ) + ) + seqs = [_make_seq(400), _make_seq(400)] + for seq in seqs: + sched.add(seq) + batch, by_id = sched.schedule() + assert all(seq.id in batch.req_ids for seq in seqs) + sched._checkpoint_holds.update( + {req_id: _chunk_end(batch, req_id) for req_id in batch.req_ids} + ) + sched.block_manager.hash_blocks = MagicMock(side_effect=RuntimeError("boom")) + with pytest.raises(RuntimeError, match="boom"): + sched.register_prefill_hashes(batch, by_id.values()) + assert not sched._checkpoint_holds diff --git a/tests/test_pp_aux_relay.py b/tests/test_pp_aux_relay.py new file mode 100644 index 0000000000..0528561cd3 --- /dev/null +++ b/tests/test_pp_aux_relay.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Which PP stage owns which aux tap, and what each stage hooks. + +Not CPU-only: `pp_aux_relay` imports `spec_decode.drafter`, which imports the +model loader and so the real aiter. Stubbing that chain would mean stubbing +the loader's whole quant surface, so this file skips on the non-GPU runner. +""" + +import contextlib +import types +from unittest import mock + +import pytest +import torch +from torch import nn + +pytest.importorskip( + "atom.spec_decode.pp_aux_relay", + reason="requires AITER to import the drafter chain", + exc_type=ImportError, +) + +from atom.models.utils import PPMissingLayer, get_pp_indices +from atom.spec_decode.drafter import AuxCaptureSpec +from atom.spec_decode.pp_aux_relay import PPAuxRelay + + +class _FakeTarget(nn.Module): + """A PP slice of a target: layers outside the stage are `PPMissingLayer`.""" + + def __init__(self, n_layers, resident_range, with_embed): + super().__init__() + lo, hi = resident_range + self.layers = nn.ModuleList( + nn.Identity() if lo <= i < hi else PPMissingLayer() for i in range(n_layers) + ) + # Like the real models: off the first stage the embedding is missing. + self.embed_tokens = nn.Identity() if with_embed else PPMissingLayer() + + +_CTX = types.SimpleNamespace(is_draft=False, ubatch_token_offset=0) + + +@pytest.fixture(autouse=True) +def _forward_context(): + _CTX.is_draft, _CTX.ubatch_token_offset = False, 0 + fc = types.SimpleNamespace(context=_CTX) + with mock.patch("atom.spec_decode.drafter.get_forward_context", return_value=fc): + yield _CTX + + +def _pp_group(group): + # The relay reads the stage; the shared residency rule reads it for `-1`. + stack = contextlib.ExitStack() + for mod in ("pp_aux_relay", "drafter"): + stack.enter_context( + mock.patch(f"atom.spec_decode.{mod}.get_pp_group", return_value=group) + ) + return stack + + +def _relay(layer_ids, n_layers, resident_range, *, first, last, hidden=8): + target = _FakeTarget(n_layers, resident_range, with_embed=first) + spec = AuxCaptureSpec( + layer_ids=tuple(layer_ids), hidden_size=hidden, extract=lambda out, mod: out + ) + group = types.SimpleNamespace(is_first_rank=first, is_last_rank=last) + with _pp_group(group): + relay = PPAuxRelay(spec, target) + relay.arm(target) + return relay, target + + +def test_relay_embedding_tap_hooks_the_embedding_not_the_last_layer(): + # Regression: `layers[-1]` is the LAST decoder layer, a PPMissingLayer on + # the first stage, so the hook silently never fired and the relayed buffer + # stayed zero. + relay, target = _relay([-1, 2], 8, (0, 4), first=True, last=False) + assert relay.own_ids == (-1, 2) + assert target.embed_tokens._forward_hooks, "embedding tap was not hooked" + assert not target.layers[-1]._forward_hooks + assert not target.layers[-1]._forward_pre_hooks + + +def _run_target(relay, target, n, hidden): + """A target forward on this stage: each own tap's module outputs `lid`.""" + for lid in relay.own_ids: + module = target.embed_tokens if lid == -1 else target.layers[lid] + module(torch.full((n, hidden), float(lid))) + + +def test_relay_embedding_tap_captures_through_the_hook(): + relay, target = _relay([-1, 2], 8, (0, 4), first=True, last=False) + rows = torch.arange(2 * 8, dtype=torch.float32).reshape(2, 8) + target.embed_tokens(rows) + target.layers[2](torch.zeros(2, 8)) + torch.testing.assert_close(relay.pack(None, 2)[:, :8], rows) + + +def test_relay_embedding_tap_ignores_capture_input(): + # -1 taps the embedding OUTPUT even when the spec captures layer inputs. + _, target = _relay([-1, 2], 8, (0, 4), first=True, last=False) + assert not target.embed_tokens._forward_pre_hooks + + +def test_relay_embedding_tap_is_incoming_on_a_later_stage(): + relay, target = _relay([-1, 2, 6], 8, (4, 8), first=False, last=False) + assert relay.own_ids == (6,) + assert relay.incoming_ids == (-1, 2) + assert relay.outgoing_ids == (-1, 2, 6) + assert not target.embed_tokens._forward_hooks + + +@pytest.mark.parametrize("pp_size", [2, 4, 8]) +def test_relay_tap_split_partitions_the_layer_ids(pp_size): + layer_ids = (-1, 2, 23, 47, 71, 89) + n_layers, seen = 93, [] + for rank in range(pp_size): + lo, hi = get_pp_indices(n_layers, rank, pp_size) + relay, _ = _relay( + layer_ids, + n_layers, + (lo, hi), + first=rank == 0, + last=rank == pp_size - 1, + ) + seen.extend(relay.own_ids) + # What a stage forwards is exactly what it was handed plus its own. + expected = () if rank == pp_size - 1 else relay.incoming_ids + relay.own_ids + assert relay.outgoing_ids == expected + assert tuple(seen) == layer_ids, "every tap is owned by exactly one stage" + + +@pytest.mark.parametrize("pp_size", [2, 4, 8]) +@pytest.mark.parametrize("layer_ids", [(-1, 2, 23, 47, 71, 89), (-1, 2, 89)]) +def test_relay_round_trip_hands_the_drafter_each_tap_in_order(pp_size, layer_ids): + # Stages that own no tap must still hand the last stage each tap's rows, + # in `layer_ids` order. + n_layers, hidden, n = 93, 8, 3 + payload, drafter = None, mock.Mock() + for rank in range(pp_size): + lo, hi = get_pp_indices(n_layers, rank, pp_size) + last = rank == pp_size - 1 + relay, target = _relay( + layer_ids, n_layers, (lo, hi), first=rank == 0, last=last, hidden=hidden + ) + if last: + relay.absorb(payload, drafter) + break + _run_target(relay, target, n, hidden) + payload = relay.pack(payload, n) + assert set(relay.incoming_ids) == {lid for lid in layer_ids if lid < lo} + if not relay.incoming_ids: + drafter.take_relayed_aux.assert_not_called() + return + (block,) = drafter.take_relayed_aux.call_args.args + for col, lid in enumerate(relay.incoming_ids): + torch.testing.assert_close( + block[:, col * hidden : (col + 1) * hidden], + torch.full((n, hidden), float(lid)), + ) + + +def test_aux_for_serves_the_relayed_block_uncopied_before_own_taps(): + from atom.spec_decode.drafter import Drafter + + hidden, n = 8, 3 + block = torch.randn(4, 2 * hidden) + own = torch.randn(4, hidden) + drafter = types.SimpleNamespace( + _captures_aux=True, + _aux_buffers=[torch.zeros(4, hidden), torch.zeros(4, hidden), own], + _relayed_aux=block, + ) + with mock.patch("atom.spec_decode.drafter.get_pcp_world_size", return_value=1): + aux = Drafter.aux_for(drafter, torch.zeros(n, hidden)) + assert len(aux) == 2 + assert aux[0].data_ptr() == block.data_ptr() and aux[0].shape == (n, 2 * hidden) + torch.testing.assert_close(aux[1], own[:n]) + + +def test_last_stage_arms_only_its_own_taps(): + from atom.spec_decode.drafter import Drafter + + hidden = 8 + # The last stage holds an embedding copy for its drafter, yet the `-1` + # tap still arrives through the relay. + target = _FakeTarget(8, (4, 8), with_embed=True) + spec = AuxCaptureSpec( + layer_ids=(-1, 2, 5), hidden_size=hidden, extract=lambda out, mod: out + ) + drafter = types.SimpleNamespace( + _aux_capture_spec=lambda: spec, + max_num_tokens=4, + device="cpu", + dtype=torch.float32, + ) + drafter._make_aux_hook = lambda i, e: Drafter._make_aux_hook(drafter, i, e) + group = types.SimpleNamespace(is_first_rank=False, is_last_rank=True) + with _pp_group(group): + Drafter.arm_aux_capture(drafter, target) + assert [b.shape[0] for b in drafter._aux_buffers] == [0, 0, 4] + assert not target.embed_tokens._forward_hooks + assert not target.layers[2]._forward_hooks + assert target.layers[5]._forward_hooks + + +@pytest.mark.parametrize("layer_ids", [(47, -1, 89, 2), (-1, 2, 2), (-1, 2, 93)]) +def test_relay_rejects_the_ids_the_drafter_rejects(layer_ids): + with pytest.raises(ValueError, match="strictly increasing|out of range"): + _relay(layer_ids, 93, (0, 47), first=True, last=False) + + +def test_relay_skips_draft_forwards(_forward_context): + relay, target = _relay([-1, 2], 8, (0, 4), first=True, last=False) + _forward_context.is_draft = True + _run_target(relay, target, 3, 8) + assert relay._send is None + + +def test_relay_rejects_a_payload_of_the_wrong_width(): + relay, _ = _relay([-1, 2, 6], 8, (4, 8), first=False, last=True) + with pytest.raises(RuntimeError, match="expected 16 aux columns"): + relay.absorb(torch.zeros(3, 8), mock.Mock()) + + +def test_relay_send_is_sized_to_the_forward_rows(): + # Nothing resident: each forward's first own tap allocates its own rows. + relay, target = _relay([-1, 2, 6], 8, (4, 8), first=False, last=False) + assert relay._send is None + for n in (3, 5, 2): + _run_target(relay, target, n, 8) + sent = relay.pack(torch.ones(n, 16), n) + assert sent.shape == (n, 24) + torch.testing.assert_close(sent[:, :16], torch.ones(n, 16)) + torch.testing.assert_close(sent[:, 16:], torch.full((n, 8), 6.0)) + assert relay._send is None, "pack must hand the buffer over" + + +@pytest.mark.parametrize("world_size", [1, 2]) +def test_build_takes_the_drafter_class_spec_and_arms_only_under_pp(world_size): + target = _FakeTarget(8, (0, 4), with_embed=True) + spec = AuxCaptureSpec(layer_ids=(-1, 2), hidden_size=8, extract=lambda o, m: o) + # Every stage, drafter or not, takes the spec from the drafter class. + drafter_cls = types.SimpleNamespace(aux_capture_spec=lambda config: spec) + config = types.SimpleNamespace(speculative_config=object()) + group = types.SimpleNamespace( + world_size=world_size, is_first_rank=True, is_last_rank=False + ) + with ( + _pp_group(group), + mock.patch( + "atom.spec_decode.pp_aux_relay.drafter_class", return_value=drafter_cls + ), + ): + relay = PPAuxRelay.build(config, target) + if world_size == 1: + assert relay is None + return + assert relay.spec is spec and relay.own_ids == (-1, 2) + assert target.embed_tokens._forward_hooks, "build must arm the relay" + + +def test_relay_hook_compiles_without_an_in_place_update(): + # Regression: `_filled.add` inside the compiled backbone made Dynamo replay + # the set via `update`, which trips the cudagraph bytecode hook's depyf + # import on images without depyf (PP2 server warmup crashed). + from torch._dynamo.convert_frame import register_bytecode_hook + + names = [] + handle = register_bytecode_hook(lambda old, new: names.extend(new.co_names)) + try: + relay, target = _relay([-1, 2], 8, (0, 4), first=True, last=False) + fwd = torch.compile( + lambda x: target.layers[2](target.embed_tokens(x) + 1), backend="eager" + ) + for _ in range(2): + fwd(torch.randn(3, 8)) + assert relay.pack(None, 3).shape == (3, 16) + finally: + handle.remove() + torch._dynamo.reset() + assert "update" not in names diff --git a/tests/test_pp_config_refusals.py b/tests/test_pp_config_refusals.py new file mode 100644 index 0000000000..39cfdef68a --- /dev/null +++ b/tests/test_pp_config_refusals.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""What `Config` refuses under PP: TBO, P/D consumers, and speculation an +architecture does not declare in `registry.pp_speculative_methods`.""" + +from types import SimpleNamespace +from unittest import mock + +import pytest + +from atom.config import Config +from atom.models.registry import pp_speculative_methods + + +def _cfg(pp=2, arch="KimiK3ForConditionalGeneration", method="dspark", **kw): + spec = SimpleNamespace(method=method, use_dspark=lambda: method == "dspark") + return SimpleNamespace( + pipeline_parallel_size=pp, + hf_config=SimpleNamespace(architectures=[arch]), + speculative_config=kw.pop("spec", spec), + dspark=SimpleNamespace(confidence_schedule=kw.pop("confidence", False)), + kv_transfer_config=kw.pop("kv", None), + enable_tbo=kw.pop("tbo", False), + ) + + +def _check(cfg): + Config._check_pp_support(cfg) + + +def test_kimi_k3_dspark_under_pp_is_allowed(): + _check(_cfg()) + + +def test_pp1_is_never_refused(): + _check(_cfg(pp=1, arch="Qwen3NextForCausalLM", method="mtp")) + + +def test_pp_without_speculation_is_allowed_for_any_model(): + _check(_cfg(arch="Qwen3NextForCausalLM", spec=None)) + + +@pytest.mark.parametrize( + "arch, method", + [ + ("KimiK3ForConditionalGeneration", "eagle3"), + ("Qwen3NextForCausalLM", "mtp"), + ("MiniMaxM3ForCausalLM", "dspark"), + ], +) +def test_unwired_spec_pairs_under_pp_are_refused(arch, method): + with pytest.raises(ValueError, match="pp_speculative_methods declares"): + _check(_cfg(arch=arch, method=method)) + + +def test_a_declared_architecture_is_allowed_without_a_config_edit(): + with mock.patch.dict(pp_speculative_methods, {"Qwen3NextForCausalLM": ("mtp",)}): + _check(_cfg(arch="Qwen3NextForCausalLM", method="mtp")) + + +def test_confidence_schedule_under_pp_is_refused(): + with pytest.raises(ValueError, match="confidence_schedule"): + _check(_cfg(confidence=True)) + + +def test_tbo_under_pp_is_refused(): + with pytest.raises(ValueError, match="TBO"): + _check(_cfg(spec=None, tbo=True)) + + +def test_kv_consumer_under_pp_is_refused(): + kv = {"kv_connector": "mooncake", "kv_role": "kv_consumer"} + with pytest.raises(ValueError, match="kv_consumer"): + _check(_cfg(spec=None, kv=kv)) + + +def test_kv_producer_under_pp_drafts_with_any_method(): + # P/D prefill with `--method mtp` on PP (GLM-5.2's Atomesh cells): the + # producer never speculates, so the Kimi-K3 + DSpark gate does not apply. + kv = {"kv_connector": "mooncake", "kv_role": "kv_producer"} + _check(_cfg(arch="Glm5ForCausalLM", method="mtp", kv=kv)) diff --git a/tests/test_pp_draft_input_ids.py b/tests/test_pp_draft_input_ids.py new file mode 100644 index 0000000000..0a125886a0 --- /dev/null +++ b/tests/test_pp_draft_input_ids.py @@ -0,0 +1,45 @@ +"""PP verify: each row's drafts land right after its anchor in `input_ids`.""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +# On the module, not on "aiter" -- see tests/test_verify_ctx_lens.py. +runner_module = pytest.importorskip( + "atom.model_engine.model_runner", + reason="requires AITER to import model_runner", + exc_type=ImportError, +) + + +def test_drafts_fill_only_their_rows_columns(): + # Rows of 3, 1 (no drafts) and 2 tokens; spec rows padded to width 2. + lens = np.array([3, 1, 2], dtype=np.int32) + cu = np.array([0, 3, 4, 6], dtype=np.int32) + anchors = np.array([100, 0, 0, 101, 102, 0], dtype=np.int32) + spec = np.array([[10, 11], [-1, -1], [20, -1]], dtype=np.int32) + batch = SimpleNamespace( + scheduled_tokens=anchors, + total_tokens_num_prefill=0, + total_tokens_num_decode=6, + total_seqs_num_prefill=0, + num_spec_step=2, + scheduled_spec_decode_tokens=spec, + num_rejected=None, + num_bonus=None, + ) + staged = SimpleNamespace(np=np.full(8, -7, dtype=np.int32)) + processor = SimpleNamespace( + is_deferred_out=False, + is_pipeline_parallel=True, + is_pd_producer=False, + use_spec=True, + input_ids=staged, + runner=SimpleNamespace( + attn_metadata_builder=SimpleNamespace(decode_spans=lambda b: (3, lens, cu)) + ), + _publish_input_ids=lambda n, group: staged.np[:n].copy(), + ) + ids, _ = runner_module.tokenIDProcessor._stage_input_ids(processor, batch, 3, None) + np.testing.assert_array_equal(ids, [100, 10, 11, 101, 102, 20]) diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 226c56dea0..aaba8b6b30 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -1773,10 +1773,10 @@ def test_fewest_tokens_left_in_steps(self, scheduler, seq_factory): scheduler.running = deque([far, near]) assert scheduler._next_decode_completion_steps() == 7 scheduler.mtp_k = 3 # up to 4 tokens a step - scheduler.spec_decode_local = True + scheduler.use_spec = True assert scheduler._next_decode_completion_steps() == 2 - # drafting only (PP): this engine verifies none, a token a step - scheduler.spec_decode_local = False + # not speculating: a token a step + scheduler.use_spec = False assert scheduler._next_decode_completion_steps() == 7 diff --git a/tests/test_scheduler_metrics.py b/tests/test_scheduler_metrics.py index 74ae3c211f..9292a56172 100644 --- a/tests/test_scheduler_metrics.py +++ b/tests/test_scheduler_metrics.py @@ -568,7 +568,7 @@ def test_prefill_context_records_full_prompt_once_and_chunk_batch_totals(monkeyp _is_preemptable=lambda seq: True, total_preemptions=0, _partial_prefill_count=1, - spec_decode_local=False, + use_spec=False, _connector_release_stalled_save=lambda seq: None, block_manager=SimpleNamespace(deallocate=lambda seq: None), waiting=deque(), @@ -638,7 +638,7 @@ def test_decode_request_context_gauge_records_first_dispatch_once(monkeypatch): owner = SimpleNamespace( _is_preemptable=lambda seq: True, total_preemptions=0, - spec_decode_local=False, + use_spec=False, _connector_release_stalled_save=lambda seq: None, block_manager=SimpleNamespace(deallocate=lambda seq: None), waiting=deque(), diff --git a/tests/test_scheduler_mtp_max_tokens.py b/tests/test_scheduler_mtp_max_tokens.py index 4d9608b1fb..e44c72df33 100644 --- a/tests/test_scheduler_mtp_max_tokens.py +++ b/tests/test_scheduler_mtp_max_tokens.py @@ -49,6 +49,8 @@ def _accept_all(sched, seq, tokens, stream_queue): num_rejected=np.asarray([0]), num_bonus=np.asarray([0]), draft_token_ids=np.asarray([[20, 21, 22]]), + # pp=1 output is deferred (drafts + lagging provisional run). + is_deferred_out=True, logprobs=( { seq.id: [ @@ -170,6 +172,130 @@ def test_earlier_cap_wins_over_later_eos(self, seq_factory): assert emitted.output_tokens == [10, 11] +class TestUndeferredSpecPlaceholders: + """PP output is undeferred: every step leaves ONE run of `mtp_k + 1`. + + `verify_ctx_lens` shifts by 1 under undeferred output only because of this: + the verify window starts `mtp_k + 1` placeholders back from the end, no + matter how many drafts the previous step rejected. + """ + + mtp_k = 3 + + def _scheduler(self, **config): + sched = Scheduler( + MockConfig( + pipeline_parallel_size=2, + **config, + speculative_config=SimpleNamespace( + num_speculative_tokens=self.mtp_k, + use_dspark=lambda: True, + ), + num_kvcache_blocks=64, + ) + ) + assert not sched.is_deferred_out and sched.use_spec + return sched + + def _step(self, sched, seq, tokens, num_rejected, stream_queue=None): + batch, seqs = sched.schedule() + return sched.postprocess( + list(seqs.values()), + ScheduledBatchOutput( + req_ids=[seq.id], + token_ids=[tuple(tokens)], + num_rejected=np.asarray([num_rejected]), + num_bonus=np.asarray([0]), + draft_token_ids=np.asarray([[7] * self.mtp_k]), + is_deferred_out=False, + ), + stream_output_queue=stream_queue, + batch=batch, + ) + + def _assert_one_full_run(self, sched, seq): + width = self.mtp_k + 1 + tail = seq.token_ids[-width - 1 :] + assert list(tail[1:]) == [sched.eos_token_id] * width + assert tail[0] != sched.eos_token_id + assert seq.num_placeholder_tokens == width + + def _prefilled(self, sched, seq_factory, max_tokens=64): + seq = seq_factory( + [1, 2, 3, 4], + sampling_params=SamplingParams(max_tokens=max_tokens, ignore_eos=True), + ) + sched.add(seq) + # The prefill step scores no drafts: one sampled token. + self._step(sched, seq, [100], num_rejected=0) + self._assert_one_full_run(sched, seq) + return seq + + @pytest.mark.parametrize("num_rejected", range(mtp_k + 1)) + def test_every_reject_count_leaves_one_full_run(self, seq_factory, num_rejected): + sched = self._scheduler() + seq = self._prefilled(sched, seq_factory) + accepted = self.mtp_k + 1 - num_rejected + real_before = seq.num_tokens - seq.num_placeholder_tokens + self._step(sched, seq, range(200, 200 + accepted), num_rejected) + self._assert_one_full_run(sched, seq) + assert seq.num_tokens - seq.num_placeholder_tokens == real_before + accepted + + def test_mixed_steps_never_grow_or_shrink_the_run(self, seq_factory): + sched = self._scheduler() + seq = self._prefilled(sched, seq_factory) + for step, num_rejected in enumerate([0, self.mtp_k, 1, self.mtp_k, 2]): + accepted = self.mtp_k + 1 - num_rejected + self._step(sched, seq, range(10 * step, 10 * step + accepted), num_rejected) + self._assert_one_full_run(sched, seq) + + def test_cap_truncates_an_all_accepted_step(self, seq_factory): + # 1 token from prefill + 4 accepted against max_tokens=3. + sched = self._scheduler() + seq = self._prefilled(sched, seq_factory, max_tokens=3) + stream_queue = mock.Mock() + finished = self._step( + sched, seq, [10, 11, 12, 13], num_rejected=0, stream_queue=stream_queue + ) + assert len(finished) == 1 + assert finished[0].leave_reason == "max_tokens" + assert finished[0].num_completion_tokens == 3 + emitted = stream_queue.put_nowait.call_args.args[0][0][1] + assert emitted.output_tokens == [10, 11] + + def test_middle_chunk_gets_no_placeholders_once_its_final_chunk_is_scheduled( + self, seq_factory + ): + # pp=2: the final chunk is scheduled while the middle one is in flight, + # flipping the live `is_partial_prefill`. The middle chunk's + # postprocess must still leave the seq without placeholders. + sched = self._scheduler(max_num_batched_tokens=16) + assert sched.advance_on_schedule + seq = seq_factory( + list(range(1, 33)), + sampling_params=SamplingParams(max_tokens=64, ignore_eos=True), + ) + sched.add(seq) + middle, middle_seqs = sched.schedule() + assert not middle.is_final_chunk[0] + sched.schedule() + assert not seq.is_partial_prefill + sched.postprocess( + list(middle_seqs.values()), + ScheduledBatchOutput( + req_ids=[seq.id], + token_ids=[(5,)], + num_rejected=np.asarray([0]), + num_bonus=np.asarray([0]), + draft_token_ids=np.asarray([[7] * self.mtp_k]), + is_deferred_out=False, + ), + batch=middle, + ) + assert seq.num_placeholder_tokens == 0 + assert seq.num_tokens == 32 + + @pytest.mark.parametrize("injected", [False, True]) def test_bad_logprobs_do_not_terminate_sibling_requests(seq_factory, injected): sched = Scheduler(MockConfig(num_kvcache_blocks=64)) diff --git a/tests/test_spec_deferred_status_queue.py b/tests/test_spec_deferred_status_queue.py index ac6ec82314..88119eb19a 100644 --- a/tests/test_spec_deferred_status_queue.py +++ b/tests/test_spec_deferred_status_queue.py @@ -17,7 +17,7 @@ _kv_config_has_producer, tokenIDProcessor, ) -from atom.model_engine.scheduler import ScheduledBatch +from atom.model_engine.scheduler import ScheduledBatch, defers_output from atom.spec_decode.status import SpecStatusQueue from atom.utils.selector import Family @@ -165,6 +165,26 @@ def _processor_for(*, kv_config: dict, family: Family) -> tokenIDProcessor: return tokenIDProcessor(runner, max_num_batched_tokens=8) +@pytest.mark.parametrize("pp_size", [1, 2]) +def test_the_scheduler_resolves_the_runners_deferral(pp_size): + """The scheduler has no runner: it reads the family off the config, and + must land where the runner does for a recurrent-state producer.""" + config = SimpleNamespace( + pipeline_parallel_size=pp_size, + kv_transfer_config={"kv_role": "kv_producer"}, + hf_config=object(), + ) + with ( + mock.patch("atom.model_engine.scheduler.get_hf_text_config"), + mock.patch( + "atom.model_engine.scheduler.attn_family", return_value=Family.KIMI_MLA + ), + ): + assert defers_output(config) is False + config.kv_transfer_config = {} + assert defers_output(config) is (pp_size == 1) + + def test_a_consumer_keeps_deferred_output(): """The other polarity of the sweep below, which fixes the role to producer. diff --git a/tests/test_speculative_sampling.py b/tests/test_speculative_sampling.py index cf33af309c..7f49d31fd5 100644 --- a/tests/test_speculative_sampling.py +++ b/tests/test_speculative_sampling.py @@ -188,6 +188,9 @@ def test_target_prefixes_cover_every_ragged_acceptance_length(stochastic): assert output[i].tolist() == expected assert counts.tolist() == accepts assert rejected.tolist() == [5 - a for a in accepts] + # The undeferred runner's counts come from the row length. + emitted = (output != -1).sum(1) + assert torch.equal(emitted - 1, counts) and torch.equal(6 - emitted, rejected) # The anchor is the row the last emitted token was sampled at. starts = [sum(w + 1 for w in widths[:i]) for i in range(len(widths))] assert anchors.tolist() == [s + a for s, a in zip(starts, accepts)] diff --git a/tests/test_verify_ctx_lens.py b/tests/test_verify_ctx_lens.py new file mode 100644 index 0000000000..83b9cf0a97 --- /dev/null +++ b/tests/test_verify_ctx_lens.py @@ -0,0 +1,101 @@ +"""The verify window's positions must be the ones the scheduler staged. + +Attention builders rebuild the window from +`tokenIDProcessor.verify_ctx_lens(context_lens)`. Token ids equal their index, +so staged tokens are their true positions. +""" + +from types import SimpleNamespace + +import numpy as np +import pytest + +# On the module, not on "aiter" -- see tests/test_dspark.py for why a name-only +# guard passes on the non-GPU runner and then aborts collection. +runner_module = pytest.importorskip( + "atom.model_engine.model_runner", + reason="requires AITER to import model_runner", + exc_type=ImportError, +) +from atom.model_engine.scheduler import ScheduledBatch +from atom.model_engine.sequence import SequenceType +from atom.model_ops.attentions.token_layout.decode import decode_positions + +MTP_K = 7 +NUM = MTP_K + 1 # anchor + drafts + + +def _verify_ctx_lens(context_lens, num_rejected, is_deferred_out): + processor = SimpleNamespace( + num_rejected=num_rejected, is_deferred_out=is_deferred_out + ) + return runner_module.tokenIDProcessor.verify_ctx_lens(processor, context_lens) + + +def _decode_batch(seq_factory, length, num_rejected, is_deferred_out): + seq = seq_factory(list(range(length))) + seq.type = SequenceType.DECODE + seq.num_rejected = num_rejected + if not is_deferred_out: + # postprocess refills the undeferred tail to one full window + seq.num_placeholder_tokens = NUM + return ScheduledBatch( + {seq.id: seq}, + [NUM], + NUM, + total_tokens_num_decode=NUM, + total_seqs_num=1, + total_seqs_num_decode=1, + num_spec_step=MTP_K, + is_deferred_out=is_deferred_out, + ) + + +@pytest.mark.parametrize("is_deferred_out", [True, False]) +@pytest.mark.parametrize("num_rejected", range(MTP_K + 1)) +def test_builder_positions_match_staged_window( + seq_factory, is_deferred_out, num_rejected +): + batch = _decode_batch(seq_factory, 162, num_rejected, is_deferred_out) + context_lens = np.asarray(batch.context_lens, dtype=np.int32) + ctx = _verify_ctx_lens(context_lens, batch.num_rejected, is_deferred_out) + np.testing.assert_array_equal(decode_positions(ctx, NUM), batch.scheduled_tokens) + + +@pytest.mark.parametrize("num_rejected", [0, 3, 6]) +def test_undeferred_window_is_not_shifted_by_num_rejected(seq_factory, num_rejected): + # The deferred shift applied to undeferred output lands the window + # `num_rejected - 1` slots early. + batch = _decode_batch(seq_factory, 162, num_rejected, is_deferred_out=False) + stale = decode_positions(batch.context_lens - batch.num_rejected, NUM) + assert stale[0] == batch.scheduled_tokens[0] - (num_rejected - 1) + + +def test_no_shift_without_counts(): + ctx = np.array([10, 20], dtype=np.int32) + for deferred in (True, False): + assert _verify_ctx_lens(ctx, None, deferred) is ctx + + +@pytest.mark.parametrize("num_rejected", [0, 1, MTP_K]) +def test_undeferred_shift_is_one_whatever_was_rejected(num_rejected): + # The undeferred window is fixed by the placeholder run postprocess left, + # so it does not move with how many drafts the sampler rejected. + ctx = np.full(4, 100, dtype=np.int32) + rejected = np.full(4, num_rejected, dtype=np.int64) + np.testing.assert_array_equal(_verify_ctx_lens(ctx, rejected, False), ctx - 1) + + +def test_deferred_shift_is_the_rejected_count_per_row(): + ctx = np.array([100, 100, 100], dtype=np.int32) + rejected = np.array([0, 3, MTP_K], dtype=np.int64) + out = _verify_ctx_lens(ctx, rejected, True) + np.testing.assert_array_equal(out, ctx - rejected) + np.testing.assert_array_equal(ctx, 100) # the caller's array is not mutated + + +@pytest.mark.parametrize("deferred", [True, False]) +def test_keeps_the_callers_dtype(deferred): + ctx = np.array([50, 60], dtype=np.int32) + out = _verify_ctx_lens(ctx, np.array([1, 2], dtype=np.int64), deferred) + assert out.dtype == np.int32