sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/srt/speculative/{ngram_utils.py → ngram_info.py}

@@ -7,6 +7,8 @@ from typing import Optional, Tuple
 import torch
 import triton
 
+from sglang.srt.server_args import get_global_server_args
+
 logger = logging.getLogger(__name__)
 
 from dataclasses import dataclass
@@ -16,10 +18,11 @@ import torch.nn.functional as F
 from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.sampler import apply_custom_logit_processor
-from sglang.srt.managers.schedule_batch import (
-    ScheduleBatch,
+from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.mem_cache.common import (
+    alloc_paged_token_slots_extend,
+    alloc_token_slots,
     get_last_loc,
-    global_server_args_dict,
 )
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
@@ -74,7 +77,10 @@ class NgramVerifyInput(SpecInput):
         batch.input_ids = self.draft_token
 
         if page_size == 1:
-            batch.out_cache_loc = batch.alloc_token_slots(len(batch.input_ids))
+            batch.out_cache_loc = alloc_token_slots(
+                batch.tree_cache,
+                len(batch.input_ids),
+            )
             end_offset = batch.seq_lens + self.draft_token_num
         else:
             # TODO(lsyin): add prefix lens cpu here to support page size > 1
@@ -87,7 +93,8 @@ class NgramVerifyInput(SpecInput):
                 batch.req_pool_indices,
                 prefix_lens,
             )
