lightseekorg--tokenspeed/python/tokenspeed/runtime/execution/drafter/eagle.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import torch
from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax
from typing_extensions import override

from tokenspeed.runtime.execution.cache_loc_kernel import (
    compute_out_cache_loc_uniform,
)
from tokenspeed.runtime.execution.context import ForwardContext
from tokenspeed.runtime.execution.drafter.base import BaseDrafter
from tokenspeed.runtime.execution.forward_batch_info import (
    CaptureHiddenMode,
    ForwardMode,
)
from tokenspeed.runtime.multimodal.inputs import maybe_substitute_mm_pad
from tokenspeed.runtime.utils.nvtx import nvtx_range

DsaTopKState = tuple[Any | None, Any | None]

if TYPE_CHECKING:
    from tokenspeed.runtime.execution.input_buffer import InputBuffers
    from tokenspeed.runtime.execution.model_runner import ModelRunner
    from tokenspeed.runtime.execution.runtime_states import RuntimeStates
    from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend
    from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool
    from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput


def _advance_draft_forward_metadata_if_supported(attn_backend, seq_lens) -> None:
    advance = getattr(attn_backend, "advance_draft_forward_metadata", None)
    if advance is not None:
        advance(seq_lens)


@dataclass
class EagleDraftInput:
    input_num_tokens: int
    num_extends: int
    forward_mode: ForwardMode
    base_model_output: torch.Tensor  # [bs]
    accept_lengths: torch.Tensor  # [bs]
    base_out_hidden_states: torch.Tensor
    global_num_tokens: list[int] | None = None
    global_bs: list[int] | None = None
    all_decode_or_idle: bool = False
    dsa_topk: DsaTopKState = (None, None)


