sglang 0.1.18__py3-none-any.whl → 0.1.19__py3-none-any.whl

sglang/srt/layers/radix_attention.py

@@ -2,10 +2,10 @@
 
 import numpy as np
 import torch
+from flashinfer.cascade import merge_state
 from torch import nn
 
 from sglang.global_config import global_config
-from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
 from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
@@ -13,18 +13,22 @@ from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetada
 
 class RadixAttention(nn.Module):
     def __init__(
-        self, num_heads: int, head_dim: int, scaling: float, num_kv_heads: int,
-        layer_id: int, logit_cap: int = -1
+        self,
+        num_heads: int,
+        head_dim: int,
+        scaling: float,
+        num_kv_heads: int,
+        layer_id: int,
+        logit_cap: int = -1,
     ):
         super().__init__()
         self.tp_q_head_num = num_heads
         self.tp_k_head_num = num_kv_heads
         self.tp_v_head_num = num_kv_heads
         self.head_dim = head_dim
+        self.scaling = scaling
         self.layer_id = layer_id
 
-        assert np.allclose(scaling, 1.0 / (head_dim**0.5))
-
         from sglang.srt.managers.controller.model_runner import global_server_args_dict
 
         if not global_server_args_dict.get("disable_flashinfer", False):
@@ -32,29 +36,17 @@ class RadixAttention(nn.Module):
             self.extend_forward = self.prefill_forward_flashinfer
             self.decode_forward = self.decode_forward_flashinfer
             # flashinfer now accepts float logit_cap argument
-            self.logit_cap = logit_cap if logit_cap > 0 else 0
+            self.logit_cap = logit_cap if logit_cap is not None and logit_cap > 0 else 0
         else:
             self.prefill_forward = self.prefill_forward_triton
             self.extend_forward = self.extend_forward_triton
             self.decode_forward = self.decode_forward_triton
-            self.logit_cap = logit_cap
+            self.logit_cap = logit_cap if logit_cap is not None else 0
 
     def prefill_forward_triton(self, q, k, v, input_metadata: InputMetadata):
-        o = torch.empty_like(q)
-
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num, self.head_dim),
-            k,
-            v,
-            o.view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.start_loc,
-            input_metadata.seq_lens,
-            input_metadata.max_seq_len,
-            self.logit_cap,
-        )
-        self.store_kv_cache(k, v, input_metadata)
-
-        return o
+        # In SGLang, we call both the typical "prefill" and "prefill with cache" as "extend".
+        # See the extend_forward_xxx functions.
+        raise NotImplementedError()
 
     def extend_forward_triton(self, q, k, v, input_metadata: InputMetadata):
         o = torch.empty_like(q)
@@ -75,7 +67,8 @@ class RadixAttention(nn.Module):
             input_metadata.extend_seq_lens,
             input_metadata.max_seq_len,
             input_metadata.max_extend_len,
-            self.logit_cap,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
@@ -96,18 +89,19 @@ class RadixAttention(nn.Module):
             input_metadata.max_seq_len,
             input_metadata.other_kv_index,
             input_metadata.total_num_tokens,
-            self.logit_cap,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
 
     def prefill_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        self.store_kv_cache(k, v, input_metadata)
-
         o1, s1 = input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
             v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
+            causal=True,
+            sm_scale=self.scaling,
             logits_soft_cap=self.logit_cap,
         )
 
@@ -118,12 +112,14 @@ class RadixAttention(nn.Module):
                 q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
                 input_metadata.token_to_kv_pool.kv_data[self.layer_id],
                 causal=False,
+                sm_scale=self.scaling,
                 logits_soft_cap=self.logit_cap,
             )
 
-            from flashinfer.cascade import merge_state
             o, _ = merge_state(o1, s1, o2, s2)
 
+        self.store_kv_cache(k, v, input_metadata)
+
         if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
             torch.cuda.synchronize()
 
@@ -135,6 +131,7 @@ class RadixAttention(nn.Module):
         o = input_metadata.flashinfer_decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            sm_scale=self.scaling,
             logits_soft_cap=self.logit_cap,
         )
 

sglang/srt/layers/token_attention.py

@@ -176,6 +176,7 @@ def _token_att_m_fwd(
     B_Start_Loc,
     B_Seqlen,
     max_len_in_batch,
+    sm_scale,
     logit_cap,
 ):
     BLOCK = 32
