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

sglang/llama3_eval.py

@@ -0,0 +1,316 @@
+# Adapt from https://github.com/fw-ai/llm_eval_meta
+
+import argparse
+import asyncio
+import os
+import pickle
+import re
+import shutil
+from collections import defaultdict
+from dataclasses import dataclass
+
+import httpx
+import numpy as np
+import openai
+import transformers
+from datasets import load_dataset
+from openai import AsyncOpenAI
+from tqdm import tqdm
+
+# Mapping providers to their clients and models
+provider_to_models = {
+    "b10": {
+        "8b": "meta-llama/Llama-3.1-8B-Instruct",
+        "70b": "meta-llama/Llama-3.1-70B-Instruct",
+        "405b": "meta-llama/Llama-3.1-405B-Instruct",
+    },
+    "oai": {
+        "8b": "meta-llama/Llama-3.1-8B-Instruct",
+        "70b": "meta-llama/Llama-3.1-70B-Instruct",
+        "405b": "meta-llama/Llama-3.1-405B-Instruct",
+    },
+    "sgl": {
+        "8b": "meta-llama/Llama-3.1-8B-Instruct",
+        "70b": "meta-llama/Llama-3.1-70B-Instruct",
+        "405b": "meta-llama/Llama-3.1-405B-Instruct",
+    },
+}
+
+
+async def fetch_responses(
+    client, prompt, semaphore, index, provider, model_size, output_dir, max_tokens
+):
+    output_file = os.path.join(output_dir, f"response_{index}.pkl")
+    if os.path.exists(output_file):
+        print(f"File {output_file} already exists, skipping.")
+        return
+
+    async with semaphore:
+        response = await client.completions.create(
+            model=provider_to_models[provider][model_size],
+            prompt=prompt,
+            temperature=0.0,
+            max_tokens=max_tokens,
+        )
+        if isinstance(response, openai.BadRequestError):
+            with open(output_file, "wb") as f:
+                pickle.dump("bad_response", f)
+        assert isinstance(response, openai.types.completion.Completion)
+        # Save response to a file
+        with open(output_file, "wb") as f:
+            pickle.dump(response, f)
+
+
+TASK_TO_MAX_TOKENS = {
+    "evals__mmlu__details": 1,
+    "evals__mmlu__0_shot__cot__details": 1024,
+    # Official meta uses 1024, but a small % (.05) of questions are answered correctly after relaxing
+    "evals__mmlu_pro__details": 2048,
+    "evals__gsm8k__details": 1024,
+}
+
+TASK_TO_EVAL_SET = {
+    "mmlu": "evals__mmlu__details",
+    "mmlu_cot": "evals__mmlu__0_shot__cot__details",
+    "mmlu_pro": "evals__mmlu_pro__details",
+    "gsm8k": "evals__gsm8k__details",
+}
+
+
+class CustomAsyncHTTPXClient(httpx.AsyncClient):
+    async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response:
+        request.url = httpx.URL(
+            f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict"
+        )
+        return await super().send(request, *args, **kwargs)
+
+
+def get_client(provider):
+    if provider not in "b10":
+        if os.getenv("OPENAI_API_KEY") == None:
+            os.environ["OPENAI_API_KEY"] = "EMPTY"
+    return {
+        "oai": AsyncOpenAI(base_url="http://127.0.0.1:8000/v1/"),
+        "b10": AsyncOpenAI(
+            api_key=f"Api-Key {os.getenv('OPENAI_API_KEY')}",
+            base_url=f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict",
+            http_client=CustomAsyncHTTPXClient(),
+        ),
+        "sgl": AsyncOpenAI(base_url="http://127.0.0.1:30000/v1/"),
+    }[provider]
+
+
+# Define the benchmark function
+async def benchmark(args):
+    ds = load_dataset(
+        "meta-llama/Llama-3.1-405B-Instruct-evals",
+        f"Llama-3.1-405B-Instruct-{TASK_TO_EVAL_SET[args.task]}",
+    )
+    semaphore = asyncio.Semaphore(args.concurrency)  # Limit to 16 concurrent tasks
+
+    if args.num_examples is None:
+        args.num_examples = len(ds["latest"]["input_final_prompts"])
+    prompts = ds["latest"]["input_final_prompts"][: args.num_examples]
+
+    # Create the output directory if it does not exist
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    tasks = []
+    # Create the tasks with tqdm progress bar
+    max_tokens = TASK_TO_MAX_TOKENS[TASK_TO_EVAL_SET[args.task]]
+    client = get_client(args.provider)
+    for idx, prompt in enumerate(tqdm(prompts, desc="Creating tasks")):
+        tasks.append(
+            asyncio.create_task(
+                fetch_responses(
+                    client,
+                    f"<|begin_of_text|>{prompt[0]}",
+                    semaphore,
+                    idx,
+                    args.provider,
+                    args.model_size,
+                    args.output_dir,
+                    max_tokens=max_tokens,
+                )
+            )
+        )
+
+    # Run the tasks with tqdm progress bar
+    for future in tqdm(
+        asyncio.as_completed(tasks), total=len(tasks), desc="Processing tasks"
+    ):
+        await future
+
+
+def get_mmlu_answer(response):
+    if response is not None:
+        return response.choices[0].text.lstrip().rstrip().upper().replace(".", "")
+    return None
+
+
+def get_mmlu_cot_answer(response):
+    pattern = r"The best answer is (.+)\.?"
+    match = re.search(pattern, response.choices[0].text)
+    if match:
+        return match.group(1).replace(".", "").replace("*", "")
+
+    pattern = r"the best answer is (.+)\.?"
+    match = re.search(pattern, response.choices[0].text)
+    if match:
+        return match.group(1).replace(".", "")
+
+    pattern = r"The correct answer is (.+)\.?"
+    match = re.search(pattern, response.choices[0].text)
+    if match:
+        return match.group(1).replace(".", "")
+
+    pattern = r"the correct answer is (.+)\.?"
+    match = re.search(pattern, response.choices[0].text)
+    if match:
+        return match.group(1).replace(".", "")
+
+
+def get_answer_gsm8k(response):
+    pattern = r"The final answer is (.+)\.?"
+    match = re.search(pattern, response.choices[0].text)
+    if match:
+        s = match.group(1)
+        for ok_symbol in ["%", "$"]:
+            s = s.replace(ok_symbol, "")
+        return s
+
+
+TASK_TO_ANSWER_EXTRACTOR = {
+    "evals__mmlu__details": get_mmlu_answer,
+    "evals__mmlu__0_shot__cot__details": get_mmlu_cot_answer,
+    "evals__gsm8k__details": get_answer_gsm8k,
+    "evals__mmlu_pro__details": get_mmlu_cot_answer,
+}
+
+
+def get_dataset_from_task(task, response_path, model_size):
+    ds_405b = load_dataset(
+        f"meta-llama/Llama-3.1-405B-Instruct-evals",
+        f"Llama-3.1-405B-Instruct-{task}",
+    )
+    ds_405b_hash_order = [x[0] for x in ds_405b["latest"]["input_final_prompts_hash"]]
+
+    if "70b" in model_size or "8b" in model_size:
+        if "70" in model_size:
+            ref_model_ds = load_dataset(
+                f"meta-llama/Llama-3.1-70B-Instruct-evals",
+                f"Llama-3.1-70B-Instruct-{task}",
+            )
+        else:
+            ref_model_ds = load_dataset(
+                f"meta-llama/Llama-3.1-8B-Instruct-evals",
+                f"Llama-3.1-8B-Instruct-{task}",
+            )
+
+        hash_to_row = {}
+        for row in ref_model_ds["latest"]:
+            hash_to_row[row["input_final_prompts_hash"][0]] = row
+        reordered_rows = []
+        for prompt_hash in ds_405b_hash_order:
+            reordered_rows.append(hash_to_row[prompt_hash])
+        ref_model_ds["latest"] = reordered_rows
+        return ref_model_ds
+
+    return ds_405b
+
+
+def analyze(task, response_path, model_size):
+    ds = get_dataset_from_task(task, response_path, model_size)
+
+    responses = []
+    total = len(ds["latest"])
+
+    for i in range(0, total):
+        response = pickle.load(
+            open(os.path.join(response_path, f"response_{i}.pkl"), "rb")
+        )
+        responses.append(response)
+
+    @dataclass
+    class Stats:
+        correct: int = 0
+        total: int = 0
+        meta_correct: int = 0
+
+        average: float = None
+
+    subtask_name_to_stats = defaultdict(lambda: Stats())
+
+    for response, ds_row in zip(responses, ds["latest"]):
+        model_answer = TASK_TO_ANSWER_EXTRACTOR[task](response)
+
+        subtask = ds_row["subtask_name"]
+
+        is_eval_correct = model_answer in ds_row["input_correct_responses"]
+        if is_eval_correct:
+            subtask_name_to_stats[subtask].correct += 1
+
+        if ds_row["is_correct"]:
+            subtask_name_to_stats[subtask].meta_correct += 1
+
+        subtask_name_to_stats[subtask].total += 1
+
+    micro_stats = Stats()
+    for subtask, stats in subtask_name_to_stats.items():
+        stats.average = stats.correct / stats.total
+        stats.meta_average = stats.meta_correct / stats.total
+
+        micro_stats.correct += stats.correct
+        micro_stats.total += stats.total
+        micro_stats.meta_correct += stats.meta_correct
+
+    micro_stats.average = micro_stats.correct / micro_stats.total
+    micro_stats.meta_average = micro_stats.meta_correct / micro_stats.total
+
+    print("Macro average", np.mean([x.average for x in subtask_name_to_stats.values()]))
+    print(
+        "Meta Macro average",
+        np.mean([x.meta_average for x in subtask_name_to_stats.values()]),
+    )
+    print("Micro average", micro_stats.average)
+    print("Meta Micro average", micro_stats.meta_average)
+
+
+# Entry point for the script
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Script to run model with specified parameters."
+    )
+    parser.add_argument(
+        "--model-size",
+        type=str,
+        default="8b",
+        help="Size of the model (e.g., 8b or 70b)",
+    )
+    parser.add_argument(
+        "--provider",
+        type=str,
+        default="sgl",
+        help="Provider name (e.g., sgl, oai, b10)",
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        required=True,
+        help="Task (e.g., mmlu, mmlu_cot, mmlu_pro, gsm8k)",
+    )
+    parser.add_argument(
+        "--num-examples", type=int, default=None, help="Number of examples to process"
+    )
+    parser.add_argument("--concurrency", type=int, default=16)
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="tmp-output-dir",
+        help="Directory to save responses",
+    )
+
+    args = parser.parse_args()
+    asyncio.run(benchmark(args))
+    analyze(TASK_TO_EVAL_SET[args.task], args.output_dir, args.model_size)
+    shutil.rmtree("tmp-output-dir", ignore_errors=True)

