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

sglang/srt/layers/logits_processor.py

@@ -1,3 +1,8 @@
+"""Logits processing."""
+
+import dataclasses
+from typing import List, Union
+
 import torch
 from torch import nn
 from vllm.distributed import (
@@ -8,6 +13,45 @@ from vllm.distributed import (
 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
+
+
+@dataclasses.dataclass
+class LogitsMetadata:
+    forward_mode: ForwardMode
+    extend_seq_lens: torch.Tensor
+    extend_start_loc: torch.Tensor
+
+    # For logprobs
+    return_logprob: bool
+    top_logprobs_nums: List[int]
+
+    @classmethod
+    def from_input_metadata(cls, input_metadata: InputMetadata):
+        return cls(
+            forward_mode=input_metadata.forward_mode,
+            extend_seq_lens=input_metadata.extend_seq_lens,
+            extend_start_loc=input_metadata.extend_start_loc,
+            return_logprob=input_metadata.return_logprob,
+            top_logprobs_nums=input_metadata.top_logprobs_nums,
+        )
+
+
 class LogitsProcessor(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -15,14 +59,14 @@ class LogitsProcessor(nn.Module):
         self.tp_size = get_tensor_model_parallel_world_size()
 
     def _get_normalized_prompt_logprobs(
-        self, prefill_token_logprobs, input_metadata: InputMetadata
+        self, prefill_token_logprobs, logits_metadata: LogitsMetadata
     ):
         logprobs_cumsum = torch.cumsum(
             prefill_token_logprobs, dim=0, dtype=torch.float32
         )
 
-        start = input_metadata.extend_start_loc.clone()
-        end = start + input_metadata.extend_seq_lens - 2
+        start = logits_metadata.extend_start_loc.clone()
+        end = start + logits_metadata.extend_seq_lens - 2
         start.clamp_(min=0, max=prefill_token_logprobs.shape[0] - 1)
         end.clamp_(min=0, max=prefill_token_logprobs.shape[0] - 1)
         sum_logp = (
@@ -31,16 +75,17 @@ class LogitsProcessor(nn.Module):
             + prefill_token_logprobs[start]
         )
         normalized_prompt_logprobs = sum_logp / (
-            (input_metadata.extend_seq_lens - 1).clamp(min=1)
+            (logits_metadata.extend_seq_lens - 1).clamp(min=1)
         )
 
         return normalized_prompt_logprobs
 
-    def _get_top_logprobs(self, all_logprobs, input_metadata: InputMetadata):
-        if input_metadata.forward_mode == ForwardMode.DECODE:
+    def _get_top_logprobs(self, all_logprobs, logits_metadata: LogitsMetadata):
+        # TODO: vectorize the code below
+        if logits_metadata.forward_mode == ForwardMode.DECODE:
             decode_top_logprobs = []
             for i in range(all_logprobs.shape[0]):
-                k = input_metadata.top_logprobs_nums[i]
+                k = logits_metadata.top_logprobs_nums[i]
                 t = all_logprobs[i].topk(k)
                 v_cpu = t.values.tolist()
                 p_cpu = t.indices.tolist()
@@ -49,14 +94,13 @@ 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.tolist()
+            extend_seq_lens_cpu = logits_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]
+                k = logits_metadata.top_logprobs_nums[i]
                 t = all_logprobs[pt : pt + extend_seq_len].topk(k)
                 vs_cpu = t.values.tolist()
                 ps_cpu = t.indices.tolist()
@@ -68,19 +112,26 @@ class LogitsProcessor(nn.Module):
 
             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:
-            last_index = (
-                torch.cumsum(input_metadata.extend_seq_lens, dim=0, dtype=torch.long)
-                - 1
-            )
+    def forward(
+        self,
+        input_ids,
+        hidden_states,
+        weight,
+        logits_metadata: Union[LogitsMetadata, InputMetadata],
+    ):
+        if isinstance(logits_metadata, InputMetadata):
+            logits_metadata = LogitsMetadata.from_input_metadata(logits_metadata)
+        assert isinstance(logits_metadata, LogitsMetadata)
 
-        # Get the last hidden states and last logits
-        if input_metadata.forward_mode == ForwardMode.DECODE:
+        # Get the last hidden states and last logits for the next token prediction
+        if logits_metadata.forward_mode == ForwardMode.DECODE:
+            last_index = None
             last_hidden = hidden_states
         else:
+            last_index = (
+                torch.cumsum(logits_metadata.extend_seq_lens, dim=0, dtype=torch.long)
+                - 1
+            )
             last_hidden = hidden_states[last_index]
 
         last_logits = torch.matmul(last_hidden, weight.T)
@@ -88,13 +139,24 @@ class LogitsProcessor(nn.Module):
             last_logits = tensor_model_parallel_all_gather(last_logits)
         last_logits = last_logits[:, : self.config.vocab_size]
 
+        if hasattr(self.config, "final_logit_softcapping"):
+            last_logits /= self.config.final_logit_softcapping
+            last_logits = torch.tanh(last_logits)
+            last_logits *= self.config.final_logit_softcapping
+
         # 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)
+        if not logits_metadata.return_logprob:
+            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:
+            if logits_metadata.forward_mode == ForwardMode.DECODE:
                 all_logits = last_logits
             else:
                 all_logits = torch.matmul(hidden_states, weight.T)
@@ -106,25 +168,25 @@ class LogitsProcessor(nn.Module):
             del all_logits
             all_logprobs[:] = torch.nn.functional.log_softmax(all_logprobs, dim=-1)
 
-            return_top_logprob = any(x > 0 for x in input_metadata.top_logprobs_nums)
+            # Get the logprob of top-k tokens
+            return_top_logprob = any(x > 0 for x in logits_metadata.top_logprobs_nums)
             if return_top_logprob:
                 prefill_top_logprobs, decode_top_logprobs = self._get_top_logprobs(
-                    all_logprobs, input_metadata
+                    all_logprobs, logits_metadata
                 )
             else:
                 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,
+            if logits_metadata.forward_mode == ForwardMode.DECODE:
+                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.
@@ -135,14 +197,16 @@ class LogitsProcessor(nn.Module):
                 ]
 
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
-                    prefill_token_logprobs, input_metadata
+                    prefill_token_logprobs, logits_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,
                 )
 
 

