sglang 0.2.9.post1__py3-none-any.whl → 0.2.11__py3-none-any.whl

sglang/srt/managers/tokenizer_manager.py

@@ -153,8 +153,9 @@ class TokenizerManager:
     async def _handle_single_request(
         self, obj, request, index=None, is_cache_for_prefill=False
     ):
-        if not is_cache_for_prefill:
-            not_use_index = not (index is not None)
+        if not is_cache_for_prefill:  # The normal case with a single prompt
+            not_use_index = index is None
+
             rid = obj.rid if not_use_index else obj.rid[index]
             input_text = obj.text if not_use_index else obj.text[index]
             input_ids = (
@@ -182,14 +183,27 @@ class TokenizerManager:
             top_logprobs_num = (
                 obj.top_logprobs_num if not_use_index else obj.top_logprobs_num[index]
             )
-        else:
-            if isinstance(obj.text, list):
-                input_text = obj.text[index]
-                rid = obj.rid[index]
+        else:  # A prefill request to cache the common prompt for parallel sampling
+            if obj.text is not None:
+                if isinstance(obj.text, list):
+                    input_text = obj.text[index]
+                    rid = obj.rid[index]
+                else:
+                    input_text = obj.text
+                    rid = obj.rid[0]
+                input_ids = self.tokenizer.encode(input_text)
             else:
-                input_text = obj.text
-                rid = obj.rid[0]
-            input_ids = self.tokenizer.encode(input_text)
+                input_text = None
+                if isinstance(obj.input_ids, list) and isinstance(
+                    obj.input_ids[0], list
+                ):
+                    # when obj["input_ids"] is List[List[int]]
+                    input_ids = obj.input_ids[index]
+                    rid = obj.rid[index]
+                else:
+                    input_ids = obj.input_ids
+                    rid = obj.rid[0]
+
             sampling_params = SamplingParams(**obj.sampling_params[0])
             sampling_params.max_new_tokens = 0
             pixel_values, image_hash, image_size = await self._get_pixel_values(
@@ -240,11 +254,11 @@ class TokenizerManager:
                 ):
                     if input_id_result is not None:
                         input_id_result.append(input_id)
-                    pass
-            if len(input_id_result) > 1 and input_id_result is not None:
+            if input_id_result is not None and len(input_id_result) > 1:
                 obj.input_ids = input_id_result
             elif input_id_result is not None:
                 obj.input_ids = input_id_result[0]
+
         # First send out all requests
         for i in range(batch_size):
             for j in range(parallel_sample_num):
@@ -264,11 +278,12 @@ class TokenizerManager:
                         input_text = None
                         input_ids = obj.input_ids[i]
                 else:
+                    assert obj.input_ids is not None
                     if batch_size == 1:
-                        input_text = obj.text
+                        input_text = None
                         input_ids = obj.input_ids
                     else:
-                        input_text = obj.text[i]
+                        input_text = None
                         input_ids = obj.input_ids[i]
                 sampling_params = self._get_sampling_params(obj.sampling_params[index])
                 pixel_values, image_hash, image_size = await self._get_pixel_values(
@@ -293,7 +308,6 @@ class TokenizerManager:
                 event = asyncio.Event()
                 state = ReqState([], False, event)
                 self.rid_to_state[rid] = state
-
         # Then wait for all responses
         output_list = []
         for i in range(batch_size):
@@ -326,7 +340,6 @@ class TokenizerManager:
                 )
                 assert state.finished
                 del self.rid_to_state[rid]
-
         yield output_list
 
     def _validate_input_length(self, input_ids: List[int]):
@@ -375,8 +388,13 @@ class TokenizerManager:
                 obj.return_text_in_logprobs,
             )
 
+            # Log requests
             if self.server_args.log_requests and state.finished:
-                logger.info(f"in={obj.text}, out={out}")
+                if obj.text is None:
+                    in_obj = {"text": self.tokenizer.decode(obj.input_ids)}
+                else:
+                    in_obj = {"text": obj.text}
+                logger.info(f"in={in_obj}, out={out}")
 
             state.out_list = []
             if state.finished:

sglang/srt/managers/tp_worker.py

@@ -39,13 +39,13 @@ from sglang.srt.managers.policy_scheduler import PolicyScheduler
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
     BaseFinishReason,
-    Batch,
-    ForwardMode,
     Req,
