sglang 0.4.2.post3__py3-none-any.whl → 0.4.2.post4__py3-none-any.whl

sglang/srt/lora/triton_ops/gate_up_lora_b.py

@@ -0,0 +1,170 @@
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.lora.utils import LoRABatchInfo
+
+
+@triton.jit
+def _gate_up_lora_b_kernel(
+    # Pointers to matrices
+    x,
+    weights,
+    output,
+    # Parameters of size
+    K,  # K = R
+    output_dim,
+    # Strides
+    x_stride_0,
+    x_stride_1,
+    w_stride_0,
+    w_stride_1,
+    w_stride_2,
+    output_stride_0,
+    output_stride_1,
+    # Information on sequence lengths and weight id
+    seg_lens,
+    seg_indptr,
+    weight_indices,
+    # Meta parameters
+    BLOCK_S: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    # For fused output scaling and adding
+    fuse_scaling_add,
+    scaling,
+):
+    # This kernel packs 2 sgemms (gate/up) into a single kernel.
+
+    # x: (s, 2 * K), s is the sum of sequence lengths, K equals to lora rank
+    # weights: (num_lora, 2 * output_dim, K)
+    # output: (s, 2 * output_dim)
+    # output_dim >> K
+
+    # Current block computes sequence with batch_id,
+    # which starts from row seg_start of x with length seg_len.
+    # gate_up_id decides which of gate or up (0: gate, 1: up)
+    batch_id = tl.program_id(axis=2)
+    gate_up_id = tl.program_id(axis=1)
+    pid = tl.program_id(axis=0)
+    seg_len = tl.load(seg_lens + batch_id)
+    w_index = tl.load(weight_indices + batch_id)
+    seg_start = tl.load(seg_indptr + batch_id)
+    n_start = gate_up_id * output_dim  # offset on output dim
+
+    # The tile in output matrix will have (pid_s, pid_n) as id
+    num_pid_n = tl.cdiv(output_dim, BLOCK_N)
+    pid_s = pid // num_pid_n
+    pid_n = pid % num_pid_n
+
+    # Create pointers for the first block of x and weights
+    # The pointers will be advanced as we move in the K direction
+    # and accumulate
+    s_offset = tl.arange(0, BLOCK_S) + pid_s * BLOCK_S
+    n_offset = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
+    k_offset = tl.arange(0, BLOCK_K)
+
+    x_ptrs = (x + seg_start * x_stride_0 + (gate_up_id * K) * x_stride_1) + (
+        s_offset[:, None] * x_stride_0 + k_offset[None, :] * x_stride_1
+    )
+    w_ptrs = (weights + w_index * w_stride_0 + n_start * w_stride_1) + (
+        k_offset[:, None] * w_stride_2 + n_offset[None, :] * w_stride_1
+    )
+
+    # Iteate to compute the block in output matrix
+    partial_sum = tl.zeros((BLOCK_S, BLOCK_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_K)):
+        x_tile = tl.load(
+            x_ptrs,
+            mask=(s_offset[:, None] < seg_len)
+            and (k_offset[None, :] < K - k * BLOCK_K),
+            other=0.0,
+        )
+        w_tile = tl.load(
+            w_ptrs,
+            mask=(k_offset[:, None] < K - k * BLOCK_K)
+            and (n_offset[None, :] < output_dim),
+            other=0.0,
+        )
+        partial_sum += tl.dot(x_tile, w_tile)
+
+        x_ptrs += BLOCK_K * x_stride_1
+        w_ptrs += BLOCK_K * w_stride_2
+
+    # Store result to output matrix
+    partial_sum *= scaling
+    partial_sum = partial_sum.to(x.dtype.element_ty)
+    output_ptr = (output + seg_start * output_stride_0 + n_start * output_stride_1) + (
+        s_offset[:, None] * output_stride_0 + n_offset[None, :] * output_stride_1
+    )
+    output_mask = (s_offset[:, None] < seg_len) and (n_offset[None, :] < output_dim)
+    if fuse_scaling_add:
+        partial_sum += tl.load(output_ptr, mask=output_mask)
+    tl.store(output_ptr, partial_sum, mask=output_mask)
+
+
+def gate_up_lora_b_fwd(
+    x: torch.Tensor,
+    gate_up_lora_b: torch.Tensor,
+    batch_info: LoRABatchInfo,
+    output_dim: int,
+    base_output: torch.Tensor = None,
+    scaling: float = 1.0,
+) -> torch.Tensor:
+
+    # x: (s, 2 * r)
+    # gate_up_lora_b: (num_lora, 2 * output_dim, r)
+    # output: (s, 2 * output_dim)
+
+    # Compute lora_output with shape (s, output_dim) as follows:
+    # lora_output[:, :output_dim] = sgemm(x[:, :r], gate_up_lora_b[:, :output_dim, :])
+    # lora_output[:, output_dim:]
+    #      = sgemm(x[:, r:], gate_up_lora_b[:, output_dim:, :])
+
+    # Get dims
+    s = x.shape[0]
+    input_dim = x.shape[1]
+    r = gate_up_lora_b.shape[-1]
+    assert input_dim == 2 * r
+
+    BLOCK_S = 16
+    BLOCK_R = 16
+    BLOCK_OUT = 64
+
+    grid_b = (
+        triton.cdiv(batch_info.max_len, BLOCK_S) * triton.cdiv(output_dim, BLOCK_OUT),
+        2,  # this dimension decides current block computes on gate or up proj
+        batch_info.bs,
+    )
+
+    if base_output is None:
+        output = torch.empty((s, 2 * output_dim), device=x.device, dtype=x.dtype)
+        fuse_scaling_add = False
+    else:
+        output = base_output
+        fuse_scaling_add = True
+
+    _gate_up_lora_b_kernel[grid_b](
+        x,
+        gate_up_lora_b,
+        output,
+        r,
+        output_dim,
+        x.stride(0),
+        x.stride(1),
+        gate_up_lora_b.stride(0),
+        gate_up_lora_b.stride(1),
+        gate_up_lora_b.stride(2),
+        output.stride(0),
+        output.stride(1),
+        batch_info.seg_lens,
+        batch_info.seg_indptr,
+        batch_info.weight_indices,
+        BLOCK_S,
+        BLOCK_OUT,
+        BLOCK_R,
+        fuse_scaling_add,
+        scaling,
+    )
+
+    return output

sglang/srt/lora/triton_ops/qkv_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 @triton.jit
@@ -108,7 +108,7 @@ def _qkv_lora_b_kernel(
 def qkv_lora_b_fwd(
     x: torch.Tensor,
     qkv_lora_b: torch.Tensor,
-    batch_info: LoraBatchInfo,
+    batch_info: LoRABatchInfo,
     output_offset: torch.Tensor,
     max_qkv_out_dim: int,
     base_output: torch.Tensor = None,
@@ -123,11 +123,11 @@ def qkv_lora_b_fwd(
     # output: (s, output_dim_q + 2 * output_dim_kv)
 
     # Compute lora_output with shape (s, output_dim) as follows:
-    # lora_output[:, :output_dim_q] = sgemm(lora_output_a[:, :r], )
+    # lora_output[:, :output_dim_q] = sgemm(x[:, :r], qkv_lora_b[:, :outptu_dim_q, :])
     # lora_output[:, output_dim_q: output_dim_q + output_dim_kv]
-    #      = sgemm(lora_output_a[:, r: 2 * r], kv_lora_b[0])
+    #      = sgemm(x[:, r: 2 * r], qkv_lora_b[:, outptu_dim_q: output_dim_q + output_dim_kv, :])
     # lora_output[:, output_dim_q + output_dim_kv: ]
-    #      = sgemm(lora_output_a[:, 2 * r: 3 * r], kv_lora_b[1])
+    #      = sgemm(x[:, 2 * r: , qkv_lora_b[:, output_dim_q + output_dim_kv: , :])
 
     # Get dims
     s = x.shape[0]

sglang/srt/lora/triton_ops/sgemm_lora_a.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 @triton.jit
@@ -91,7 +91,7 @@ def _sgemm_lora_a_kernel(
 
 
 def sgemm_lora_a_fwd(
-    x: torch.Tensor, weights: torch.Tensor, batch_info: LoraBatchInfo
+    x: torch.Tensor, weights: torch.Tensor, batch_info: LoRABatchInfo
 ) -> torch.Tensor:
     # x: (s, input_dim)
     # weights: (num_lora, r, input_dim)

sglang/srt/lora/triton_ops/sgemm_lora_b.py

@@ -2,7 +2,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.lora.lora import LoraBatchInfo
+from sglang.srt.lora.utils import LoRABatchInfo
 
 
 @triton.jit
@@ -98,7 +98,7 @@ def _sgemm_lora_b_kernel(
 def sgemm_lora_b_fwd(
     x: torch.Tensor,
     weights: torch.Tensor,
-    batch_info: LoraBatchInfo,
+    batch_info: LoRABatchInfo,
     base_output: torch.Tensor = None,
     scaling: float = 1.0,
 ) -> torch.Tensor:

sglang/srt/lora/utils.py

@@ -0,0 +1,141 @@
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import Optional, Set, Tuple
+
+import torch
+
+from sglang.srt.hf_transformers_utils import AutoConfig
+
+
+@dataclass
+class LoRABatchInfo:
+    # Batch size
+    bs: int
+
+    # Lengths of each sequence in shape (bs,)
+    seg_lens: torch.Tensor
+
+    # Indice pointers of each sequence in shape (bs + 1, )
+    seg_indptr: torch.Tensor
+
+    # Maximum sequence length of current batch
+    max_len: int
+
+    # The index of lora adapter used by each sequence, in shape (bs,)
+    weight_indices: torch.Tensor
+
+
+class LoRAType(Enum):
+    LORA_A = 0
+    LORA_B = 1
+
+
+def get_layer_id(name: str) -> int:
+    """
+    Extract integer id of layer from its name in string.
+    """
+    match = re.search(r"layers\.(\d+)\.", name)
+    if match is None:
+        return None
+    return int(match.group(1))
+
+
+def get_customized_names_from_hf_names(
+    hf_module_names: Set[str], base_model: torch.nn.Module
+) -> Set[str]:
+    """
+    This function takes in a set of huggingface style module names:
+         e.g., {"k_proj", "q_proj", "v_proj", "o_proj"}
+    and outputs a set of module names of customized sglang layers:
+         e.g., {"qkv_proj", "o_proj"}
+    """
+    if hasattr(base_model, "get_module_name"):
+        return {base_model.get_module_name(name) for name in hf_module_names}
+    else:
+        """
+        Fallback solution of mapping from config module name to module name in model class.
+        Please check if it aligns with your base model.
+        Please implement the function in the model class if it is not.
+        You can reference this function in llama.py.
+        """
+        params_mapping = {
+            "q_proj": "qkv_proj",
+            "k_proj": "qkv_proj",
+            "v_proj": "qkv_proj",
+            "gate_proj": "gate_up_proj",
+            "up_proj": "gate_up_proj",
+        }
+        return {params_mapping.get(name, name) for name in hf_module_names}
+
+
+def get_hidden_dim(
+    module_name: str, config: AutoConfig, base_model: torch.nn.Module
+) -> Tuple[int]:
+    """
+    Given a module_name (might be a stacked name), return the hidden dims of modules's input and output.
+    """
+
+    if hasattr(base_model, "get_hidden_dim"):
+        return base_model.get_hidden_dim(module_name)
+    else:
+        """
+        WARNING: get_hidden_dim() is not defined,
+        which is used to get the hidden dim for different lora modules
+        Use the default one, but please check if it is correct for your model.
+        Please implement the function in the model class if it is not.
+        You can reference this function in llama.py.
+        """
+        if module_name in ["q_proj", "o_proj", "qkv_proj"]:
+            return config.hidden_size, config.hidden_size
+        elif module_name in ["kv_proj"]:
+            return config.hidden_size, config.hidden_size // (
+                config.num_attention_heads // config.num_key_value_heads
+            )
+        elif module_name == "gate_up_proj":
+            return config.hidden_size, config.intermediate_size
+        elif module_name == "down_proj":
+            return config.intermediate_size, config.hidden_size
+        else:
+            raise NotImplementedError()
+
+
+def get_stacked_name(name: str) -> Tuple[str]:
+    """
+    Mapping a target module name to (stacked name for Lora A, stacked name for Lora B)
+    """
+    params_mapping = {
+        "q_proj": ("qkv_proj", "q_proj"),
+        "k_proj": ("qkv_proj", "kv_proj"),
+        "v_proj": ("qkv_proj", "kv_proj"),
+        "gate_proj": ("gate_up_proj", "gate_up_proj"),
+        "up_proj": ("gate_up_proj", "gate_up_proj"),
+    }
+    return params_mapping.get(name, (name, name))
+
+
+def get_stacked_multiply(module_name: str) -> int:
+    """
+    Mapping a lora module name to its magnification at output dimension
+    """
+    stacked_rank = {
+        "qkv_proj": 3,
+        "kv_proj": 2,
+        "gate_up_proj": 2,
+    }
+    return stacked_rank[module_name] if module_name in stacked_rank else 1
+
+
+def get_weight_name(
+    target_name: str, lora_weight_names: Set[Tuple[str]], lora_type: LoRAType
+) -> Optional[str]:
+    """
+    target_name is name of a given module,
+    lora_weight_names is a set of lora stacked name pairs (see get_stacked_name method above)
+    If there is a weight name in lora_weight_names that can match target_name, return this name
+    Else return None
+    """
+    idx = 0 if lora_type == LoRAType.LORA_A else 1
+    for weight_name_pair in lora_weight_names:
+        if weight_name_pair[idx] in target_name:
+            return weight_name_pair[idx]

sglang/srt/model_executor/cuda_graph_runner.py

@@ -237,6 +237,7 @@ class CudaGraphRunner:
                 "1. disable cuda graph by --disable-cuda-graph\n"
                 "2. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
                 "3. disable torch compile by not using --enable-torch-compile\n"
+                "4. set --cuda-graph-max-bs to a smaller value (e.g., 32)\n"
                 "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
             )
 
@@ -462,8 +463,11 @@ class CudaGraphRunner:
                     ),
                     positions=None,
                     retrive_index=None,
+                    retrive_next_token=None,
+                    retrive_next_sibling=None,
                     retrive_cum_len=None,
                     draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
+                    spec_steps=self.model_runner.server_args.speculative_num_steps,
                     capture_hidden_mode=CaptureHiddenMode.FULL,
                 )
 

sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -85,6 +85,7 @@ class EAGLEDraftCudaGraphRunner:
                 "1. disable cuda graph by --disable-cuda-graph\n"
                 "2. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
                 "3. disable torch compile by not using --enable-torch-compile\n"
+                "4. specify --dtype to the same dtype (e.g. bfloat16)\n"
                 "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
             )
 

sglang/srt/speculative/eagle_utils.py

@@ -4,6 +4,7 @@ import dataclasses
 from typing import TYPE_CHECKING, List
 
 import torch
+import torch.nn.functional as F
 import triton
 import triton.language as tl
 
@@ -11,7 +12,14 @@ from sglang.srt.layers.attention.flashinfer_backend import (
     create_flashinfer_kv_indices_triton,
 )
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
-from sglang.srt.speculative.build_eagle_tree import build_tree_kernel
+from sglang.srt.speculative.build_eagle_tree import (
+    build_tree_kernel,
+    build_tree_kernel_efficient,
+)
+from sglang.srt.utils import is_cuda_available
+
+if is_cuda_available():
+    from sgl_kernel import tree_speculative_sampling_target_only
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -160,8 +168,11 @@ class EagleVerifyInput:
     custom_mask: torch.Tensor
     positions: torch.Tensor
     retrive_index: torch.Tensor
+    retrive_next_token: torch.Tensor
+    retrive_next_sibling: torch.Tensor
     retrive_cum_len: torch.Tensor
     draft_token_num: int
+    spec_steps: int
     capture_hidden_mode: CaptureHiddenMode
 
     @classmethod
@@ -175,10 +186,45 @@ class EagleVerifyInput:
         seq_lens_sum: int,
         topk: int,
         spec_steps: int,
-        num_verify_token: int,
+        num_verify_tokens: int,
+        is_all_greedy: bool,
     ):
-        tree_mask, position, retrive_index, retrive_cum_len, draft_tokens = (
-            build_tree_kernel(
+        if is_all_greedy:
+            tree_mask, position, retrive_index, retrive_cum_len, draft_tokens = (
+                build_tree_kernel(
+                    verified_id,
+                    score_list,  # b, n, topk; n= 1 + (num_steps-1) * self.topk
+                    token_list,
+                    parents_list,
+                    seq_lens,
+                    seq_lens_sum,
+                    topk,
+                    spec_steps,
+                    num_verify_tokens,
+                )
+            )
+
+            return cls(
+                draft_tokens,
+                tree_mask,
+                position,
+                retrive_index,
+                None,
+                None,
+                retrive_cum_len,
+                num_verify_tokens,
+                spec_steps,
+                CaptureHiddenMode.FULL,
+            )
+        else:
+            (
+                tree_mask,
+                position,
+                retrive_index,
+                retrive_next_token,
+                retrive_next_sibling,
+                draft_tokens,
+            ) = build_tree_kernel_efficient(
                 verified_id,
                 score_list,
                 token_list,
@@ -187,18 +233,21 @@ class EagleVerifyInput:
                 seq_lens_sum,
                 topk,
                 spec_steps,
-                num_verify_token,
+                num_verify_tokens,
+            )
+
+            return cls(
+                draft_tokens,
+                tree_mask,
+                position,
+                retrive_index,
+                retrive_next_token,
+                retrive_next_sibling,
+                None,
+                num_verify_tokens,
+                spec_steps,
+                CaptureHiddenMode.FULL,
             )
-        )
-        return cls(
-            draft_tokens,
-            tree_mask,
-            position,
-            retrive_index,
-            retrive_cum_len,
-            num_verify_token,
-            CaptureHiddenMode.FULL,
-        )
 
     def prepare_for_verify(self, batch: ScheduleBatch):
         batch.input_ids = self.draft_token
@@ -313,12 +362,6 @@ class EagleVerifyInput:
                 uniform_samples=coins,
                 target_probs=target_probs,
                 draft_probs=draft_probs,
-                threshold_single=global_server_args_dict[
-                    "speculative_accept_threshold_single"
-                ],
-                threshold_acc=global_server_args_dict[
-                    "speculative_accept_threshold_acc"
-                ],
                 deterministic=True,
             )
 

sglang/srt/speculative/eagle_worker.py

@@ -185,6 +185,7 @@ class EAGLEWorker(TpModelWorker):
             self.topk,
             self.speculative_num_steps,
             self.server_args.speculative_num_draft_tokens,
+            batch.sampling_info.is_all_greedy,
         )
 
         # Free cache locations

sglang/version.py

@@ -1 +1 @@
-__version__ = "0.4.2.post3"
+__version__ = "0.4.2.post4"

{sglang-0.4.2.post3.dist-info → sglang-0.4.2.post4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: sglang
-Version: 0.4.2.post3
+Version: 0.4.2.post4
 Summary: SGLang is yet another fast serving framework for large language models and vision language models.
 License:                                  Apache License
                                    Version 2.0, January 2004
@@ -239,11 +239,11 @@ Requires-Dist: xgrammar>=0.1.10; extra == "runtime-common"
 Provides-Extra: srt
 Requires-Dist: sglang[runtime_common]; extra == "srt"
 Requires-Dist: cuda-python; extra == "srt"
-Requires-Dist: sgl-kernel>=0.0.3.post2; extra == "srt"
+Requires-Dist: sgl-kernel>=0.0.3.post3; extra == "srt"
 Requires-Dist: torch; extra == "srt"
-Requires-Dist: vllm==0.6.4.post1; extra == "srt"
+Requires-Dist: vllm<=0.7.2,>=0.6.4.post1; extra == "srt"
 Requires-Dist: flashinfer_python>=0.2.0.post2; extra == "srt"
-Requires-Dist: outlines<0.1.0,>=0.0.44; extra == "srt"
+Requires-Dist: outlines<=0.1.11,>=0.0.44; extra == "srt"
 Provides-Extra: srt-hip
 Requires-Dist: sglang[runtime_common]; extra == "srt-hip"
 Requires-Dist: torch; extra == "srt-hip"