[Core] in batch prefix caching by delay scheduling

python/sglang/benchmark_prefix.py

@@ -0,0 +1,155 @@
+# Benchmark with lots of common prefixes. Used to benchmark prefix caching performance.
+# python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000
+
+import random
+import string
+import time
+
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+import sglang as sgl
+from sglang import set_default_backend
+from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
+
+
+def generate_random_string(token_length: int) -> str:
+    random_string = "".join(
+        random.choices(string.ascii_letters + string.digits, k=token_length * 100)
+    )
+    tokenized_output = tokenizer.encode(random_string, add_special_tokens=False)[
+        :token_length
+    ]
+
+    if len(tokenized_output) < token_length:
+        tokenized_output = tokenized_output + [tokenizer.pad_token_id] * (
+            token_length - len(tokenized_output)
+        )
+
+    decoded_string = tokenizer.decode(tokenized_output, skip_special_tokens=False)
+    return decoded_string
+
+
+def generate_unique_prefix(base_text, index):
+    return str(index) + base_text[len(str(index)) :]
+
+
+@sgl.function
+def text_qa(s, question, gen_len):
+    s += "Q: " + question + "\n"
+    s += "A:" + sgl.gen("answer", stop="\n", temperature=0, max_tokens=gen_len)
+
+
+def prepare_prompts(num_prefix, num_samples_per_prefix, prefix_length, suffix_length):
+    base_prefix = generate_random_string(prefix_length)
+
+    tot_input_len = 0
+    all_prompts = []
+    for i in tqdm(range(num_prefix), desc="prepare prompts"):
+        unique_prefix = generate_unique_prefix(base_prefix, i)
+        prompt_list = []
+        for j in range(num_samples_per_prefix):
+            suffix = generate_random_string(suffix_length)
+            prompt = unique_prefix + suffix
+            prompt_list.append(prompt)
+            tot_input_len += len(tokenizer.encode(prompt))
+        all_prompts.append(prompt_list)
+    return all_prompts, tot_input_len
+
+
+def test_prefix_caching(all_prompts, num_prefix, num_samples_per_prefix, gen_len):
+    tot_time = 0
+    tot_output_len = 0
+    for i in tqdm(range(num_prefix), desc="prefix batches"):
+        tic = time.time()
+        output = text_qa.run_batch(
+            list(zip(all_prompts[i], [gen_len] * len(all_prompts[i]))),
+            progress_bar=False,
+        )
+        tot_time += time.time() - tic
+
+        for i, out in enumerate(output):
+            answer = out["answer"]
+            output_token_length = len(tokenizer.encode(answer))
+            tot_output_len += output_token_length
+
+    print(f"Total output token length: {tot_output_len}")
+    rps = num_prefix * num_samples_per_prefix / tot_time
+    tps = (tot_input_len + tot_output_len) / tot_time
+    print(f"Throughput: {rps:.2f} requests/s, {tps:.2f} tokens/s")
+
+
+def test_prefix_caching_hint(all_prompts, num_prefix, num_samples_per_prefix, gen_len):
+    tot_time = 0
+    tot_output_len = 0
+    for i in tqdm(range(num_prefix), desc="prefix batches"):
+        tic = time.time()
+        text_qa.run_batch(
+            list(zip(all_prompts[i][:1], [gen_len])),
+        )
+        output = text_qa.run_batch(
+            list(zip(all_prompts[i], [gen_len] * len(all_prompts[i]))),
+            progress_bar=False,
+        )
+        tot_time += time.time() - tic
+
+        for i, out in enumerate(output):
+            answer = out["answer"]
+            output_token_length = len(tokenizer.encode(answer))
+            tot_output_len += output_token_length
+
+    print(f"Total output token length: {tot_output_len}")
+    rps = num_prefix * num_samples_per_prefix / tot_time
+    tps = (tot_input_len + tot_output_len) / tot_time
+    print(f"Throughput: {rps:.2f} requests/s, {tps:.2f} tokens/s")
+
+
+def test_prefix_caching_send_all(
+    all_prompts, num_prefix, num_samples_per_prefix, gen_len
+):
+    all_prompts = [x for prompt_list in all_prompts for x in prompt_list]
+
+    # run
+    tic = time.time()
+    output = text_qa.run_batch(
+        list(zip(all_prompts, [gen_len] * len(all_prompts))), progress_bar=True
+    )
+    tot_time = time.time() - tic
+
+    # print
+    tot_output_len = 0
+    for i, out in enumerate(output):
+        answer = out["answer"]
+        output_token_length = len(tokenizer.encode(answer))
+        tot_output_len += output_token_length
+
+    print(f"Total output token length: {tot_output_len}")
+    rps = num_prefix * num_samples_per_prefix / tot_time
+    tps = (tot_input_len + tot_output_len) / tot_time
+    print(f"Throughput: {rps:.2f} requests/s, {tps:.2f} tokens/s")
+
+
+if __name__ == "__main__":
+    tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
+    backend = RuntimeEndpoint("http://127.0.0.1:30000")
+    set_default_backend(backend)
+
+    random.seed(0)
+    num_prefix = 10
+    num_samples_per_prefix = 32
+    prefix_length = 1024
+    suffix_length = 128
+    gen_len = 1
+    all_prompts, tot_input_len = prepare_prompts(
+        num_prefix, num_samples_per_prefix, prefix_length, suffix_length
+    )
+    print(f"Total input token length: {tot_input_len}")
+
+    backend.flush_cache()
+    test_prefix_caching(all_prompts, num_prefix, num_samples_per_prefix, gen_len)
+    backend.flush_cache()
+    test_prefix_caching_hint(all_prompts, num_prefix, num_samples_per_prefix, gen_len)
+    backend.flush_cache()
+    test_prefix_caching_send_all(
+        all_prompts, num_prefix, num_samples_per_prefix, gen_len
+    )

