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

sglang/bench_offline_throughput.py

@@ -331,6 +331,7 @@ def throughput_test(
         extra_request_body=extra_request_body,
         profile=bench_args.profile,
     )
+    backend.shutdown()
 
     if bench_args.result_filename:
         with open(bench_args.result_filename, "a") as fout:

sglang/bench_serving.py

@@ -897,6 +897,7 @@ async def benchmark(
     else:
         raise ValueError(f"Unknown backend: {backend}")
 
+    # Limit concurrency
     # From https://github.com/vllm-project/vllm/pull/9390
     semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None
 
@@ -906,6 +907,7 @@ async def benchmark(
         async with semaphore:
             return await request_func(request_func_input=request_func_input, pbar=pbar)
 
+    # Warmup
     print("Starting initial single prompt test run...")
     test_prompt, test_prompt_len, test_output_len = input_requests[0]
     test_input = RequestFuncInput(
@@ -924,11 +926,15 @@ async def benchmark(
             f"are correctly specified. Error: {test_output.error}"
         )
     else:
-        requests.post(base_url + "/flush_cache")
         print("Initial test run completed. Starting main benchmark run...")
 
-    time.sleep(1.5)
+    # Flush cache
+    if "sglang" in backend:
+        requests.post(base_url + "/flush_cache")
+
+    time.sleep(1.0)
 
+    # Start profiler
     if profile:
         print("Starting profiler...")
         profile_output = await async_request_profile(
@@ -939,6 +945,7 @@ async def benchmark(
 
     pbar = None if disable_tqdm else tqdm(total=len(input_requests))
 
+    # Run all requests
     benchmark_start_time = time.perf_counter()
     tasks: List[asyncio.Task] = []
     async for request in get_request(input_requests, request_rate):
@@ -959,6 +966,7 @@ async def benchmark(
         )
     outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks)
 
+    # Stop profiler
     if profile:
         print("Stopping profiler...")
         profile_output = await async_request_profile(api_url=base_url + "/stop_profile")
@@ -968,8 +976,8 @@ async def benchmark(
     if pbar is not None:
         pbar.close()
 
+    # Compute metrics and print results
     benchmark_duration = time.perf_counter() - benchmark_start_time
-
     metrics, output_lens = calculate_metrics(
         input_requests=input_requests,
         outputs=outputs,

sglang/lang/backend/openai.py

@@ -366,6 +366,11 @@ class OpenAI(BaseBackend):
 def openai_completion(
     client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
 ):
+    # if "ebnf" is in kwargs, warn and remove
+    if "ebnf" in kwargs:
+        warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.")
+        del kwargs["ebnf"]
+
     for attempt in range(retries):
         try:
             if is_chat:
@@ -398,6 +403,11 @@ def openai_completion(
 def openai_completion_stream(
     client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
 ):
+    # if "ebnf" is in kwargs, warn and remove
+    if "ebnf" in kwargs:
+        warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.")
+        del kwargs["ebnf"]
+
     for attempt in range(retries):
         try:
             if is_chat:

sglang/srt/configs/model_config.py

@@ -15,7 +15,7 @@
 import json
 import logging
 from enum import IntEnum, auto
-from typing import List, Optional, Union
+from typing import List, Optional, Set, Union
 
 import torch
 from transformers import PretrainedConfig
@@ -47,6 +47,7 @@ class ModelConfig:
         self.model_path = model_path
         self.revision = revision
         self.quantization = quantization
+
         # Parse args
         self.model_override_args = json.loads(model_override_args)
         self.hf_config = get_config(
@@ -130,7 +131,8 @@ class ModelConfig:
         # Veirfy quantization
         self._verify_quantization()
 
-        # Multimodel attrs
+        # Cache attributes
+        self.hf_eos_token_id = self.get_hf_eos_token_id()
         self.image_token_id = getattr(self.hf_config, "image_token_id", None)
 
     # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289
@@ -271,6 +273,13 @@ class ModelConfig:
                     self.quantization,
                 )
 
+    def get_hf_eos_token_id(self) -> Optional[Set[int]]:
+        eos_ids = getattr(self.hf_config, "eos_token_id", None)
+        if eos_ids:
+            # it can be either int or list of int
+            eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
+        return eos_ids
+
 
 def get_hf_text_config(config: PretrainedConfig):
     """Get the "sub" config relevant to llm for multi modal models.

sglang/srt/constrained/xgrammar_backend.py

@@ -126,6 +126,12 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
                     f"Skip invalid json_schema: json_schema={key_string}, {e=}"
                 )
                 return None
+        elif key_type == "ebnf":
+            try:
+                ctx = self.grammar_compiler.compile_grammar(key_string)
+            except RuntimeError as e:
+                logging.warning(f"Skip invalid ebnf: ebnf={key_string}, {e=}")
+                return None
         elif key_type == "regex":
             logger.warning(
                 "regex hasn't been supported by xgrammar yet. This is skipped."

sglang/srt/layers/attention/__init__.py

@@ -2,7 +2,6 @@ from abc import ABC, abstractmethod
 from typing import Optional
 
 import torch
-from torch import nn
 
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch

sglang/srt/layers/attention/flashinfer_backend.py

@@ -8,8 +8,9 @@ Each backend supports two operators: extend (i.e. prefill with cached prefix) an
 """
 
 import os
+from dataclasses import dataclass
 from enum import Enum, auto
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Union
 
 import torch
 import triton
@@ -38,12 +39,25 @@ class WrapperDispatch(Enum):
     CROSS_ATTENTION = auto()
 
 
+@dataclass
+class DecodeMetadata:
+    decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper]
+
+
+@dataclass
+class PrefillMetadata:
+    prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper]
+    use_ragged: bool
+    extend_no_prefix: bool
+
+
 class FlashInferAttnBackend(AttentionBackend):
     """Flashinfer attention kernels."""
 
     def __init__(self, model_runner: ModelRunner):
         super().__init__()
 
+        # Parse constants
         self.decode_use_tensor_cores = should_use_tensor_core(
             kv_cache_dtype=model_runner.kv_cache_dtype,
             num_attention_heads=model_runner.model_config.num_attention_heads
@@ -52,7 +66,6 @@ class FlashInferAttnBackend(AttentionBackend):
                 model_runner.tp_size
             ),
         )
-
         self.max_context_len = model_runner.model_config.context_len
 
         assert not (
@@ -120,8 +133,8 @@ class FlashInferAttnBackend(AttentionBackend):
         )
 
         # Other metadata
-        self.forward_metadata = None
-        self.cuda_graph_metadata = {}
+        self.forward_metadata: Union[PrefillMetadata, DecodeMetadata] = None
+        self.decode_cuda_graph_metadata = {}
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         if forward_batch.forward_mode.is_decode():
@@ -129,10 +142,10 @@ class FlashInferAttnBackend(AttentionBackend):
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
                 forward_batch.seq_lens_sum,
-                decode_wrappers=None,
+                decode_wrappers=self.decode_wrappers,
                 encoder_lens=forward_batch.encoder_lens,
             )
-            self.forward_metadata = (self.decode_wrappers,)
+            self.forward_metadata = DecodeMetadata(self.decode_wrappers)
         else:
             prefix_lens = forward_batch.extend_prefix_lens
 
@@ -149,11 +162,13 @@ class FlashInferAttnBackend(AttentionBackend):
                 forward_batch.seq_lens,
                 forward_batch.seq_lens_sum,
                 prefix_lens,
+                prefill_wrappers=self.prefill_wrappers_paged,
                 use_ragged=use_ragged,
                 encoder_lens=forward_batch.encoder_lens,
             )
-
-            self.forward_metadata = (use_ragged, extend_no_prefix)
+            self.forward_metadata = PrefillMetadata(
+                self.prefill_wrappers_paged, use_ragged, extend_no_prefix
+            )
 
     def init_cuda_graph_state(self, max_bs: int):
         cuda_graph_kv_indices = torch.zeros(
@@ -194,8 +209,8 @@ class FlashInferAttnBackend(AttentionBackend):
             decode_wrappers=decode_wrappers,
             encoder_lens=encoder_lens,
         )
-        self.cuda_graph_metadata[bs] = decode_wrappers
-        self.forward_metadata = (decode_wrappers,)
+        self.decode_cuda_graph_metadata[bs] = decode_wrappers
+        self.forward_metadata = DecodeMetadata(decode_wrappers)
 
     def init_forward_metadata_replay_cuda_graph(
         self,
@@ -209,7 +224,7 @@ class FlashInferAttnBackend(AttentionBackend):
             req_pool_indices[:bs],
             seq_lens[:bs],
             seq_lens_sum,
-            decode_wrappers=self.cuda_graph_metadata[bs],
+            decode_wrappers=self.decode_cuda_graph_metadata[bs],
             encoder_lens=encoder_lens[:bs] if encoder_lens is not None else None,
         )
 
@@ -225,18 +240,16 @@ class FlashInferAttnBackend(AttentionBackend):
         forward_batch: ForwardBatch,
         save_kv_cache=True,
     ):
-        prefill_wrapper_paged = self.prefill_wrappers_paged[
+        prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[
             self._get_wrapper_idx(layer)
         ]
-
-        use_ragged, extend_no_prefix = self.forward_metadata
         cache_loc = (
             forward_batch.out_cache_loc
             if not layer.is_cross_attention
             else forward_batch.encoder_out_cache_loc
         )
 
-        if not use_ragged:
+        if not self.forward_metadata.use_ragged:
             if k is not None:
                 assert v is not None
                 if save_kv_cache:
@@ -260,7 +273,7 @@ class FlashInferAttnBackend(AttentionBackend):
                 logits_soft_cap=layer.logit_cap,
             )
 
-            if extend_no_prefix:
+            if self.forward_metadata.extend_no_prefix:
                 o = o1
             else:
                 o2, s2 = prefill_wrapper_paged.forward_return_lse(
@@ -287,7 +300,9 @@ class FlashInferAttnBackend(AttentionBackend):
         forward_batch: ForwardBatch,
         save_kv_cache=True,
     ):
-        decode_wrapper = self.forward_metadata[0][self._get_wrapper_idx(layer)]
+        decode_wrapper = self.forward_metadata.decode_wrappers[
+            self._get_wrapper_idx(layer)
+        ]
         cache_loc = (
             forward_batch.out_cache_loc
             if not layer.is_cross_attention
@@ -322,7 +337,7 @@ class FlashInferAttnBackend(AttentionBackend):
 
 class FlashInferIndicesUpdaterDecode:
     def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
-        # Constants
+        # Parse Constants
         self.num_qo_heads = (
             model_runner.model_config.num_attention_heads // model_runner.tp_size
         )
@@ -340,9 +355,8 @@ class FlashInferIndicesUpdaterDecode:
         self.kv_indptr = attn_backend.kv_indptr
         self.kv_last_page_len = attn_backend.kv_last_page_len
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
-        self.decode_wrappers = attn_backend.decode_wrappers
 
-        # Dispatch
+        # Dispatch the update function
         if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
             self.update = self.update_sliding_window
         elif self.attn_backend.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
@@ -356,7 +370,7 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers: List,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: torch.Tensor,
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
@@ -367,7 +381,7 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers: List,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: torch.Tensor,
     ):
         decode_wrappers = decode_wrappers or self.decode_wrappers
@@ -385,11 +399,9 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers: List,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: torch.Tensor,
     ):