-            batch.out_cache_loc = batch.alloc_paged_token_slots_extend(
+            batch.out_cache_loc = alloc_paged_token_slots_extend(
+                batch.tree_cache,
                 prefix_lens,
                 prefix_lens_cpu,
                 end_offset,
@@ -345,10 +352,8 @@ class NgramVerifyInput(SpecInput):
             uniform_samples_for_final_sampling=coins_for_final_sampling,
             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"],
+            threshold_single=get_global_server_args().speculative_accept_threshold_single,
+            threshold_acc=get_global_server_args().speculative_accept_threshold_acc,
             deterministic=True,
         )
 

sglang/srt/speculative/ngram_worker.py

@@ -6,11 +6,12 @@ import torch
 from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
 
 from sglang.srt.managers.schedule_batch import ScheduleBatch
+from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
-from sglang.srt.model_executor.forward_batch_info import ForwardBatchOutput, ForwardMode
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.cpp_ngram.ngram_cache import NgramCache
-from sglang.srt.speculative.ngram_utils import NgramVerifyInput
+from sglang.srt.speculative.ngram_info import NgramVerifyInput
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 
 logger = logging.getLogger(__name__)
@@ -207,18 +208,18 @@ class NGRAMWorker:
             batch_tokens.append(put_ids)
         self.ngram_cache.batch_put(batch_tokens)
 
-    def forward_batch_generation(self, batch: ScheduleBatch) -> ForwardBatchOutput:
+    def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult:
         self._prepare_for_speculative_decoding(batch)
         model_worker_batch = batch.get_model_worker_batch()
         num_accepted_tokens = 0
 
         if model_worker_batch.forward_mode.is_target_verify():
-            forward_batch_output = self.target_worker.forward_batch_generation(
+            batch_result = self.target_worker.forward_batch_generation(
                 model_worker_batch, is_verify=True
             )
             logits_output, can_run_cuda_graph = (
-                forward_batch_output.logits_output,
-                forward_batch_output.can_run_cuda_graph,
+                batch_result.logits_output,
+                batch_result.can_run_cuda_graph,
             )
             verify_input = model_worker_batch.spec_info
             logits_output, next_token_ids, num_accepted_tokens = verify_input.verify(
@@ -228,16 +229,16 @@ class NGRAMWorker:
             batch.forward_mode = ForwardMode.DECODE
 
         else:
-            forward_batch_output = self.target_worker.forward_batch_generation(
+            batch_result = self.target_worker.forward_batch_generation(
                 model_worker_batch
             )
             logits_output, next_token_ids, can_run_cuda_graph = (
-                forward_batch_output.logits_output,
-                forward_batch_output.next_token_ids,
-                forward_batch_output.can_run_cuda_graph,
+                batch_result.logits_output,
+                batch_result.next_token_ids,
+                batch_result.can_run_cuda_graph,
             )
 
-        return ForwardBatchOutput(
+        return GenerationBatchResult(
             logits_output=logits_output,
             next_token_ids=next_token_ids,
             num_accepted_tokens=num_accepted_tokens,

sglang/srt/speculative/spec_info.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from enum import IntEnum, auto
+from functools import lru_cache
 from typing import List, Tuple
 
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch
@@ -27,6 +28,7 @@ class SpeculativeAlgorithm(IntEnum):
     def is_ngram(self):
         return self == SpeculativeAlgorithm.NGRAM
 
+    @lru_cache(maxsize=None)
     @staticmethod
     def from_string(name: str):
         name_map = {

sglang/srt/speculative/spec_utils.py

@@ -3,24 +3,33 @@ from __future__ import annotations
 import logging
 import os
 import time
+from contextlib import contextmanager
 from typing import TYPE_CHECKING, List
 
 import torch
 import triton
 import triton.language as tl
+from huggingface_hub import snapshot_download
 
 from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
+from sglang.srt.distributed.parallel_state import (
+    GroupCoordinator,
+    patch_tensor_parallel_group,
+)
 from sglang.srt.environ import envs
+from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.utils import is_cuda, is_hip
 
+if TYPE_CHECKING:
+    from sglang.srt.speculative.eagle_info import EagleVerifyInput
+
+
 if is_cuda():
     from sgl_kernel import fast_topk
 elif is_hip():
     from sgl_kernel import fast_topk
 
-if TYPE_CHECKING:
-    from sglang.srt.speculative.eagle_info import EagleVerifyInput
 
 logger = logging.getLogger(__name__)
 
@@ -436,7 +445,7 @@ def select_top_k_tokens(
     return input_ids, hidden_states, scores, tree_info
 
 
-def _generate_simulated_accept_index(
+def generate_simulated_accept_index(
     accept_index,
     predict,
     accept_length,
@@ -604,3 +613,29 @@ def generate_token_bitmask(
 
     verify_input.grammar = grammar
     return allocate_token_bitmask
+
+
+def load_token_map(token_map_path: str) -> List[int]:
+    if not os.path.exists(token_map_path):
+        cache_dir = snapshot_download(
+            os.path.dirname(token_map_path),
+            ignore_patterns=["*.bin", "*.safetensors"],
+        )
+        token_map_path = os.path.join(cache_dir, os.path.basename(token_map_path))
+    hot_token_id = torch.load(token_map_path, weights_only=True)
+    return torch.tensor(hot_token_id, dtype=torch.int64)
+
+
+@contextmanager
+def draft_tp_context(tp_group: GroupCoordinator):
+    # Draft model doesn't use dp and has its own tp group.
+    # We disable mscclpp now because it doesn't support 2 comm groups.
+    with patch_tensor_parallel_group(tp_group):
+        yield
+
+
+def detect_nan(logits_output: LogitsProcessorOutput):
+    logits = logits_output.next_token_logits
+    if torch.any(torch.isnan(logits)):
+        logger.error("Detected errors during sampling! NaN in the logits.")
+        raise ValueError("Detected errors during sampling! NaN in the logits.")

sglang/srt/speculative/standalone_worker.py

@@ -1,29 +1,20 @@
 import logging
-from contextlib import contextmanager
 from typing import Optional
 
 import torch
 
-from sglang.srt.distributed import GroupCoordinator, patch_tensor_parallel_group
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.server_args import ServerArgs
-from sglang.srt.speculative.eagle_worker import EAGLEWorker, load_token_map
+from sglang.srt.speculative.eagle_worker import EAGLEWorker
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import draft_tp_context, load_token_map
 from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda
 
 if is_cuda():
-    from sgl_kernel import segment_packbits
+    from sgl_kernel import segment_packbits  # noqa: F401
 
 logger = logging.getLogger(__name__)
-RETURN_ORIGINAL_LOGPROB = get_bool_env_var("RETURN_ORIGINAL_LOGPROB")
-
-
-@contextmanager
-def draft_tp_context(tp_group: GroupCoordinator):
-    # Draft model doesn't use dp and has its own tp group.
-    # We disable mscclpp now because it doesn't support 2 comm groups.
-    with patch_tensor_parallel_group(tp_group):
-        yield
+SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
 
 
 class StandaloneWorker(EAGLEWorker):
@@ -51,7 +42,6 @@ class StandaloneWorker(EAGLEWorker):
         self.speculative_algorithm = SpeculativeAlgorithm.from_string(
             server_args.speculative_algorithm
         )
-        self.padded_static_len = -1
 
         # Override the context length of the draft model to be the same as the target model.
         server_args.context_length = target_worker.model_runner.model_config.context_len

sglang/srt/tokenizer/tiktoken_tokenizer.py

@@ -133,9 +133,9 @@ class TiktokenTokenizer:
         )
         return self.encode(ret) if tokenize else ret
 
-    def __call__(self, text, **kwargs):
+    def __call__(self, text: List[str], **kwargs):
         return {
-            "input_ids": self.encode(text),
+            "input_ids": [self.encode(x) for x in text],
         }
 
     def init_xgrammar(self):

sglang/srt/two_batch_overlap.py

@@ -4,10 +4,11 @@ import copy
 import dataclasses
 import logging
 from dataclasses import replace
-from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence
 
 import torch
 
+from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.communicator import (
     CommunicateContext,
@@ -20,9 +21,11 @@ from sglang.srt.layers.moe import (
     get_tbo_token_distribution_threshold,
     is_tbo_enabled,
 )
-from sglang.srt.layers.moe.token_dispatcher import DeepEPDispatcher
-from sglang.srt.layers.quantization import deep_gemm_wrapper
-from sglang.srt.managers.schedule_batch import ScheduleBatch, global_server_args_dict
+from sglang.srt.layers.moe.token_dispatcher import (
+    DeepEPDispatcher,
+    MooncakeEPDispatcher,
+)
+from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import (
     ForwardBatch,
     ForwardMode,
@@ -30,12 +33,13 @@ from sglang.srt.model_executor.forward_batch_info import (
 )
 from sglang.srt.operations import execute_operations, execute_overlapped_operations
 from sglang.srt.operations_strategy import OperationsStrategy
-from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import BumpAllocator, empty_context, get_bool_env_var, is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.layers.moe.token_dispatcher import DispatchOutput
+    from sglang.srt.speculative.eagle_info import EagleVerifyInput
 
 _is_hip = is_hip()
 
@@ -153,7 +157,7 @@ def _update_device_and_sum_field_from_cpu_field(
         cpu_value
         if isinstance(cpu_value, torch.Tensor)
         else torch.tensor(cpu_value, dtype=old_device_value.dtype)
-    ).to(device=global_server_args_dict["device"], non_blocking=True)
+    ).to(device=get_global_server_args().device, non_blocking=True)
     setattr(batch, device_field, new_device_value)
 
     if sum_field is not None:
@@ -362,7 +366,7 @@ class TboDPAttentionPreparer:
     ):
 
         deepep_mode = get_deepep_mode()
-        enable_deepep_moe = get_moe_a2a_backend().is_deepep()
+        enable_a2a_moe = not get_moe_a2a_backend().is_none()
         enable_two_batch_overlap = is_tbo_enabled()
 
         self.enable_two_batch_overlap = enable_two_batch_overlap
@@ -391,7 +395,7 @@ class TboDPAttentionPreparer:
                     local_batch.forward_mode.is_extend()
                     and not local_batch.forward_mode.is_target_verify()
                 )
-                and enable_deepep_moe
+                and enable_a2a_moe
                 and (resolved_deepep_mode.is_low_latency())
             )
         else:
@@ -582,7 +586,7 @@ class TboForwardBatchPreparer:
             sum_field=None,
         )
         _, child_b.extend_start_loc = compute_position(
-            global_server_args_dict["attention_backend"],
+            get_global_server_args().attention_backend,
             child_b.extend_prefix_lens,
             child_b.extend_seq_lens,
             child_b.extend_num_tokens,
@@ -667,6 +671,7 @@ class TboForwardBatchPreparer:
             "can_run_dp_cuda_graph",
             "dp_padding_mode",
             "global_forward_mode",
+            "is_prefill_only",
             "spec_algorithm",
             "capture_hidden_mode",
             "padded_static_len",
@@ -686,7 +691,7 @@ class TboForwardBatchPreparer:
 
         # TODO improve, e.g. unify w/ `init_raw`
         if (
-            global_server_args_dict["moe_dense_tp_size"] == 1
+            get_global_server_args().moe_dense_tp_size == 1
             and batch.global_dp_buffer_len is not None
         ):
             sum_len = end_token_index - start_token_index
@@ -754,7 +759,7 @@ class TboForwardBatchPreparer:
         value_a = min(tbo_split_token_index, num_token_non_padded)
         value_b = max(0, num_token_non_padded - tbo_split_token_index)
         return torch.tensor([value_a, value_b], dtype=torch.int32).to(
-            device=global_server_args_dict["device"], non_blocking=True
+            device=get_global_server_args().device, non_blocking=True
         )
 
     @classmethod
@@ -966,9 +971,14 @@ def _model_forward_tbo_merge_outputs(output_a, output_b):
 class MaybeTboDeepEPDispatcher:
     def __init__(self, **kwargs):
         num_inner_dispatchers = 2 if is_tbo_enabled() else 1
-        self._inners = [
-            DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
-        ]
+        if get_moe_a2a_backend().is_deepep():
+            self._inners = [
+                DeepEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
+            ]
+        elif get_moe_a2a_backend().is_mooncake():
+            self._inners = [
+                MooncakeEPDispatcher(**kwargs) for _ in range(num_inner_dispatchers)
+            ]
 
     def _execute(self, name, tbo_subbatch_index: Optional[int] = None, **kwargs):
         return getattr(self._inners[tbo_subbatch_index or 0], name)(**kwargs)
@@ -990,3 +1000,7 @@ class MaybeTboDeepEPDispatcher:
 
     def combine_b(self, **kwargs):
         return self._execute("combine_b", **kwargs)
+
+    def set_quant_config(self, quant_config: dict):
+        for inner in self._inners:
+            inner.set_quant_config(quant_config)

sglang/srt/utils/__init__.py

@@ -1,2 +1,2 @@
 # Temporarily do this to avoid changing all imports in the repo
-from .common import *
+from sglang.srt.utils.common import *

sglang/srt/{bench_utils.py → utils/bench_utils.py}

@@ -1,4 +1,5 @@
 import os
+import re
 import sys
 from contextlib import nullcontext
 
@@ -108,7 +109,8 @@ def bench_kineto(
     if not with_multiple_kernels:
         for name in kernel_names:
             assert (
-                sum([name in line for line in prof_lines]) == 1
+                sum([int(re.search(name, line) is not None) for line in prof_lines])
+                == 1
             ), f"Errors of the kernel {name} in the profiling table (table: {prof_lines})"
 
     # Save chrome traces
@@ -122,7 +124,7 @@ def bench_kineto(
         total_time = 0
         total_num = 0
         for line in prof_lines:
-            if name in line:
+            if re.search(name, line) is not None:
                 time_str = line.split()[-2]
                 num_str = line.split()[-1]
                 for unit, scale in units.items():