class Eagle(BaseDrafter):
    """
    Draft model runner that implements the Eagle/Eagle3 algorithm.
    """

    def __init__(
        self,
        spec_num_tokens: int,
        spec_num_steps: int,
        page_size: int,
        draft_model_runner: ModelRunner,
        req_to_page: torch.Tensor,
        attn_backend: AttentionBackend | None = None,
        token_to_kv_pool: BaseTokenToKVPool | None = None,
        runtime_states: RuntimeStates | None = None,
        input_buffers: InputBuffers | None = None,
        vocab_size: int | None = None,
    ) -> None:

        super().__init__(
            spec_num_tokens,
            spec_num_steps,
            draft_model_runner,
            runtime_states=runtime_states,
            input_buffers=input_buffers,
            page_size=page_size,
            req_to_page=req_to_page,
            attn_backend=attn_backend,
            token_to_kv_pool=token_to_kv_pool,
            vocab_size=vocab_size,
        )

        self.device = draft_model_runner.device
        hot_token_ids = draft_model_runner.model.get_hot_token_id()

        if hot_token_ids is not None:
            self.hot_token_ids = hot_token_ids.to(self.device)
        else:
            self.hot_token_ids = None

        # For constructing fallback global_num_tokens during CUDA graph capture.
        self.dp_size = draft_model_runner.mapping.attn.dp_size
        self.world_size = draft_model_runner.mapping.world_size

        # Drafter-owned alias source for the draft attn backend; advanced in
        # place during multi-step decode.
        self.draft_seq_lens_buf = torch.zeros_like(self.input_buffers.seq_lens_buf)

        # Persistent output buffer for the draft step's compute_out_cache_loc.
        self.draft_out_cache_loc_buf = torch.empty(
            (self.input_buffers.max_bs * (spec_num_steps - 1),),
            dtype=torch.int32,
            device=self.device,
        )

        # Precomputed `arange(max_bs) * spec_num_tokens - 1`
        # gather_ids = gather_ids_offsets + accept_lengths
        self.padded_gather_ids_offsets_buf = (
            torch.arange(
                self.input_buffers.max_bs, dtype=torch.int64, device=self.device
            )
            * spec_num_tokens
            - 1
        )

        # VLM placeholder id plumbed by ModelExecutor; None for text-only targets.
        self.mm_pad_substitute_id: int | None = None
        hf_config = getattr(draft_model_runner.model_config, "hf_config", None)
        self._dsa_reuse_mtp_topk = bool(
            getattr(hf_config, "index_share_for_mtp_iteration", False)
        )

    def _accepted_output_indices(
        self,
        accept_lengths: torch.Tensor,
        row_count: int,
        *,
        base_offset: int = 0,
    ) -> torch.Tensor:
        """Return safe flat output-token indices for each decode request.

        ``accept_lengths`` is the number of tokens that may be committed.  When
        the context-length cap reduces a row to 0 there is no real newly
        committed output token, but the drafter still runs to preserve graph
        shape.  Use the row's first verify output as a valid dummy source rather
        than producing ``row * N - 1``.
        """
        safe_accept_lengths = (
            accept_lengths[:row_count].to(torch.int64).clamp(1, self.spec_num_tokens)
        )
        return (
            self.padded_gather_ids_offsets_buf[:row_count]
            + safe_accept_lengths
            + base_offset
        )

    def set_mm_pad_substitute_id(self, token_id: int) -> None:
        self.mm_pad_substitute_id = token_id

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _attach_dsa_topk(
        self,
        ctx: ForwardContext,
        dsa_topk: DsaTopKState,
    ) -> None:
        if not self._dsa_reuse_mtp_topk:
            return
        ctx.dsa_prefill_topk, ctx.dsa_decode_topk = dsa_topk

    def _extract_dsa_topk(
        self,
        ctx: ForwardContext,
        dsa_topk: DsaTopKState,
    ) -> DsaTopKState:
        if not self._dsa_reuse_mtp_topk:
            return dsa_topk
        return ctx.dsa_prefill_topk, ctx.dsa_decode_topk

    def _map_hot(self, ids: torch.Tensor) -> torch.Tensor:
        """Map token ids through hot_token_ids if available, otherwise return as-is."""
        return self.hot_token_ids[ids] if self.hot_token_ids is not None else ids

    def _get_first_step_input(
        self,
        draft_input: EagleDraftInput,
        bs: int,
        input_num_tokens: int,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Returns (input_ids, gather_ids) for the first draft step.

        The first-step input shape matches the base model's: ragged
        ``[prefill_part || decode_part]`` under MIXED, full prefill chunks
        under EXTEND, ``base_model_output`` directly under DECODE.
        """
        num_extends = draft_input.num_extends
        num_decodes = bs - num_extends
        if num_extends > 0:
            num_decode_tokens = num_decodes * self.spec_num_tokens
            num_prefill_tokens = input_num_tokens - num_decode_tokens

            input_ids = self.input_buffers.shifted_prefill_ids_buf[:input_num_tokens]
            unpadded_input_lengths = self.input_buffers.input_lengths_buf[:bs]
            if num_decodes > 0:
                input_ids[num_prefill_tokens:].copy_(
                    draft_input.base_model_output[num_extends:]
                )
                unpadded_input_lengths[num_extends:].copy_(
                    draft_input.accept_lengths[num_extends:]
                )

            last_indices = unpadded_input_lengths[:num_extends].cumsum(0) - 1
            last_input_ids = input_ids[last_indices]
            input_ids[last_indices] = torch.where(
                last_input_ids == -1,
                draft_input.base_model_output[:num_extends],
                last_input_ids,
            )

            gather_ids = last_indices
            if num_decodes > 0:
                gather_ids = torch.cat(
                    [
                        gather_ids,
                        self._accepted_output_indices(
                            draft_input.accept_lengths[num_extends:],
                            num_decodes,
                            base_offset=num_prefill_tokens,
                        ),
                    ]
                )
        else:
            input_ids = draft_input.base_model_output
            gather_ids = self._accepted_output_indices(
                draft_input.accept_lengths,
                bs,
            )

        return input_ids, gather_ids

    @nvtx_range("draft_first_step", color="purple")
    def _run_first_step(
        self,
        bs: int,
        draft_input: EagleDraftInput,
    ) -> tuple[LogitsProcessorOutput, DsaTopKState]:

        buffers = self.input_buffers
        forward_mode = draft_input.forward_mode

        input_ids, gather_ids = self._get_first_step_input(
            draft_input, bs, draft_input.input_num_tokens
        )
        input_ids = maybe_substitute_mm_pad(input_ids, self.mm_pad_substitute_id)
        draft_model = self.draft_model_runner.model
        input_num_tokens = draft_input.input_num_tokens

        ctx = ForwardContext(
            attn_backend=self.attn_backend,
            token_to_kv_pool=self.token_to_kv_pool,
            req_to_page=self.req_to_page,
            bs=bs,
            num_extends=draft_input.num_extends,
            input_num_tokens=input_num_tokens,
            forward_mode=forward_mode,
            capture_hidden_mode=CaptureHiddenMode.LAST,
            gather_ids=gather_ids,
            global_num_tokens=draft_input.global_num_tokens,
            global_bs=draft_input.global_bs,
            all_decode_or_idle=draft_input.all_decode_or_idle,
            draft_seq_lens_buf=self.draft_seq_lens_buf,
            accept_lengths=draft_input.accept_lengths,
        )

        dsa_topk = draft_input.dsa_topk
        prepare_dsa_topk = getattr(draft_model, "prepare_dsa_topk_for_mtp_decode", None)
        compute_dsa_topk_first_step = bool(
            getattr(draft_model, "compute_dsa_topk_first_step", False)
        )
        if compute_dsa_topk_first_step:
            # GLM NextN has its own indexer weights. Compute first-step top-k
            # in the draft model, then select rows used by later MTP steps.
            dsa_topk = (None, None)
        elif draft_input.num_extends == 0 and prepare_dsa_topk is not None:
            dsa_topk = prepare_dsa_topk(dsa_topk, gather_ids)
        else:
            dsa_topk = (None, None)
        self._attach_dsa_topk(ctx, dsa_topk)

        logits_output = self.draft_model_runner.forward(
            ctx=ctx,
            input_ids=input_ids,
            positions=buffers.positions_buf[:input_num_tokens],
            out_cache_loc=buffers.out_cache_loc_buf[:input_num_tokens],
            captured_hidden_states=draft_input.base_out_hidden_states,
            spec_step_idx=0,
        )
        dsa_topk = self._extract_dsa_topk(ctx, dsa_topk)
        if compute_dsa_topk_first_step and prepare_dsa_topk is not None:
            dsa_topk = prepare_dsa_topk(
                dsa_topk,
                gather_ids,
                num_prefill_rows=draft_input.num_extends,
            )
        return logits_output, dsa_topk

    @nvtx_range("draft_multi_step", color="purple")
    def _run_multi_step_decode(
        self,
        bs: int,
        draft_ids: torch.Tensor,
        next_tokens: torch.Tensor,
        logits_output: LogitsProcessorOutput,
        draft_input: EagleDraftInput,
        dsa_topk: DsaTopKState,
    ) -> None:
        num_extends = draft_input.num_extends
        num_decodes = bs - num_extends
        req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs]
        cache_start = self.input_buffers.seq_lens_buf[:bs].clone()
        # Step 1's write position uses vc+accept_length after target verify so
        # rotary/cache metadata stay on the accepted prefix, not rejected tail.
        if num_decodes > 0:
            cache_start[num_extends:] = (
                self.runtime_states.valid_cache_lengths.index_select(
                    0, req_pool_indices[num_extends:]
                )
                + draft_input.accept_lengths[num_extends:]
            )

        # Write cache slots for steps 1..N-1.
        cache_locs = self.draft_out_cache_loc_buf[: bs * (self.spec_num_steps - 1)]
        compute_out_cache_loc_uniform(
            out_cache_loc_ptr=cache_locs,
            req_pool_indices=req_pool_indices,
            uniform_input_length=self.spec_num_steps - 1,
            cache_start=cache_start,
            req_to_pages=self.req_to_page,
            page_size=self.page_size,
        )
        cache_locs = cache_locs.view(bs, self.spec_num_steps - 1)
        # +1 is the kernel's read-inclusive convention; advanced per iter.
        draft_seq_lens = self.draft_seq_lens_buf[:bs]
        torch.add(cache_start, 1, out=draft_seq_lens)

        positions = cache_start.clone()

        for i in range(1, self.spec_num_steps):
            # make a ctx every time model runner forward
            # Multi-step decode is pure DECODE mode: one token per request.
            # global_num_tokens must reflect each rank's batch size, not the
            # target model's total tokens (which may be bs * spec_num_tokens).
            global_num_tokens = draft_input.global_num_tokens

            if self.dp_size > 1:
                if draft_input.global_bs is not None:
                    global_num_tokens = draft_input.global_bs
                else:
                    # CUDA graph capture path: uniform batch size across ranks.
                    global_num_tokens = [bs] * self.world_size

            ctx = ForwardContext(
                bs=bs,
                num_extends=0,
                attn_backend=self.attn_backend,
                token_to_kv_pool=self.token_to_kv_pool,
                req_to_page=self.req_to_page,
                input_num_tokens=bs,
                forward_mode=ForwardMode.DECODE,
                capture_hidden_mode=CaptureHiddenMode.LAST,
                global_num_tokens=global_num_tokens,
                global_bs=draft_input.global_bs,
                all_decode_or_idle=draft_input.all_decode_or_idle,
            )
            self._attach_dsa_topk(ctx, dsa_topk)

            out_cache_loc = cache_locs[:, i - 1].contiguous()
            # Keep attention metadata on the accepted prefix; rejected verify
            # tail slots may still contain stale draft KV.
            _advance_draft_forward_metadata_if_supported(
                ctx.attn_backend,
                draft_seq_lens,
            )

            with nvtx_range("draft_forward", color="red"):
                logits_output = self.draft_model_runner.forward(
                    ctx=ctx,
                    input_ids=self._map_hot(draft_ids),
                    positions=positions,
                    out_cache_loc=out_cache_loc,
                    captured_hidden_states=logits_output.hidden_states,
                    spec_step_idx=i,
                )
                dsa_topk = self._extract_dsa_topk(ctx, dsa_topk)

            with nvtx_range("draft_sample", color="yellow"):
                if logits_output.next_token_ids is not None:
                    draft_ids = logits_output.next_token_ids
                else:
                    draft_ids = sampling_argmax(logits_output.next_token_logits)
                # Column 0 holds last_verified_ids; drafter writes step `i` into column `i + 1`.
                next_tokens[:, i + 1] = self._map_hot(draft_ids)
                if i + 1 < self.spec_num_steps:
                    positions.add_(1)
                    draft_seq_lens.add_(1)

    # ------------------------------------------------------------------
    # Public entry point (type-based dispatch from ModelExecutor)
    # ------------------------------------------------------------------

    @override
    def get_candidates(
        self,
        base_ctx: ForwardContext,
    ) -> torch.Tensor | None:
        num_extends = base_ctx.num_extends
        num_decodes = base_ctx.bs - num_extends
        if num_decodes == 0:
            return None

        num_decode_tokens = num_decodes * self.spec_num_tokens
        num_prefill_tokens = base_ctx.input_num_tokens - num_decode_tokens
        return self.input_buffers.input_ids_buf[
            num_prefill_tokens : base_ctx.input_num_tokens
        ].reshape(num_decodes, self.spec_num_tokens)

    @override
    def draft(
        self,
        draft_input: EagleDraftInput,
    ) -> torch.Tensor:

        bs = draft_input.accept_lengths.shape[0]

        # Layout: column 0 holds the last verified id (the base model's accepted token);
        # columns 1..spec_num_steps hold the drafter's speculative tokens.
        next_tokens = torch.empty(
            (bs, self.spec_num_steps + 1),
            dtype=torch.int32,
            device=self.device,
        )

        # Last verified id per request → next_tokens[:, 0].
        num_extends = draft_input.num_extends
        num_decodes = bs - num_extends
        if num_extends > 0:
            next_tokens[:num_extends, 0] = draft_input.base_model_output[:num_extends]
        if num_decodes > 0:
            indices = self._accepted_output_indices(
                draft_input.accept_lengths[num_extends:],
                num_decodes,
            )
            if num_extends > 0:
                indices.add_(num_extends)
            torch.index_select(
                draft_input.base_model_output,
                0,
                indices,
                out=next_tokens[num_extends:, 0],
            )
        if self.spec_num_steps > 0:
            next_tokens[:, 1:] = next_tokens[:, :1]

        # Seed the draft attn backend's aliased seq_lens for the first step.
        self.draft_seq_lens_buf[:bs].copy_(self.input_buffers.seq_lens_buf[:bs])

        # First draft step. LogitsProcessor prunes `[num_prefill_tokens + num_decodes * spec_num_tokens, ...]`
        # down to `[bs, ...]`, so logits/hidden_states arrive here already aligned to one row per request.
        logits_output, dsa_topk = self._run_first_step(bs, draft_input)

        if logits_output.next_token_ids is not None:
            draft_ids = logits_output.next_token_ids
        else:
            draft_ids = sampling_argmax(logits_output.next_token_logits)
        next_tokens[:, 1] = self._map_hot(draft_ids)

        if self.spec_num_steps <= 1:
            return next_tokens

        if self.input_buffers.all_extends_mid_chunk and self.dp_size == 1:
            # Skip multi-step when the whole batch is mid-chunk EXTEND:
            # no request completes a target-side speculative verification
            # after this forward, so any speculative tokens would be discarded.
            #
            # In DP we still run, because peer ranks may have completing
            # extends or decodes; diverging here would desync the drafter's
            # dense-TP / MoE-EP collectives (NCCL hang or RSAG mismatch).
            return next_tokens

        # Draft step 2+ (multi-step decode).
        # Multi-step decode operates on full bs; drop the [num_extends:]
        # slice that step 0 may have set up for MIXED target. No-op on
        # backends that fill separate prefill/decode metadata at init
        # time.
        with self.attn_backend.override_num_extends(0):
            self._run_multi_step_decode(
                bs,
                draft_ids,
                next_tokens,
                logits_output,
                draft_input,
                dsa_topk,
            )
        return next_tokens

    @override
    @nvtx_range("drafter", color="purple")
    def run(
        self,
        base_ctx: ForwardContext,
        logits_output: LogitsProcessorOutput,
        output_tokens: torch.Tensor,
        accept_lengths: torch.Tensor,
    ) -> torch.Tensor:

        draft_input = EagleDraftInput(
            input_num_tokens=base_ctx.input_num_tokens,
            num_extends=base_ctx.num_extends,
            forward_mode=base_ctx.forward_mode,
            base_model_output=output_tokens,
            accept_lengths=accept_lengths,
            base_out_hidden_states=logits_output.hidden_states,
            global_num_tokens=base_ctx.global_num_tokens,
            global_bs=base_ctx.global_bs,
            all_decode_or_idle=base_ctx.all_decode_or_idle,
            dsa_topk=(base_ctx.dsa_prefill_topk, base_ctx.dsa_decode_topk),
        )

        # next_tokens layout: column 0 = last verified id, columns 1.. = drafter tokens.
        return self.draft(draft_input)