sglang 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

sglang/srt/layers/attention_backend.py

@@ -1,474 +0,0 @@
-from __future__ import annotations
-
-"""
-Support different attention backends.
-Now there are two backends: FlashInfer and Triton.
-FlashInfer is faster and Triton is easier to customize.
-Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
-"""
-
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING
-
-import torch
-import torch.nn as nn
-
-from sglang.global_config import global_config
-from sglang.srt.layers.flashinfer_utils import update_flashinfer_indices
-from sglang.srt.managers.schedule_batch import ScheduleBatch, global_server_args_dict
-from sglang.srt.model_executor.forward_batch_info import ForwardMode, InputMetadata
-from sglang.srt.utils import is_hip
-
-if TYPE_CHECKING:
-    from sglang.srt.model_executor.model_runner import ModelRunner
-
-# ROCm: flashinfer available later
-if not is_hip():
-    from flashinfer import (
-        BatchDecodeWithPagedKVCacheWrapper,
-        BatchPrefillWithPagedKVCacheWrapper,
-        BatchPrefillWithRaggedKVCacheWrapper,
-    )
-    from flashinfer.cascade import merge_state
-    from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
-
-
-class AttentionBackend(ABC):
-    """The base class of attention backends"""
-
-    @abstractmethod
-    def init_forward_metadata(
-        self, batch: ScheduleBatch, input_metadata: InputMetadata
-    ):
-        """Init the metadata for a forward pass."""
-        raise NotImplementedError()
-
-    def init_cuda_graph_state(self, max_bs: int):
-        """Init the global shared states for cuda graph."""
-        raise NotImplementedError()
-
-    def init_forward_metadata_capture_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
-    ):
-        """Init the metadata for a forward pass for capturing a cuda graph."""
-        raise NotImplementedError()
-
-    def init_forward_metadata_replay_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
-    ):
-        """Init the metadata for a forward pass for replying a cuda graph."""
-        raise NotImplementedError()
-
-    def get_cuda_graph_seq_len_fill_value(self):
-        """Get the fill value for padded seq lens. Typically, it is 0 or 1."""
-        raise NotImplementedError()
-
-    def forward(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        """Run forward on an attention layer."""
-        if input_metadata.forward_mode.is_decode():
-            return self.forward_decode(q, k, v, layer, input_metadata)
-        else:
-            return self.forward_extend(q, k, v, layer, input_metadata)
-
-    def forward_decode(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        """Run a forward for decode."""
-        raise NotImplementedError()
-
-    def forward_extend(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        """Run a forward for extend."""
-        raise NotImplementedError()
-
-
-class FlashInferAttnBackend(AttentionBackend):
-    """Flashinfer attention kernels."""
-
-    def __init__(self, model_runner: ModelRunner):
-        super().__init__()
-        self.model_runner = model_runner
-
-        if not _grouped_size_compiled_for_decode_kernels(
-            model_runner.model_config.num_attention_heads // model_runner.tp_size,
-            model_runner.model_config.get_num_kv_heads(model_runner.tp_size),
-        ):
-            self.decode_use_tensor_cores = True
-        else:
-            self.decode_use_tensor_cores = False
-
-        self.workspace_buffer = torch.empty(
-            global_config.flashinfer_workspace_size,
-            dtype=torch.uint8,
-            device="cuda",
-        )
-
-        if model_runner.sliding_window_size is None:
-            self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
-                self.workspace_buffer, "NHD"
-            )
-            self.prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
-                self.workspace_buffer, "NHD"
-            )
-            self.decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-                self.workspace_buffer,
-                "NHD",
-                use_tensor_cores=self.decode_use_tensor_cores,
-            )
-        else:
-            # Two wrappers: one for sliding window attention and one for full attention.
-            # Using two wrappers is unnecessary in the current PR, but are prepared for future PRs
-            self.prefill_wrapper_ragged = None
-            self.prefill_wrapper_paged = []
-            self.decode_wrapper = []
-            for _ in range(2):
-                self.prefill_wrapper_paged.append(
-                    BatchPrefillWithPagedKVCacheWrapper(self.workspace_buffer, "NHD")
-                )
-                self.decode_wrapper.append(
-                    BatchDecodeWithPagedKVCacheWrapper(
-                        self.workspace_buffer,
-                        "NHD",
-                        use_tensor_cores=self.decode_use_tensor_cores,
-                    )
-                )
-
-        self.forward_metadata = None
-        self.cuda_graph_metadata = {}
-
-    def init_forward_metadata(
-        self, batch: ScheduleBatch, input_metadata: InputMetadata
-    ):
-        if input_metadata.forward_mode.is_decode():
-            prefix_lens = None
-            use_ragged = False
-            total_num_tokens = None
-        else:
-            prefix_lens = input_metadata.extend_prefix_lens
-
-            # Some heuristics to check whether to use ragged forward
-            use_ragged = False
-            if (
-                torch.sum(input_metadata.seq_lens).item() >= 4096
-                and self.model_runner.sliding_window_size is None
-            ):
-                use_ragged = True
-
-            total_num_tokens = torch.sum(input_metadata.seq_lens).item()
-
-        update_flashinfer_indices(
-            input_metadata.forward_mode,
-            self.model_runner,
-            input_metadata.req_pool_indices,
-            input_metadata.seq_lens,
-            prefix_lens,
-            use_ragged=use_ragged,
-        )
-
-        self.forward_metadata = (use_ragged, total_num_tokens, self.decode_wrapper)
-
-    def init_cuda_graph_state(self, max_bs: int):
-        self.cuda_graph_kv_indptr = torch.zeros(
-            (max_bs + 1,), dtype=torch.int32, device="cuda"
-        )
-        self.cuda_graph_kv_indices = torch.zeros(
-            (max_bs * self.model_runner.model_config.context_len,),
-            dtype=torch.int32,
-            device="cuda",
-        )
-        self.cuda_graph_kv_last_page_len = torch.ones(
-            (max_bs,), dtype=torch.int32, device="cuda"
-        )
-
-        if self.model_runner.sliding_window_size is not None:
-            self.cuda_graph_kv_indptr = [
-                self.cuda_graph_kv_indptr,
-                self.cuda_graph_kv_indptr.clone(),
-            ]
-            self.cuda_graph_kv_indices = [
-                self.cuda_graph_kv_indices,
-                self.cuda_graph_kv_indices.clone(),
-            ]
-
-    def init_forward_metadata_capture_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
-    ):
-        if self.model_runner.sliding_window_size is None:
-            decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-                self.workspace_buffer,
-                "NHD",
-                use_cuda_graph=True,
-                use_tensor_cores=self.decode_use_tensor_cores,
-                paged_kv_indptr_buffer=self.cuda_graph_kv_indptr[: bs + 1],
-                paged_kv_indices_buffer=self.cuda_graph_kv_indices,
-                paged_kv_last_page_len_buffer=self.cuda_graph_kv_last_page_len[:bs],
-            )
-        else:
-            decode_wrapper = []
-            for i in range(2):
-                decode_wrapper.append(
-                    BatchDecodeWithPagedKVCacheWrapper(
-                        self.workspace_buffer,
-                        "NHD",
-                        use_cuda_graph=True,
-                        use_tensor_cores=self.decode_use_tensor_cores,
-                        paged_kv_indptr_buffer=self.cuda_graph_kv_indptr[i][: bs + 1],
-                        paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
-                        paged_kv_last_page_len_buffer=self.cuda_graph_kv_last_page_len[
-                            :bs
-                        ],
-                    )
-                )
-
-        update_flashinfer_indices(
-            ForwardMode.DECODE,
-            self.model_runner,
-            req_pool_indices,
-            seq_lens,
-            None,
-            decode_wrapper,
-        )
-
-        self.cuda_graph_metadata[bs] = decode_wrapper
-
-        self.forward_metadata = (False, None, decode_wrapper)
-
-    def init_forward_metadata_replay_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
-    ):
-        update_flashinfer_indices(
-            ForwardMode.DECODE,
-            self.model_runner,
-            req_pool_indices[:bs],
-            seq_lens[:bs],
-            None,
-            self.cuda_graph_metadata[bs],
-        )
-
-    def get_cuda_graph_seq_len_fill_value(self):
-        return 0
-
-    def forward_extend(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        if not isinstance(self.prefill_wrapper_paged, list):
-            prefill_wrapper_paged = self.prefill_wrapper_paged
-        else:
-            if layer.sliding_window_size != -1:
-                prefill_wrapper_paged = self.prefill_wrapper_paged[0]
-            else:
-                prefill_wrapper_paged = self.prefill_wrapper_paged[1]
-
-        use_ragged, total_num_tokens, decode_wrapper = self.forward_metadata
-
-        if not use_ragged:
-            if k is not None:
-                assert v is not None
-                input_metadata.token_to_kv_pool.set_kv_buffer(
-                    layer.layer_id, input_metadata.out_cache_loc, k, v
-                )
-            o = prefill_wrapper_paged.forward(
-                q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                input_metadata.token_to_kv_pool.get_kv_buffer(layer.layer_id),
-                causal=True,
-                sm_scale=layer.scaling,
-                window_left=layer.sliding_window_size,
-                logits_soft_cap=layer.logit_cap,
-            )
-        else:
-            o1, s1 = self.prefill_wrapper_ragged.forward_return_lse(
-                q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                k.contiguous().view(-1, layer.tp_k_head_num, layer.head_dim),
-                v.contiguous().view(-1, layer.tp_v_head_num, layer.head_dim),
-                causal=True,
-                sm_scale=layer.scaling,
-                logits_soft_cap=layer.logit_cap,
-            )
-
-            if input_metadata.extend_no_prefix:
-                o = o1
-            else:
-                o2, s2 = prefill_wrapper_paged.forward_return_lse(
-                    q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                    input_metadata.token_to_kv_pool.get_kv_buffer(layer.layer_id),
-                    causal=False,
-                    sm_scale=layer.scaling,
-                    logits_soft_cap=layer.logit_cap,
-                )
-
-                o, _ = merge_state(o1, s1, o2, s2)
-
-            input_metadata.token_to_kv_pool.set_kv_buffer(
-                layer.layer_id, input_metadata.out_cache_loc, k, v
-            )
-
-        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
-
-    def forward_decode(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        use_ragged, total_num_tokens, decode_wrapper = self.forward_metadata
-
-        if isinstance(decode_wrapper, list):
-            if layer.sliding_window_size != -1:
-                decode_wrapper = decode_wrapper[0]
-            else:
-                decode_wrapper = decode_wrapper[1]
-
-        if k is not None:
-            assert v is not None
-            input_metadata.token_to_kv_pool.set_kv_buffer(
-                layer.layer_id, input_metadata.out_cache_loc, k, v
-            )
-
-        o = decode_wrapper.forward(
-            q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-            input_metadata.token_to_kv_pool.get_kv_buffer(layer.layer_id),
-            sm_scale=layer.scaling,
-            logits_soft_cap=layer.logit_cap,
-        )
-
-        return o.view(-1, layer.tp_q_head_num * layer.head_dim)
-
-
-class TritonAttnBackend(AttentionBackend):
-    def __init__(self, model_runner: ModelRunner):
-        # Lazy import to avoid the initialization of cuda context
-        from sglang.srt.layers.triton_attention.decode_attention import (
-            decode_attention_fwd,
-        )
-        from sglang.srt.layers.triton_attention.extend_attention import (
-            extend_attention_fwd,
-        )
-
-        super().__init__()
-
-        self.decode_attention_fwd = decode_attention_fwd
-        self.extend_attention_fwd = extend_attention_fwd
-        self.num_head = (
-            model_runner.model_config.num_attention_heads // model_runner.tp_size
-        )
-
-        if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
-            self.reduce_dtype = torch.float32
-        else:
-            self.reduce_dtype = torch.float16
-
-        self.forward_metadata = None
-
-        self.cuda_graph_max_seq_len = model_runner.model_config.context_len
-
-    def init_forward_metadata(
-        self, batch: ScheduleBatch, input_metadata: InputMetadata
-    ):
-        """Init auxiliary variables for triton attention backend."""
-
-        if input_metadata.forward_mode.is_decode():
-            start_loc = torch.zeros_like(input_metadata.seq_lens, dtype=torch.int32)
-            start_loc[1:] = torch.cumsum(input_metadata.seq_lens[:-1], dim=0)
-
-            total_num_tokens = torch.sum(input_metadata.seq_lens).item()
-            attn_logits = torch.empty(
-                (self.num_head, total_num_tokens),
-                dtype=self.reduce_dtype,
-                device="cuda",
-            )
-
-            max_seq_len = torch.max(input_metadata.seq_lens).item()
-            max_extend_len = None
-        else:
-            start_loc = attn_logits = max_seq_len = None
-            prefix_lens = torch.tensor(batch.prefix_lens_cpu, device="cuda")
-            max_extend_len = torch.max(input_metadata.seq_lens - prefix_lens).item()
-
-        self.forward_metadata = start_loc, attn_logits, max_seq_len, max_extend_len
-
-    def init_cuda_graph_state(self, max_bs: int):
-        self.cuda_graph_max_total_num_tokens = max_bs * self.cuda_graph_max_seq_len
-
-        self.cuda_graph_start_loc = torch.zeros(
-            (max_bs,), dtype=torch.int32, device="cuda"
-        )
-        self.cuda_graph_attn_logits = torch.empty(
-            (
-                self.num_head,
-                self.cuda_graph_max_total_num_tokens,
-            ),
-            dtype=self.reduce_dtype,
-            device="cuda",
-        )
-
-    def init_forward_metadata_capture_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
-    ):
-        self.forward_metadata = (
-            self.cuda_graph_start_loc,
-            self.cuda_graph_attn_logits,
-            self.cuda_graph_max_seq_len,
-            None,
-        )
-
-    def init_forward_metadata_replay_cuda_graph(
-        self, bs: int, req_pool_indices, seq_lens
-    ):
-        self.cuda_graph_start_loc.zero_()
-        self.cuda_graph_start_loc[1:bs] = torch.cumsum(seq_lens[: bs - 1], dim=0)
-
-    def get_cuda_graph_seq_len_fill_value(self):
-        return 1
-
-    def forward_extend(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        # TODO: reuse the buffer across layers
-        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)
-
-        input_metadata.token_to_kv_pool.set_kv_buffer(
-            layer.layer_id, input_metadata.out_cache_loc, k, v
-        )
-
-        start_loc, attn_logits, max_seq_len, max_extend_len = self.forward_metadata
-        self.extend_attention_fwd(
-            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-            k.contiguous(),
-            v.contiguous(),
-            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
-            input_metadata.token_to_kv_pool.get_key_buffer(layer.layer_id),
-            input_metadata.token_to_kv_pool.get_value_buffer(layer.layer_id),
-            input_metadata.req_to_token_pool.req_to_token,
-            input_metadata.req_pool_indices,
-            input_metadata.seq_lens,
-            input_metadata.extend_seq_lens,
-            input_metadata.extend_start_loc,
-            max_extend_len,
-            layer.scaling,
-            layer.logit_cap,
-        )
-        return o
-
-    def forward_decode(self, q, k, v, layer: nn.Module, input_metadata: InputMetadata):
-        # During torch.compile, there is a bug in rotary_emb that causes the
-        # output value to have a 3D tensor shape. This reshapes the output correctly.
-        q = q.reshape(-1, layer.tp_q_head_num * layer.qk_head_dim)
-
-        # TODO: reuse the buffer across layers
-        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)
-
-        start_loc, attn_logits, max_seq_len, max_extend_len = self.forward_metadata
-
-        input_metadata.token_to_kv_pool.set_kv_buffer(
-            layer.layer_id, input_metadata.out_cache_loc, k, v
-        )
-
-        self.decode_attention_fwd(
-            q.view(-1, layer.tp_q_head_num, layer.qk_head_dim),
-            input_metadata.token_to_kv_pool.get_key_buffer(layer.layer_id),
-            input_metadata.token_to_kv_pool.get_value_buffer(layer.layer_id),
-            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
-            input_metadata.req_to_token_pool.req_to_token,
-            input_metadata.req_pool_indices,
-            start_loc,
-            input_metadata.seq_lens,
-            attn_logits,
-            max_seq_len,
-            layer.scaling,
-            layer.logit_cap,
-        )
-        return o