@@ -183,7 +184,6 @@ def _token_att_m_fwd(
     Lq, Lk = q.shape[-1], k_buffer.shape[-1]
     assert Lq == Lk
     assert Lk in {16, 32, 64, 128, 256}
-    sm_scale = 1.0 / (Lk**0.5)
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
@@ -317,6 +317,7 @@ def token_attention_fwd(
     max_len_in_batch,
     other_kv_index,
     total_num_tokens,
+    sm_scale=None,
     logit_cap=-1,
     att_m=None,
 ):
@@ -324,6 +325,7 @@ def token_attention_fwd(
         att_m = torch.empty(
             (q.shape[-2], total_num_tokens), dtype=REDUCE_TORCH_TYPE, device="cuda"
         )
+    sm_scale = 1.0 / (Lq**0.5) if sm_scale is None else sm_scale
 
     _token_att_m_fwd(
         q,
@@ -334,6 +336,7 @@ def token_attention_fwd(
         b_start_loc,
         b_seq_len,
         max_len_in_batch,
+        sm_scale,
         logit_cap,
     )
     _token_softmax_reducev_fwd(

sglang/srt/managers/controller/infer_batch.py

@@ -3,7 +3,7 @@
 import warnings
 from dataclasses import dataclass
 from enum import IntEnum, auto
-from typing import List
+from typing import List, Union
 
 import numpy as np
 import torch
@@ -31,7 +31,7 @@ class BaseFinishReason:
 
 
 class FINISH_MATCHED_TOKEN(BaseFinishReason):
-    def __init__(self, matched: int | List[int]):
+    def __init__(self, matched: Union[int, List[int]]):
         super().__init__()
         self.matched = matched
 

sglang/srt/managers/controller/manager_single.py

@@ -99,4 +99,4 @@ def start_controller_process(
     except Exception:
         logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
     finally:
-        kill_parent_process()
+        kill_parent_process()

sglang/srt/managers/controller/model_runner.py

@@ -127,7 +127,7 @@ class InputMetadata:
                 num_qo_heads,
                 num_kv_heads,
                 head_dim,
-                1
+                1,
             )
         else:
             self.flashinfer_decode_wrapper.end_forward()
@@ -140,7 +140,7 @@ class InputMetadata:
                 head_dim,
                 1,
                 pos_encoding_mode="NONE",
-                data_type=self.token_to_kv_pool.kv_data[0].dtype
+                data_type=self.token_to_kv_pool.kv_data[0].dtype,
             )
 
     def init_extend_args(self):
@@ -228,7 +228,7 @@ class InputMetadata:
             ret.init_flashinfer_args(
                 model_runner.model_config.num_attention_heads // tp_size,
                 model_runner.model_config.get_num_kv_heads(tp_size),
-                model_runner.model_config.head_dim
+                model_runner.model_config.head_dim,
             )
 
         return ret
@@ -259,7 +259,10 @@ class ModelRunner:
         logger.info(f"[gpu_id={self.gpu_id}] Set cuda device.")
         torch.cuda.set_device(self.gpu_id)
         logger.info(f"[gpu_id={self.gpu_id}] Init nccl begin.")
-        monkey_patch_vllm_p2p_access_check(self.gpu_id)
+
+        if not server_args.enable_p2p_check:
+            monkey_patch_vllm_p2p_access_check(self.gpu_id)
+
         if server_args.nccl_init_addr:
             nccl_init_method = f"tcp://{server_args.nccl_init_addr}"
         else:
@@ -269,7 +272,7 @@ class ModelRunner:
             world_size=self.tp_size,
             rank=self.tp_rank,
             local_rank=self.gpu_id,
-            distributed_init_method=nccl_init_method
+            distributed_init_method=nccl_init_method,
         )
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
         total_gpu_memory = get_available_gpu_memory(
@@ -323,7 +326,7 @@ class ModelRunner:
             device_config=device_config,
             load_config=load_config,
             lora_config=None,
-            vision_language_config=None,
+            multimodal_config=None,
             parallel_config=None,
             scheduler_config=None,
             cache_config=None,
@@ -341,7 +344,13 @@ class ModelRunner:
         )
         head_dim = self.model_config.head_dim
         head_num = self.model_config.get_num_kv_heads(self.tp_size)
-        cell_size = head_num * head_dim * self.model_config.num_hidden_layers * 2 * torch._utils._element_size(self.dtype)
+        cell_size = (
+            head_num
+            * head_dim
+            * self.model_config.num_hidden_layers
+            * 2
+            * torch._utils._element_size(self.dtype)
+        )
         rest_memory = available_gpu_memory - total_gpu_memory * (
             1 - self.mem_fraction_static
         )
@@ -384,33 +393,36 @@ class ModelRunner:
     def init_flash_infer(self):
         if not global_server_args_dict.get("disable_flashinfer", False):
             from flashinfer import (
-                BatchPrefillWithRaggedKVCacheWrapper,
-                BatchPrefillWithPagedKVCacheWrapper,
                 BatchDecodeWithPagedKVCacheWrapper,
+                BatchPrefillWithPagedKVCacheWrapper,
+                BatchPrefillWithRaggedKVCacheWrapper,
             )
             from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
 
             if not _grouped_size_compiled_for_decode_kernels(
                 self.model_config.num_attention_heads // self.tp_size,
-                self.model_config.get_num_kv_heads(self.tp_size)):
+                self.model_config.get_num_kv_heads(self.tp_size),
+            ):
                 use_tensor_cores = True
             else:
                 use_tensor_cores = False
 
             workspace_buffers = torch.empty(
-                3, 96 * 1024 * 1024, dtype=torch.uint8, device="cuda"
+                2, 96 * 1024 * 1024, dtype=torch.uint8, device="cuda"
             )
-            self.flashinfer_prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
-                workspace_buffers[0], "NHD"
+            self.flashinfer_prefill_wrapper_ragged = (
+                BatchPrefillWithRaggedKVCacheWrapper(workspace_buffers[0], "NHD")
             )
             self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
                 workspace_buffers[1], "NHD"
             )
             self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
