Fix prefill OOM error in the case of large page size (sgl-project#5081)

Author: xiezhq-hermann

python/sglang/srt/managers/schedule_policy.py

@@ -455,7 +455,10 @@ def add_one_req(
         total_tokens = req.extend_input_len + min(
             req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS_ESTIMATION
         )
-        input_tokens = req.extend_input_len
+        input_tokens = (
+            -(-req.extend_input_len // self.tree_cache.page_size)
+            * self.tree_cache.page_size
+        )
         prefix_len = len(req.prefix_indices)
 
         if total_tokens >= self.rem_total_tokens:
@@ -477,7 +480,10 @@ def add_one_req(
                     req.last_node_global, req.prefix_indices
                 )
                 req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices)
-                input_tokens = req.extend_input_len
+                input_tokens = (
+                    -(-req.extend_input_len // self.tree_cache.page_size)
+                    * self.tree_cache.page_size
+                )
                 prefix_len = len(req.prefix_indices)
 
             if self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens:

python/sglang/srt/managers/scheduler.py

@@ -507,7 +507,11 @@ def init_memory_pool_and_cache(self):
             self.tree_cache = ChunkCache(
                 req_to_token_pool=self.req_to_token_pool,
                 token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+<<<<<<< HEAD
                 token_to_kv_pool_allocator_local=self.token_to_kv_pool_allocator_local,
+=======
+                page_size=self.page_size,
+>>>>>>> f8e46093 (Fix prefill OOM error in the case of large page size (#5081))
             )
         else:
             if self.enable_hierarchical_cache:

python/sglang/srt/mem_cache/chunk_cache.py

@@ -24,10 +24,12 @@ def __init__(
         self,
         req_to_token_pool: ReqToTokenPool,
         token_to_kv_pool_allocator: TokenToKVPoolAllocator,
+        page_size: int,
         token_to_kv_pool_allocator_local: TokenToKVPoolAllocator = None,
     ):
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.page_size = page_size
         self.token_to_kv_pool_allocator_local = token_to_kv_pool_allocator_local
 
     def reset(self):