-        decode_wrappers = decode_wrappers or self.decode_wrappers
-
         for wrapper_id in range(2):
             if wrapper_id == 0:
                 # Sliding window attention
@@ -419,11 +431,9 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers: List,
+        decode_wrappers: List[BatchDecodeWithPagedKVCacheWrapper],
         encoder_lens: torch.Tensor,
     ):
-        decode_wrappers = decode_wrappers or self.decode_wrappers
-
         for wrapper_id in range(2):
             if wrapper_id == 0:
                 # Normal attention
@@ -446,7 +456,7 @@ class FlashInferIndicesUpdaterDecode:
 
     def call_begin_forward(
         self,
-        wrapper,
+        wrapper: BatchDecodeWithPagedKVCacheWrapper,
         req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,
@@ -486,7 +496,7 @@ class FlashInferIndicesUpdaterDecode:
 
 class FlashInferIndicesUpdaterPrefill:
     def __init__(self, model_runner: ModelRunner, attn_backend: AttentionBackend):
-        # Constants
+        # Parse Constants
         self.num_qo_heads = (
             model_runner.model_config.num_attention_heads // model_runner.tp_size
         )
@@ -505,10 +515,9 @@ class FlashInferIndicesUpdaterPrefill:
         self.kv_last_page_len = attn_backend.kv_last_page_len
         self.qo_indptr = attn_backend.qo_indptr
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
-        self.wrapper_ragged = attn_backend.prefill_wrapper_ragged
-        self.wrappers_paged = attn_backend.prefill_wrappers_paged
+        self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
 
-        # Dispatch
+        # Dispatch the update function
         if self.attn_backend.dispatch_reason == WrapperDispatch.SLIDING_WINDOW:
             self.update = self.update_sliding_window
         elif self.attn_backend.dispatch_reason == WrapperDispatch.CROSS_ATTENTION:
@@ -523,6 +532,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: torch.Tensor,
     ):
