sglang 0.4.5.post3__py3-none-any.whl → 0.4.6__py3-none-any.whl

sglang/srt/disaggregation/prefill.py

@@ -176,17 +176,25 @@ class SchedulerDisaggregationPrefillMixin:
     """
 
     @torch.no_grad()
-    def event_loop_normal_disagg_prefill(self):
+    def event_loop_normal_disagg_prefill(self: Scheduler):
         """A normal scheduler loop for prefill worker in disaggregation mode."""
 
         while True:
             recv_reqs = self.recv_requests()
             self.process_input_requests(recv_reqs)
             self.waiting_queue.extend(
-                self.disagg_prefill_pending_queue.pop_bootstrapped()
+                self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
             )
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
+
+            # Handle DP attention
+            if (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            ):
+                batch, _ = self.prepare_dp_attn_batch(batch)
+
             self.cur_batch = batch
 
             if batch:
@@ -206,17 +214,25 @@ class SchedulerDisaggregationPrefillMixin:
             self.running_batch.batch_is_full = False
 
     @torch.no_grad()
-    def event_loop_overlap_disagg_prefill(self):
+    def event_loop_overlap_disagg_prefill(self: Scheduler):
         self.result_queue = deque()
 
         while True:
             recv_reqs = self.recv_requests()
             self.process_input_requests(recv_reqs)
             self.waiting_queue.extend(
-                self.disagg_prefill_pending_queue.pop_bootstrapped()
+                self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
             )
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
+
+            # Handle DP attention
+            if (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            ):
+                batch, _ = self.prepare_dp_attn_batch(batch)
+
             self.cur_batch = batch
 
             if batch:
@@ -310,7 +326,7 @@ class SchedulerDisaggregationPrefillMixin:
                 raise Exception("Transferring failed")
 
         for req in done_reqs:
-            self.disagg_prefill_pending_queue.req_to_metadata_buffer_idx_allocator.free(
+            self.disagg_prefill_bootstrap_queue.req_to_metadata_buffer_idx_allocator.free(
                 req.metadata_buffer_index
             )
 
@@ -326,9 +342,8 @@ class SchedulerDisaggregationPrefillMixin:
                 # only finished requests to running_batch.
                 self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
                 self.tree_cache.cache_unfinished_req(self.chunked_req)
-                if (
-                    self.enable_overlap
-                ):  # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
+                if self.enable_overlap:
+                    # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
                     self.chunked_req.tmp_end_idx = min(
                         len(self.chunked_req.fill_ids),
                         len(self.chunked_req.origin_input_ids),
@@ -374,7 +389,7 @@ class SchedulerDisaggregationPrefillMixin:
             .numpy()
         )
         if last_chunk is True:
-            self.disagg_prefill_pending_queue.store_prefill_results(
+            self.disagg_prefill_bootstrap_queue.store_prefill_results(
                 req.metadata_buffer_index, token_id
             )
         page_indices = kv_to_page_indices(kv_indices, page_size)

sglang/srt/entrypoints/http_server.py

@@ -84,6 +84,7 @@ from sglang.srt.utils import (
     add_api_key_middleware,
     add_prometheus_middleware,
     delete_directory,
+    get_bool_env_var,
     kill_process_tree,
     set_uvicorn_logging_configs,
 )
@@ -126,7 +127,10 @@ async def lifespan(fast_api_app: FastAPI):
 
 
 # Fast API
-app = FastAPI(lifespan=lifespan)
+app = FastAPI(
+    lifespan=lifespan,
+    openapi_url=None if get_bool_env_var("DISABLE_OPENAPI_DOC") else "/openapi.json",
+)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -277,7 +281,9 @@ async def generate_from_file_request(file: UploadFile, request: Request):
     )
 
     try:
-        ret = await _global_state.generate_request(obj, request).__anext__()
+        ret = await _global_state.tokenizer_manager.generate_request(
+            obj, request
+        ).__anext__()
         return ret
     except ValueError as e:
         logger.error(f"Error: {e}")

sglang/srt/function_call_parser.py

@@ -491,6 +491,7 @@ class DeepSeekV3Detector(BaseFormatDetector):
         self.eot_token = "<｜tool▁calls▁end｜>"
         self.func_call_regex = r"<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>"
         self.func_detail_regex = r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)\n```<｜tool▁call▁end｜>"
+        self._last_arguments = ""
 
     def has_tool_call(self, text: str) -> bool:
         """Check if the text contains a deepseek format tool call."""
@@ -528,13 +529,84 @@ class DeepSeekV3Detector(BaseFormatDetector):
 
     def structure_info(self) -> _GetInfoFunc:
         return lambda name: StructureInfo(
-            begin="<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>"
-            + name
-            + "\n```json\n",
-            end="\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜>",
-            trigger="<｜tool▁calls▁begin｜>",
+            begin=">" + name + "\n```json\n",
+            end="\n```<",
+            trigger=">" + name + "\n```json\n",
         )
 
+    def parse_streaming_increment(
+        self, new_text: str, tools: List[Tool]
+    ) -> StreamingParseResult:
+        """
+        Streaming incremental parsing tool calls for DeepSeekV3 format.
+        """
+        self._buffer += new_text
+        current_text = self._buffer
+
+        if self.bot_token not in current_text:
+            self._buffer = ""
+            for e_token in [self.eot_token, "```", "<｜tool▁call▁end｜>"]:
+                if e_token in new_text:
+                    new_text = new_text.replace(e_token, "")
+            return StreamingParseResult(normal_text=new_text)
+
+        if not hasattr(self, "_tool_indices"):
+            self._tool_indices = {
+                tool.function.name: i
+                for i, tool in enumerate(tools)
+                if tool.function and tool.function.name
+            }
+
+        calls: list[ToolCallItem] = []
+        try:
+            partial_match = re.search(
+                pattern=r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)",
+                string=current_text,
+                flags=re.DOTALL,
+            )
+            if partial_match:
+                func_name = partial_match.group(2).strip()
+                func_args_raw = partial_match.group(3).strip()
+
+                if not self.current_tool_name_sent:
+                    calls.append(
+                        ToolCallItem(
+                            tool_index=self._tool_indices.get(func_name, 0),
+                            name=func_name,
+                            parameters="",
+                        )
+                    )
+                    self.current_tool_name_sent = True
+                else:
+                    argument_diff = (
+                        func_args_raw[len(self._last_arguments) :]
+                        if func_args_raw.startswith(self._last_arguments)
+                        else func_args_raw
+                    )
+
+                    if argument_diff:
+                        calls.append(
+                            ToolCallItem(
+                                tool_index=self._tool_indices.get(func_name, 0),
+                                name=None,
+                                parameters=argument_diff,
+                            )
+                        )
+                        self._last_arguments += argument_diff
+
+                    if _is_complete_json(func_args_raw):
+                        result = StreamingParseResult(normal_text="", calls=calls)
+                        self._buffer = ""
+                        self._last_arguments = ""
+                        self.current_tool_name_sent = False
+                        return result
+
+            return StreamingParseResult(normal_text="", calls=calls)
+
+        except Exception as e:
+            logger.error(f"Error in parse_streaming_increment: {e}")
+            return StreamingParseResult(normal_text=current_text)
+
 
 class MultiFormatParser:
     def __init__(self, detectors: List[BaseFormatDetector]):

sglang/srt/layers/attention/base_attn_backend.py

@@ -62,6 +62,7 @@ class AttentionBackend(ABC):
         layer: RadixAttention,
         forward_batch: ForwardBatch,
         save_kv_cache: bool = True,
+        **kwargs,
     ):
         """Run forward on an attention layer."""
         if forward_batch.forward_mode.is_decode():
@@ -72,6 +73,7 @@ class AttentionBackend(ABC):
                 layer,
                 forward_batch,
                 save_kv_cache=save_kv_cache,
+                **kwargs,
             )
         else:
             return self.forward_extend(
@@ -81,6 +83,7 @@ class AttentionBackend(ABC):
                 layer,
                 forward_batch,
                 save_kv_cache=save_kv_cache,
+                **kwargs,
             )
 
     def forward_decode(

sglang/srt/layers/attention/flashattention_backend.py

@@ -623,6 +623,9 @@ class FlashAttentionBackend(AttentionBackend):
         layer: RadixAttention,
         forward_batch: ForwardBatch,
         save_kv_cache=True,
+        # For multi-head latent attention
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
     ):
         if k is not None:
             assert v is not None
@@ -637,11 +640,11 @@ class FlashAttentionBackend(AttentionBackend):
                         layer, cache_loc, k, v, layer.k_scale, layer.v_scale
                     )
                 else:
-                    forward_batch.token_to_kv_pool.set_kv_buffer(
+                    forward_batch.token_to_kv_pool.set_mla_kv_buffer(
                         layer,
                         cache_loc,
                         k,
-                        v,
+                        k_rope,
                     )
 
         # Use precomputed metadata across all layers
@@ -815,9 +818,15 @@ class FlashAttentionBackend(AttentionBackend):
                 c_kv_cache = c_kv.view(
                     -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
                 )
-                q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
-                q_nope = q_all[:, :, : layer.v_head_dim]
-                q_rope = q_all[:, :, layer.v_head_dim :]
+                if q_rope is not None:
+                    q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                    q_rope = q_rope.view(
+                        -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                    )
+                else:
+                    q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                    q_nope = q_all[:, :, : layer.v_head_dim]
+                    q_rope = q_all[:, :, layer.v_head_dim :]
 
                 result = flash_attn_with_kvcache(
                     q=q_rope,
@@ -877,6 +886,9 @@ class FlashAttentionBackend(AttentionBackend):
         layer: RadixAttention,
         forward_batch: ForwardBatch,
         save_kv_cache=True,
+        # For multi-head latent attention
+        q_rope: Optional[torch.Tensor] = None,
+        k_rope: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         if k is not None:
             assert v is not None
@@ -891,11 +903,11 @@ class FlashAttentionBackend(AttentionBackend):
                         layer, cache_loc, k, v, layer.k_scale, layer.v_scale
                     )
                 else:
-                    forward_batch.token_to_kv_pool.set_kv_buffer(
+                    forward_batch.token_to_kv_pool.set_mla_kv_buffer(
                         layer,
                         cache_loc,
                         k,
-                        v,
+                        k_rope,
                     )
 
         # Use precomputed metadata across all layers
@@ -1047,9 +1059,15 @@ class FlashAttentionBackend(AttentionBackend):
                 -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
             )
 
-            q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
-            q_nope = q_all[:, :, : layer.v_head_dim]
-            q_rope = q_all[:, :, layer.v_head_dim :]
+            if q_rope is not None:
+                q_nope = q.view(-1, layer.tp_q_head_num, layer.v_head_dim)
+                q_rope = q_rope.view(
+                    -1, layer.tp_q_head_num, layer.head_dim - layer.v_head_dim
+                )
+            else:
+                q_all = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
+                q_nope = q_all[:, :, : layer.v_head_dim]
+                q_rope = q_all[:, :, layer.v_head_dim :]
             max_seqlen_q = metadata.max_seq_len_q
 
             result = flash_attn_with_kvcache(

sglang/srt/layers/attention/flashmla_backend.py

@@ -68,9 +68,6 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
         self.num_q_heads = (
             model_runner.model_config.num_attention_heads // get_attention_tp_size()
         )
-        self.num_kv_heads = model_runner.model_config.get_num_kv_heads(
-            get_attention_tp_size()
-        )
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.num_local_heads = (
             model_runner.model_config.num_attention_heads // get_attention_tp_size()
@@ -111,8 +108,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
                 )
                 mla_metadata, num_splits = get_mla_metadata(
                     forward_batch.seq_lens.to(torch.int32),
-                    Q_LEN * self.num_q_heads // self.num_kv_heads,
-                    self.num_kv_heads,
+                    Q_LEN * self.num_q_heads,
+                    1,
                 )
                 self.forward_metadata = FlashMLADecodeMetadata(
                     mla_metadata,
@@ -141,8 +138,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
 
         self.cuda_graph_mla_metadata, self.cuda_graph_num_splits = get_mla_metadata(
             torch.ones(max_bs, dtype=torch.int32, device=cuda_graph_kv_indices.device),
-            Q_LEN * self.num_q_heads // self.num_kv_heads,
-            self.num_kv_heads,
+            Q_LEN * self.num_q_heads,
+            1,
         )
         self.cuda_graph_kv_indices = cuda_graph_kv_indices
 
@@ -171,8 +168,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
                 )
                 mla_metadata, num_splits = get_mla_metadata(
                     seq_lens.to(torch.int32),
-                    Q_LEN * self.num_q_heads // self.num_kv_heads,
-                    self.num_kv_heads,
+                    Q_LEN * self.num_q_heads,
+                    1,
                 )
                 self.cuda_graph_mla_metadata.copy_(mla_metadata)
                 self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)
@@ -221,8 +218,8 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
             )
             mla_metadata, num_splits = get_mla_metadata(
                 seq_lens.to(torch.int32),
-                Q_LEN * self.num_q_heads // self.num_kv_heads,
-                self.num_kv_heads,
+                Q_LEN * self.num_q_heads,
+                1,
             )
             self.cuda_graph_mla_metadata.copy_(mla_metadata)
             self.cuda_graph_num_splits[: bs + 1].copy_(num_splits)

sglang/srt/layers/attention/vision.py

@@ -271,6 +271,8 @@ class VisionSdpaAttention(nn.Module):
         Returns:
              [b * s, h, head_size]
         """
