sgl-project--sglang/python/sglang/srt/model_executor/runner_utils/buffers.py

# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Static-buffer dataclasses used by the CUDA graph runners.

DecodeInputBuffers backs the decode-phase capture/replay path.
PrefillInputBuffers backs the prefill-phase capture/replay path.

Both subclass ForwardInputBuffers so that buffer-pool sharing works
the same way as for non-cuda-graph forward paths.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import torch

from sglang.srt.environ import envs
from sglang.srt.model_executor.forward_batch_info import (
    ForwardBatch,
    NgramEmbeddingInfo,
    PPProxyTensors,
    compute_local_num_token_non_padded,
)
from sglang.srt.model_executor.input_buffers import ForwardInputBuffers

_has_foreach_copy = hasattr(torch, "_foreach_copy_")


def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
    """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs."""

    def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
        if _has_foreach_copy:
            torch._foreach_copy_(dsts, srcs)
        else:
            for dst, src in zip(dsts, srcs):
                dst.copy_(src)

    groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
    for dst, src in zip(dsts, srcs):
        key = (dst.dtype, src.dtype)
        if key not in groups:
            groups[key] = ([], [])
        groups[key][0].append(dst)
        groups[key][1].append(src)
    for group_dsts, group_srcs in groups.values():
        foreach_copy(group_dsts, group_srcs)


@dataclass
class DecodeInputBuffers(ForwardInputBuffers):
    input_ids: torch.Tensor
    input_embeds: torch.Tensor
    req_pool_indices: torch.Tensor
    seq_lens: torch.Tensor
    seq_lens_cpu: torch.Tensor
    out_cache_loc: torch.Tensor
    positions: torch.Tensor
    mrope_positions: torch.Tensor
    num_token_non_padded: torch.Tensor
    custom_mask: torch.Tensor
    next_token_logits_buffer: torch.Tensor
    mamba_track_indices: Optional[torch.Tensor]
    mamba_track_mask: Optional[torch.Tensor]
    global_num_tokens_gpu: torch.Tensor
    global_num_tokens_for_logprob_gpu: torch.Tensor
    encoder_lens: Optional[torch.Tensor]
    pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
    ngram_embedding_info: Optional[NgramEmbeddingInfo]
    rids_int: Optional[torch.Tensor]
    bootstrap_room_ids_int: Optional[torch.Tensor]

    @classmethod
    def create(
        cls,
        *,
        device: torch.device,
        max_bs: int,
        max_num_token: int,
        hidden_size: int,
        next_token_logits_buffer: torch.Tensor,
        dtype: torch.dtype,
        dp_size: int,
        pp_size: int,
        is_encoder_decoder: bool,
        require_mlp_tp_gather: bool,
        seq_len_fill_value: int,
        encoder_len_fill_value: int,
        num_tokens_per_bs: int,
        cache_loc_dtype: torch.dtype,
        enable_mamba_track: bool,
        ne_token_table: Optional[torch.Tensor] = None,
        hc_hidden_size: Optional[int] = None,
        pp_proxy_topk_size: Optional[int] = None,
    ) -> DecodeInputBuffers:
        with torch.device(device):
            input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
            input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
            req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
            seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int64)
            out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
            positions = torch.zeros((max_num_token,), dtype=torch.int64)
            mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
            num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
            custom_mask = torch.ones(
                (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
                dtype=torch.bool,
            )
            mamba_track_indices = (
                torch.zeros((max_bs,), dtype=torch.int64)
                if enable_mamba_track
                else None
            )
            mamba_track_mask = (
                torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
            )

            if pp_size > 1:
                is_mhc = hc_hidden_size is not None
                hs = hc_hidden_size if is_mhc else hidden_size
                pp_proxy_tensors = {
                    "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
                }
                if not is_mhc:
                    pp_proxy_tensors["residual"] = torch.zeros(
                        (max_bs, hidden_size), dtype=dtype
                    )
                if pp_proxy_topk_size is not None:
                    pp_proxy_tensors["topk_indices"] = torch.zeros(
                        (max_num_token, pp_proxy_topk_size), dtype=torch.int32
                    )
            else:
                pp_proxy_tensors = None

            if is_encoder_decoder:
                encoder_lens = torch.full(
                    (max_bs,), encoder_len_fill_value, dtype=torch.int32
                )
            else:
                encoder_lens = None

            if require_mlp_tp_gather:
                global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
                global_num_tokens_for_logprob_gpu = torch.zeros(
                    (dp_size,), dtype=torch.int32
                )
            else:
                global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
                global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)

            ngram_embedding_info = (
                NgramEmbeddingInfo(
                    token_table=ne_token_table,
                    column_starts=torch.zeros([max_bs], dtype=torch.int32),
                    req_lens=torch.ones([max_bs], dtype=torch.int32),
                    out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
                    out_req_lens=torch.ones([max_bs], dtype=torch.int32),
                    skip_token_table_update=torch.zeros([max_bs], dtype=torch.bool),
                )
                if ne_token_table is not None
                else None
            )

            if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
                rids_int = torch.zeros((max_bs,), dtype=torch.int64)
                bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
            else:
                rids_int = None
                bootstrap_room_ids_int = None

        seq_lens_cpu = torch.full(
            (max_bs,),
            seq_len_fill_value,
            dtype=torch.int64,
            device="cpu",
        )

        return cls(
            input_ids=input_ids,
            input_embeds=input_embeds,
            req_pool_indices=req_pool_indices,
            seq_lens=seq_lens,
            seq_lens_cpu=seq_lens_cpu,
            out_cache_loc=out_cache_loc,
            positions=positions,
            mrope_positions=mrope_positions,
            num_token_non_padded=num_token_non_padded,
            custom_mask=custom_mask,
            next_token_logits_buffer=next_token_logits_buffer,
            mamba_track_indices=mamba_track_indices,
            mamba_track_mask=mamba_track_mask,
            encoder_lens=encoder_lens,
            global_num_tokens_gpu=global_num_tokens_gpu,
            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
            pp_proxy_tensors=pp_proxy_tensors,
            ngram_embedding_info=ngram_embedding_info,
            rids_int=rids_int,
            bootstrap_room_ids_int=bootstrap_room_ids_int,
        )

    def populate_from_forward_batch(
        self,
        *,
        forward_batch: ForwardBatch,
        raw_bs: int,
        raw_num_token: int,
        bs: int,
        seq_len_fill_value: int,
        require_gathered_buffer: bool,
        num_tokens_per_bs: int,
        dsa_enable_prefill_cp: bool,
        enable_num_token_non_padded_flag: bool,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ):
        if bs != raw_bs:
            self.seq_lens.fill_(seq_len_fill_value)
            self.out_cache_loc.zero_()
            if self.mamba_track_indices is not None:
                self.mamba_track_indices.zero_()
            if self.mamba_track_mask is not None:
                self.mamba_track_mask.fill_(False)

        # Build batched copy lists for all GPU tensors.
        dsts = [
            self.input_ids[:raw_num_token],
            self.req_pool_indices[:raw_bs],
            self.seq_lens[:raw_bs],
            self.out_cache_loc[:raw_num_token],
            self.positions[:raw_num_token],
        ]
        srcs = [
            forward_batch.input_ids,
            forward_batch.req_pool_indices,
            forward_batch.seq_lens,
            forward_batch.out_cache_loc,
            forward_batch.positions,
        ]

        if self.ngram_embedding_info is not None:
            ngram_embedding_info = forward_batch.ngram_embedding_info
            self.ngram_embedding_info.column_starts[:raw_bs].copy_(
                ngram_embedding_info.column_starts
            )
            self.ngram_embedding_info.req_lens[:raw_bs].copy_(
                ngram_embedding_info.req_lens
            )

        if (
            self.mamba_track_indices is not None
            and forward_batch.mamba_track_indices is not None
        ):
            dsts.append(self.mamba_track_indices[:raw_bs])
            srcs.append(forward_batch.mamba_track_indices)
        if (
            self.mamba_track_mask is not None
            and forward_batch.mamba_track_mask is not None
        ):
            dsts.append(self.mamba_track_mask[:raw_bs])
            srcs.append(forward_batch.mamba_track_mask)

        if self.encoder_lens is not None and forward_batch.encoder_lens is not None:
            dsts.append(self.encoder_lens[:raw_bs])
            srcs.append(forward_batch.encoder_lens)

        if forward_batch.mrope_positions is not None:
            dsts.append(self.mrope_positions[:, :raw_num_token])
            srcs.append(forward_batch.mrope_positions)

        if self.rids_int is not None and forward_batch.rids_int is not None:
            dsts.append(self.rids_int[:raw_bs])
            srcs.append(forward_batch.rids_int)
        if (
            self.bootstrap_room_ids_int is not None
            and forward_batch.bootstrap_room_ids_int is not None
        ):
            dsts.append(self.bootstrap_room_ids_int[:raw_bs])
            srcs.append(forward_batch.bootstrap_room_ids_int)

        if require_gathered_buffer:
            self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
            self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)

        if enable_num_token_non_padded_flag:
            if require_gathered_buffer and not dsa_enable_prefill_cp:
                num_tokens_per_dp = bs * num_tokens_per_bs
                local = compute_local_num_token_non_padded(
                    global_num_token_non_padded=forward_batch.num_token_non_padded,
                    num_tokens_per_dp=num_tokens_per_dp,
                )
                dsts.append(self.num_token_non_padded)
                srcs.append(local)
            else:
                dsts.append(self.num_token_non_padded)
                srcs.append(forward_batch.num_token_non_padded)

        # Pipeline-parallel proxy tensors.
        if pp_proxy_tensors is not None and self.pp_proxy_tensors is not None:
            for key, buf in self.pp_proxy_tensors.items():
                src = pp_proxy_tensors.tensors[key]
                dim = src.shape[0]
                dsts.append(buf[:dim])
                srcs.append(src)

        # Batch all GPU copies, grouped by dtype pair.
        _grouped_foreach_copy_(dsts, srcs)

        if forward_batch.seq_lens_cpu is not None:
            if bs != raw_bs:
                self.seq_lens_cpu.fill_(seq_len_fill_value)
            self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)


