sglang 0.1.20__py3-none-any.whl → 0.1.22__py3-none-any.whl

sglang/srt/layers/radix_attention.py

@@ -7,8 +7,8 @@ from torch import nn
 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.infer_batch import global_server_args_dict
 from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
+from sglang.srt.server import global_server_args_dict
 
 
 class RadixAttention(nn.Module):
@@ -136,7 +136,33 @@ class RadixAttention(nn.Module):
             return self.decode_forward(q, k, v, input_metadata)
 
     def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
-        key_buffer = input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id)
-        key_buffer[input_metadata.out_cache_loc] = cache_k
-        value_buffer = input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id)
-        value_buffer[input_metadata.out_cache_loc] = cache_v
+        kv_cache = input_metadata.token_to_kv_pool.kv_data[self.layer_id]
+        _store_kv_cache(cache_k, cache_v, kv_cache, input_metadata.out_cache_loc)
+
+
+try:
+
+    @torch.library.custom_op("mylib::store_kv_cache", mutates_args={"kv_cache"})
+    def _store_kv_cache(
+        k: torch.Tensor,
+        v: torch.Tensor,
+        kv_cache: torch.Tensor,
+        cache_loc: torch.Tensor,
+    ) -> None:
+        kv_cache[cache_loc, 0] = k
+        kv_cache[cache_loc, 1] = v
+
+    @_store_kv_cache.register_fake
+    def _(k, v, kv_cache, cache_loc):
+        pass
+
+except:
+
+    def _store_kv_cache(
+        k: torch.Tensor,
+        v: torch.Tensor,
+        kv_cache: torch.Tensor,
+        cache_loc: torch.Tensor,
+    ) -> None:
+        kv_cache[cache_loc, 0] = k
+        kv_cache[cache_loc, 1] = v

sglang/srt/layers/token_attention.py

@@ -5,8 +5,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.managers.controller.model_runner import global_server_args_dict
-from sglang.srt.utils import wrap_kernel_launcher
+from sglang.srt.server import global_server_args_dict
 
 if global_server_args_dict.get("attention_reduce_in_fp32", False):
     REDUCE_TRITON_TYPE = tl.float32