@@ -535,6 +545,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: torch.Tensor,
     ):
@@ -546,8 +557,8 @@ class FlashInferIndicesUpdaterPrefill:
             paged_kernel_lens_sum = seq_lens_sum
 
         self.call_begin_forward(
-            self.wrapper_ragged,
-            self.wrappers_paged[0],
+            self.prefill_wrapper_ragged,
+            prefill_wrappers[0],
             req_pool_indices,
             paged_kernel_lens,
             paged_kernel_lens_sum,
@@ -565,6 +576,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: torch.Tensor,
     ):
@@ -584,8 +596,8 @@ class FlashInferIndicesUpdaterPrefill:
             kv_start_idx = seq_lens - paged_kernel_lens
 
             self.call_begin_forward(
-                self.wrapper_ragged,
-                self.wrappers_paged[wrapper_id],
+                self.prefill_wrapper_ragged,
+                prefill_wrappers[wrapper_id],
                 req_pool_indices,
                 paged_kernel_lens,
                 paged_kernel_lens_sum,
@@ -603,6 +615,7 @@ class FlashInferIndicesUpdaterPrefill:
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
         prefix_lens: torch.Tensor,
+        prefill_wrappers: List[BatchPrefillWithPagedKVCacheWrapper],
         use_ragged: bool,
         encoder_lens: torch.Tensor,
     ):
