sgl-project--sglang/python/sglang/srt/layers/attention/intel_amx_backend.py

from __future__ import annotations

from typing import TYPE_CHECKING

import torch

from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.srt.model_executor.forward_batch_info import ForwardBatch

if TYPE_CHECKING:
    from sglang.srt.layers.radix_attention import RadixAttention
    from sglang.srt.model_executor.model_runner import ModelRunner


class IntelAMXAttnBackend(AttentionBackend):
    def __init__(self, model_runner: ModelRunner):
        import sgl_kernel  # noqa: F401

        super().__init__()
        self.forward_metadata = None
        self.extend_metadata = None
        self.draft_decode_metadata = None
        self.device = model_runner.device
        # Pool refs — captured at construction so they survive deletion of the
        # corresponding ForwardBatch fields.
        self.req_to_token_pool = model_runner.req_to_token_pool
        self.token_to_kv_pool = model_runner.token_to_kv_pool
        self.max_context_len = model_runner.model_config.context_len

        # full->SWA translated out_cache_loc, computed once per forward (the only
        # set_kv_buffer is in eager forward_extend; decode writes KV in-kernel).
        self.use_sliding_window_kv_pool = (
            isinstance(self.token_to_kv_pool, SWAKVPool)
            and self.token_to_kv_pool.swa_layer_nums > 0
        )
        self.swa_out_cache_loc = None

        self.num_head = (
            model_runner.model_config.num_attention_heads // model_runner.tp_size
        )

        # [NB]: `layer_id` set to 0 for qwen3-next models, as not all attn layers require kv pool
        # using "full_attention_layer_id_mapping" to map which layer needs kv pool
        layer_id = 0
        if hasattr(model_runner.token_to_kv_pool, "full_attention_layer_id_mapping"):
            layer_id = [*model_runner.token_to_kv_pool.full_attention_layer_id_mapping][
                0
            ]
        self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(
            layer_id
        ).shape[-1]
        self.decode_attention_fwd = torch.ops.sgl_kernel.decode_attention_cpu
        self.extend_attention_fwd = torch.ops.sgl_kernel.extend_attention_cpu

        # Number of KV splits used by decode_attention_cpu; attn_logits is
        # sized [bs, num_head, num_kv_splits, v_head_dim + 1] to match.
        self.num_kv_splits = 8

        # speculative decoding params
        self.num_draft_tokens = model_runner.server_args.speculative_num_draft_tokens

    def _build_extend_metadata(self, forward_batch: ForwardBatch):
        """Resolve (seq_lens, extend_seq_lens, extend_start_loc, tree_mask) for
        forward_extend, once per forward pass.

        In TARGET_VERIFY mode the batch carries no extend_* fields, so they are
        derived from spec_info (mirrors the CUDA unified path in
        triton_backend.py); each request extends by exactly num_draft_tokens
        tokens. Outside spec decoding the fields are passed through.
        """
        bs = forward_batch.batch_size
        seq_lens = forward_batch.seq_lens
        tree_mask = None

        if forward_batch.forward_mode.is_target_verify():
            spec_info = forward_batch.spec_info
            if spec_info is None:
                raise RuntimeError(
                    "spec_info is unset in TARGET_VERIFY mode; the extend_* "
                    "metadata can only be derived from spec_info for "
                    "speculative verify batches."
                )
            num_draft_tokens = spec_info.draft_token_num
            extend_seq_lens = torch.full(
                (bs,), num_draft_tokens, dtype=torch.int32, device=self.device
            )
            # Uniform extend lengths: start locations form a plain range.
            extend_start_loc = torch.arange(
                0,
                bs * num_draft_tokens,
                num_draft_tokens,
                dtype=torch.int32,
                device=self.device,
            )
            seq_lens = forward_batch.seq_lens + num_draft_tokens
            # Speculative verify with a token tree: each draft token may only
            # attend to its ancestors among the draft tokens (the committed
            # prefix stays fully visible).
            #
            # NOTE: unlike triton_backend.py, which forwards spec_info.custom_mask
            # unconditionally, the mask is gated on tree_topk here. tree_topk == 1
            # means the draft tokens form a simple chain whose visibility is
            # exactly the kernel's built-in causal masking, and skipping the explicit
            # mask lets extend_attention_cpu take its faster mask-free path. EAGLE
            # has tree_topk == topk (> 1 for real trees); NGRAM has tree_topk == -1
            # (irregular tree); both need the mask.
            if spec_info.tree_topk != 1:
                custom_mask = spec_info.custom_mask
                if custom_mask is not None and custom_mask.numel() > 0:
                    tree_mask = custom_mask
        else:
            extend_seq_lens = forward_batch.extend_seq_lens
            extend_start_loc = forward_batch.extend_start_loc

        return seq_lens, extend_seq_lens, extend_start_loc, tree_mask

    def init_forward_metadata(self, forward_batch: ForwardBatch):
        """Init the metadata for a forward pass."""

        bs = forward_batch.batch_size
        attn_logits = torch.zeros(
            (
                bs,
                self.num_head,
                self.num_kv_splits,
                self.v_head_dim + 1,
            ),
            dtype=torch.float32,
            device=self.device,
        )
        if forward_batch.forward_mode.is_decode_or_idle():
            max_extend_len = None
            self.extend_metadata = None
        elif forward_batch.forward_mode.is_target_verify():
            max_extend_len = self.num_draft_tokens
            self.extend_metadata = self._build_extend_metadata(forward_batch)
        else:
            max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
            self.extend_metadata = self._build_extend_metadata(forward_batch)
        self.forward_metadata = (attn_logits, max_extend_len)

        if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
            self.swa_out_cache_loc = (
                self.token_to_kv_pool.translate_loc_from_full_to_swa(
                    forward_batch.out_cache_loc
                )
            )
        else:
            self.swa_out_cache_loc = None

    def get_cpu_graph_seq_len_fill_value(self):
        return 1

    def init_forward_metadata_capture_cpu_graph(
        self,
        bs: int,
        num_tokens: int,
        req_pool_indices: torch.Tensor,
        seq_lens: torch.Tensor,
        encoder_lens,
        forward_mode,
        spec_info,
    ):
        attn_logits = torch.zeros(
            (
                bs,
                self.num_head,
                self.num_kv_splits,
                self.v_head_dim + 1,
            ),
            dtype=torch.float32,
            device=self.device,
        )
        max_extend_len = None
        self.forward_metadata = (attn_logits, max_extend_len)
        self.extend_metadata = None

    def init_cpu_graph_state(self, max_bs: int, max_num_tokens: int):
        pass

    def forward_extend(
        self,
        q,
        k,
        v,
        layer: RadixAttention,
        forward_batch: ForwardBatch,
        save_kv_cache=True,
        sinks=None,
    ):
        if layer.qk_head_dim != layer.v_head_dim:
            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
        else:
            o = torch.empty_like(q)
        cache_loc = (
            forward_batch.out_cache_loc
            if not layer.is_cross_attention
            else forward_batch.encoder_out_cache_loc
        )
        if save_kv_cache and k is not None and v is not None:
            # Cross-attention never writes to the SWA pool, so only thread the
            # full->SWA location for non-cross-attention layers.
            swa_loc = None if layer.is_cross_attention else self.swa_out_cache_loc
            self.token_to_kv_pool.set_kv_buffer(
                layer, KVWriteLoc(cache_loc, swa_loc), k, v
            )

        # Precomputed once per forward pass in init_forward_metadata (spec
        # verify batches carry no extend_* fields; see _build_extend_metadata).
        seq_lens, extend_seq_lens, extend_start_loc, tree_mask = self.extend_metadata

        _, max_extend_len = self.forward_metadata
        self.extend_attention_fwd(
            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
            k,
            v,
            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
            self.token_to_kv_pool.get_key_buffer(layer.layer_id),
            self.token_to_kv_pool.get_value_buffer(layer.layer_id),
            self.req_to_token_pool.req_to_token,
            forward_batch.req_pool_indices,
            seq_lens,
            extend_seq_lens,
            extend_start_loc,
            max_extend_len,
            layer.scaling,
            layer.logit_cap,
            layer.is_cross_attention,
            layer.sliding_window_size + 1,
            forward_batch.encoder_lens,
            sinks,
            tree_mask,
        )
        return o

    def forward_decode(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        layer: RadixAttention,
        forward_batch: ForwardBatch,
        save_kv_cache=True,
        sinks=None,
    ):
        attn_logits, _ = self.forward_metadata

        if self.draft_decode_metadata is not None:
            req_to_token, seq_lens, req_pool_indices = self.draft_decode_metadata
        else:
            req_to_token = self.req_to_token_pool.req_to_token
            req_pool_indices = forward_batch.req_pool_indices
            seq_lens = forward_batch.seq_lens

        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)

        if layer.qk_head_dim != layer.v_head_dim:
            o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
        else:
            o = torch.empty_like(q)
        cache_loc = (
            forward_batch.out_cache_loc
            if not layer.is_cross_attention
            else forward_batch.encoder_out_cache_loc
        )
        self.decode_attention_fwd(
            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
            self.token_to_kv_pool.get_key_buffer(layer.layer_id),
            self.token_to_kv_pool.get_value_buffer(layer.layer_id),
            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
            k,
            v,
            cache_loc,
            attn_logits,
            req_to_token,
            req_pool_indices,
            seq_lens,
            layer.scaling,
            layer.logit_cap,
            layer.is_cross_attention,
            layer.sliding_window_size + 1,
            forward_batch.encoder_lens,
            sinks,
        )
        return o

    def support_triton(self):
        return False