+    ScheduleBatch,
 )
 from sglang.srt.mem_cache.chunk_cache import ChunkCache
 from sglang.srt.mem_cache.radix_cache import RadixCache
 from sglang.srt.model_config import ModelConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import (
@@ -172,7 +172,7 @@ class ModelTpServer:
 
         # Init running status
         self.waiting_queue: List[Req] = []
-        self.running_batch: Batch = None
+        self.running_batch: ScheduleBatch = None
         self.out_pyobjs = []
         self.decode_forward_ct = 0
         self.stream_interval = server_args.stream_interval
@@ -200,7 +200,6 @@ class ModelTpServer:
         )
         self.new_token_ratio = self.min_new_token_ratio
         self.new_token_ratio_decay = global_config.new_token_ratio_decay
-        self.new_token_ratio_recovery = global_config.new_token_ratio_recovery
 
     def exposed_step(self, recv_reqs):
         try:
@@ -290,10 +289,10 @@ class ModelTpServer:
                 "KV cache pool leak detected!"
             )
 
-        if self.req_to_token_pool.can_use_mem_size != self.req_to_token_pool.size:
+        if len(self.req_to_token_pool.free_slots) != self.req_to_token_pool.size:
             warnings.warn(
                 "Warning: "
-                f"available req slots={self.req_to_token_pool.can_use_mem_size}, "
+                f"available req slots={len(self.req_to_token_pool.free_slots)}, "
                 f"total slots={self.req_to_token_pool.size}\n"
                 "Memory pool leak detected!"
             )
@@ -353,7 +352,7 @@ class ModelTpServer:
         )
         self.waiting_queue.append(req)
 
-    def get_new_prefill_batch(self) -> Optional[Batch]:
+    def get_new_prefill_batch(self) -> Optional[ScheduleBatch]:
         # TODO(lsyin): organize this function
         running_bs = (
             len(self.running_batch.reqs) if self.running_batch is not None else 0
@@ -364,12 +363,13 @@ class ModelTpServer:
         # Compute matched prefix length
         for req in self.waiting_queue:
             req.input_ids = req.origin_input_ids + req.output_ids
+            try_match_ids = req.input_ids
+            if req.return_logprob:
+                try_match_ids = req.input_ids[: req.logprob_start_len]
+            # NOTE: the prefix_indices must always be aligned with last_node
             prefix_indices, last_node = self.tree_cache.match_prefix(
-                rid=req.rid,
-                key=req.input_ids,
+                rid=req.rid, key=try_match_ids
             )
-            if req.return_logprob:
-                prefix_indices = prefix_indices[: req.logprob_start_len]
             req.extend_input_len = len(req.input_ids) - len(prefix_indices)
             req.prefix_indices = prefix_indices
             req.last_node = last_node
@@ -525,7 +525,7 @@ class ModelTpServer:
             )
 
         # Return the new batch