+        if self.flatten_batch:
+            assert bsz == 1, "flatten_batch is True, bsz must be 1"
 
         s = q.shape[0] // bsz
 

sglang/srt/layers/layernorm.py

@@ -22,8 +22,6 @@ import torch.nn as nn
 from sglang.srt.custom_op import CustomOp
 from sglang.srt.utils import is_cuda, is_hip
 
-logger = logging.getLogger(__name__)
-
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 
@@ -36,19 +34,9 @@ if _is_cuda:
     )
 
 if _is_hip:
+    from vllm._custom_ops import fused_add_rms_norm, rms_norm
 
-    from aiter.ops.rmsnorm import rms_norm, rmsnorm2d_fwd_with_add
-
-    rmsnorm = rms_norm
-
-    def fused_add_rmsnorm(
-        x: torch.Tensor,
-        residual: torch.Tensor,
-        w: torch.Tensor,
-        eps: float,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        rmsnorm2d_fwd_with_add(x, x, residual, residual, w, eps)
-        return x, residual
+logger = logging.getLogger(__name__)
 
 
 class RMSNorm(CustomOp):
@@ -61,23 +49,49 @@ class RMSNorm(CustomOp):
         self.weight = nn.Parameter(torch.ones(hidden_size))
         self.variance_epsilon = eps
 
+    def forward(self, *args, **kwargs):
+        if torch.compiler.is_compiling():
+            return self.forward_native(*args, **kwargs)
+        if _is_cuda:
+            return self.forward_cuda(*args, **kwargs)
+        elif _is_hip:
+            return self.forward_hip(*args, **kwargs)
+        else:
+            return self.forward_native(*args, **kwargs)
+
     def forward_cuda(
         self,
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-
         if residual is not None:
             fused_add_rmsnorm(x, residual, self.weight.data, self.variance_epsilon)
             return x, residual
         out = rmsnorm(x, self.weight.data, self.variance_epsilon)
         return out
 
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if not x.is_contiguous():
+            # NOTE: Romove this if aiter kernel supports discontinuous input
+            x = x.contiguous()
+        if residual is not None:
+            fused_add_rms_norm(x, residual, self.weight.data, self.variance_epsilon)
+            return x, residual
+        out = torch.empty_like(x)
+        rms_norm(out, x, self.weight.data, self.variance_epsilon)
+        return out
+
     def forward_native(
         self,
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if not x.is_contiguous():
+            x = x.contiguous()
         orig_dtype = x.dtype
         x = x.to(torch.float32)
         if residual is not None:
@@ -103,6 +117,14 @@ class GemmaRMSNorm(CustomOp):
         self.weight = nn.Parameter(torch.zeros(hidden_size))
         self.variance_epsilon = eps
 
+    def forward(self, *args, **kwargs):
+        if torch.compiler.is_compiling():
+            return self.forward_native(*args, **kwargs)
+        if _is_cuda:
+            return self.forward_cuda(*args, **kwargs)
+        else:
+            return self.forward_native(*args, **kwargs)
+
     def forward_native(
         self,
         x: torch.Tensor,
@@ -156,6 +178,6 @@ class Gemma3RMSNorm(nn.Module):
 
 if not (_is_cuda or _is_hip):
     logger.info(
-        "sgl-kernel is not available on Non-NV platforms. Fallback to other kernel libraries."
+        "sgl-kernel layernorm implementation is not available on current platform. Fallback to other kernel libraries."
     )
     from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm

sglang/srt/layers/logits_processor.py

@@ -335,13 +335,13 @@ class LogitsProcessor(nn.Module):
                     aux_pruned_states = torch.cat(aux_pruned_states, dim=-1)
                     hidden_states_to_store = (
                         aux_pruned_states[sample_indices]
-                        if sample_indices
+                        if sample_indices is not None
                         else aux_pruned_states
                     )
                 else:
                     hidden_states_to_store = (
                         pruned_states[sample_indices]
-                        if sample_indices
+                        if sample_indices is not None
                         else pruned_states
                     )
             else:

sglang/srt/layers/moe/fused_moe_native.py

@@ -8,6 +8,7 @@ from typing import Callable, Optional
 import torch
 from torch.nn import functional as F
 
+from sglang.srt.layers.activation import GeluAndMul, SiluAndMul
 from sglang.srt.layers.moe.topk import select_experts
 
 
@@ -30,7 +31,7 @@ def fused_moe_forward_native(
 ) -> torch.Tensor:
 
     if apply_router_weight_on_input:
-        raise NotImplementedError
+        raise NotImplementedError()
 
     topk_weights, topk_ids = select_experts(
         hidden_states=x,
@@ -75,9 +76,6 @@ def moe_forward_native(
     activation: str = "silu",
     routed_scaling_factor: Optional[float] = None,
 ) -> torch.Tensor:
-
-    from sglang.srt.layers.activation import GeluAndMul, SiluAndMul
-
     topk_weights, topk_ids = select_experts(
         hidden_states=x,
         router_logits=router_logits,