@@ -162,10 +161,6 @@ def _fwd_kernel_stage2(
     tl.store(out_ptrs, acc)
 
 
-cached_kernel_stage1 = None
-cached_kernel_stage2 = None
-
-
 def _token_att_m_fwd(
     q,
     k_buffer,
@@ -194,28 +189,6 @@ def _token_att_m_fwd(
     else:
         num_warps = 2
 
-    global cached_kernel_stage1
-    if cached_kernel_stage1:
-        cached_kernel_stage1(
-            grid,
-            num_warps,
-            q,
-            k_buffer,
-            sm_scale,
-            Req_to_tokens,
-            B_req_idx,
-            B_Start_Loc,
-            B_Seqlen,
-            att_out,
-            Req_to_tokens.stride(0),
-            q.stride(0),
-            q.stride(1),
-            k_buffer.stride(0),
-            k_buffer.stride(1),
-            att_out.stride(0),
-        )
-        return
-
     _fwd_kernel_stage1[grid](
         q,
         k_buffer,
@@ -238,7 +211,6 @@ def _token_att_m_fwd(
         num_warps=num_warps,
         num_stages=1,
     )
-    cached_kernel_stage1 = wrap_kernel_launcher(_fwd_kernel_stage1)
 
 
 def _token_softmax_reducev_fwd(
@@ -257,27 +229,6 @@ def _token_softmax_reducev_fwd(
 
     num_warps = 1
 
-    global cached_kernel_stage2
-    if cached_kernel_stage2:
-        cached_kernel_stage2(
-            grid,
-            num_warps,
-            logics,
-            v_buffer,
-            o,
-            req_to_tokens,
-            b_req_idx,
-            b_start_loc,
-            b_seq_len,
-            logics.stride(0),
-            v_buffer.stride(0),
-            v_buffer.stride(1),
-            o.stride(0),
-            o.stride(1),
-            req_to_tokens.stride(0),
-        )
-        return
-
     _fwd_kernel_stage2[grid](
         logics,
         v_buffer,
@@ -298,7 +249,6 @@ def _token_softmax_reducev_fwd(
         num_warps=num_warps,
         num_stages=3,
     )
-    cached_kernel_stage2 = wrap_kernel_launcher(_fwd_kernel_stage2)
 
 
 def token_attention_fwd(

sglang/srt/managers/controller/cuda_graph_runner.py

@@ -3,12 +3,16 @@
 import bisect
 
 import torch
+from flashinfer import BatchDecodeWithPagedKVCacheWrapper
+from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 from vllm.distributed.parallel_state import graph_capture
 
-from sglang.global_config import global_config
 from sglang.srt.layers.logits_processor import LogitProcessorOutput
 from sglang.srt.managers.controller.infer_batch import (
-    Batch, ForwardMode, InputMetadata, init_flashinfer_args
+    Batch,
+    ForwardMode,
+    InputMetadata,
+    init_flashinfer_args,
 )
 
 
@@ -24,18 +28,28 @@ class CudaGraphRunner:
         # Common inputs
         self.max_bs = max_batch_size_to_capture
         self.input_ids = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+        self.req_pool_indices = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
         self.seq_lens = torch.ones((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.position_ids_offsets = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
-        self.out_cache_loc = torch.zeros((self.max_bs,), dtype=torch.int32, device="cuda")
+        self.position_ids_offsets = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
+        self.out_cache_loc = torch.zeros(
+            (self.max_bs,), dtype=torch.int32, device="cuda"
+        )
 
         # FlashInfer inputs
-        self.flashinfer_workspace_buffer = self.model_runner.flashinfer_workspace_buffers[0]
+        self.flashinfer_workspace_buffer = (
+            self.model_runner.flashinfer_workspace_buffers[0]
+        )
         self.flashinfer_kv_indptr = torch.zeros(
             (self.max_bs + 1,), dtype=torch.int32, device="cuda"
         )
         self.flashinfer_kv_indices = torch.zeros(
-            (self.max_bs * model_runner.model_config.context_len,), dtype=torch.int32, device="cuda"
+            (self.max_bs * model_runner.model_config.context_len,),
+            dtype=torch.int32,
+            device="cuda",
         )
         self.flashinfer_kv_last_page_len = torch.ones(
             (self.max_bs,), dtype=torch.int32, device="cuda"
@@ -49,16 +63,18 @@ class CudaGraphRunner:
         with graph_capture() as graph_capture_context:
             self.stream = graph_capture_context.stream
             for bs in batch_size_list:
-                graph, input_buffers, output_buffers, flashinfer_handler = self.capture_one_batch_size(bs)
+                (
+                    graph,
+                    input_buffers,
+                    output_buffers,
+                    flashinfer_handler,
+                ) = self.capture_one_batch_size(bs)
                 self.graphs[bs] = graph
                 self.input_buffers[bs] = input_buffers
                 self.output_buffers[bs] = output_buffers
                 self.flashinfer_handlers[bs] = flashinfer_handler
 
     def capture_one_batch_size(self, bs):
-        from flashinfer import BatchDecodeWithPagedKVCacheWrapper
-        from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
-
         graph = torch.cuda.CUDAGraph()
         stream = self.stream
 
@@ -71,17 +87,19 @@ class CudaGraphRunner:
 
         # FlashInfer inputs
         if not _grouped_size_compiled_for_decode_kernels(
-            self.model_runner.model_config.num_attention_heads // self.model_runner.tp_size,
+            self.model_runner.model_config.num_attention_heads
+            // self.model_runner.tp_size,
             self.model_runner.model_config.get_num_kv_heads(self.model_runner.tp_size),
         ):
             use_tensor_cores = True
         else:
             use_tensor_cores = False
         flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-            self.flashinfer_workspace_buffer, "NHD",
+            self.flashinfer_workspace_buffer,
+            "NHD",
             use_cuda_graph=True,
             use_tensor_cores=use_tensor_cores,
-            paged_kv_indptr_buffer=self.flashinfer_kv_indptr[:bs+1],
+            paged_kv_indptr_buffer=self.flashinfer_kv_indptr[: bs + 1],
             paged_kv_indices_buffer=self.flashinfer_kv_indices,
             paged_kv_last_page_len_buffer=self.flashinfer_kv_last_page_len[:bs],
         )
@@ -132,8 +150,8 @@ class CudaGraphRunner:
         index = bisect.bisect_left(self.batch_size_list, raw_bs)
         bs = self.batch_size_list[index]
         if bs != raw_bs:
-            self.seq_lens.zero_()
-            self.position_ids_offsets.fill_(1)
+            self.seq_lens.fill_(1)
+            self.position_ids_offsets.zero_()
             self.out_cache_loc.zero_()
 
         # Common inputs
@@ -163,10 +181,18 @@ class CudaGraphRunner:
         else:
             output = LogitProcessorOutput(
                 next_token_logits=output.next_token_logits[:raw_bs],
-                next_token_logprobs=output.next_token_logprobs[:raw_bs] if output.next_token_logprobs is not None else None,
+                next_token_logprobs=(
+                    output.next_token_logprobs[:raw_bs]
+                    if output.next_token_logprobs is not None
+                    else None
+                ),
                 normalized_prompt_logprobs=None,
                 prefill_token_logprobs=None,
                 prefill_top_logprobs=None,
-                decode_top_logprobs=output.decode_top_logprobs[:raw_bs] if output.decode_top_logprobs is not None else None,
+                decode_top_logprobs=(
+                    output.decode_top_logprobs[:raw_bs]
+                    if output.decode_top_logprobs is not None
+                    else None
+                ),
             )
         return output

sglang/srt/managers/controller/infer_batch.py

@@ -7,6 +7,7 @@ from typing import List, Union
 
 import numpy as np
 import torch
+from flashinfer.sampling import top_k_top_p_sampling_from_probs
 
 from sglang.srt.constrained import RegexGuide
 from sglang.srt.constrained.jump_forward import JumpForwardMap
@@ -15,9 +16,6 @@ from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
 
 INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
 
-# Store some global server args
-global_server_args_dict = {}
-
 
 class ForwardMode(IntEnum):
     # Prefill a new sequence. This is deprecated now. "EXTEND" covers this case.
@@ -84,6 +82,15 @@ class Req:
         self.input_ids = None  # input_ids = origin_input_ids + output_ids
 
         # For incremental decoding
+        # ----- | --------- read_ids -------|
+        # ----- |   surr_ids  |
+        # xxxxx | xxxxxxxxxxx | xxxxxxxxxxx |
+        # ----- ^ ----------- ^ ----------- ^
+        # ----- 1 ----------- 2 ----------- 3
+        # 1: surr_offset
+        # 2: read_offset
+        # 3: last token
+        self.vid = 0  # version id to sync decode status with in detokenizer_manager
         self.decoded_text = ""
         self.surr_offset = None  # Surrounding offset to defeat the cleanup algorithm
         self.read_offset = None
@@ -134,7 +141,7 @@ class Req:
         return self.finished_reason is not None
 
     # Based on https://github.com/vllm-project/vllm/blob/7a64d24aad69e4d2548aa0bf528d9fe63428ab01/vllm/transformers_utils/detokenizer.py#L194-L313
-    def init_detokenize_incrementally(self):
+    def init_incremental_detokenize(self):
         first_iter = self.surr_offset is None or self.read_offset is None
 
         if first_iter:
@@ -144,13 +151,11 @@ class Req:
             )
 
         all_ids = self.origin_input_ids_unpadded + self.output_ids
-        surr_ids = all_ids[self.surr_offset : self.read_offset]
-        read_ids = all_ids[self.surr_offset :]
-
-        return surr_ids, read_ids, len(all_ids)
+        return all_ids[self.surr_offset :], self.read_offset - self.surr_offset
 
-    def detokenize_incrementally(self, inplace: bool = True):
-        surr_ids, read_ids, num_all_tokens = self.init_detokenize_incrementally()
+    def get_next_inc_detokenization(self):
+        read_ids, read_offset = self.init_incremental_detokenize()
+        surr_ids = read_ids[:read_offset]
 
         surr_text = self.tokenizer.decode(
             surr_ids,
@@ -164,19 +169,10 @@ class Req:
         )
 
         if len(new_text) > len(surr_text) and not new_text.endswith("�"):
-            new_text = new_text[len(surr_text) :]
-            if inplace:
-                self.decoded_text += new_text
-                self.surr_offset = self.read_offset
-                self.read_offset = num_all_tokens
-
-            return True, new_text
+            return True, new_text[len(surr_text) :]
 
         return False, ""
 
-    def max_new_tokens(self):
-        return self.sampling_params.max_new_tokens
-
     def check_finished(self):
         if self.finished():
             return
@@ -275,6 +271,7 @@ class Batch:
     prefix_lens: torch.Tensor = None
     position_ids_offsets: torch.Tensor = None
     out_cache_loc: torch.Tensor = None
+    extend_num_tokens: int = None
 
     # For processing logprobs
     return_logprob: bool = False
@@ -285,10 +282,6 @@ class Batch:
     image_sizes: List[List[int]] = None
     image_offsets: List[int] = None
 
-    # Other arguments for control
-    output_ids: torch.Tensor = None
-    extend_num_tokens: int = None
-
     # Batched sampling params
     temperatures: torch.Tensor = None
     top_ps: torch.Tensor = None
@@ -330,6 +323,13 @@ class Batch:
         seq_lens = []
 
         req_pool_indices = self.req_to_token_pool.alloc(bs)
+
+        if req_pool_indices is None:
+            raise RuntimeError(
+                "Out of memory. "
+                "Please set a smaller number for `--max-running-requests`."
+            )
+
         req_pool_indices_cpu = req_pool_indices.cpu().numpy()
         for i in range(bs):
             flatten_input_ids.extend(input_ids[i])
@@ -352,7 +352,7 @@ class Batch:
         extend_num_tokens = seq_lens.sum() - prefix_lens.sum()
         out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
         if out_cache_loc is None:
-            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.dec_refs)
+            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.free)
             out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
 
             if out_cache_loc is None:
@@ -401,10 +401,10 @@ class Batch:
         ).view(-1, 1)
         self.top_ps = torch.tensor(
             [r.sampling_params.top_p for r in reqs], dtype=torch.float, device=device
-        ).view(-1, 1)
+        )
         self.top_ks = torch.tensor(
             [r.sampling_params.top_k for r in reqs], dtype=torch.int, device=device
-        ).view(-1, 1)
+        )
         self.frequency_penalties = torch.tensor(
             [r.sampling_params.frequency_penalty for r in reqs],
             dtype=torch.float,
@@ -422,7 +422,7 @@ class Batch:
         if self.token_to_kv_pool.available_size() >= bs:
             return True
 
-        self.tree_cache.evict(bs, self.token_to_kv_pool.dec_refs)
+        self.tree_cache.evict(bs, self.token_to_kv_pool.free)
 
         if self.token_to_kv_pool.available_size() >= bs:
             return True
@@ -453,7 +453,7 @@ class Batch:
             token_indices = self.req_to_token_pool.req_to_token[
                 req_pool_indices_cpu[idx]
             ][last_uncached_pos : seq_lens_cpu[idx]]
-            self.token_to_kv_pool.dec_refs(token_indices)
+            self.token_to_kv_pool.free(token_indices)
 
             # release the last node
             self.tree_cache.dec_lock_ref(req.last_node)
@@ -502,7 +502,7 @@ class Batch:
                     cur_output_ids = req.output_ids
 
                     req.output_ids.extend(suffix_ids)
-                    decode_res, new_text = req.detokenize_incrementally(inplace=False)
+                    decode_res, new_text = req.get_next_inc_detokenization()
                     if not decode_res:
                         req.output_ids = cur_output_ids
                         continue
@@ -521,6 +521,9 @@ class Batch:
                         req.output_ids = cur_output_ids
                         continue
 
+                    # The decode status has diverged from detokenizer_manager
+                    req.vid += 1
+
                     # insert the old request into tree_cache
                     if req_pool_indices_cpu is None:
                         req_pool_indices_cpu = self.req_pool_indices.tolist()
@@ -596,8 +599,7 @@ class Batch:
             "logit_bias",
         ]:
             self_val = getattr(self, item, None)
-            # logit_bias can be None
-            if self_val is not None:
+            if self_val is not None:  # logit_bias can be None
                 setattr(self, item, self_val[new_indices])
 
     def merge(self, other: "Batch"):
@@ -663,18 +665,21 @@ class Batch:
 
         # TODO(lmzheng): apply penalty
         probs = torch.softmax(logits, dim=-1)
-        probs_sort, probs_idx = _top_p_top_k(probs, self.top_ps, self.top_ks)
-        try:
-            sampled_index = torch.multinomial(probs_sort, num_samples=1)
-        except RuntimeError as e:
-            warnings.warn(f"Ignore errors in sampling: {e}")
-            sampled_index = torch.ones(probs_sort.shape[:-1] + (1,), dtype=torch.int64, device=probs.device)
-        batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(
-            -1
+
+        max_top_k_round, batch_size = 32, probs.shape[0]
+        uniform_samples = torch.rand((max_top_k_round, batch_size), device=probs.device)
+        batch_next_token_ids, _ = top_k_top_p_sampling_from_probs(
+            probs, uniform_samples, self.top_ks, self.top_ps
         )
-        batch_next_token_probs = torch.gather(
-            probs_sort, dim=1, index=sampled_index
-        ).view(-1)
+
+        # FIXME: this is a temporary fix for the illegal token ids
+        illegal_mask = torch.logical_or(
+            batch_next_token_ids < 0, batch_next_token_ids >= probs.shape[-1]
+        )
+        if torch.any(illegal_mask):
+            warnings.warn("Illegal sampled token ids")
+            probs = probs.masked_fill(torch.isnan(probs), 0.0)
+            batch_next_token_ids = torch.argmax(probs, dim=-1)
 
         if has_regex:
             batch_next_token_ids_cpu = batch_next_token_ids.cpu().numpy()
@@ -684,18 +689,7 @@ class Batch:
                         req.regex_fsm_state, batch_next_token_ids_cpu[i]
                     )
 
-        return batch_next_token_ids, batch_next_token_probs
-
-
-def _top_p_top_k(probs: torch.Tensor, top_ps: torch.Tensor, top_ks: torch.Tensor):
-    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
-    probs_sum = torch.cumsum(probs_sort, dim=-1)
-    probs_sort[(probs_sum - probs_sort) > top_ps] = 0.0
-    probs_sort[
-        torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1) >= top_ks
-    ] = 0.0
-    probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
-    return probs_sort, probs_idx
+        return batch_next_token_ids
 
 
 @dataclass
@@ -749,8 +743,14 @@ class InputMetadata:
         skip_flashinfer_init=False,
     ):
         if not skip_flashinfer_init and not model_runner.server_args.disable_flashinfer:
-            init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens, prefix_lens,
-                                 model_runner.flashinfer_decode_wrapper)
+            init_flashinfer_args(
+                forward_mode,
+                model_runner,
+                req_pool_indices,
+                seq_lens,
+                prefix_lens,
+                model_runner.flashinfer_decode_wrapper,
+            )
 
         batch_size = len(req_pool_indices)
 
@@ -807,16 +807,25 @@ class InputMetadata:
         )
 
         if model_runner.server_args.disable_flashinfer:
-            (ret.triton_max_seq_len,
-             ret.triton_max_extend_len,
-             ret.triton_start_loc,
-             ret.triton_prefix_lens) = init_triton_args(forward_mode, seq_lens, prefix_lens)
+            (
+                ret.triton_max_seq_len,
+                ret.triton_max_extend_len,
+                ret.triton_start_loc,
+                ret.triton_prefix_lens,
+            ) = init_triton_args(forward_mode, seq_lens, prefix_lens)
 
         return ret
 
 
-def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens, prefix_lens,
-                         flashinfer_decode_wrapper):
+def init_flashinfer_args(
+    forward_mode,
+    model_runner,
+    req_pool_indices,
+    seq_lens,
+    prefix_lens,
+    flashinfer_decode_wrapper,
+):
+    """Init auxiliary variables for FlashInfer attention backend."""
     num_qo_heads = model_runner.model_config.num_attention_heads // model_runner.tp_size
     num_kv_heads = model_runner.model_config.get_num_kv_heads(model_runner.tp_size)
     head_dim = model_runner.model_config.head_dim
@@ -827,9 +836,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
     else:
         paged_kernel_lens = prefix_lens
 
-    kv_indptr = torch.zeros(
-        (batch_size + 1,), dtype=torch.int32, device="cuda"
-    )
+    kv_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
     kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
     req_pool_indices_cpu = req_pool_indices.cpu().numpy()
     paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
@@ -842,9 +849,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
         ],
         dim=0,
     ).contiguous()
-    kv_last_page_len = torch.ones(
-        (batch_size,), dtype=torch.int32, device="cuda"
-    )
+    kv_last_page_len = torch.ones((batch_size,), dtype=torch.int32, device="cuda")
 
     if forward_mode == ForwardMode.DECODE:
         flashinfer_decode_wrapper.end_forward()
@@ -859,9 +864,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
         )
     else:
         # extend part
-        qo_indptr = torch.zeros(
-            (batch_size + 1,), dtype=torch.int32, device="cuda"
-        )
+        qo_indptr = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda")
         qo_indptr[1:] = torch.cumsum(seq_lens - prefix_lens, dim=0)
 
         model_runner.flashinfer_prefill_wrapper_ragged.end_forward()
@@ -888,6 +891,7 @@ def init_flashinfer_args(forward_mode, model_runner, req_pool_indices, seq_lens,
 
 
 def init_triton_args(forward_mode, seq_lens, prefix_lens):
+    """Init auxiliary variables for triton attention backend."""
     batch_size = len(seq_lens)
     max_seq_len = int(torch.max(seq_lens))
     start_loc = torch.zeros((batch_size,), dtype=torch.int32, device="cuda")