sglang/srt/constrained/outlines_backend.py

@@ -42,6 +42,7 @@ class OutlinesGrammar(BaseGrammarObject):
         self.guide = guide
         self.jump_forward_map = jump_forward_map
         self.state = 0
+        self.finished = False
 
     def accept_token(self, token: int):
         self.state = self.guide.get_next_state(self.state, token)
@@ -84,6 +85,10 @@ class OutlinesGrammar(BaseGrammarObject):
     ) -> torch.Tensor:
         return torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device)
 
+    @staticmethod
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        return vocab_mask
+
     def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
         tokens = torch.tensor(
             self.guide.get_next_instruction(self.state).tokens, dtype=torch.int64

sglang/srt/constrained/xgrammar_backend.py

@@ -45,6 +45,7 @@ class XGrammarGrammar(BaseGrammarObject):
         self.matcher = matcher
         self.vocab_size = vocab_size
         self.ctx = ctx
+        self.finished = False
 
     def accept_token(self, token: int):
         assert self.matcher.accept_token(token)
@@ -85,12 +86,11 @@ class XGrammarGrammar(BaseGrammarObject):
         self.matcher.fill_next_token_bitmask(vocab_mask, idx)
 
     @staticmethod
-    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
-        if vocab_mask.device.type != logits.device.type:
-            # vocab_mask must then be on the same device as logits
-            # when applying the token bitmask, so we check and move if needed
-            vocab_mask = vocab_mask.to(logits.device)
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        return vocab_mask.to(device, non_blocking=True)
 
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
         apply_token_bitmask_inplace(logits, vocab_mask)
 
     def copy(self):
@@ -117,7 +117,10 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         key_type, key_string = key
         if key_type == "json":
             try:
-                ctx = self.grammar_compiler.compile_json_schema(schema=key_string)
+                if key_string == "$$ANY$$":
+                    ctx = self.grammar_compiler.compile_builtin_json_grammar()
+                else:
+                    ctx = self.grammar_compiler.compile_json_schema(schema=key_string)
             except RuntimeError as e:
                 logging.warning(
                     f"Skip invalid json_schema: json_schema={key_string}, {e=}"

sglang/srt/layers/attention/__init__.py

@@ -52,12 +52,13 @@ class AttentionBackend(ABC):
         v: torch.Tensor,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
     ):
         """Run forward on an attention layer."""
         if forward_batch.forward_mode.is_decode():
-            return self.forward_decode(q, k, v, layer, forward_batch)
+            return self.forward_decode(q, k, v, layer, forward_batch, save_kv_cache)
         else:
-            return self.forward_extend(q, k, v, layer, forward_batch)
+            return self.forward_extend(q, k, v, layer, forward_batch, save_kv_cache)
 
     def forward_decode(
         self,
@@ -66,6 +67,7 @@ class AttentionBackend(ABC):
         v: torch.Tensor,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
     ):
         """Run a forward for decode."""
         raise NotImplementedError()
@@ -77,6 +79,7 @@ class AttentionBackend(ABC):
         v: torch.Tensor,
         layer: RadixAttention,
         forward_batch: ForwardBatch,
+        save_kv_cache: bool = True,
     ):
         """Run a forward for extend."""
         raise NotImplementedError()

