sglang 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

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,110 @@
+"""A data parallel worker thread."""
+import asyncio
+import logging
+import queue
+import threading
+from typing import List, Callable
+
+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