lightseekorg--tokenspeed/python/tokenspeed/runtime/layers/attention/backends/base.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

import inspect
from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import TYPE_CHECKING

import torch

from tokenspeed.runtime.execution.breakable_cuda_graph import break_point

if TYPE_CHECKING:
    from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
    from tokenspeed.runtime.layers.attention.configs.base import BaseAttnConfig
    from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool
    from tokenspeed.runtime.layers.paged_attention import PagedAttention
    from tokenspeed.runtime.pd.utils import StepCounter


def init_backend_cuda_graph_state(
    backend: "AttentionBackend",
    max_bs: int,
    seq_lens_buf: torch.Tensor,
    **extras,
) -> None:
    """Call ``backend.init_cuda_graph_state`` with only the kwargs its
    signature accepts (VAR_KEYWORD accepts all of them).

    Signature-probe instead of try/except TypeError: paged_cache_group_specs
    is load-bearing for the state shed, so a TypeError raised from inside the
    backend's body must propagate rather than silently retry without specs.

    Shared by the cuda-graph wrapper and by composite backends (hybrid) that
    forward to user-selectable sub-backends with possibly narrow signatures.
    """
    params = inspect.signature(backend.init_cuda_graph_state).parameters
    if not any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()):
        extras = {k: v for k, v in extras.items() if k in params}
    backend.init_cuda_graph_state(max_bs, seq_lens_buf, **extras)