sglang/srt/layers/attention/double_sparsity_backend.py

@@ -165,7 +165,13 @@ class DoubleSparseAttnBackend(AttentionBackend):
         return 1
 
     def forward_extend(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         # TODO: reuse the buffer across layers
         if layer.qk_head_dim != layer.v_head_dim:
@@ -181,9 +187,10 @@ class DoubleSparseAttnBackend(AttentionBackend):
             .expand(k.shape[0], -1, -1),
         )
 
-        forward_batch.token_to_kv_pool.set_kv_buffer(
-            layer, forward_batch.out_cache_loc, k, v, k_label
-        )
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v, k_label
+            )
 
         (
             start_loc,
@@ -212,7 +219,13 @@ class DoubleSparseAttnBackend(AttentionBackend):
         return o
 
     def forward_decode(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         # During torch.compile, there is a bug in rotary_emb that causes the
         # output value to have a 3D tensor shape. This reshapes the output correctly.
@@ -242,9 +255,10 @@ class DoubleSparseAttnBackend(AttentionBackend):
             .expand(k.shape[0], -1, -1),
         )
 
-        forward_batch.token_to_kv_pool.set_kv_buffer(
-            layer, forward_batch.out_cache_loc, k, v, k_label
-        )
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v, k_label
+            )
 
         # NOTE(Andy) shouldn't be used when max_len_in_batch < heavy_token_num
         #            and set a minimum value for sparse_decode

