[Spec] Dissolve EagleDraftInputV2Mixin so spec-info dataclasses hold data only (#29220)

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

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:

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)
 

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
 

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,
 )

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.

python/sglang/srt/speculative/eagle_info_v2.py

@@ -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,
-        )

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,
+    )

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,

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.

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,

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,

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)

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")

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.",
         )