sglang/srt/managers/controller_multi.py

@@ -1,207 +0,0 @@
-"""
-Copyright 2023-2024 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.
-"""
-
-"""
-A controller that manages multiple data parallel workers.
-Each data parallel worker can manage multiple tensor parallel workers.
-"""
-
-import dataclasses
-import logging
-import multiprocessing
-from enum import Enum, auto
-
-import numpy as np
-import zmq
-
-from sglang.srt.managers.controller_single import (
-    start_controller_process as start_controller_process_single,
-)
-from sglang.srt.managers.io_struct import (
-    AbortReq,
-    FlushCacheReq,
-    TokenizedGenerateReqInput,
-)
-from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import configure_logger, kill_parent_process
-from sglang.utils import get_exception_traceback
-
-logger = logging.getLogger(__name__)
-
-
-class LoadBalanceMethod(Enum):
-    """Load balance method."""
-
-    ROUND_ROBIN = auto()
-    SHORTEST_QUEUE = auto()
-
-    @classmethod
-    def from_str(cls, method: str):
-        method = method.upper()
-        try:
-            return cls[method]
-        except KeyError as exc:
-            raise ValueError(f"Invalid load balance method: {method}") from exc
-
-
-@dataclasses.dataclass
-class WorkerHandle:
-    """Store the handle of a data parallel worker."""
-
-    proc: multiprocessing.Process
-    queue: multiprocessing.Queue
-
-
-class ControllerMulti:
-    """A controller that manages multiple data parallel workers."""
-
-    def __init__(
-        self,
-        server_args: ServerArgs,
-        port_args: PortArgs,
-    ):
-        # Parse args
-        self.server_args = server_args
-        self.port_args = port_args
-        self.load_balance_method = LoadBalanceMethod.from_str(
-            server_args.load_balance_method
-        )
-
-        # Init communication
-        context = zmq.Context()
-        self.recv_from_tokenizer = context.socket(zmq.PULL)
-        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.controller_port}")
-
-        # Dispatch method
-        self.round_robin_counter = 0
-        dispatch_lookup = {
-            LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
-            LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler,
-        }
-        self.dispatching = dispatch_lookup[self.load_balance_method]
-
-        # Start data parallel workers
-        self.workers = []
-        for i in range(server_args.dp_size):
-            self.start_dp_worker(i)
-
-    def start_dp_worker(self, dp_worker_id: int):
-        tp_size = self.server_args.tp_size
-
-        pipe_controller_reader, pipe_controller_writer = multiprocessing.Pipe(
-            duplex=False
-        )
-
-        gpu_ids = list(range(dp_worker_id * tp_size, (dp_worker_id + 1) * tp_size))
-        queue = multiprocessing.Queue()
-        proc = multiprocessing.Process(
-            target=start_controller_process_single,
-            args=(
-                self.server_args,
-                self.port_args,
-                pipe_controller_writer,
-                True,
-                gpu_ids,
-                dp_worker_id,
-                queue,
-            ),
-        )
-        proc.start()
-
-        controller_init_state = pipe_controller_reader.recv()
-        if controller_init_state != "init ok":
-            raise RuntimeError(
-                f"Initialization failed. controller_init_state: {controller_init_state}"
-            )
-        self.workers.append(
-            WorkerHandle(
-                proc=proc,
-                queue=queue,
-            )
-        )
-
-    def round_robin_scheduler(self, input_requests):
-        for r in input_requests:
-            self.workers[self.round_robin_counter].queue.put(r)
-            self.round_robin_counter = (self.round_robin_counter + 1) % len(
-                self.workers
-            )
-
-    def shortest_queue_scheduler(self, input_requests):
-        for r in input_requests:
-            queue_sizes = [worker.queue.qsize() for worker in self.workers]
-            wid = np.argmin(queue_sizes)
-            self.workers[wid].queue.put(r)
-
-    def loop_for_forward(self):
-        while True:
-            recv_reqs = self.recv_requests()
-            self.dispatching(recv_reqs)
-
-    def recv_requests(self):
-        recv_reqs = []
-
-        while True:
-            try:
-                recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
-            except zmq.ZMQError:
-                break
-
-            if isinstance(recv_req, FlushCacheReq):
-                # TODO(lsyin): apply more specific flushCacheReq
-                for worker in self.workers:
-                    worker.queue.put(recv_req)
-            elif isinstance(recv_req, AbortReq):
-                in_queue = False
-                for i, req in enumerate(recv_reqs):
-                    if req.rid == recv_req.rid:
-                        recv_reqs[i] = recv_req
-                        in_queue = True
-                        break
-                if not in_queue:
-                    # Send abort req to all TP groups
-                    for worker in self.workers:
-                        worker.queue.put(recv_req)
-            elif isinstance(recv_req, TokenizedGenerateReqInput):
-                recv_reqs.append(recv_req)
-            else:
-                logger.error(f"Invalid object: {recv_req}")
-
-        return recv_reqs
-
-
-def start_controller_process(
-    server_args: ServerArgs,
-    port_args: PortArgs,
-    pipe_writer,
-):
-    """Start a controller process."""
-
-    configure_logger(server_args)
-
-    try:
-        controller = ControllerMulti(server_args, port_args)
-    except Exception:
-        pipe_writer.send(get_exception_traceback())
-        raise
-
-    pipe_writer.send("init ok")
-
-    try:
-        controller.loop_for_forward()
-    except Exception:
-        logger.error("Exception in ControllerMulti:\n" + get_exception_traceback())
-    finally:
-        kill_parent_process()