sglang/srt/layers/attention/flashinfer_backend.py

@@ -221,7 +221,13 @@ class FlashInferAttnBackend(AttentionBackend):
         return 0
 
     def forward_extend(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         prefill_wrapper_paged = self.prefill_wrappers_paged[
             self._get_wrapper_idx(layer)
@@ -237,7 +243,8 @@ class FlashInferAttnBackend(AttentionBackend):
         if not use_ragged:
             if k is not None:
                 assert v is not None
-                forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
+                if save_kv_cache:
+                    forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
 
             o = prefill_wrapper_paged.forward(
                 q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
@@ -270,12 +277,19 @@ class FlashInferAttnBackend(AttentionBackend):
 
                 o, _ = merge_state(o1, s1, o2, s2)
 
-            forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
 
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)
 
     def forward_decode(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         decode_wrapper = self.forward_metadata[0][self._get_wrapper_idx(layer)]
         cache_loc = (
@@ -286,7 +300,8 @@ class FlashInferAttnBackend(AttentionBackend):
 
         if k is not None:
             assert v is not None
-            forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
+            if save_kv_cache:
+                forward_batch.token_to_kv_pool.set_kv_buffer(layer, cache_loc, k, v)
 
         o = decode_wrapper.forward(
             q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
@@ -663,6 +678,7 @@ class FlashInferIndicesUpdaterPrefill:
                 self.num_qo_heads,
                 self.num_kv_heads,
                 self.head_dim,
+                q_data_type=self.q_data_type,
             )
 
         # cached part
@@ -676,6 +692,7 @@ class FlashInferIndicesUpdaterPrefill:
             self.num_kv_heads,
             self.head_dim,
             1,
+            q_data_type=self.q_data_type,
         )
 
 

sglang/srt/layers/attention/torch_native_backend.py

@@ -216,16 +216,23 @@ class TorchNativeAttnBackend(AttentionBackend):
         return output
 
     def forward_extend(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         if layer.qk_head_dim != layer.v_head_dim:
             o = q.new_empty((q.shape[0], layer.tp_q_head_num * layer.v_head_dim))
         else:
             o = torch.empty_like(q)
 
-        forward_batch.token_to_kv_pool.set_kv_buffer(
-            layer, forward_batch.out_cache_loc, k, v
-        )
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
 
         use_gqa = layer.tp_q_head_num != layer.tp_k_head_num
 
@@ -249,7 +256,13 @@ class TorchNativeAttnBackend(AttentionBackend):
         return o
 
     def forward_decode(
-        self, q, k, v, layer: RadixAttention, forward_batch: ForwardBatch
+        self,
+        q,
+        k,
+        v,
+        layer: RadixAttention,
+        forward_batch: ForwardBatch,
+        save_kv_cache=True,
     ):
         # During torch.compile, there is a bug in rotary_emb that causes the
         # output value to have a 3D tensor shape. This reshapes the output correctly.
@@ -260,9 +273,10 @@ class TorchNativeAttnBackend(AttentionBackend):
         else:
             o = torch.empty_like(q)
 
-        forward_batch.token_to_kv_pool.set_kv_buffer(
-            layer, forward_batch.out_cache_loc, k, v
-        )
+        if save_kv_cache:
+            forward_batch.token_to_kv_pool.set_kv_buffer(
+                layer, forward_batch.out_cache_loc, k, v
+            )
 
         use_gqa = layer.tp_q_head_num != layer.tp_k_head_num