sglang 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

sglang/srt/layers/logits_processor.py

@@ -1,3 +1,8 @@
+"""Logits processing."""
+
+import dataclasses
+from typing import List
+
 import torch
 from torch import nn
 from vllm.distributed import (
@@ -5,7 +10,25 @@ from vllm.distributed import (
     tensor_model_parallel_all_gather,
 )
 
-from sglang.srt.managers.router.model_runner import ForwardMode, InputMetadata
+from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
+
+
+@dataclasses.dataclass
+class LogitProcessorOutput:
+    # The logits of the next tokens.       shape: [#seq, vocab_size]
+    next_token_logits: torch.Tensor
+    # The logprobs of the next tokens.     shape: [#seq, vocab_size]
+    next_token_logprobs: torch.Tensor
+
+    # The normlaized logprobs of prompts.  shape: [#seq]
+    normalized_prompt_logprobs: torch.Tensor
+    # The logprobs of prefill tokens.      shape: [#token, vocab_size]
+    prefill_token_logprobs: torch.Tensor
+
+    # The logprob and id of the top-k tokens in prefill positions.  shape [#seq, #token, k] of Tuple(logprob, token_id)
+    prefill_top_logprobs: List
+    # The logprob and id of the top-k tokens in decode positions.   shape [#seq, #token, k] of Tuple(logprob, token_id)
+    decode_top_logprobs: List
 
 
 class LogitsProcessor(nn.Module):
@@ -37,6 +60,7 @@ class LogitsProcessor(nn.Module):
         return normalized_prompt_logprobs
 
     def _get_top_logprobs(self, all_logprobs, input_metadata: InputMetadata):
+        # TODO: vectorize the code below
         if input_metadata.forward_mode == ForwardMode.DECODE:
             decode_top_logprobs = []
             for i in range(all_logprobs.shape[0]):
@@ -49,37 +73,34 @@ class LogitsProcessor(nn.Module):
         else:
             prefill_top_logprobs, decode_top_logprobs = [], []
             pt = 0
-            # NOTE: the GPU-CPU overhead can be reduced
-            extend_seq_lens_cpu = input_metadata.extend_seq_lens.cpu().numpy()
-            for i in range(len(extend_seq_lens_cpu)):
-                if extend_seq_lens_cpu[i] == 0:
+            extend_seq_lens_cpu = input_metadata.extend_seq_lens.tolist()
+            for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
+                if extend_seq_len == 0:
                     prefill_top_logprobs.append([])
                     decode_top_logprobs.append([])
                     continue
                 k = input_metadata.top_logprobs_nums[i]
-                t = all_logprobs[pt : pt + extend_seq_lens_cpu[i]].topk(k)
+                t = all_logprobs[pt : pt + extend_seq_len].topk(k)
                 vs_cpu = t.values.tolist()
                 ps_cpu = t.indices.tolist()
                 prefill_top_logprobs.append(
                     [list(zip(vs_cpu[j], ps_cpu[j])) for j in range(len(vs_cpu) - 1)]
                 )
                 decode_top_logprobs.append(list(zip(vs_cpu[-1], ps_cpu[-1])))
-                pt += extend_seq_lens_cpu[i]
+                pt += extend_seq_len
+
             return prefill_top_logprobs, decode_top_logprobs
 
     def forward(self, input_ids, hidden_states, weight, input_metadata: InputMetadata):
-        # Get last index for next token prediction, except for DECODE mode.
-        last_index = None
-        if input_metadata.forward_mode != ForwardMode.DECODE:
+        # Get the last hidden states and last logits for the next token prediction
+        if input_metadata.forward_mode == ForwardMode.DECODE:
+            last_index = None
+            last_hidden = hidden_states
+        else:
             last_index = (
                 torch.cumsum(input_metadata.extend_seq_lens, dim=0, dtype=torch.long)
                 - 1
             )
-
-        # Get the last hidden states and last logits
-        if input_metadata.forward_mode == ForwardMode.DECODE:
-            last_hidden = hidden_states
-        else:
             last_hidden = hidden_states[last_index]
 
         last_logits = torch.matmul(last_hidden, weight.T)
@@ -89,8 +110,14 @@ class LogitsProcessor(nn.Module):
 
         # Return only last_logits if logprob is not requested
         if not input_metadata.return_logprob:
-            hidden_states = None
-            return last_logits, (None, None, None, None, None)
+            return LogitProcessorOutput(
+                next_token_logits=last_logits,
+                next_token_logprobs=None,
+                normalized_prompt_logprobs=None,
+                prefill_token_logprobs=None,
+                prefill_top_logprobs=None,
+                decode_top_logprobs=None,
+            )
         else:
             # When logprob is requested, compute the logits for all tokens.
             if input_metadata.forward_mode == ForwardMode.DECODE:
@@ -105,6 +132,7 @@ class LogitsProcessor(nn.Module):
             del all_logits
             all_logprobs[:] = torch.nn.functional.log_softmax(all_logprobs, dim=-1)
 
+            # Get the logprob of top-k tokens
             return_top_logprob = any(x > 0 for x in input_metadata.top_logprobs_nums)
             if return_top_logprob:
                 prefill_top_logprobs, decode_top_logprobs = self._get_top_logprobs(
@@ -114,16 +142,15 @@ class LogitsProcessor(nn.Module):
                 prefill_top_logprobs = decode_top_logprobs = None
 
             if input_metadata.forward_mode == ForwardMode.DECODE:
-                last_logprobs = all_logprobs
-                return last_logits, (
-                    None,
-                    None,
-                    None,
-                    decode_top_logprobs,
-                    last_logprobs,
+                return LogitProcessorOutput(
+                    next_token_logits=last_logits,
+                    next_token_logprobs=all_logprobs,
+                    normalized_prompt_logprobs=None,
+                    prefill_token_logprobs=None,
+                    prefill_top_logprobs=None,
+                    decode_top_logprobs=decode_top_logprobs,
                 )
             else:
-                # Compute the logprobs for the last token of each request.
                 last_logprobs = all_logprobs[last_index]
 
                 # Compute the logprobs and normalized logprobs for the prefill tokens.
@@ -136,16 +163,18 @@ class LogitsProcessor(nn.Module):
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
                     prefill_token_logprobs, input_metadata
                 )
-                return last_logits, (
-                    prefill_token_logprobs,
-                    normalized_prompt_logprobs,
-                    prefill_top_logprobs,
-                    decode_top_logprobs,
-                    last_logprobs,
+
+                return LogitProcessorOutput(
+                    next_token_logits=last_logits,
+                    next_token_logprobs=last_logprobs,
+                    normalized_prompt_logprobs=normalized_prompt_logprobs,
+                    prefill_token_logprobs=prefill_token_logprobs,
+                    prefill_top_logprobs=prefill_top_logprobs,
+                    decode_top_logprobs=decode_top_logprobs,
                 )
 
 
-if __name__ == "__main__":
+def test():
     all_logprobs = torch.tensor(
         #       s                     s                s
         [[0, 1, 2, 3], [1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6], [4, 5, 6, 7]],
@@ -173,3 +202,7 @@ if __name__ == "__main__":
     print("start", start)
     print("end", end)
     print("sum_logp", sum_logp)
+
+
+if __name__ == "__main__":
+    test()

sglang/srt/layers/radix_attention.py

@@ -1,14 +1,21 @@
+"""Radix attention."""
+
+import numpy as np
 import torch
 from torch import nn
 
+from sglang.global_config import global_config
 from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
-from sglang.srt.managers.router.model_runner import ForwardMode, InputMetadata
+from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
 
 
 class RadixAttention(nn.Module):
-    def __init__(self, num_heads, head_dim, scaling, num_kv_heads, layer_id):
+    def __init__(
+        self, num_heads: int, head_dim: int, scaling: float, num_kv_heads: int,
+        layer_id: int, logit_cap: int = -1
+    ):
         super().__init__()
         self.tp_q_head_num = num_heads
         self.tp_k_head_num = num_kv_heads
@@ -16,16 +23,21 @@ class RadixAttention(nn.Module):
         self.head_dim = head_dim
         self.layer_id = layer_id
 
-        from sglang.srt.managers.router.model_runner import global_server_args_dict
+        assert np.allclose(scaling, 1.0 / (head_dim**0.5))
+
+        from sglang.srt.managers.controller.model_runner import global_server_args_dict
 
-        if global_server_args_dict.get("enable_flashinfer", False):
+        if not global_server_args_dict.get("disable_flashinfer", False):
             self.prefill_forward = self.prefill_forward_flashinfer
             self.extend_forward = self.prefill_forward_flashinfer
             self.decode_forward = self.decode_forward_flashinfer
+            # flashinfer now accepts float logit_cap argument
+            self.logit_cap = logit_cap if logit_cap > 0 else 0
         else:
             self.prefill_forward = self.prefill_forward_triton
             self.extend_forward = self.extend_forward_triton
             self.decode_forward = self.decode_forward_triton
+            self.logit_cap = logit_cap
 
     def prefill_forward_triton(self, q, k, v, input_metadata: InputMetadata):
         o = torch.empty_like(q)
@@ -38,6 +50,7 @@ class RadixAttention(nn.Module):
             input_metadata.start_loc,
             input_metadata.seq_lens,
             input_metadata.max_seq_len,
+            self.logit_cap,
         )
         self.store_kv_cache(k, v, input_metadata)
 
@@ -62,6 +75,7 @@ class RadixAttention(nn.Module):
             input_metadata.extend_seq_lens,
             input_metadata.max_seq_len,
             input_metadata.max_extend_len,
+            self.logit_cap,
         )
 
         return o
@@ -82,6 +96,7 @@ class RadixAttention(nn.Module):
             input_metadata.max_seq_len,
             input_metadata.other_kv_index,
             input_metadata.total_num_tokens,
+            self.logit_cap,
         )
 
         return o
@@ -89,19 +104,38 @@ class RadixAttention(nn.Module):
     def prefill_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
         self.store_kv_cache(k, v, input_metadata)
 
-        o = input_metadata.prefill_wrapper.forward(
+        o1, s1 = input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
+            v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
+            logits_soft_cap=self.logit_cap,
         )
 
+        if input_metadata.no_prefix:
+            o = o1
+        else:
+            o2, s2 = input_metadata.flashinfer_prefill_wrapper_paged.forward_return_lse(
+                q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
+                input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+                causal=False,
+                logits_soft_cap=self.logit_cap,
+            )
+
+            from flashinfer.cascade import merge_state
+            o, _ = merge_state(o1, s1, o2, s2)
+
+        if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
+            torch.cuda.synchronize()
+
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
     def decode_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
         self.store_kv_cache(k, v, input_metadata)
 
-        o = input_metadata.decode_wrapper.forward(
+        o = input_metadata.flashinfer_decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            logits_soft_cap=self.logit_cap,
         )
 
         return o.view(-1, self.tp_q_head_num * self.head_dim)

sglang/srt/layers/token_attention.py

@@ -5,7 +5,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.managers.router.model_runner import global_server_args_dict
+from sglang.srt.managers.controller.model_runner import global_server_args_dict
 from sglang.srt.utils import wrap_kernel_launcher
 
 if global_server_args_dict.get("attention_reduce_in_fp32", False):
@@ -16,6 +16,12 @@ else:
     REDUCE_TORCH_TYPE = torch.float16
 
 
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
 @triton.jit
 def _fwd_kernel_stage1(
     Q,
@@ -35,6 +41,7 @@ def _fwd_kernel_stage1(
     kv_group_num: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    logit_cap: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -77,6 +84,10 @@ def _fwd_kernel_stage1(
         ).to(REDUCE_TRITON_TYPE)
         att_value = tl.sum(q[None, :] * k, 1)
         att_value *= sm_scale
+
+        if logit_cap > 0:
+            att_value = logit_cap * tanh(att_value / logit_cap)
+
         off_o = cur_head * att_stride_h + (cur_batch_in_all_start_index + offs_n)
         tl.store(Att_Out + off_o, att_value, mask=offs_n_new < cur_batch_end_index)
 
@@ -165,6 +176,7 @@ def _token_att_m_fwd(
     B_Start_Loc,
     B_Seqlen,
     max_len_in_batch,
+    logit_cap,
 ):
     BLOCK = 32
     # shape constraints
@@ -223,6 +235,7 @@ def _token_att_m_fwd(
         kv_group_num=kv_group_num,
         BLOCK_DMODEL=Lk,
         BLOCK_N=BLOCK,
+        logit_cap=logit_cap,
         num_warps=num_warps,
         num_stages=1,
     )
@@ -304,6 +317,7 @@ def token_attention_fwd(
     max_len_in_batch,
     other_kv_index,
     total_num_tokens,
+    logit_cap=-1,
     att_m=None,
 ):
     if att_m is None:
@@ -320,6 +334,7 @@ def token_attention_fwd(
         b_start_loc,
         b_seq_len,
         max_len_in_batch,
+        logit_cap,
     )
     _token_softmax_reducev_fwd(
         att_m,

sglang/srt/managers/controller/dp_worker.py

@@ -0,0 +1,113 @@
+"""A data parallel worker thread."""
+
+import asyncio
+import logging
+import queue
+import threading
+from typing import Callable, List
+
+import uvloop
+import zmq
+
+from sglang.global_config import global_config
+from sglang.srt.managers.controller.tp_worker import ModelTpClient
+from sglang.srt.managers.io_struct import BatchTokenIDOut
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import kill_parent_process
+from sglang.utils import get_exception_traceback
+
+logger = logging.getLogger("srt.controller")
+CHECKING_INTERVAL = 5
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+
+class DataParallelWorkerThread(threading.Thread):
+    def __init__(
+        self,
+        worker_id: int,
+        request_queue: queue.Queue,
+        detokenizer_port: int,
+        step_func: Callable,
+    ):
+        super(DataParallelWorkerThread, self).__init__()
+        self.worker_id = worker_id
+        self.request_queue = request_queue
+        self.liveness = True
+        self.request_dependency_delay = global_config.request_dependency_delay
+
+        context = zmq.asyncio.Context()
+        self.send_to_detokenizer = context.socket(zmq.PUSH)
+        self.send_to_detokenizer.connect(f"tcp://127.0.0.1:{detokenizer_port}")
+
+        self.step = step_func
+
+    async def loop_for_forward(self):
+        while self.liveness:
+            requests = []
+            while not self.request_queue.empty():
+                requests.append(self.request_queue.get())
+
+            out_pyobjs: List[BatchTokenIDOut] = []
+            try:
+                out_pyobjs = await self.step(requests)
+            except Exception:
+                for r in requests:
+                    self.request_queue.put(r)
+                logger.error(
+                    f"Worker thread {self.worker_id}: "
+                    f"failed to get back from Model Server\n"
+                    f"{get_exception_traceback()}"
+                )
+                self.liveness = False
+                # Crash the whole server when there are any errors.
+                # TODO(lianmin): make this an option.
+                kill_parent_process()
+                return
+
+            for obj in out_pyobjs:
+                self.send_to_detokenizer.send_pyobj(obj)
+
+            # async sleep for receiving the subsequent request and avoiding cache miss
+            if len(out_pyobjs) != 0:
+                has_finished = any(
+                    [obj.finished_reason is not None for obj in out_pyobjs]
+                )
+                if has_finished:
+                    await asyncio.sleep(self.request_dependency_delay)
+            await asyncio.sleep(global_config.wait_for_new_request_delay)
+
+    async def monitoring(self):
+        while True:
+            await asyncio.sleep(CHECKING_INTERVAL)
+            # can plug in monitoring logic here
+
+    def run(self):
+        logger.info(f"DataParallelWorkerThread {self.worker_id} start")
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.create_task(self.monitoring())
+        loop.run_until_complete(self.loop_for_forward())
+
+
+def start_data_parallel_worker(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    model_overide_args,
+    gpu_ids: List[int],
+    worker_id: int,
+):
+    model_tp_client = ModelTpClient(
+        gpu_ids,
+        server_args,
+        port_args.model_port_args[worker_id],
+        model_overide_args,
+    )
+    worker_thread = DataParallelWorkerThread(
+        worker_id=worker_id,
+        request_queue=queue.Queue(),
+        detokenizer_port=port_args.detokenizer_port,
+        step_func=model_tp_client.step,
+    )
+    worker_thread.start()
+    return worker_thread