-        new_batch = Batch.init_new(
+        new_batch = ScheduleBatch.init_new(
             can_run_list,
             self.req_to_token_pool,
             self.token_to_kv_pool,
@@ -534,7 +534,7 @@ class ModelTpServer:
         self.waiting_queue = [x for x in self.waiting_queue if x not in can_run_list]
         return new_batch
 
-    def forward_prefill_batch(self, batch: Batch):
+    def forward_prefill_batch(self, batch: ScheduleBatch):
         # Build batch tensors
         batch.prepare_for_extend(
             self.model_config.vocab_size, self.int_token_logit_bias
@@ -623,14 +623,13 @@ class ModelTpServer:
                 )
             req.output_top_logprobs.append(output.output_top_logprobs[i])
 
-    def cache_filled_batch(self, batch: Batch):
-        req_pool_indices_cpu = batch.req_pool_indices.cpu().numpy()
+    def cache_filled_batch(self, batch: ScheduleBatch):
         for i, req in enumerate(batch.reqs):
             new_prefix_indices, new_last_node = self.tree_cache.cache_req(
                 rid=req.rid,
                 token_ids=tuple(req.input_ids),
                 last_uncached_pos=len(req.prefix_indices),
-                req_pool_idx=req_pool_indices_cpu[i],
+                req_pool_idx=req.req_pool_idx,
                 del_in_memory_pool=False,
                 old_last_node=req.last_node,
             )
@@ -638,9 +637,9 @@ class ModelTpServer:
 
             if req is self.current_inflight_req:
                 # inflight request would get a new req idx
-                self.req_to_token_pool.free(int(req_pool_indices_cpu[i]))
+                self.req_to_token_pool.free(req.req_pool_idx)
 
-    def forward_decode_batch(self, batch: Batch):
+    def forward_decode_batch(self, batch: ScheduleBatch):
         # Check if decode out of memory
         if not batch.check_decode_mem():
             old_ratio = self.new_token_ratio
@@ -699,7 +698,7 @@ class ModelTpServer:
 
         self.handle_finished_requests(batch)
 
-    def handle_finished_requests(self, batch: Batch):
+    def handle_finished_requests(self, batch: ScheduleBatch):
         output_rids = []
         output_vids = []
         decoded_texts = []
@@ -781,14 +780,13 @@ class ModelTpServer:
         # Remove finished reqs
         if finished_indices:
             # Update radix cache
-            req_pool_indices_cpu = batch.req_pool_indices.tolist()
             for i in finished_indices:
                 req = batch.reqs[i]
                 self.tree_cache.cache_req(
                     rid=req.rid,
                     token_ids=tuple(req.origin_input_ids + req.output_ids)[:-1],
                     last_uncached_pos=len(req.prefix_indices),
-                    req_pool_idx=req_pool_indices_cpu[i],
+                    req_pool_idx=req.req_pool_idx,
                 )
 
                 self.tree_cache.dec_lock_ref(req.last_node)
@@ -799,7 +797,7 @@ class ModelTpServer:
             else:
                 batch.reqs = []
 
-    def filter_out_inflight(self, batch: Batch):
+    def filter_out_inflight(self, batch: ScheduleBatch):
         # TODO(lsyin): reduce the overhead, make a special version for this
         if self.current_inflight_req is None:
             return

sglang/srt/mem_cache/memory_pool.py

@@ -16,6 +16,7 @@ limitations under the License.
 """Memory pool."""
 
 import logging
+from typing import List
 
 import torch
 
@@ -27,62 +28,42 @@ class ReqToTokenPool:
 
     def __init__(self, size: int, max_context_len: int):
         self.size = size
-        self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
+        self.free_slots = list(range(size))
         self.req_to_token = torch.empty(
             (size, max_context_len), dtype=torch.int32, device="cuda"
         )
-        self.can_use_mem_size = size
 
-    def alloc(self, need_size: int):
-        if need_size > self.can_use_mem_size:
+    def alloc(self, need_size: int) -> List[int]:
+        if need_size > len(self.free_slots):
             return None
 
-        select_index = (
-            torch.nonzero(self.mem_state).squeeze(1)[:need_size].to(torch.int32)
-        )
-        self.mem_state[select_index] = False
-        self.can_use_mem_size -= need_size
+        select_index = self.free_slots[:need_size]
+        self.free_slots = self.free_slots[need_size:]
 
         return select_index
 
     def free(self, free_index):
-        self.mem_state[free_index] = True
         if isinstance(free_index, (int,)):
-            self.can_use_mem_size += 1
+            self.free_slots.append(free_index)
         else:
-            self.can_use_mem_size += free_index.shape[0]
+            self.free_slots.extend(free_index)
 
     def clear(self):
-        self.mem_state.fill_(True)
-        self.can_use_mem_size = len(self.mem_state)
+        self.free_slots = list(range(self.size))
 
 
-class TokenToKVPool:
+class BaseTokenToKVPool:
     """A memory pool that maps a token to its kv cache locations"""
 
     def __init__(
         self,
         size: int,
-        dtype: torch.dtype,
-        head_num: int,
-        head_dim: int,
-        layer_num: int,
     ):
         self.size = size
 
         # We also add one slot. This slot is used for writing dummy output from padded tokens.
         self.mem_state = torch.ones((self.size + 1,), dtype=torch.bool, device="cuda")
 
-        # [size, head_num, head_dim] for each layer
-        self.k_buffer = [
-            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device="cuda")
-            for _ in range(layer_num)
-        ]
-        self.v_buffer = [
-            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device="cuda")
-            for _ in range(layer_num)
-        ]
-
         # Prefetch buffer
         self.prefetch_buffer = torch.empty(0, device="cuda", dtype=torch.int32)
         self.prefetch_chunk_size = 512
@@ -90,15 +71,6 @@ class TokenToKVPool:
         self.can_use_mem_size = self.size
         self.clear()
 