sglang/srt/layers/radix_attention.py

@@ -1,52 +1,52 @@
-import torch
+"""Radix attention."""
+
 import numpy as np
+import torch
+from flashinfer.cascade import merge_state
 from torch import nn
 
-from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
+from sglang.global_config import global_config
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
 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, logit_cap=-1):
+    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
         self.tp_v_head_num = num_kv_heads
         self.head_dim = head_dim
+        self.scaling = scaling
         self.layer_id = layer_id
-        self.logit_cap = logit_cap
-
-        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 is not None and 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 if logit_cap is not None else 0
 
     def prefill_forward_triton(self, q, k, v, input_metadata: InputMetadata):
-        o = torch.empty_like(q)
-
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num, self.head_dim),
-            k,
-            v,
-            o.view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.start_loc,
-            input_metadata.seq_lens,
-            input_metadata.max_seq_len,
-            self.logit_cap,
-        )
-        self.store_kv_cache(k, v, input_metadata)
-
-        return o
+        # In SGLang, we call both the typical "prefill" and "prefill with cache" as "extend".
+        # See the extend_forward_xxx functions.
+        raise NotImplementedError()
 
     def extend_forward_triton(self, q, k, v, input_metadata: InputMetadata):
         o = torch.empty_like(q)
@@ -67,7 +67,8 @@ class RadixAttention(nn.Module):
             input_metadata.extend_seq_lens,
             input_metadata.max_seq_len,
             input_metadata.max_extend_len,
-            self.logit_cap,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
@@ -88,27 +89,50 @@ class RadixAttention(nn.Module):
             input_metadata.max_seq_len,
             input_metadata.other_kv_index,
             input_metadata.total_num_tokens,
-            self.logit_cap,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
 
     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),
+            causal=True,
+            sm_scale=self.scaling,
+            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,
+                sm_scale=self.scaling,
+                logits_soft_cap=self.logit_cap,
+            )
+
+            o, _ = merge_state(o1, s1, o2, s2)
+
+        self.store_kv_cache(k, v, input_metadata)
+
+        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],
+            sm_scale=self.scaling,
+            logits_soft_cap=self.logit_cap,
         )
 
         return o.view(-1, self.tp_q_head_num * self.head_dim)

sglang/srt/layers/token_attention.py

@@ -176,6 +176,7 @@ def _token_att_m_fwd(
     B_Start_Loc,
     B_Seqlen,
     max_len_in_batch,
+    sm_scale,
     logit_cap,
 ):
     BLOCK = 32
@@ -183,7 +184,6 @@ def _token_att_m_fwd(
     Lq, Lk = q.shape[-1], k_buffer.shape[-1]
     assert Lq == Lk
     assert Lk in {16, 32, 64, 128, 256}
-    sm_scale = 1.0 / (Lk**0.5)
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
@@ -317,6 +317,7 @@ def token_attention_fwd(
     max_len_in_batch,
     other_kv_index,
     total_num_tokens,
+    sm_scale=None,
     logit_cap=-1,
     att_m=None,
 ):
@@ -324,6 +325,7 @@ def token_attention_fwd(
         att_m = torch.empty(
             (q.shape[-2], total_num_tokens), dtype=REDUCE_TORCH_TYPE, device="cuda"
         )
+    sm_scale = 1.0 / (Lq**0.5) if sm_scale is None else sm_scale
 
     _token_att_m_fwd(
         q,
@@ -334,6 +336,7 @@ def token_attention_fwd(
         b_start_loc,
         b_seq_len,
         max_len_in_batch,
+        sm_scale,
         logit_cap,
     )
     _token_softmax_reducev_fwd(

sglang/srt/managers/controller/dp_worker.py

@@ -1,9 +1,10 @@
 """A data parallel worker thread."""
+
 import asyncio
 import logging
 import queue
 import threading
-from typing import List, Callable
+from typing import Callable, List
 
 import uvloop
 import zmq
@@ -69,7 +70,9 @@ class DataParallelWorkerThread(threading.Thread):
 
             # 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])
+                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)
@@ -107,4 +110,4 @@ def start_data_parallel_worker(
         step_func=model_tp_client.step,
     )
     worker_thread.start()
-    return worker_thread
+    return worker_thread