python/sglang/lang/backend/runtime_endpoint.py

@@ -55,6 +55,7 @@ def flush_cache(self):
             self.base_url + "/flush_cache",
             api_key=self.api_key,
             verify=self.verify,
+            method="POST",
         )
         self._assert_success(res)
 

python/sglang/srt/managers/schedule_batch.py

@@ -256,6 +256,7 @@ def __init__(
 
         # Prefix info
         self.prefix_indices = []
+        # Tokens to run prefill. input_tokens - shared_prefix_tokens.
         self.extend_input_len = 0
         self.last_node = None
 
@@ -316,6 +317,7 @@ def finished(self) -> bool:
     def init_next_round_input(self, tree_cache: Optional[BasePrefixCache] = None):
         self.fill_ids = self.origin_input_ids + self.output_ids
         if tree_cache is not None:
+            # tree cache is None if the prefix is not computed with tree cache.
             self.prefix_indices, self.last_node = tree_cache.match_prefix(
                 rid=self.rid, key=self.adjust_max_prefix_ids()
             )

python/sglang/srt/managers/schedule_policy.py

@@ -20,9 +20,11 @@
 from enum import Enum, auto
 from typing import Dict, List, Optional
 
+import torch
+
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.radix_cache import TreeNode
+from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
 
 # Clip the estimation of max_new_tokens for the request whose max_new_tokens is very large.
 # This can prevent the server from being too conservative.
@@ -32,6 +34,13 @@
     os.environ.get("SGLANG_CLIP_MAX_NEW_TOKENS_ESTIMATION", "4096")
 )
 
+# The threshold to apply in-batch prefix caching.
+# If we use too small value, in-batch prefix caching cannot be used. E.g.,
+# imagine "the" prefix.
+IN_BATCH_PREFIX_CACHING_THRESHOLD = int(
+    os.environ.get("SGLANG_IN_BATCH_PREFIX_CACHING_THRESHOLD", "32")
+)
+
 
 class SchedulePolicy:
     def __init__(self, policy: str, tree_cache: BasePrefixCache):
@@ -51,18 +60,50 @@ def calc_priority(self, waiting_queue: List[Req]):
 
         # Compute matched prefix length
         prefix_computed = False
+        # rid to deprioritize in the current run.
+        temporary_deprioritized = set()
         if policy == "lpm" or policy == "dfs-weight":
+            # It is used to find the matching prefix for in-batch prefix caching.
+            temp_radix = RadixCache(None, None, False)
             for r in waiting_queue:
+                prefix_ids = r.adjust_max_prefix_ids()
                 # NOTE: the prefix_indices must always be aligned with last_node
                 r.prefix_indices, r.last_node = self.tree_cache.match_prefix(
-                    rid=r.rid, key=r.adjust_max_prefix_ids()
+                    rid=r.rid, key=prefix_ids
                 )
 