-                workspace_buffers[2], "NHD", use_tensor_cores=use_tensor_cores
+                workspace_buffers[0], "NHD", use_tensor_cores=use_tensor_cores
             )
         else:
-            self.flashinfer_prefill_wrapper_ragged = self.flashinfer_prefill_wrapper_paged = None
+            self.flashinfer_prefill_wrapper_ragged = (
+                self.flashinfer_prefill_wrapper_paged
+            ) = None
             self.flashinfer_decode_wrapper = None
 
     @torch.inference_mode()

sglang/srt/managers/controller/tp_worker.py

@@ -34,11 +34,11 @@ from sglang.srt.managers.io_struct import (
 from sglang.srt.model_config import ModelConfig
 from sglang.srt.server_args import ModelPortArgs, ServerArgs
 from sglang.srt.utils import (
+    connect_rpyc_service,
     get_int_token_logit_bias,
     is_multimodal_model,
     set_random_seed,
     start_rpyc_service_process,
-    connect_rpyc_service,
     suppress_other_loggers,
 )
 from sglang.utils import get_exception_traceback
@@ -368,9 +368,11 @@ class ModelTpServer:
             if (
                 req.extend_input_len + req.max_new_tokens() + new_batch_total_tokens
                 < available_size
-                and (req.extend_input_len + new_batch_input_tokens
-                <= self.max_prefill_tokens
-                or len(can_run_list) == 0)
+                and (
+                    req.extend_input_len + new_batch_input_tokens
+                    <= self.max_prefill_tokens
+                    or len(can_run_list) == 0
+                )
             ):
                 delta = self.tree_cache.inc_lock_ref(req.last_node)
                 available_size += delta
@@ -452,7 +454,9 @@ class ModelTpServer:
                     next_token_ids,
                 ].tolist()
                 output.prefill_token_logprobs = output.prefill_token_logprobs.tolist()
-                output.normalized_prompt_logprobs = output.normalized_prompt_logprobs.tolist()
+                output.normalized_prompt_logprobs = (
+                    output.normalized_prompt_logprobs.tolist()
+                )
 
             next_token_ids = next_token_ids.tolist()
         else:
@@ -582,7 +586,9 @@ class ModelTpServer:
             req.check_finished()
 
             if req.return_logprob:
-                req.decode_token_logprobs.append((next_token_logprobs[i], next_token_id))
+                req.decode_token_logprobs.append(
+                    (next_token_logprobs[i], next_token_id)
+                )
                 if req.top_logprobs_num > 0:
                     req.decode_top_logprobs.append(output.decode_top_logprobs[i])
 
@@ -759,16 +765,27 @@ class ModelTpClient:
             with ThreadPoolExecutor(self.tp_size) as executor:
                 # Launch model processes
                 if server_args.nnodes == 1:
-                    self.procs = list(executor.map(
-                        lambda args: start_rpyc_service_process(*args),
-                        [(ModelTpService, p) for p in model_port_args.model_tp_ports],
-                    ))
+                    self.procs = list(
+                        executor.map(
+                            lambda args: start_rpyc_service_process(*args),
+                            [
+                                (ModelTpService, p)
+                                for p in model_port_args.model_tp_ports
+                            ],
+                        )
+                    )
                     addrs = [("localhost", p) for p in model_port_args.model_tp_ports]
                 else:
-                    addrs = [(ip, port) for ip, port in zip(model_port_args.model_tp_ips, model_port_args.model_tp_ports)]
-
-                self.model_services = list(executor.map(
-                    lambda args: connect_rpyc_service(*args), addrs))
+                    addrs = [
+                        (ip, port)
+                        for ip, port in zip(
+                            model_port_args.model_tp_ips, model_port_args.model_tp_ports
+                        )
+                    ]
+
+                self.model_services = list(
+                    executor.map(lambda args: connect_rpyc_service(*args), addrs)
+                )
 
                 # Init model
                 def init_model(i):

sglang/srt/managers/detokenizer_manager.py

@@ -11,7 +11,7 @@ from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.controller.infer_batch import FINISH_MATCHED_STR
 from sglang.srt.managers.io_struct import BatchStrOut, BatchTokenIDOut
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.utils import get_exception_traceback, graceful_registry
+from sglang.utils import find_printable_text, get_exception_traceback, graceful_registry
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -57,6 +57,8 @@ class DetokenizerManager:
             output_strs = []
             for i in range(len(recv_obj.rids)):
                 new_text = read_texts[i][len(surr_texts[i]) :]
+                if recv_obj.finished_reason[i] is None:
+                    new_text = find_printable_text(new_text)
                 output_strs.append(recv_obj.decoded_texts[i] + new_text)
 
                 if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
@@ -67,7 +69,7 @@ class DetokenizerManager:
             self.send_to_tokenizer.send_pyobj(
                 BatchStrOut(
                     rids=recv_obj.rids,
-                    output_str=output_strs,
+                    output_strs=output_strs,
                     meta_info=recv_obj.meta_info,
                     finished_reason=recv_obj.finished_reason,
                 )

sglang/srt/managers/io_struct.py

@@ -122,7 +122,7 @@ class BatchTokenIDOut:
 @dataclass
 class BatchStrOut:
     rids: List[str]
-    output_str: List[str]
+    output_strs: List[str]
     meta_info: List[Dict]
     finished_reason: List[BaseFinishReason]
 

sglang/srt/managers/tokenizer_manager.py

@@ -316,7 +316,7 @@ class TokenizerManager:
 
                 recv_obj.meta_info[i]["id"] = rid
                 out_dict = {
-                    "text": recv_obj.output_str[i],
+                    "text": recv_obj.output_strs[i],
                     "meta_info": recv_obj.meta_info[i],
                 }
                 state.out_list.append(out_dict)
@@ -333,17 +333,18 @@ class TokenizerManager:
             ret["meta_info"]["decode_token_logprobs"] = self.detokenize_logprob_tokens(
                 ret["meta_info"]["decode_token_logprobs"], return_text_in_logprobs
             )
-        if top_logprobs_num > 0:
-            ret["meta_info"][
-                "prefill_top_logprobs"
-            ] = self.detokenize_top_logprobs_tokens(
-                ret["meta_info"]["prefill_top_logprobs"], return_text_in_logprobs
-            )
-            ret["meta_info"][
-                "decode_top_logprobs"
-            ] = self.detokenize_top_logprobs_tokens(
-                ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
-            )
+
+            if top_logprobs_num > 0:
+                ret["meta_info"][
+                    "prefill_top_logprobs"
+                ] = self.detokenize_top_logprobs_tokens(
+                    ret["meta_info"]["prefill_top_logprobs"], return_text_in_logprobs
+                )
+                ret["meta_info"][
+                    "decode_top_logprobs"
+                ] = self.detokenize_top_logprobs_tokens(
+                    ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
+                )
         return ret
 
     def detokenize_logprob_tokens(self, token_logprobs, decode_to_text):
@@ -383,7 +384,7 @@ def get_pixel_values(
     try:
         processor = processor or global_processor
         image, image_size = load_image(image_data)
-        if image_size != None:
+        if image_size is not None:
             image_hash = hash(image_data)
             pixel_values = processor.image_processor(image)["pixel_values"]
             for _ in range(len(pixel_values)):

sglang/srt/model_config.py

@@ -115,6 +115,12 @@ def get_hf_text_config(config: PretrainedConfig):
     """Get the "sub" config relevant to llm for multi modal models.
     No op for pure text models.
     """
+    class_name = config.architectures[0]
+    if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"):
+        # We support non-hf version of llava models, so we do not want to
+        # read the wrong values from the unused default text_config.
+        return config
+
     if hasattr(config, "text_config"):
         # The code operates under the assumption that text_config should have
         # `num_attention_heads` (among others). Assert here to fail early