@dataclass
class PrefillInputBuffers(ForwardInputBuffers):
    input_ids: torch.Tensor
    out_cache_loc: torch.Tensor
    num_token_non_padded: torch.Tensor
    mamba_track_indices: Optional[torch.Tensor]
    mamba_track_mask: Optional[torch.Tensor]
    mamba_track_seqlens: Optional[torch.Tensor]
    positions: torch.Tensor
    input_embeds: Optional[torch.Tensor]
    mrope_positions: Optional[torch.Tensor]

    @classmethod
    def create(
        cls,
        *,
        device: torch.device,
        max_bs: int,
        max_num_tokens: int,
        cache_loc_dtype: torch.dtype,
        is_multimodal: bool,
        hidden_size: int,
        dtype: torch.dtype,
        enable_mamba_track: bool,
    ) -> PrefillInputBuffers:
        with torch.device(device):
            input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
            out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
            num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
            mamba_track_indices = (
                torch.zeros((max_bs,), dtype=torch.int64)
                if enable_mamba_track
                else None
            )
            mamba_track_mask = (
                torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
            )
            mamba_track_seqlens = (
                torch.zeros((max_bs,), dtype=torch.int32)
                if enable_mamba_track
                else None
            )
            positions = torch.zeros((max_num_tokens,), dtype=torch.int64)

            if is_multimodal:
                input_embeds = torch.zeros((max_num_tokens, hidden_size), dtype=dtype)
                mrope_positions = torch.zeros((3, max_num_tokens), dtype=torch.int64)
            else:
                input_embeds = None
                mrope_positions = None

        return cls(
            input_ids=input_ids,
            out_cache_loc=out_cache_loc,
            num_token_non_padded=num_token_non_padded,
            mamba_track_indices=mamba_track_indices,
            mamba_track_mask=mamba_track_mask,
            mamba_track_seqlens=mamba_track_seqlens,
            positions=positions,
            input_embeds=input_embeds,
            mrope_positions=mrope_positions,
        )

    def populate_from_forward_batch(
        self,
        *,
        forward_batch: ForwardBatch,
        raw_num_tokens: int,
        static_num_tokens: int,
        is_multimodal: bool,
    ) -> None:
        """Copy serving-batch values into static buffers and zero out
        the padding region between raw_num_tokens and
        static_num_tokens.
        """
        if static_num_tokens != raw_num_tokens:
            self.out_cache_loc.zero_()
            self.input_ids[raw_num_tokens:static_num_tokens].zero_()
            self.positions[raw_num_tokens:static_num_tokens].zero_()
            if is_multimodal:
                self.input_embeds[raw_num_tokens:static_num_tokens].zero_()
            if forward_batch.mrope_positions is not None:
                self.mrope_positions[:, raw_num_tokens:static_num_tokens].zero_()

        bs = forward_batch.batch_size

        self.input_ids[:raw_num_tokens].copy_(forward_batch.input_ids)
        self.positions[:raw_num_tokens].copy_(forward_batch.positions)
        self.out_cache_loc[:raw_num_tokens].copy_(forward_batch.out_cache_loc)

        if (
            self.mamba_track_indices is not None
            and forward_batch.mamba_track_indices is not None
        ):
            self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
        if (
            self.mamba_track_mask is not None
            and forward_batch.mamba_track_mask is not None
        ):
            self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
        if (
            self.mamba_track_seqlens is not None
            and forward_batch.mamba_track_seqlens is not None
        ):
            self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)

        if forward_batch.mrope_positions is not None:
            self.mrope_positions[:, :raw_num_tokens].copy_(
                forward_batch.mrope_positions
            )