@@ -619,8 +632,8 @@ class FlashInferIndicesUpdaterPrefill:
                 paged_kernel_lens_sum = paged_kernel_lens.sum().item()
 
             self.call_begin_forward(
-                self.wrapper_ragged,
-                self.wrappers_paged[wrapper_id],
+                self.prefill_wrapper_ragged,
+                prefill_wrappers[wrapper_id],
                 req_pool_indices,
                 paged_kernel_lens,
                 paged_kernel_lens_sum,
@@ -634,8 +647,8 @@ class FlashInferIndicesUpdaterPrefill:
 
     def call_begin_forward(
         self,
-        wrapper_ragged,
-        wrapper_paged,
+        wrapper_ragged: BatchPrefillWithRaggedKVCacheWrapper,
+        wrapper_paged: BatchPrefillWithPagedKVCacheWrapper,
         req_pool_indices: torch.Tensor,
         paged_kernel_lens: torch.Tensor,
         paged_kernel_lens_sum: int,

sglang/srt/layers/attention/triton_ops/extend_attention.py

@@ -292,27 +292,33 @@ def extend_attention_fwd(
         BLOCK_DPE = 0
     BLOCK_DV = triton.next_power_of_2(Lv)
 
-    if is_cuda_available and CUDA_CAPABILITY[0] >= 9:
-        if Lq <= 256:
-            BLOCK_M, BLOCK_N = (128, 64)
-        else:
-            BLOCK_M, BLOCK_N = (32, 64)
-    elif is_cuda_available and CUDA_CAPABILITY[0] >= 8:
-        if Lq <= 128:
-            BLOCK_M, BLOCK_N = (128, 128)
-        elif Lq <= 256:
-            BLOCK_M, BLOCK_N = (64, 64)
-        else:
-            BLOCK_M, BLOCK_N = (32, 64)
+    if is_hip_:
+        BLOCK_M, BLOCK_N = (64, 64)
+        num_warps = 4
+
     else:
-        BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
+        if is_cuda_available and CUDA_CAPABILITY[0] >= 9:
+            if Lq <= 256:
+                BLOCK_M, BLOCK_N = (128, 64)
+            else:
+                BLOCK_M, BLOCK_N = (32, 64)
+        elif is_cuda_available and CUDA_CAPABILITY[0] >= 8:
+            if Lq <= 128:
+                BLOCK_M, BLOCK_N = (128, 128)
+            elif Lq <= 256:
+                BLOCK_M, BLOCK_N = (64, 64)
+            else:
+                BLOCK_M, BLOCK_N = (32, 64)
+        else:
+            BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
+
+        num_warps = 4 if Lk <= 64 else 8
 
     sm_scale = sm_scale or 1.0 / (Lq**0.5)
     batch_size, head_num = b_seq_len.shape[0], q_extend.shape[1]
     kv_group_num = q_extend.shape[1] // k_extend.shape[1]
 
     grid = (batch_size, head_num, triton.cdiv(max_len_extend, BLOCK_M))
-    num_warps = 4 if Lk <= 64 else 8
     num_stages = 1
 
     extra_kargs = {}

sglang/srt/layers/logits_processor.py

@@ -24,7 +24,11 @@ from vllm.distributed import (
 )
 
 from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+)
 
 
 @dataclasses.dataclass
@@ -46,6 +50,10 @@ class LogitsProcessorOutput:
     output_top_logprobs_val: List = None
     output_top_logprobs_idx: List = None
 
+    # Used by speculative decoding (EAGLE)
+    # The output of transformer layers
+    hidden_states: Optional[torch.Tensor] = None
+
 
 @dataclasses.dataclass
 class LogitsMetadata:
@@ -61,6 +69,8 @@ class LogitsMetadata:
     extend_logprob_start_lens_cpu: Optional[List[int]] = None
     extend_logprob_pruned_lens_cpu: Optional[List[int]] = None
 
+    capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL
+
     @classmethod
     def from_forward_batch(cls, forward_batch: ForwardBatch):
         extend_logprob_pruned_lens_cpu = None
@@ -78,6 +88,11 @@ class LogitsMetadata:
         else:
             return_top_logprob = False
 
+        if forward_batch.spec_info:
+            capture_hidden_mode = forward_batch.spec_info.capture_hidden_mode
+        else:
+            capture_hidden_mode = CaptureHiddenMode.NULL
+
         return cls(
             forward_mode=forward_batch.forward_mode,
             top_logprobs_nums=forward_batch.top_logprobs_nums,
@@ -87,6 +102,7 @@ class LogitsMetadata:
             extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
             extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
             extend_logprob_pruned_lens_cpu=extend_logprob_pruned_lens_cpu,
+            capture_hidden_mode=capture_hidden_mode,
         )
 
 
@@ -116,7 +132,10 @@ class LogitsProcessor(nn.Module):
         assert isinstance(logits_metadata, LogitsMetadata)
 
         # Get the last hidden states and last logits for the next token prediction
-        if logits_metadata.forward_mode.is_decode():
+        if (
+            logits_metadata.forward_mode.is_decode()
+            or logits_metadata.forward_mode.is_target_verify()
+        ):
             last_index = None
             last_hidden = hidden_states
         else:
@@ -137,6 +156,15 @@ class LogitsProcessor(nn.Module):
         if not logits_metadata.return_logprob:
             return LogitsProcessorOutput(
                 next_token_logits=last_logits,
+                hidden_states=(
+                    hidden_states
+                    if logits_metadata.capture_hidden_mode.is_full()
+                    else (
+                        last_hidden
+                        if logits_metadata.capture_hidden_mode.is_last()
+                        else None
+                    )
+                ),
             )
         else:
             last_logprobs = self.compute_temp_top_p_normalized_logprobs(