sglang 0.4.9__py3-none-any.whl → 0.4.9.post1__py3-none-any.whl

sglang/srt/managers/cache_controller.py

@@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
-import concurrent.futures
 import logging
 import math
 import threading
@@ -169,12 +168,23 @@ class HiCacheController:
         page_size: int,
         load_cache_event: threading.Event = None,
         write_policy: str = "write_through_selective",
+        io_backend: str = "",
     ):
         self.mem_pool_device_allocator = token_to_kv_pool_allocator
         self.mem_pool_device = token_to_kv_pool_allocator.get_kvcache()
         self.mem_pool_host = mem_pool_host
         self.write_policy = write_policy
         self.page_size = page_size
+        # using kernel for small page KV cache transfer and DMA for large pages
+        if not io_backend:
+            IO_BACKEND_PAGE_SIZE_THRESHOLD = 64
+            self.io_backend = (
+                "direct"
+                if self.page_size >= IO_BACKEND_PAGE_SIZE_THRESHOLD
+                else "kernel"
+            )
+        else:
+            self.io_backend = io_backend
 
         self.load_cache_event = load_cache_event
         self.layer_done_counter = LayerDoneCounter(self.mem_pool_device.layer_num)
@@ -203,12 +213,7 @@ class HiCacheController:
         self.load_stream = torch.cuda.Stream()
 
         self.write_thread = threading.Thread(
-            target=(
-                self.write_thread_func_buffer
-                if self.page_size == 1
-                else self.write_thread_func_direct
-            ),
-            daemon=True,
+            target=self.write_thread_func_direct, daemon=True
         )
         self.load_thread = threading.Thread(
             target=self.load_thread_func_layer_by_layer, daemon=True
@@ -229,12 +234,7 @@ class HiCacheController:
         self.ack_load_queue.queue.clear()
 
         self.write_thread = threading.Thread(
-            target=(
-                self.write_thread_func_buffer
-                if self.page_size == 1
-                else self.write_thread_func_direct
-            ),
-            daemon=True,
+            target=self.write_thread_func_direct, daemon=True
         )
         self.load_thread = threading.Thread(
             target=self.load_thread_func_layer_by_layer, daemon=True
@@ -281,6 +281,15 @@ class HiCacheController:
         )
         return device_indices
 
+    def move_indices(self, host_indices, device_indices):
+        # move indices to GPU if using kernels, to host if using direct indexing
+        if self.io_backend == "kernel":
+            return host_indices.to(self.mem_pool_device.device), device_indices
+        elif self.io_backend == "direct":
+            return host_indices, device_indices.cpu()
+        else:
+            raise ValueError(f"Unsupported io backend")
+
     def write_thread_func_direct(self):
         """
         Directly write through KV caches to host memory without buffering.
@@ -289,10 +298,14 @@ class HiCacheController:
         while not self.stop_event.is_set():
             try:
                 operation = self.write_queue.get(block=True, timeout=1)
-                self.mem_pool_host.write_page_all_layers(
-                    operation.host_indices,
-                    operation.device_indices,
-                    self.mem_pool_device,
+                host_indices, device_indices = self.move_indices(
+                    operation.host_indices, operation.device_indices
+                )
+                self.mem_pool_device.backup_to_host_all_layer(
+                    self.mem_pool_host,
+                    host_indices,
+                    device_indices,
+                    self.io_backend,
                 )
                 self.write_stream.synchronize()
                 self.mem_pool_host.complete_io(operation.host_indices)
@@ -304,27 +317,6 @@ class HiCacheController:
             except Exception as e:
                 logger.error(e)
 
-    def load_thread_func_direct(self):
-        """
-        Directly load KV caches from host memory to device memory without buffering.
-        """
-        torch.cuda.set_stream(self.load_stream)
-        while not self.stop_event.is_set():
-            try:
-                operation = self.load_queue.get(block=True, timeout=1)
-                operation.data = self.mem_pool_host.get_flat_data(
-                    operation.host_indices
-                )
-                self.mem_pool_device.transfer(operation.device_indices, operation.data)
-                self.mem_pool_host.complete_io(operation.host_indices)
-                for node_id in operation.node_ids:
-                    if node_id != 0:
-                        self.ack_load_queue.put(node_id)
-            except Empty:
-                continue
-            except Exception as e:
-                logger.error(e)
-
     def load_thread_func_layer_by_layer(self):
         """
         Load KV caches from host memory to device memory layer by layer.
@@ -349,22 +341,18 @@ class HiCacheController:
 
             # start layer-wise KV cache transfer from CPU to GPU
             self.layer_done_counter.reset()
+            host_indices, device_indices = self.move_indices(
+                batch_operation.host_indices, batch_operation.device_indices
+            )
             for i in range(self.mem_pool_host.layer_num):
-                if self.page_size == 1:
-                    flat_data = self.mem_pool_host.get_flat_data_by_layer(
-                        batch_operation.host_indices, i
-                    )
-                    self.mem_pool_device.transfer_per_layer(
-                        batch_operation.device_indices, flat_data, i
-                    )
-                else:
-                    self.mem_pool_host.load_page_per_layer(
-                        batch_operation.host_indices,
-                        batch_operation.device_indices,
-                        self.mem_pool_device,
-                        i,
-                    )
-                    self.load_stream.synchronize()
+                self.mem_pool_device.load_from_host_per_layer(
+                    self.mem_pool_host,
+                    host_indices,
+                    device_indices,
+                    i,
+                    self.io_backend,
+                )
+                self.load_stream.synchronize()
                 self.layer_done_counter.increment()
 
             self.mem_pool_host.complete_io(batch_operation.host_indices)
@@ -372,148 +360,6 @@ class HiCacheController:
                 if node_id != 0:
                     self.ack_load_queue.put(node_id)
 
-    def write_aux_func(self, no_wait=False):
-        """
-        Auxiliary function to prepare the buffer for write operations.
-        """
-        torch.cuda.set_stream(self.write_stream)
-
-        def _to_op(op_):
-            assert op_.device_indices.is_cuda, "Device indices should be on GPU"
-            op_.data = self.mem_pool_device.get_flat_data(op_.device_indices).to(
-                self.mem_pool_host.device
-            )
-            self.write_buffer.put(op_)
-            return op_
-
-        buffer = None
-        while not self.stop_event.is_set():
-            try:
-                operation = self.write_queue.get(block=True, timeout=1)
-                factor = (
-                    len(operation.device_indices) // self.write_buffer.max_buffer_size
-                )
-
-                if factor >= 1:
-                    if buffer is not None:
-                        _to_op(buffer)
-                        buffer = None
-
-                    if factor < 2:
-                        _to_op(operation)
-                    else:
-                        split_ops = operation.split(factor)
-                        for op_ in split_ops:
-                            _to_op(op_)
-                    continue
-
-                if buffer is None:
-                    buffer = operation
-                else:
-                    buffer.merge(operation)
-                if (
-                    no_wait
-                    or len(buffer.host_indices) >= self.write_buffer.max_buffer_size
-                    or self.write_queue.empty()
-                    or self.write_buffer.empty()
-                ):
-                    _to_op(buffer)
-                    buffer = None
-            except Empty:
-                continue
-            except Exception as e:
-                logger.error(e)
-
-    def load_aux_func(self):
-        """
-        Auxiliary function to prepare the buffer for load operations.
-        """
-
-        def _pin_op(op_, put=True):
-            op_.data = (
-                self.mem_pool_host.get_flat_data(op_.host_indices)
-                .contiguous()
-                .pin_memory()
-            )
-            if put:
-                self.load_buffer.put(op_)
-            return op_
-
-        buffer = None
-        while not self.stop_event.is_set():
-            try:
-                operation = self.load_queue.get(block=True, timeout=1)
-                factor = len(operation.host_indices) // self.load_buffer.max_buffer_size
-
-                if factor >= 1:
-                    if buffer is not None:
-                        _pin_op(buffer)
-                        buffer = None
-
-                    if factor < 2:
-                        _pin_op(operation)
-                    else:
-                        split_ops = operation.split(factor)
-                        split_args = [(op_, True) for op_ in split_ops[:-1]]
-                        split_args.append((split_ops[-1], False))
-                        # Spawn threads to pin each op concurrently
-                        with concurrent.futures.ThreadPoolExecutor() as executor:
-                            pinned_ops = list(
-                                executor.map(
-                                    lambda x: _pin_op(x[0], put=x[1]), split_args
-                                )
-                            )
-                        # preserve the order of last op to ensure correct ack
-                        self.load_buffer.put(pinned_ops[-1])
-                    continue
-
-                if buffer is None:
-                    buffer = operation
-                else:
-                    buffer.merge(operation)
-                if (
-                    len(buffer.host_indices) >= self.load_buffer.max_buffer_size
-                    or self.load_queue.empty()
-                    or self.load_buffer.empty()
-                ):
-                    _pin_op(buffer)
-                    buffer = None
-            except Empty:
-                continue
-            except Exception as e:
-                logger.error(e)
-
-    # todo (zhiqiang): double buffering to be deprecated
-    def write_thread_func_buffer(self):
-        aux_thread = threading.Thread(target=self.write_aux_func, daemon=True)
-        aux_thread.start()
-
-        while not self.stop_event.is_set():
-            operation = self.write_buffer.get()
-            if operation is None:
-                continue
-            self.mem_pool_host.assign_flat_data(operation.host_indices, operation.data)
-            self.mem_pool_host.complete_io(operation.host_indices)
-            for node_id in operation.node_ids:
-                if node_id != 0:
-                    self.ack_write_queue.put(node_id)
-        aux_thread.join()
-
-    def load_thread_func_buffer(self):
-        torch.cuda.set_stream(self.load_stream)
-        aux_thread = threading.Thread(target=self.load_aux_func, daemon=True)
-        aux_thread.start()
-        while not self.stop_event.is_set():
-            operation = self.load_buffer.get()
-            if operation is None:
-                continue
-            self.mem_pool_device.transfer(operation.device_indices, operation.data)
-            self.mem_pool_host.complete_io(operation.host_indices)
-            for node_id in operation.node_ids:
-                if node_id != 0:
-                    self.ack_load_queue.put(node_id)
-        aux_thread.join()
-
     def evict_device(
         self, device_indices: torch.Tensor, host_indices: torch.Tensor
     ) -> int:

sglang/srt/managers/io_struct.py

@@ -200,6 +200,8 @@ class GenerateReqInput:
                 self.text = [self.text]
             if self.input_ids is not None:
                 self.input_ids = [self.input_ids]
+            if self.input_embeds is not None:
+                self.input_embeds = [self.input_embeds]
 
     def _normalize_single_inputs(self):
         """Normalize inputs for a single example."""
@@ -324,7 +326,9 @@ class GenerateReqInput:
             new_rids = [f"{self.rid}_{i}" for i in range(num)]
             self.rid = new_rids
         elif isinstance(self.rid, list):
-            if len(self.rid) != num:
+            # Note: the length of rid shall be the same as the batch_size,
+            # as the rid would be expanded for parallel sampling in tokenizer_manager
+            if len(self.rid) != self.batch_size:
                 raise ValueError(
                     "The specified rids length mismatch with the batch_size for batch processing."
                 )
@@ -400,6 +404,9 @@ class GenerateReqInput:
         return GenerateReqInput(
             text=self.text[i] if self.text is not None else None,
             input_ids=self.input_ids[i] if self.input_ids is not None else None,
+            input_embeds=(
+                self.input_embeds[i] if self.input_embeds is not None else None
+            ),
             image_data=self.image_data[i],
             audio_data=self.audio_data[i],
             sampling_params=self.sampling_params[i],

sglang/srt/managers/mm_utils.py

@@ -248,7 +248,9 @@ def _get_chunked_prefill_embedding(
 ) -> Optional[torch.Tensor]:
     # Calculate embedding for each request, try to get it from cache to avoid repeated calculation
     embedding_list = []
-    for i in range(len(items_size) - 1):
+    # FIXME(Xinyuan): temporary workaround for eagle3, which may have len(items_size) > len(prefix_length)
+    max_iterations = min(len(items_size) - 1, len(prefix_length))
+    for i in range(max_iterations):
         if items_size[i] == items_size[i + 1]:
             continue
         embedding_items_per_req = embedding_items[items_size[i] : items_size[i + 1]]
@@ -269,7 +271,7 @@ def _get_chunked_prefill_embedding(
         embedding_per_req_chunk, _, end_index = get_embedding_chunk(
             embedding=embedding_per_req,
             extend_prefix_len=prefix_length[i],
-            extend_seq_len=extend_length[i],
+            extend_seq_len=extend_length[i] if i < len(extend_length) else 0,
             items_offset=items_offset,
         )
         # remove this item from cache if chunk reaches to the end

sglang/srt/managers/schedule_batch.py

@@ -101,6 +101,7 @@ GLOBAL_SERVER_ARGS_KEYS = [
     "triton_attention_reduce_in_fp32",
     "num_reserved_decode_tokens",
     "weight_loader_disable_mmap",
+    "enable_triton_kernel_moe",
 ]
 
 # Put some global args for easy access
@@ -842,7 +843,6 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     global_num_tokens_for_logprob: Optional[List[int]] = None
     is_extend_in_batch: bool = False
     can_run_dp_cuda_graph: bool = False
-    is_extend_in_batch: bool = False
     tbo_split_seq_index: Optional[int] = None
     global_forward_mode: Optional[ForwardMode] = None
 

sglang/srt/managers/scheduler.py

@@ -13,6 +13,7 @@
 # ==============================================================================
 """A scheduler that manages a tensor parallel GPU worker."""
 
+import datetime
 import faulthandler
 import logging
 import os
@@ -590,6 +591,12 @@ class Scheduler(
                     hicache_ratio=server_args.hicache_ratio,
                     hicache_size=server_args.hicache_size,
                     hicache_write_policy=server_args.hicache_write_policy,
+                    hicache_io_backend=(
+                        "direct"
+                        if server_args.attention_backend
+                        == "fa3"  # hot fix for incompatibility
+                        else server_args.hicache_io_backend
+                    ),
                 )
                 self.tp_worker.register_hicache_layer_transfer_counter(
                     self.tree_cache.cache_controller.layer_done_counter
@@ -1313,10 +1320,12 @@ class Scheduler(
             f += f"#unbootstrapped-req: {len(self.disagg_prefill_bootstrap_queue.queue)}, "
             f += f"#queue-req: {len(self.waiting_queue)}, "
             f += f"#transferring-req: {len(self.disagg_prefill_inflight_queue)}, "
-            f += f"input throughput (token/s): {self.last_input_throughput:.2f} "
+            f += f"input throughput (token/s): {self.last_input_throughput:.2f}, "
         else:
             f += f"#running-req: {running_bs}, "
-            f += f"#queue-req: {len(self.waiting_queue)}"
+            f += f"#queue-req: {len(self.waiting_queue)}, "
+
+        f += f"timestamp: {datetime.datetime.now().isoformat()}"
 
         logger.info(f)
 
@@ -1378,7 +1387,8 @@ class Scheduler(
         msg += (
             f"cuda graph: {can_run_cuda_graph}, "
             f"gen throughput (token/s): {self.last_gen_throughput:.2f}, "
-            f"#queue-req: {len(self.waiting_queue)}"
+            f"#queue-req: {len(self.waiting_queue)}, "
+            f"timestamp: {datetime.datetime.now().isoformat()}"
         )
 
         logger.info(msg)
@@ -2333,9 +2343,8 @@ class Scheduler(
 
     def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
         tags = recv_req.tags
-        import subprocess
 
-        if tags is None:
+        if tags is None or len(tags) == 0:
             tags = [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]
 
         if GPU_MEMORY_TYPE_KV_CACHE in tags:
@@ -2346,17 +2355,20 @@ class Scheduler(
             self.stashed_model_static_state = _export_static_state(
                 self.tp_worker.worker.model_runner.model
             )
+            torch.distributed.barrier(self.tp_cpu_group)
             self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
 
         return ReleaseMemoryOccupationReqOutput()
 
     def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput):
         tags = recv_req.tags
+
         if tags is None or len(tags) == 0:
             tags = [GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE]
 
         if GPU_MEMORY_TYPE_WEIGHTS in tags:
             self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
+            torch.distributed.barrier(self.tp_cpu_group)
             _import_static_state(
                 self.tp_worker.worker.model_runner.model,
                 self.stashed_model_static_state,

sglang/srt/mem_cache/hiradix_cache.py

@@ -34,6 +34,7 @@ class HiRadixCache(RadixCache):
         hicache_ratio: float,
         hicache_size: int,
         hicache_write_policy: str,
+        hicache_io_backend: str,
     ):
         self.kv_cache = token_to_kv_pool_allocator.get_kvcache()
         if isinstance(self.kv_cache, MHATokenToKVPool):
@@ -56,6 +57,7 @@ class HiRadixCache(RadixCache):
             page_size,
             load_cache_event=self.load_cache_event,
             write_policy=hicache_write_policy,
+            io_backend=hicache_io_backend,
         )
 
         # record the nodes with ongoing write through