-    def get_key_buffer(self, layer_id: int):
-        return self.k_buffer[layer_id]
-
-    def get_value_buffer(self, layer_id: int):
-        return self.v_buffer[layer_id]
-
-    def get_kv_buffer(self, layer_id: int):
-        return self.k_buffer[layer_id], self.v_buffer[layer_id]
-
     def available_size(self):
         return self.can_use_mem_size + len(self.prefetch_buffer)
 
@@ -139,3 +111,67 @@ class TokenToKVPool:
 
         # We also add one slot. This slot is used for writing dummy output from padded tokens.
         self.mem_state[0] = False
+
+
+class MHATokenToKVPool(BaseTokenToKVPool):
+
+    def __init__(
+        self,
+        size: int,
+        dtype: torch.dtype,
+        head_num: int,
+        head_dim: int,
+        layer_num: int,
+    ):
+        super().__init__(size)
+
+        # [size, head_num, head_dim] for each layer
+        self.k_buffer = [
+            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device="cuda")
+            for _ in range(layer_num)
+        ]
+        self.v_buffer = [
+            torch.empty((size + 1, head_num, head_dim), dtype=dtype, device="cuda")
+            for _ in range(layer_num)
+        ]
+
+    def get_key_buffer(self, layer_id: int):
+        return self.k_buffer[layer_id]
+
+    def get_value_buffer(self, layer_id: int):
+        return self.v_buffer[layer_id]
+
+    def get_kv_buffer(self, layer_id: int):
+        return self.k_buffer[layer_id], self.v_buffer[layer_id]
+
+
+class MLATokenToKVPool(BaseTokenToKVPool):
+
+    def __init__(
+        self,
+        size: int,
+        dtype: torch.dtype,
+        kv_lora_rank: int,
+        qk_rope_head_dim: int,
+        layer_num: int,
+    ):
+        super().__init__(size)
+
+        self.kv_lora_rank = kv_lora_rank
+        self.kv_buffer = [
+            torch.empty(
+                (size + 1, 1, kv_lora_rank + qk_rope_head_dim),
+                dtype=dtype,
+                device="cuda",
+            )
+            for _ in range(layer_num)
+        ]
+
+    def get_key_buffer(self, layer_id: int):
+        return self.kv_buffer[layer_id]
+
+    def get_value_buffer(self, layer_id: int):
+        return self.kv_buffer[layer_id][..., : self.kv_lora_rank]
+
+    def get_kv_buffer(self, layer_id: int):
+        return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id)

sglang/srt/model_config.py

@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 """
 
+from enum import IntEnum, auto
 from typing import Optional
 
 from transformers import PretrainedConfig
@@ -20,6 +21,11 @@ from transformers import PretrainedConfig
 from sglang.srt.hf_transformers_utils import get_config, get_context_length
 
 
+class AttentionArch(IntEnum):
+    MLA = auto()
+    MHA = auto()
+
+
 class ModelConfig:
     def __init__(
         self,
@@ -55,6 +61,11 @@ class ModelConfig:
         # FIXME: temporary special judge for deepseek v2 MLA architecture
         if "DeepseekV2ForCausalLM" in self.hf_config.architectures:
             self.head_dim = 256
+            self.attention_arch = AttentionArch.MLA
+            self.kv_lora_rank = self.hf_config.kv_lora_rank
+            self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim
+        else:
+            self.attention_arch = AttentionArch.MHA
 
         self.num_attention_heads = self.hf_config.num_attention_heads
         self.num_key_value_heads = getattr(self.hf_config, "num_key_value_heads", None)

sglang/srt/model_executor/cuda_graph_runner.py

@@ -29,8 +29,8 @@ from sglang.srt.layers.logits_processor import (
     LogitsMetadata,
     LogitsProcessor,
 )
-from sglang.srt.managers.schedule_batch import (
-    Batch,
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.model_executor.forward_batch_info import (
     ForwardMode,
     InputMetadata,
     init_flashinfer_args,
@@ -202,7 +202,7 @@ class CudaGraphRunner:
         self.graph_memory_pool = graph.pool()
         return graph, None, out, flashinfer_decode_wrapper
 
-    def replay(self, batch: Batch):
+    def replay(self, batch: ScheduleBatch):
         assert batch.out_cache_loc is not None
         raw_bs = len(batch.reqs)