class IntelAMXMultiStepDraftBackend:
    """
    Wrap multiple intel amx attention backends as one for multiple consecutive
    draft decoding steps.
    """

    def __init__(
        self,
        model_runner: ModelRunner,
        topk: int,
        speculative_num_steps: int,
    ):
        from sgl_kernel import build_draft_decode_metadata_cpu

        self.build_draft_decode_metadata = build_draft_decode_metadata_cpu
        self.topk = topk
        self.speculative_num_steps = speculative_num_steps
        self.attn_backends: list[IntelAMXAttnBackend] = []
        for _ in range(self.speculative_num_steps - 1):
            self.attn_backends.append(IntelAMXAttnBackend(model_runner))
        self.device = model_runner.device
        self.pool_len = model_runner.req_to_token_pool.req_to_token.shape[1]

    def init_forward_metadata(self, forward_batch: ForwardBatch):
        num_seqs = forward_batch.batch_size
        topk = self.topk
        bs = num_seqs * topk
        num_steps = self.speculative_num_steps
        req_to_token = self.attn_backends[0].req_to_token_pool.req_to_token
        seq_lens = forward_batch.seq_lens
        pool_len = self.pool_len
        num_head = self.attn_backends[0].num_head
        v_head_dim = self.attn_backends[0].v_head_dim
        device = self.device

        # Build expanded req_to_token via C++ kernel
        req_to_token_draft = self.build_draft_decode_metadata(
            req_to_token,
            forward_batch.req_pool_indices,
            seq_lens,
            topk,
            num_steps,
            pool_len,
        )

        req_pool_indices_expanded = torch.arange(bs, dtype=torch.int64, device=device)

        num_kv_splits = self.attn_backends[0].num_kv_splits
        for step in range(num_steps - 1):
            # Each candidate sees prefix + (step + 1) draft tokens.
            seq_lens_expanded = seq_lens.repeat_interleave(topk) + step + 1
            attn_logits = torch.zeros(
                (bs, num_head, num_kv_splits, v_head_dim + 1),
                dtype=torch.float32,
                device=device,
            )
            self.attn_backends[step].forward_metadata = (attn_logits, None)
            self.attn_backends[step].draft_decode_metadata = (
                req_to_token_draft,
                seq_lens_expanded,
                req_pool_indices_expanded,
            )