+                # NOTE(sang): This logic is for In-batch prefix caching;
+                # If there are more than 1 request that have small matching prefix from
+                # existing cache, but all those requests share the same prefix, we prefer
+                # to schedule only one of them so that we can increase the cache hit rate.
+                # We prefer to set IN_BATCH_PREFIX_CACHING_THRESHOLD > 0 because too small
+                # threshold means we cannot use in-batch prefix caching for short prefixes.
+                # It is kind of common when the engine is long running (e.g., imagine "the").
+                if len(r.prefix_indices) <= IN_BATCH_PREFIX_CACHING_THRESHOLD:
+                    in_batch_matching_prefixes, _ = temp_radix.match_prefix(
+                        rid=r.rid, key=prefix_ids
+                    )
+                    if (
+                        len(in_batch_matching_prefixes)
+                        >= IN_BATCH_PREFIX_CACHING_THRESHOLD
+                    ):
+                        temporary_deprioritized.add(r.rid)
+                    else:
+                        temp_radix.insert(prefix_ids, torch.tensor(prefix_ids))
+
             prefix_computed = True
 
         if policy == "lpm":
             # Longest Prefix Match
-            waiting_queue.sort(key=lambda x: -len(x.prefix_indices))
+            def get_priority(r: Req):
+                score = 0
+                if r.rid in temporary_deprioritized:
+                    score = float("inf")
+                else:
+                    score = -len(r.prefix_indices)
+                return score
+
+            waiting_queue.sort(key=get_priority)
         elif policy == "fcfs":
             # first come first serve
             pass

python/sglang/srt/managers/scheduler.py

@@ -713,7 +713,7 @@ def check_memory(self):
             if crash_on_warnings():
                 raise ValueError(msg)
 
-    def get_next_batch_to_run(self):
+    def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
         # Merge the prefill batch into the running batch
         if self.last_batch and self.last_batch.forward_mode.is_extend():
             if self.being_chunked_req:

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Callable
+from typing import Callable, List, Tuple
 
 
 class BasePrefixCache(ABC):
@@ -10,7 +10,7 @@ def reset(self):
         pass
 
     @abstractmethod
-    def match_prefix(self, **kwargs):
+    def match_prefix(self, **kwargs) -> Tuple[List[int], int]:
         pass
 
     @abstractmethod

python/sglang/srt/mem_cache/chunk_cache.py

@@ -2,7 +2,7 @@
 
 """Cache for chunked prefill, used when RadixCache is disabled."""
 
-from typing import TYPE_CHECKING, Callable, List, Optional
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
 
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
@@ -30,7 +30,7 @@ def __init__(
     def reset(self):
         self.entries = {}
 
-    def match_prefix(self, rid: int, key: List[int]):
+    def match_prefix(self, rid: int, key: List[int]) -> Tuple[List[int], int]:
         if rid not in self.entries:
             return [], None
 

python/sglang/srt/mem_cache/radix_cache.py

@@ -22,7 +22,7 @@
 import heapq
 import time
 from collections import defaultdict
-from typing import TYPE_CHECKING, Callable, List, Optional
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
 
 import torch
 
@@ -76,7 +76,17 @@ def reset(self):
         self.root_node.lock_ref = 1
         self.evictable_size_ = 0
 
-    def match_prefix(self, key: List, **kwargs):
+    def match_prefix(self, key: List[int], **kwargs) -> Tuple[torch.Tensor, int]:
+        """Find the matching prefix from the radix tree.
+        Args:
+            key: A list of token IDs to find a matching prefix.
+        Returns:
+            A tuple of a tensor of matching prefix token IDs and
+            the last node that contains the prefix values. Note that
+            this API can modify the internal state of the Radix tree.
+            The last node create a new child if the prefix is shorter
+            than the last node's value.
+        """
         if self.disable:
             return [], self.root_node
 

python/sglang/utils.py

@@ -79,7 +79,14 @@ def status_code(self):
         return self.resp.status
 
 
-def http_request(url, json=None, stream=False, api_key=None, verify=None):
+def http_request(
+    url,
+    json=None,
+    stream=False,
+    api_key=None,
+    verify=None,
+    method: Optional[str] = None,
+):
     """A faster version of requests.post with low-level urllib API."""
     headers = {"Content-Type": "application/json; charset=utf-8"}
 
@@ -90,7 +97,7 @@ def http_request(url, json=None, stream=False, api_key=None, verify=None):
     if stream:
         return requests.post(url, json=json, stream=True, headers=headers)
     else:
-        req = urllib.request.Request(url, headers=headers)
+        req = urllib.request.Request(url, headers=headers, method=method)
         if json is None:
             data = None
         else: