From c7734e687120afb46e1b99c44891184d54ea34a4 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 24 Jun 2026 18:08:30 -0700 Subject: [PATCH] [Spec] Dissolve `EagleDraftInputV2Mixin` so spec-info dataclasses hold data only (#29220) --- .../srt/layers/attention/triton_backend.py | 2 +- python/sglang/srt/managers/schedule_batch.py | 12 +- python/sglang/srt/mem_cache/common.py | 2 +- .../srt/model_executor/pool_configurator.py | 2 +- .../srt/speculative/dflash_worker_v2.py | 2 +- python/sglang/srt/speculative/eagle_info.py | 3 +- .../sglang/srt/speculative/eagle_info_v2.py | 107 ------------------ python/sglang/srt/speculative/eagle_utils.py | 86 ++++++++++++++ .../sglang/srt/speculative/eagle_worker_v2.py | 2 +- .../speculative/frozen_kv_mtp_worker_v2.py | 4 +- .../multi_layer_eagle_worker_v2.py | 2 +- python/sglang/srt/speculative/ngram_info.py | 3 +- python/sglang/srt/speculative/spec_utils.py | 12 ++ .../speculative_draft_extend_runner.py | 4 +- .../spec/test_decode_bookkeeping_ownership.py | 13 ++- 15 files changed, 122 insertions(+), 134 deletions(-) delete mode 100644 python/sglang/srt/speculative/eagle_info_v2.py diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 3f0bab194..fa765787a 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/python/sglang/srt/layers/attention/triton_backend.py @@ -543,7 +543,7 @@ class TritonAttnBackend(AttentionBackend): dtype=torch.int32, device=self.device, ) - # DRAFT_EXTEND_V2: seq_lens = prefix + extend (bumped by eagle_info_v2). + # DRAFT_EXTEND_V2: seq_lens = prefix + extend (bumped on the draft-extend path). # Triton extend kernel receives extend K/V as separate tensors, so # kv_indptr/kv_indices must cover only the prefix portion. # extend_seq_lens_tensor is only attached to spec_info at real diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index cb2823dfd..65348246c 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -120,7 +120,6 @@ if TYPE_CHECKING: from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator from sglang.srt.managers.scheduler_components.metrics_reporter import PrefillStats from sglang.srt.session.session_controller import Session - from sglang.srt.speculative.eagle_info import EagleDraftInput from sglang.srt.speculative.spec_info import SpecInput, SpeculativeAlgorithm INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5 @@ -2436,7 +2435,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return self._new_tokens_required_next_decode_spec_v2(requests, page_size) def _new_tokens_required_next_decode_spec_v2(self, requests, page_size): - """Tight estimate matching eagle_info_v2.prepare_for_decode allocation.""" + """Tight estimate matching eagle_utils.eagle_prepare_for_decode allocation.""" reserve = get_alloc_reserve_per_decode() total = 0 for r in requests: @@ -2613,7 +2612,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): def prepare_for_decode(self): self.forward_mode = ForwardMode.DECODE - bs = len(self.reqs) # Decode embeds the last output token via embed_tokens; clear the stale # prefill-time tensor so it doesn't leak into ForwardBatch. self.input_embeds = None @@ -2623,10 +2621,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.attn_cp_metadata = None if not self.spec_algorithm.is_none(): - # Spec decoding: the draft input owns decode preparation - # (allocation, pre-claim, seq-lens bookkeeping). - draft_input: EagleDraftInput = self.spec_info - draft_input.prepare_for_decode(self) + # Spec decoding owns decode preparation (allocation, seq-lens bookkeeping). + from sglang.srt.speculative.spec_utils import spec_prepare_for_decode + + spec_prepare_for_decode(self) return if self.sampling_info.penalizer_orchestrator.is_required: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 5ff01b4f2..d6baf5179 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -243,7 +243,7 @@ def get_alloc_reserve_per_decode(server_args: Optional[ServerArgs] = None) -> in """KV length reserved per request at each decode step. The 2x is a double-buffer that absorbs the kv_committed_len lag in overlap - mode; see eagle_info_v2.prepare_for_decode. + mode; see eagle_utils.eagle_prepare_for_decode. """ return 2 * get_alloc_len_per_decode(server_args) diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index dfa0c5b20..6ac3a0b1d 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -367,7 +367,7 @@ class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator): decode_alloc = spec_decode_alloc_len_per_request(sa) else: # spec-v2: the overlap allocator keeps 2 * alloc_len outstanding - # (eagle_info_v2.prepare_for_decode: kv_committed_len + 2 * alloc_len). + # (eagle_utils.eagle_prepare_for_decode: kv_committed_len + 2 * alloc_len). decode_alloc = 2 * get_alloc_len_per_decode(sa) per_request = trailing_tokens + decode_alloc diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index d8670e4b9..b2255c3d7 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -31,9 +31,9 @@ from sglang.srt.speculative.dflash_utils import ( is_dflash_sampling_verify_available, parse_dflash_draft_config, ) -from sglang.srt.speculative.eagle_info_v2 import assign_extend_cache_locs_func from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func +from sglang.srt.speculative.triton_ops.cache_locs import assign_extend_cache_locs_func from sglang.srt.speculative.triton_ops.dflash_accept_bonus import ( _compute_dflash_accept_bonus_triton_unchecked, ) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 6fc815bc7..5efbc771f 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -9,7 +9,6 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode from sglang.srt.server_args import get_global_server_args -from sglang.srt.speculative.eagle_info_v2 import EagleDraftInputV2Mixin from sglang.srt.speculative.spec_info import SpecInput, SpecInputType logger = logging.getLogger(__name__) @@ -155,7 +154,7 @@ class EagleVerifyInput(SpecInput): @dataclass -class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): +class EagleDraftInput(SpecInput): # For idle stubs use `create_idle_input`, not the bare ctor: `filter_batch` # / `merge_batch` slice / cat `topk_p` / `topk_index` / `hidden_states` / # `bonus_tokens` unconditionally. diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py deleted file mode 100644 index 37137c503..000000000 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ /dev/null @@ -1,107 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -import torch - -from sglang.srt.managers.schedule_batch import ScheduleBatch -from sglang.srt.mem_cache.common import ( - alloc_paged_token_slots_extend, - alloc_token_slots, - get_alloc_reserve_per_decode, - get_last_loc, -) -from sglang.srt.speculative.triton_ops.cache_locs import ( - assign_extend_cache_locs_func as assign_extend_cache_locs_func, -) -from sglang.srt.speculative.triton_ops.eagle import ( - fill_bonus_tokens as fill_bonus_tokens, -) - -if TYPE_CHECKING: - from sglang.srt.speculative.eagle_info import ( - EagleDraftInput, - ) - - -@dataclass -class EagleDraftInputV2Mixin: - def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch): - batch.maybe_evict_swa() - - from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func - - bs = batch.batch_size() - - # Accumulate penalty - # This is a relaxed version of penalties for speculative decoding. - if batch.sampling_info.penalizer_orchestrator.is_required: - batch.cumulate_penalty_output_tokens() - - page_size = batch.token_to_kv_pool_allocator.page_size - double_alloc = get_alloc_reserve_per_decode() - - cur_kv_lens = [0] * bs - nxt_kv_lens = [0] * bs - num_needed_tokens = 0 - for i, r in enumerate(batch.reqs): - cur = r.kv_allocated_len - # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur. - # kv_committed_len is honest (bonus committed in resolve, not here), - # so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs. - nxt = max(cur, r.kv_committed_len + double_alloc) - cur_kv_lens[i] = cur - nxt_kv_lens[i] = nxt - num_needed_tokens += nxt - cur - r.kv_allocated_len = nxt - r.decode_batch_idx += 1 - - cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu") - nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu") - - # Fail fast if the page>1 + topk>1 draft over-allocation - # (get_alloc_reserve_per_decode) outgrows the req_to_token row: the write below - # would OOB and free would leak KV. The row is widened to hold it in _init_pools - # (PR #26972); fail here with a clear error, not on a later cryptic CUDA assert. - from sglang.srt.server_args import get_global_server_args - - if page_size > 1 and (get_global_server_args().speculative_eagle_topk or 1) > 1: - max_alloc_len = int(nxt_kv_lens_cpu.max()) - row_width = batch.req_to_token_pool.req_to_token.shape[1] - assert max_alloc_len <= row_width, ( - f"spec v2 page>1 topk>1 draft over-allocation ({max_alloc_len}) exceeds " - f"req_to_token row width ({row_width}); page_size={page_size}. Widen the " - f"row to hold committed + get_alloc_reserve_per_decode (PR #26972)." - ) - - # non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR - # barrier has chained to the prev forward -> host stalls a full forward. - cur_kv_lens_device = cur_kv_lens_cpu.to(device=batch.device, non_blocking=True) - nxt_kv_lens_device = nxt_kv_lens_cpu.to(device=batch.device, non_blocking=True) - if page_size == 1: - out_cache_loc = alloc_token_slots(batch.tree_cache, num_needed_tokens) - else: - last_loc = get_last_loc( - batch.req_to_token_pool.req_to_token, - batch.req_pool_indices, - cur_kv_lens_device, - ) - out_cache_loc = alloc_paged_token_slots_extend( - batch.tree_cache, - cur_kv_lens_device, - cur_kv_lens_cpu, - nxt_kv_lens_device, - nxt_kv_lens_cpu, - last_loc, - num_needed_tokens, - ) - - assign_req_to_token_pool_func( - batch.req_pool_indices, - batch.req_to_token_pool.req_to_token, - cur_kv_lens_device, - nxt_kv_lens_device, - out_cache_loc, - bs, - ) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 21de68eb8..ffca4ab41 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -6,6 +6,12 @@ from typing import TYPE_CHECKING, List, Optional import torch +from sglang.srt.mem_cache.common import ( + alloc_paged_token_slots_extend, + alloc_token_slots, + get_alloc_reserve_per_decode, + get_last_loc, +) from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu from sglang.srt.utils.async_probe import maybe_detect_oob @@ -561,3 +567,83 @@ def eagle_sample( # tensor includes the trailing/bonus token via out-of-place +1 so the # name no longer flips semantics mid-function (naming doc C2). return predict, num_correct_drafts + 1, accept_index + + +def eagle_prepare_for_decode(batch: ScheduleBatch): + batch.maybe_evict_swa() + + from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func + + bs = batch.batch_size() + + # Accumulate penalty + # This is a relaxed version of penalties for speculative decoding. + if batch.sampling_info.penalizer_orchestrator.is_required: + batch.cumulate_penalty_output_tokens() + + page_size = batch.token_to_kv_pool_allocator.page_size + double_alloc = get_alloc_reserve_per_decode() + + cur_kv_lens = [0] * bs + nxt_kv_lens = [0] * bs + num_needed_tokens = 0 + for i, r in enumerate(batch.reqs): + cur = r.kv_allocated_len + # max(cur, ...) clamps so adaptive downswitch cannot make nxt < cur. + # kv_committed_len is honest (bonus committed in resolve, not here), + # so it lags batch.seq_lens by ~1 verify in overlap; 2*alloc absorbs. + nxt = max(cur, r.kv_committed_len + double_alloc) + cur_kv_lens[i] = cur + nxt_kv_lens[i] = nxt + num_needed_tokens += nxt - cur + r.kv_allocated_len = nxt + r.decode_batch_idx += 1 + + cur_kv_lens_cpu = torch.tensor(cur_kv_lens, dtype=torch.int32, device="cpu") + nxt_kv_lens_cpu = torch.tensor(nxt_kv_lens, dtype=torch.int32, device="cpu") + + # Fail fast if the page>1 + topk>1 draft over-allocation + # (get_alloc_reserve_per_decode) outgrows the req_to_token row: the write below + # would OOB and free would leak KV. The row is widened to hold it in _init_pools + # (PR #26972); fail here with a clear error, not on a later cryptic CUDA assert. + from sglang.srt.server_args import get_global_server_args + + if page_size > 1 and (get_global_server_args().speculative_eagle_topk or 1) > 1: + max_alloc_len = int(nxt_kv_lens_cpu.max()) + row_width = batch.req_to_token_pool.req_to_token.shape[1] + assert max_alloc_len <= row_width, ( + f"spec v2 page>1 topk>1 draft over-allocation ({max_alloc_len}) exceeds " + f"req_to_token row width ({row_width}); page_size={page_size}. Widen the " + f"row to hold committed + get_alloc_reserve_per_decode (PR #26972)." + ) + + # non_blocking H2D: a blocking .to() syncs the schedule stream, which the WAR + # barrier has chained to the prev forward -> host stalls a full forward. + cur_kv_lens_device = cur_kv_lens_cpu.to(device=batch.device, non_blocking=True) + nxt_kv_lens_device = nxt_kv_lens_cpu.to(device=batch.device, non_blocking=True) + if page_size == 1: + out_cache_loc = alloc_token_slots(batch.tree_cache, num_needed_tokens) + else: + last_loc = get_last_loc( + batch.req_to_token_pool.req_to_token, + batch.req_pool_indices, + cur_kv_lens_device, + ) + out_cache_loc = alloc_paged_token_slots_extend( + batch.tree_cache, + cur_kv_lens_device, + cur_kv_lens_cpu, + nxt_kv_lens_device, + nxt_kv_lens_cpu, + last_loc, + num_needed_tokens, + ) + + assign_req_to_token_pool_func( + batch.req_pool_indices, + batch.req_to_token_pool.req_to_token, + cur_kv_lens_device, + nxt_kv_lens_device, + out_cache_loc, + bs, + ) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 23589830f..e79b54e69 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -64,7 +64,6 @@ from sglang.srt.speculative.eagle_info import ( EagleDraftInput, EagleVerifyInput, ) -from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens from sglang.srt.speculative.eagle_utils import ( TreeMaskMode, _eagle_prefill_tail_tokens, @@ -88,6 +87,7 @@ from sglang.srt.speculative.spec_utils import ( select_top_k_tokens, spec_stage_span, ) +from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens from sglang.srt.utils.async_probe import ( maybe_detect_inf, maybe_detect_nan, diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index d8d4b62ac..3f1b88a87 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -424,8 +424,8 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): assert isinstance(spec_info, FrozenKVMTPDraftInput) # NOTE: per-iter bookkeeping (penalty cumulation, maybe_evict_swa, - # decode_batch_idx tick) is done by the inherited - # EagleDraftInputV2Mixin.prepare_for_decode (scheduler-driven, see + # decode_batch_idx tick) is done by the scheduler-driven + # eagle_utils.eagle_prepare_for_decode (see # ScheduleBatch.prepare_for_decode), not here -- matching EAGLE v2. # Repeating evict/tick here would double-run them: the idx clock # gates SWA eviction timing and the SWA prefix-lock release. diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 3fb0228ae..6f4364f24 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -48,7 +48,6 @@ from sglang.srt.speculative.eagle_info import ( EagleDraftInput, EagleVerifyInput, ) -from sglang.srt.speculative.eagle_info_v2 import fill_bonus_tokens from sglang.srt.speculative.eagle_utils import ( TreeMaskMode, build_tree_kernel_efficient, @@ -66,6 +65,7 @@ from sglang.srt.speculative.spec_utils import ( record_stream_for_v2_verify, select_top_k_tokens, ) +from sglang.srt.speculative.triton_ops.eagle import fill_bonus_tokens from sglang.srt.utils import is_npu from sglang.srt.utils.async_probe import ( maybe_detect_inf, diff --git a/python/sglang/srt/speculative/ngram_info.py b/python/sglang/srt/speculative/ngram_info.py index 21136b66a..958b2cfe7 100644 --- a/python/sglang/srt/speculative/ngram_info.py +++ b/python/sglang/srt/speculative/ngram_info.py @@ -6,11 +6,10 @@ import torch from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton -from sglang.srt.speculative.eagle_info_v2 import EagleDraftInputV2Mixin from sglang.srt.speculative.spec_info import SpecInput, SpecInputType -class NgramVerifyInput(SpecInput, EagleDraftInputV2Mixin): +class NgramVerifyInput(SpecInput): def __init__( self, draft_token: torch.Tensor = None, diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 312f366f9..da3593318 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -650,3 +650,15 @@ def commit_mamba_states_after_verify( mamba_steps_to_track=mamba_steps_to_track, model=model_runner.model, ) + + +def spec_prepare_for_decode(batch: ScheduleBatch) -> None: + """eagle/ngram share a stateless free function; dflash keeps stateful + prep on its draft input -- the dispatcher routes. + """ + if batch.spec_algorithm.is_dflash(): + batch.spec_info.prepare_for_decode(batch) + else: + from sglang.srt.speculative.eagle_utils import eagle_prepare_for_decode + + eagle_prepare_for_decode(batch) diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index e9d14666e..2bf6fc5d9 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -149,7 +149,7 @@ def _make_eagle_draft_extend_v2_input(case, batch, *, device: str): def _set_draft_extend_v2_prefix_lens(batch, case, *, device: str): # Production sets seq_lens = prefix + extend before init_forward_metadata - # (eagle_info_v2.py bumps seq_lens by num_draft_tokens). Match that here. + # (the draft-extend path bumps seq_lens by num_draft_tokens). Match that here. seq_lens = tuple(p + e for p, e in zip(case.prefix_lens, case.input_lens)) batch.seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) batch.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32, device="cpu") @@ -821,7 +821,7 @@ def _set_draft_extend_v2_prefix_lens( device: str, ) -> None: # Production sets seq_lens = prefix + extend before init_forward_metadata - # (eagle_info_v2.py bumps seq_lens by num_draft_tokens). Match that here. + # (the draft-extend path bumps seq_lens by num_draft_tokens). Match that here. seq_lens = tuple(p + e for p, e in zip(case.prefix_lens, case.input_lens)) batch.seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device) batch.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32, device="cpu") diff --git a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py index c69166b8a..ba505a990 100644 --- a/test/registered/unit/spec/test_decode_bookkeeping_ownership.py +++ b/test/registered/unit/spec/test_decode_bookkeeping_ownership.py @@ -4,7 +4,8 @@ Per-request accounting state (`decode_batch_idx` / `extend_batch_idx` iter clocks, `kv_committed_len` / `kv_allocated_len` KV watermarks, `spec_verify_ct`, and the `maybe_evict_swa()` call) must only be advanced by the reviewed owner sites in _OWNER_SITES; spec-v2 draft workers must not -repeat any of them (the scheduler-driven mixin / resolve path already does). +repeat any of them (the scheduler-driven free function / resolve path already +does). A clock that runs fast fires SWA eviction in the overlap race window and releases the SWA prefix lock early; neither shows up in e2e CI or the idle leak checker, hence this AST-level guard. @@ -39,7 +40,7 @@ _EVICT_METHOD = "maybe_evict_swa" # attribute (`= 0` resets exempt) or "evict" for a `maybe_evict_swa()` call. # Any added/removed/recounted site fails until reviewed here. _SB = "managers/schedule_batch.py" -_MIXIN = ("speculative/eagle_info_v2.py", "EagleDraftInputV2Mixin.prepare_for_decode") +_EAGLE_DECODE = ("speculative/eagle_utils.py", "eagle_prepare_for_decode") _RESOLVE = ( "managers/scheduler_components/batch_result_processor.py", "SchedulerBatchResultProcessor._resolve_spec_v2_tokens", @@ -56,9 +57,9 @@ _OWNER_SITES = { ("mem_cache/common.py", "alloc_for_extend", "evict"): 1, ("mem_cache/common.py", "alloc_for_decode", "evict"): 1, # spec v2: no pre-claim; resolve commits the full accepted run uniformly. - (*_MIXIN, "decode_batch_idx"): 1, - (*_MIXIN, "evict"): 1, - (*_MIXIN, "kv_allocated_len"): 1, + (*_EAGLE_DECODE, "decode_batch_idx"): 1, + (*_EAGLE_DECODE, "evict"): 1, + (*_EAGLE_DECODE, "kv_allocated_len"): 1, (*_RESOLVE, "kv_committed_len"): 1, (*_RESOLVE, "spec_verify_ct"): 1, ( @@ -228,7 +229,7 @@ class TestDecodeBookkeepingOwnership(CustomTestCase): + "\n ".join(map(str, sorted(violations))) + "\nUnder spec v2 the iter-clock ticks, `maybe_evict_swa`, and " "KV watermark settlement are owned by the scheduler-driven " - "mixin / resolve path. Remove these from the worker.", + "free function / resolve path. Remove these from the worker.", )