sglang 0.3.4__py3-none-any.whl → 0.3.4.post2__py3-none-any.whl

sglang/srt/managers/tp_worker_overlap_thread.py

@@ -0,0 +1,209 @@
+"""
+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 tensor parallel worker."""
+
+import logging
+import threading
+import time
+from queue import Queue
+from typing import Optional
+
+import torch
+
+from sglang.srt.managers.io_struct import UpdateWeightReqInput
+from sglang.srt.managers.schedule_batch import ModelWorkerBatch
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+@torch.compile(dynamic=True)
+def resolve_future_token_ids(input_ids, future_token_ids_map):
+    input_ids[:] = torch.where(
+        input_ids < 0,
+        future_token_ids_map[torch.clamp(-input_ids, min=0)],
+        input_ids,
+    )
+
+
+class TpModelWorkerClient:
+    """A tensor parallel model worker."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        dp_rank: Optional[int],
+        nccl_port: int,
+    ):
+        # Load the model
+        self.worker = TpModelWorker(server_args, gpu_id, tp_rank, dp_rank, nccl_port)
+        self.max_running_requests = self.worker.max_running_requests
+        self.device = self.worker.device
+
+        # Init future mappings
+        self.future_token_ids_ct = 0
+        self.future_token_ids_limit = self.max_running_requests * 3
+        self.future_token_ids_map = torch.empty(
+            (self.max_running_requests * 5,), dtype=torch.int32, device=self.device
+        )
+
+        # Launch threads
+        self.input_queue = Queue()
+        self.output_queue = Queue()
+        self.forward_stream = torch.cuda.Stream()
+        self.forward_thread = threading.Thread(
+            target=self.forward_thread_func,
+        )
+        self.forward_thread.start()
+
+        self.copy_queue = Queue()
+        self.copy_thread = threading.Thread(
+            target=self.copy_thread_func,
+        )
+        self.copy_thread.start()
+
+    def get_worker_info(self):
+        return self.worker.get_worker_info()
+
+    def get_pad_input_ids_func(self):
+        return self.worker.get_pad_input_ids_func()
+
+    def get_tp_cpu_group(self):
+        return self.worker.get_tp_cpu_group()
+
+    def get_memory_pool(self):
+        return (
+            self.worker.model_runner.req_to_token_pool,
+            self.worker.model_runner.token_to_kv_pool,
+        )
+
+    def forward_thread_func(self):
+        with torch.cuda.stream(self.forward_stream):
+            self.forward_thread_func_()
+
+    @torch.inference_mode()
+    def forward_thread_func_(self):
+        while True:
+            self.has_inflight_batch = False
+            model_worker_batch, future_token_ids_ct = self.input_queue.get()
+            if not model_worker_batch:
+                break
+            self.has_inflight_batch = True
+            self.launch_event = threading.Event()
+
+            # Resolve future tokens in the input
+            input_ids = model_worker_batch.input_ids
+            resolve_future_token_ids(input_ids, self.future_token_ids_map)
+
+            # Run forward
+            logits_output, next_token_ids = self.worker.forward_batch_generation(
+                model_worker_batch
+            )
+
+            # Update the future token ids map
+            bs = len(model_worker_batch.seq_lens)
+            self.future_token_ids_map[
+                future_token_ids_ct + 1 : future_token_ids_ct + bs + 1
+            ] = next_token_ids
+
+            # Copy results to the CPU
+            if model_worker_batch.return_logprob:
+                logits_output.next_token_logprobs = logits_output.next_token_logprobs[
+                    torch.arange(len(next_token_ids), device=self.device),
+                    next_token_ids,
+                ].to("cpu", non_blocking=True)
+                if logits_output.input_token_logprobs is not None:
+                    logits_output.input_token_logprobs = (
+                        logits_output.input_token_logprobs.to("cpu", non_blocking=True)
+                    )
+                    logits_output.normalized_prompt_logprobs = (
+                        logits_output.normalized_prompt_logprobs.to(
+                            "cpu", non_blocking=True
+                        )
+                    )
+            next_token_ids = next_token_ids.to("cpu", non_blocking=True)
+            copy_event = torch.cuda.Event(blocking=True)
+            copy_event.record()
+
+            self.launch_event.set()
+            self.copy_queue.put((copy_event, logits_output, next_token_ids))
+
+    def copy_thread_func(self):
+        while True:
+            copy_event, logits_output, next_token_ids = self.copy_queue.get()
+            if not copy_event:
+                break
+            while not copy_event.query():
+                time.sleep(1e-5)
+
+            if logits_output.next_token_logprobs is not None:
+                logits_output.next_token_logprobs = (
+                    logits_output.next_token_logprobs.tolist()
+                )
+                if logits_output.input_token_logprobs is not None:
+                    logits_output.input_token_logprobs = (
+                        logits_output.input_token_logprobs.tolist()
+                    )
+                    logits_output.normalized_prompt_logprobs = (
+                        logits_output.normalized_prompt_logprobs.tolist()
+                    )
+
+            self.output_queue.put((logits_output, next_token_ids.tolist()))
+
+    def resulve_batch_result(self, bid: int):
+        logits_output, next_token_ids = self.output_queue.get()
+        if self.has_inflight_batch:
+            # Wait until the batch is launched
+            self.launch_event.wait()
+        return logits_output, next_token_ids
+
+    def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch):
+        # Push a new batch to the queue
+        self.input_queue.put((model_worker_batch.copy(), self.future_token_ids_ct))
+
+        # Allocate output future objects
+        bs = len(model_worker_batch.seq_lens)
+        future_next_token_ids = torch.arange(
+            -(self.future_token_ids_ct + 1),
+            -(self.future_token_ids_ct + 1 + bs),
+            -1,
+            dtype=torch.int32,
+            device=self.device,
+        )
+        self.future_token_ids_ct = (
+            self.future_token_ids_ct + bs
+        ) % self.future_token_ids_limit
+        return None, future_next_token_ids
+
+    def forward_batch_embedding(self, model_worker_batch: ModelWorkerBatch):
+        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
+        logits_output = self.model_runner.forward(forward_batch)
+        embeddings = logits_output.embeddings
+        return embeddings
+
+    def update_weights(self, recv_req: UpdateWeightReqInput):
+        success, message = self.model_runner.update_weights(
+            recv_req.model_path, recv_req.load_format
+        )
+        return success, message
+
+    def __delete__(self):
+        self.input_queue.put((None, None))
+        self.copy_queue.put((None, None, None))

sglang/srt/mem_cache/memory_pool.py

@@ -13,27 +13,46 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-"""Memory pool."""
+"""
+Memory pool.
+
+SGLang has two levels of memory pool.
+ReqToTokenPool maps a a request to its token locations.
+BaseTokenToKVPool maps a token location to its KV cache data.
+"""
 
 import logging
 from typing import List, Tuple, Union
 
 import torch
 
+from sglang.srt.layers.radix_attention import RadixAttention
+
 logger = logging.getLogger(__name__)
 
 
 class ReqToTokenPool:
     """A memory pool that maps a request to its token locations."""
 
-    def __init__(self, size: int, max_context_len: int, device: str):
+    def __init__(self, size: int, max_context_len: int, device: str, use_records: bool):
         self.size = size
         self.max_context_len = max_context_len
         self.device = device
-        self.req_to_token = torch.empty(
+        self.req_to_token = torch.zeros(
             (size, max_context_len), dtype=torch.int32, device=device
         )
         self.free_slots = list(range(size))
+        self.write_records = []
+        self.use_records = use_records
+
+        if self.use_records:
+            self.write = self.write_with_records
+        else:
+            self.write = self.write_without_records
+
+    def write(self, indices, values):
+        # Keep the signature for type checking. It will be assigned during runtime.
+        raise NotImplementedError()
 
     def available_size(self):
         return len(self.free_slots)
@@ -55,10 +74,27 @@ class ReqToTokenPool:
 
     def clear(self):
         self.free_slots = list(range(self.size))
+        self.write_records = []
+
+    def write_without_records(self, indices, values):
+        self.req_to_token[indices] = values
+
+    def write_with_records(self, indices, values):
+        self.req_to_token[indices] = values
+        self.write_records.append((indices, values))
+
+    def get_write_records(self):
+        ret = self.write_records
+        self.write_records = []
+        return ret
+
+    def apply_write_records(self, write_records: List[Tuple]):
+        for indices, values in write_records:
+            self.req_to_token[indices] = values
 
 
 class BaseTokenToKVPool:
-    """A memory pool that maps a token to its kv cache locations"""
+    """A memory pool that maps a token location to its kv cache data."""
 
     def __init__(
         self,
@@ -68,12 +104,12 @@ class BaseTokenToKVPool:
     ):
         self.size = size
         self.dtype = dtype
-        self.device = device
         if dtype == torch.float8_e5m2:
             # NOTE: Store as torch.uint8 because Tensor index_put is not implemented for torch.float8_e5m2
             self.store_dtype = torch.uint8
         else:
             self.store_dtype = dtype
+        self.device = device
 
         self.free_slots = None
         self.is_not_in_free_group = True
@@ -124,7 +160,7 @@ class BaseTokenToKVPool:
 
     def set_kv_buffer(
         self,
-        layer_id: int,
+        layer: RadixAttention,
         loc: torch.Tensor,
         cache_k: torch.Tensor,
         cache_v: torch.Tensor,
@@ -179,14 +215,14 @@ class MHATokenToKVPool(BaseTokenToKVPool):
 
     def set_kv_buffer(
         self,
-        layer_id: int,
+        layer: RadixAttention,
         loc: torch.Tensor,
         cache_k: torch.Tensor,
         cache_v: torch.Tensor,
     ):
+        layer_id = layer.layer_id
         if cache_k.dtype != self.dtype:
             cache_k = cache_k.to(self.dtype)
-        if cache_v.dtype != self.dtype:
             cache_v = cache_v.to(self.dtype)
         if self.store_dtype != self.dtype:
             self.k_buffer[layer_id][loc] = cache_k.view(self.store_dtype)
@@ -196,6 +232,14 @@ class MHATokenToKVPool(BaseTokenToKVPool):
             self.v_buffer[layer_id][loc] = cache_v
 
 
+# This compiled version is slower in the unit test
+# python3 -m unittest test_bench_serving.TestBenchServing.test_offline_throughput_non_stream_small_batch_size
+@torch.compile(dynamic=True)
+def copy_two_array(loc, dst_1, src_1, dst_2, src_2, dtype, store_dtype):
+    dst_1[loc] = src_1.to(dtype).view(store_dtype)
+    dst_2[loc] = src_2.to(dtype).view(store_dtype)
+
+
 class MLATokenToKVPool(BaseTokenToKVPool):
 
     def __init__(
@@ -235,11 +279,12 @@ class MLATokenToKVPool(BaseTokenToKVPool):
 
     def set_kv_buffer(
         self,
-        layer_id: int,
+        layer: RadixAttention,
         loc: torch.Tensor,
         cache_k: torch.Tensor,
         cache_v: torch.Tensor,
     ):
+        layer_id = layer.layer_id
         if cache_k.dtype != self.dtype:
             cache_k = cache_k.to(self.dtype)
         if self.store_dtype != self.dtype:
@@ -294,13 +339,14 @@ class DoubleSparseTokenToKVPool(BaseTokenToKVPool):
 
     def set_kv_buffer(
         self,
-        layer_id: int,
+        layer: RadixAttention,
         loc: torch.Tensor,
         cache_k: torch.Tensor,
         cache_v: torch.Tensor,
         cache_label: torch.Tensor,
     ):
         # NOTE(Andy): ignore the dtype check
+        layer_id = layer.layer_id
         self.k_buffer[layer_id][loc] = cache_k
         self.v_buffer[layer_id][loc] = cache_v
         self.label_buffer[layer_id][loc] = cache_label

sglang/srt/mem_cache/radix_cache.py

@@ -145,9 +145,10 @@ class RadixCache(BasePrefixCache):
         # The prefix indices could be updated, reuse it
         new_indices, new_last_node = self.match_prefix(token_ids)
         assert len(new_indices) == len(token_ids)
-        self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, len(req.prefix_indices) : len(new_indices)
-        ] = new_indices[len(req.prefix_indices) :]
+        self.req_to_token_pool.write(
+            (req.req_pool_idx, slice(len(req.prefix_indices), len(new_indices))),
+            new_indices[len(req.prefix_indices) :],
+        )
 
         self.dec_lock_ref(req.last_node)
         self.inc_lock_ref(new_last_node)

sglang/srt/model_executor/cuda_graph_runner.py

@@ -92,6 +92,11 @@ def set_torch_compile_config():
     torch._dynamo.config.accumulated_cache_size_limit = 1024
 
 
+@torch.compile(dynamic=True)
+def clamp_position(seq_lens):
+    return torch.clamp((seq_lens - 1), min=0).to(torch.int64)
+
+
 class CudaGraphRunner:
     """A CudaGraphRunner runs the forward pass of a model with cuda graph and torch.compile."""
 
@@ -105,13 +110,13 @@ class CudaGraphRunner:
         self.graph_memory_pool = None
         self.use_torch_compile = model_runner.server_args.enable_torch_compile
         self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
+        self.is_encoder_decoder = self.model_runner.model_config.is_encoder_decoder
 
         # Batch sizes to capture
         if self.model_runner.server_args.disable_cuda_graph_padding:
             self.capture_bs = list(range(1, 32)) + [64, 128]
         else:
-            self.capture_bs = [1, 2, 4] + [i * 8 for i in range(1, 21)]
-
+            self.capture_bs = [1, 2, 3, 4] + [i * 8 for i in range(1, 21)]
         self.capture_bs = [
             bs for bs in self.capture_bs if bs <= model_runner.req_to_token_pool.size
         ]
@@ -128,10 +133,14 @@ class CudaGraphRunner:
         # Attention backend
         self.max_bs = max(self.capture_bs)
         self.model_runner.attn_backend.init_cuda_graph_state(self.max_bs)
+
         self.seq_len_fill_value = (
             self.model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
         )
 
+        # FIXME(lsyin): leave it here for now, I don't know whether it is necessary
+        self.encoder_len_fill_value = 0
+
         if self.use_torch_compile:
             set_torch_compile_config()
 
@@ -143,10 +152,20 @@ class CudaGraphRunner:
                 (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32
             )
             self.out_cache_loc = torch.zeros((self.max_bs,), dtype=torch.int32)
+            self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int32)
+
+            if self.is_encoder_decoder:
+                # NOTE: encoder_lens can influence the full_text_row_masked_out_mask tensor when doing mixed batch
+                self.encoder_lens = torch.full(
+                    (self.max_bs,), self.encoder_len_fill_value, dtype=torch.int32
+                )
+            else:
+                self.encoder_lens = None
 
         # Capture
         try:
-            self.capture()
+            with self.model_capture_mode():
+                self.capture()
         except RuntimeError as e:
             raise Exception(
                 f"Capture cuda graph failed: {e}\n"
@@ -157,11 +176,32 @@ class CudaGraphRunner:
                 "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
             )
 
-    def can_run(self, batch_size: int):
-        if self.disable_padding:
-            return batch_size in self.graphs
-        else:
-            return batch_size <= self.max_bs
+    @contextmanager
+    def model_capture_mode(self):
+        if hasattr(self.model_runner.model, "capture_mode"):
+            self.model_runner.model.capture_mode = True
+
+        yield
+
+        if hasattr(self.model_runner.model, "capture_mode"):
+            self.model_runner.model.capture_mode = False
+
+    def can_run(self, forward_batch: ForwardBatch):
+        is_bs_supported = (
+            forward_batch.batch_size in self.graphs
+            if self.disable_padding
+            else forward_batch.batch_size <= self.max_bs
+        )
+
+        # NOTE: cuda graph cannot handle mixed batch (encoder_len = 0)
+        # If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph
+        # because the full_text_row_masked_out_mask tensor will always be ones
+        is_encoder_lens_supported = (
+            torch.all(forward_batch.encoder_lens > 0)
+            if self.is_encoder_decoder
+            else True
+        )
+        return is_bs_supported and is_encoder_lens_supported
 
     def capture(self):
         with graph_capture() as graph_capture_context:
@@ -188,10 +228,20 @@ class CudaGraphRunner:
         req_pool_indices = self.req_pool_indices[:bs]
         seq_lens = self.seq_lens[:bs]
         out_cache_loc = self.out_cache_loc[:bs]
+        if self.is_encoder_decoder:
+            encoder_lens = self.encoder_lens[:bs]
+        else:
+            encoder_lens = None
+
+        seq_lens_sum = seq_lens.sum().item()
+        mrope_positions = self.mrope_positions[:, :bs]
 
         # Attention backend
         self.model_runner.attn_backend.init_forward_metadata_capture_cuda_graph(
-            bs, req_pool_indices, seq_lens
+            bs,
+            req_pool_indices,
+            seq_lens,
+            encoder_lens,
         )
 
         # Run and capture
@@ -206,11 +256,15 @@ class CudaGraphRunner:
                 token_to_kv_pool=self.model_runner.token_to_kv_pool,
                 attn_backend=self.model_runner.attn_backend,
                 out_cache_loc=out_cache_loc,
+                seq_lens_sum=seq_lens_sum,
+                encoder_lens=encoder_lens,
                 return_logprob=False,
                 top_logprobs_nums=[0] * bs,
-                positions=torch.clamp((seq_lens - 1), min=0).to(torch.int64),
+                positions=clamp_position(seq_lens),
+                mrope_positions=mrope_positions,
             )
-            return forward(input_ids, forward_batch.positions, forward_batch)
+            logits_output = forward(input_ids, forward_batch.positions, forward_batch)
+            return logits_output.next_token_logits
 
         for _ in range(2):
             torch.cuda.synchronize()
@@ -241,7 +295,7 @@ class CudaGraphRunner:
         index = bisect.bisect_left(self.capture_bs, raw_bs)
         bs = self.capture_bs[index]
         if bs != raw_bs:
-            self.seq_lens.fill_(self.seq_len_fill_value)
+            self.seq_lens.fill_(1)
             self.out_cache_loc.zero_()
 
         # Common inputs
@@ -249,31 +303,32 @@ class CudaGraphRunner:
         self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
         self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
         self.out_cache_loc[:raw_bs].copy_(forward_batch.out_cache_loc)
+        if self.is_encoder_decoder:
+            self.encoder_lens[:raw_bs].copy_(forward_batch.encoder_lens)
+        if forward_batch.mrope_positions is not None:
+            self.mrope_positions[:, :raw_bs].copy_(forward_batch.mrope_positions)
 
         # Attention backend
         self.model_runner.attn_backend.init_forward_metadata_replay_cuda_graph(
-            bs, self.req_pool_indices, self.seq_lens
+            bs,
+            self.req_pool_indices,
+            self.seq_lens,
+            forward_batch.seq_lens_sum + (bs - raw_bs),
+            self.encoder_lens,
         )
 
         # Replay
         self.graphs[bs].replay()
-        logits_output = self.output_buffers[bs]
-
-        # Unpad
-        if bs != raw_bs:
-            logits_output = LogitsProcessorOutput(
-                next_token_logits=logits_output.next_token_logits[:raw_bs],
-                next_token_logprobs=None,
-                normalized_prompt_logprobs=None,
-                input_token_logprobs=None,
-                input_top_logprobs=None,
-                output_top_logprobs=None,
-            )
+        next_token_logits = self.output_buffers[bs][:raw_bs]
 
         # Extract logprobs
         if forward_batch.return_logprob:
-            logits_output.next_token_logprobs = torch.nn.functional.log_softmax(
-                logits_output.next_token_logits, dim=-1
+            next_token_logprobs = torch.nn.functional.log_softmax(
+                next_token_logits, dim=-1
+            )
+            logits_output = LogitsProcessorOutput(
+                next_token_logits=next_token_logits,
+                next_token_logprobs=next_token_logprobs,
             )
             return_top_logprob = any(x > 0 for x in forward_batch.top_logprobs_nums)
             if return_top_logprob:
@@ -282,7 +337,11 @@ class CudaGraphRunner:
                     top_logprobs_nums=forward_batch.top_logprobs_nums,
                 )
                 logits_output.output_top_logprobs = LogitsProcessor.get_top_logprobs(
-                    logits_output.next_token_logprobs, logits_metadata
+                    next_token_logprobs, logits_metadata
                 )[1]
+        else:
+            logits_output = LogitsProcessorOutput(
+                next_token_logits=next_token_logits,
+            )
 
         return logits_output