class AttentionBackend(ABC):
    """The base class of attention backends"""

    uses_paged_cache_groups: bool = False
    # Flat KV-cache per-group block tables (absolute index, null hole = 0). A
    # separate flag from uses_paged_cache_groups because the two mechanisms have
    # different hole/index semantics; a group-aware flat backend (Phase 4) sets
    # this True. Default False keeps every existing backend on today's path.
    uses_flat_cache_groups: bool = False
    # False for flat-capable backends whose spec-verify path is not wired yet.
    flat_spec_capable: bool = True
    uses_padded_decode_token_mask: bool = False

    def __init__(self, config: BaseAttnConfig) -> None:
        self.device = config.device
        self.num_qo_heads = config.num_attention_heads // config.attn_tp_size
        self.num_kv_heads = max(config.num_kv_heads // config.attn_tp_size, 1)
        self.dtype = config.dtype
        self.head_dim = config.head_dim
        self.is_draft = config.is_draft
        self.spec_num_tokens = config.speculative_num_draft_tokens
        # True when this backend's CUDA-graph block-table (kv_indices) buffer is
        # aliased to a peer backend's (e.g. a drafter sharing the target's), so
        # the replay path skips rebuilding it — the peer already populates it.
        self._block_table_aliased = False

    @contextmanager
    def override_num_extends(self, num_extends: int):
        """Temporarily override the decode-metadata slice discriminator for the
        wrapped block. Used by MLA backends to flip between drafter step 0
        (slice = [num_extends:]) and step 1+ (slice = [0:]).

        Default no-op for backends that fill separate prefill/decode metadata
        at init time.
        """
        yield

    def support_kv_cache_prewrite(
        self, forward_mode: ForwardMode | None = None
    ) -> bool:
        return False

    def select_out_cache_loc(self, layer, out_cache_loc, forward_mode=None):
        """Flat per-group write-location hook for out-of-backend KV writers
        (fused RoPE prewrite); identity for backends without flat cache
        groups (see uses_flat_cache_groups). ``forward_mode`` picks the
        metadata slot for backends that prewrite on extend as well."""
        return out_cache_loc

    @property
    def sinks_dtype(self) -> torch.dtype:
        return torch.bfloat16

    @abstractmethod
    def init_forward_metadata(self, *args, **kwargs):
        """Init the metadata for a forward pass.

        When use_cuda_graph=True the backend should use its pre-allocated
        cuda-graph buffers instead of the normal eager buffers.
        """
        raise NotImplementedError()

    def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor):
        """Init the global shared states for cuda graph. `seq_lens_buf` is
        the controller-owned per-request seq_lens; backends should reference
        (alias) it rather than copy, and must not mutate the contents."""
        raise NotImplementedError()

    def init_forward_metadata_capture_cuda_graph(
        self,
        bs: int,
        req_pool_indices: torch.Tensor,
        seq_lens: torch.Tensor,
        forward_mode: ForwardMode,
        flat_cache_group_ids: tuple[str, ...] = (),
        **kwargs,
    ):
        """Init the metadata for a forward pass for capturing a cuda graph.

        ``flat_cache_group_ids`` names the flat KV-cache groups whose page
        tables arrive at replay; a flat-capable backend (uses_flat_cache_groups)
        allocates its persistent per-group buffers from these ids — no table
        data exists at capture time. Empty tuple for non-flat backends.
        """
        raise NotImplementedError()

    def init_forward_metadata_replay_cuda_graph(
        self,
        bs: int,
        req_pool_indices: torch.Tensor,
        seq_lens: torch.Tensor,
        forward_mode: ForwardMode = None,
        req_to_page: torch.Tensor = None,
        flat_block_tables: dict[str, torch.Tensor] | None = None,
        **kwargs,
    ):
        """Update pre-allocated CUDA-graph metadata buffers in-place before replay.

        Called instead of init_forward_metadata when use_cuda_graph=True, so
        that the captured kernels (which hold pointers into the pre-allocated
        buffers) see the current batch's data without any new allocations.
        ``flat_block_tables`` carries the per-group flat page tables
        (group_id -> [>=bs, cols]) for flat-capable backends; a backend that
        captured flat buffers must be handed non-empty tables whenever bs > 0.
        Default: fall back to init_forward_metadata (correct but may not work
        for all backends that use separate cuda-graph buffer pools).
        """
        raise NotImplementedError(
            f"{type(self).__name__} must implement init_forward_metadata_replay_cuda_graph "
            "for CUDA graph support"
        )

    def configure_runtime(self, **kwargs) -> None:
        """Configure runtime state after model loading (e.g. sliding_window_size).

        Called once during ModelExecutor initialization with information that is
        not available at backend construction time.  Default: no-op.
        """
        pass

    def register_step_counter(self, step_counter: StepCounter):
        self.step_counter = step_counter

    @contextmanager
    def record_pd_cache_step(
        self,
        forward_mode: ForwardMode,
        save_kv_cache: bool,
        record_kv_cache: bool | None,
    ):
        """Anchor the PD layerwise cache-step record to the wrapped KV write.

        Records the ``StepCounter`` step before the attention call when the KV
        was pre-written (``save_kv_cache=False``) and after it otherwise, so a
        layerwise cache transfer always observes a fully written layer. See
        ``forward`` for the ``record_kv_cache`` override contract. No-op when no
        step counter is registered. Backends that own the record (e.g. the
        hybrid wrapper, which counts once per model layer across full-attn +
        mamba children) reuse this to avoid duplicating the gate logic.
        """
        if record_kv_cache is None:
            record_cache = not forward_mode.is_decode() and not forward_mode.is_idle()
        else:
            record_cache = record_kv_cache
        record_cache = record_cache and getattr(self, "step_counter", None) is not None

        if record_cache and not save_kv_cache:
            self.step_counter.record_cache()
        yield
        if record_cache and save_kv_cache:
            self.step_counter.record_cache()

    @break_point
    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer: PagedAttention,
        out_cache_loc: torch.Tensor,
        token_to_kv_pool: BaseTokenToKVPool,
        forward_mode: ForwardMode,
        bs: int,
        save_kv_cache: bool = True,
        record_kv_cache: bool | None = None,
        **kwargs,
    ):
        """Run forward on an attention layer with explicit scheduler metadata.

        ``record_kv_cache`` overrides the PD layerwise cache-step recording:
        ``None`` keeps the default (record on the EXTEND-side path), an explicit
        bool forces it so a DECODE-dispatched draft catch-up can still record.
        """
        with self.record_pd_cache_step(forward_mode, save_kv_cache, record_kv_cache):
            if forward_mode.is_decode():
                ret = self.forward_decode(
                    q,
                    k,
                    v,
                    layer,
                    out_cache_loc,
                    token_to_kv_pool,
                    bs,
                    save_kv_cache=save_kv_cache,
                    **kwargs,
                )
            else:
                ret = self.forward_extend(
                    q,
                    k,
                    v,
                    layer,
                    out_cache_loc,
                    token_to_kv_pool,
                    bs,
                    save_kv_cache=save_kv_cache,
                    forward_mode=forward_mode,
                    **kwargs,
                )
        return ret

    def forward_decode(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer: PagedAttention,
        out_cache_loc: torch.Tensor,
        token_to_kv_pool: BaseTokenToKVPool,
        bs: int,
        save_kv_cache: bool = True,
        **kwargs,
    ):
        """Run a forward for decode."""
        raise NotImplementedError()

    def forward_extend(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer: PagedAttention,
        out_cache_loc: torch.Tensor,
        token_to_kv_pool: BaseTokenToKVPool,
        bs: int,
        save_kv_cache: bool = True,
        **kwargs,
    ):
        """Run a forward for extend."""
        raise NotImplementedError()