sglang 0.3.5.post2__py3-none-any.whl → 0.3.6__py3-none-any.whl

sglang/srt/constrained/outlines_backend.py

@@ -81,9 +81,22 @@ class OutlinesGrammar(BaseGrammarObject):
     ):
         self.state = next_state
 
-    def fill_vocab_mask(self, vocab_mask: torch.Tensor):
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        return torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device)
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        tokens = torch.tensor(
+            self.guide.get_next_instruction(self.state).tokens, dtype=torch.int64
+        ).to(vocab_mask.device, non_blocking=True)
+        vocab_mask = vocab_mask[idx]
         vocab_mask.fill_(1)
-        vocab_mask[self.guide.get_next_instruction(self.state).tokens] = 0
+        vocab_mask.scatter_(0, tokens, torch.zeros_like(tokens, dtype=torch.bool))
+
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor):
+        logits.masked_fill_(vocab_mask, float("-inf"))
 
     def copy(self):
         return OutlinesGrammar(self.guide, self.jump_forward_map)

sglang/srt/constrained/xgrammar_backend.py

@@ -21,7 +21,12 @@ from typing import List, Tuple
 import torch
 
 try:
-    from xgrammar import CachedGrammarCompiler, CompiledGrammar, GrammarMatcher
+    from xgrammar import (
+        CachedGrammarCompiler,
+        CompiledGrammar,
+        GrammarMatcher,
+        TokenizerInfo,
+    )
 
     import_error = None
 except ImportError as e:
@@ -80,19 +85,23 @@ class XGrammarGrammar(BaseGrammarObject):
         for i in range(k, len(new_output_ids)):
             assert self.matcher.accept_token(new_output_ids[i])
 
-    def fill_vocab_mask(self, vocab_mask: torch.Tensor):
-        # Note that this bitmask is a bitset, not bool
-        bitmask = self.matcher.get_next_token_bitmask()
-        # Mask the tokens that are not allowed
-        vocab_mask[
-            self.matcher.get_rejected_tokens_from_bitmask(bitmask, self.vocab_size)
-        ] = 1
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        return self.matcher.allocate_token_bitmask(vocab_size, batch_size)
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        self.matcher.fill_next_token_bitmask(vocab_mask, idx)
+
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
+        GrammarMatcher.apply_token_bitmask_inplace(logits, vocab_mask)
 
     def copy(self):
         matcher = GrammarMatcher(
             self.ctx,
             max_rollback_tokens=MAX_ROLLBACK_TOKENS,
-            mask_vocab_size=self.vocab_size,
+            vocab_size=self.vocab_size,
         )
         return XGrammarGrammar(matcher, self.vocab_size, self.ctx)
 
@@ -112,7 +121,8 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
             self.grammar_cache = None
             return
 
-        self.grammar_cache = CachedGrammarCompiler(tokenizer_or_vocab=tokenizer)
+        tokenizer_info = TokenizerInfo.from_huggingface(tokenizer)
+        self.grammar_cache = CachedGrammarCompiler(tokenizer_info=tokenizer_info)
         self.vocab_size = vocab_size
 
     def init_value_impl(self, key: Tuple[str, str]) -> XGrammarGrammar:
@@ -122,9 +132,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         key_type, key_string = key
         if key_type == "json":
             try:
-                ctx = self.grammar_cache.get_compiled_grammar_for_json_schema(
-                    key_string
-                )
+                ctx = self.grammar_cache.compile_json_schema_grammar(schema=key_string)
             except RuntimeError as e:
                 logging.warning(
                     f"Skip invalid json_schema: json_schema={key_string}, {e=}"
@@ -141,7 +149,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         matcher = GrammarMatcher(
             ctx,
             max_rollback_tokens=MAX_ROLLBACK_TOKENS,
-            mask_vocab_size=self.vocab_size,
+            vocab_size=self.vocab_size,
         )
         return XGrammarGrammar(matcher, self.vocab_size, ctx)
 

sglang/srt/layers/activation.py

@@ -32,12 +32,14 @@ from vllm.distributed import (
 )
 from vllm.model_executor.custom_op import CustomOp
 
+from sglang.srt.layers.custom_op_util import register_custom_op
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.utils import set_weight_attrs
 
 logger = logging.getLogger(__name__)
 
 
+@register_custom_op("sglang_silu_and_mul")
 class SiluAndMul(CustomOp):
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:
         d = x.shape[-1] // 2
@@ -51,6 +53,7 @@ class SiluAndMul(CustomOp):
         return out
 
 
+@register_custom_op("sglang_gelu_and_mul")
 class GeluAndMul(CustomOp):
     def __init__(self, approximate="tanh"):
         super().__init__()

sglang/srt/layers/attention/flashinfer_backend.py

@@ -8,7 +8,7 @@ Each backend supports two operators: extend (i.e. prefill with cached prefix) an
 """
 
 from enum import Enum, auto
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
 
 import torch
 import triton
@@ -136,15 +136,17 @@ class FlashInferAttnBackend(AttentionBackend):
             prefix_lens = forward_batch.extend_prefix_lens
 
             # Some heuristics to check whether to use ragged forward
-            use_ragged = False
             if forward_batch.extend_num_tokens >= 4096 and self.num_wrappers == 1:
                 use_ragged = True
-
-            extend_no_prefix = not torch.any(forward_batch.extend_prefix_lens).item()
+                extend_no_prefix = not any(forward_batch.extend_prefix_lens_cpu)
+            else:
+                use_ragged = False
+                extend_no_prefix = False
 
             self.indices_updater_prefill.update(
                 forward_batch.req_pool_indices,
                 forward_batch.seq_lens,
+                forward_batch.seq_lens_sum,
                 prefix_lens,
                 use_ragged=use_ragged,
                 encoder_lens=forward_batch.encoder_lens,
@@ -314,7 +316,6 @@ class FlashInferIndicesUpdaterDecode:
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = model_runner.kv_cache_dtype
         self.q_data_type = model_runner.dtype
-        self.max_context_len = model_runner.req_to_token_pool.req_to_token.size(1)
         self.sliding_window_size = model_runner.sliding_window_size
 
         self.attn_backend = attn_backend
@@ -335,7 +336,12 @@ class FlashInferIndicesUpdaterDecode:
             self.update = self.update_single_wrapper
 
     def update(
-        self, req_pool_indices, seq_lens, seq_lens_sum, decode_wrappers, encoder_lens
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrappers: List,
+        encoder_lens: torch.Tensor,
     ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
@@ -345,8 +351,8 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers=None,
-        encoder_lens=None,
+        decode_wrappers: List,
+        encoder_lens: torch.Tensor,
     ):
         decode_wrappers = decode_wrappers or self.decode_wrappers
         self.call_begin_forward(
@@ -363,8 +369,8 @@ class FlashInferIndicesUpdaterDecode:
         req_pool_indices: torch.Tensor,
         seq_lens: torch.Tensor,
         seq_lens_sum: int,
-        decode_wrappers=None,
-        encoder_lens=None,
+        decode_wrappers: List,
+        encoder_lens: torch.Tensor,
     ):
         decode_wrappers = decode_wrappers or self.decode_wrappers
 
@@ -394,11 +400,11 @@ class FlashInferIndicesUpdaterDecode:
 
     def update_cross_attention(
         self,
-        req_pool_indices,
-        seq_lens,
-        seq_lens_sum,
-        decode_wrappers=None,
-        encoder_lens=None,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        decode_wrappers: List,
+        encoder_lens: torch.Tensor,
     ):
         decode_wrappers = decode_wrappers or self.decode_wrappers
 
@@ -425,11 +431,11 @@ class FlashInferIndicesUpdaterDecode:
     def call_begin_forward(
         self,
         wrapper,
-        req_pool_indices,
-        paged_kernel_lens,
-        paged_kernel_lens_sum,
-        kv_indptr,
-        kv_start_idx,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: int,
+        kv_indptr: torch.Tensor,
+        kv_start_idx: torch.Tensor,
     ):
         bs = len(req_pool_indices)
         kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
@@ -445,7 +451,7 @@ class FlashInferIndicesUpdaterDecode:
             kv_indptr,
             kv_start_idx,
             kv_indices,
-            self.max_context_len,
+            self.req_to_token.shape[1],
         )
 
         wrapper.end_forward()
@@ -474,7 +480,6 @@ class FlashInferIndicesUpdaterPrefill:
         self.head_dim = model_runner.model_config.head_dim
         self.data_type = model_runner.kv_cache_dtype
         self.q_data_type = model_runner.dtype
-        self.max_context_len = model_runner.req_to_token_pool.req_to_token.size(1)
         self.sliding_window_size = model_runner.sliding_window_size
 
         self.attn_backend = attn_backend
@@ -496,23 +501,40 @@ class FlashInferIndicesUpdaterPrefill:
             assert self.attn_backend.num_wrappers == 1
             self.update = self.update_single_wrapper
 
-    def update(self, req_pool_indices, seq_lens, prefix_lens, use_ragged, encoder_lens):
+    def update(
+        self,
+        req_pool_indices: torch.Tnesor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        use_ragged: bool,
+        encoder_lens: torch.Tensor,
+    ):
         # Keep the signature for type checking. It will be assigned during runtime.
         raise NotImplementedError()
 
     def update_single_wrapper(
-        self, req_pool_indices, seq_lens, prefix_lens, use_ragged, encoder_lens
+        self,
+        req_pool_indices: torch.Tnesor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        use_ragged: bool,
+        encoder_lens: torch.Tensor,
     ):
         if use_ragged:
             paged_kernel_lens = prefix_lens
+            paged_kernel_lens_sum = paged_kernel_lens.sum().item()
         else:
             paged_kernel_lens = seq_lens
+            paged_kernel_lens_sum = seq_lens_sum
 
         self.call_begin_forward(
             self.wrapper_ragged,
             self.wrappers_paged[0],
             req_pool_indices,
             paged_kernel_lens,
+            paged_kernel_lens_sum,
             seq_lens,
             prefix_lens,
             None,
@@ -522,7 +544,13 @@ class FlashInferIndicesUpdaterPrefill:
         )
 
     def update_sliding_window(
-        self, req_pool_indices, seq_lens, prefix_lens, use_ragged, encoder_lens
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        use_ragged: bool,
+        encoder_lens: torch.Tensor,
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
@@ -531,9 +559,12 @@ class FlashInferIndicesUpdaterPrefill:
                     seq_lens,
                     torch.tensor(self.sliding_window_size) + seq_lens - prefix_lens,
                 )
+                paged_kernel_lens_sum = paged_kernel_lens.sum().item()
             else:
                 # full attention
                 paged_kernel_lens = seq_lens
+                paged_kernel_lens_sum = seq_lens_sum
+
             kv_start_idx = seq_lens - paged_kernel_lens
 
             self.call_begin_forward(
@@ -541,6 +572,7 @@ class FlashInferIndicesUpdaterPrefill:
                 self.wrappers_paged[wrapper_id],
                 req_pool_indices,
                 paged_kernel_lens,
+                paged_kernel_lens_sum,
                 seq_lens,
                 prefix_lens,
                 kv_start_idx,
@@ -550,23 +582,32 @@ class FlashInferIndicesUpdaterPrefill:
             )
 
     def update_cross_attention(
-        self, req_pool_indices, seq_lens, prefix_lens, use_ragged, encoder_lens
+        self,
+        req_pool_indices: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_sum: int,
+        prefix_lens: torch.Tensor,
+        use_ragged: bool,
+        encoder_lens: torch.Tensor,
     ):
         for wrapper_id in range(2):
             if wrapper_id == 0:
                 # normal attention
                 paged_kernel_lens = seq_lens
                 kv_start_idx = encoder_lens
+                paged_kernel_lens_sum = seq_lens_sum
             else:
                 # cross attention
                 paged_kernel_lens = encoder_lens
                 kv_start_idx = torch.zeros_like(encoder_lens)
+                paged_kernel_lens_sum = paged_kernel_lens.sum().item()
 
             self.call_begin_forward(
                 self.wrapper_ragged,
                 self.wrappers_paged[wrapper_id],
                 req_pool_indices,
                 paged_kernel_lens,
+                paged_kernel_lens_sum,
                 seq_lens,
                 prefix_lens,
                 kv_start_idx,
@@ -579,19 +620,22 @@ class FlashInferIndicesUpdaterPrefill:
         self,
         wrapper_ragged,
         wrapper_paged,
-        req_pool_indices,
-        paged_kernel_lens,
-        seq_lens,
-        prefix_lens,
-        kv_start_idx,
-        kv_indptr,
-        qo_indptr,
-        use_ragged,
+        req_pool_indices: torch.Tensor,
+        paged_kernel_lens: torch.Tensor,
+        paged_kernel_lens_sum: int,
+        seq_lens: torch.Tensor,
+        prefix_lens: torch.Tensor,
+        kv_start_idx: torch.Tensor,
+        kv_indptr: torch.Tensor,
+        qo_indptr: torch.Tensor,
+        use_ragged: bool,
     ):
         bs = len(req_pool_indices)
         kv_indptr[1 : bs + 1] = torch.cumsum(paged_kernel_lens, dim=0)
         kv_indptr = kv_indptr[: bs + 1]
-        kv_indices = torch.empty(kv_indptr[-1], dtype=torch.int32, device="cuda")
+        kv_indices = torch.empty(
+            paged_kernel_lens_sum, dtype=torch.int32, device="cuda"
+        )
         create_flashinfer_kv_indices_triton[(bs,)](
             self.req_to_token,
             req_pool_indices,
@@ -599,7 +643,7 @@ class FlashInferIndicesUpdaterPrefill:
             kv_indptr,
             kv_start_idx,
             kv_indices,
-            self.max_context_len,
+            self.req_to_token.shape[1],
         )
 
         qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
@@ -638,10 +682,11 @@ def create_flashinfer_kv_indices_triton(
     kv_indptr,
     kv_start_idx,
     kv_indices_ptr,
-    max_context_len: tl.constexpr,
+    req_to_token_ptr_stride: tl.constexpr,
 ):
     BLOCK_SIZE: tl.constexpr = 512
     pid = tl.program_id(axis=0)
+
     req_pool_index = tl.load(req_pool_indices_ptr + pid)
     kv_indices_offset = tl.load(kv_indptr + pid)
 
@@ -652,15 +697,15 @@ def create_flashinfer_kv_indices_triton(
         kv_end = kv_start
     kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
 
-    req_to_token_ptr += req_pool_index * max_context_len
-    kv_indices_ptr += kv_indices_offset
-
-    ld_offset = kv_start + tl.arange(0, BLOCK_SIZE)
-    st_offset = tl.arange(0, BLOCK_SIZE)
     num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
-    for _ in range(num_loop):
-        mask = ld_offset < kv_end
-        data = tl.load(req_to_token_ptr + ld_offset, mask=mask)
-        tl.store(kv_indices_ptr + st_offset, data, mask=mask)
-        ld_offset += BLOCK_SIZE
-        st_offset += BLOCK_SIZE
+    for i in range(num_loop):
+        offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE
+        mask = offset < kv_end - kv_start
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + offset,
+            mask=mask,
+        )
+        tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)

sglang/srt/layers/attention/triton_backend.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 import torch
-import torch.nn as nn
 
 from sglang.srt.layers.attention import AttentionBackend
 from sglang.srt.managers.schedule_batch import global_server_args_dict
@@ -28,9 +27,13 @@ class TritonAttnBackend(AttentionBackend):
 
         self.decode_attention_fwd = decode_attention_fwd
         self.extend_attention_fwd = extend_attention_fwd
-        self.num_head = (
-            model_runner.model_config.num_attention_heads // model_runner.tp_size
-        )
+
+        if model_runner.server_args.enable_dp_attention:
+            self.num_head = model_runner.model_config.num_attention_heads
+        else:
+            self.num_head = (
+                model_runner.model_config.num_attention_heads // model_runner.tp_size
+            )
 
         if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
             self.reduce_dtype = torch.float32
@@ -50,7 +53,7 @@ class TritonAttnBackend(AttentionBackend):
             start_loc = torch.zeros_like(forward_batch.seq_lens, dtype=torch.int32)
             start_loc[1:] = torch.cumsum(forward_batch.seq_lens[:-1], dim=0)
 
-            total_num_tokens = torch.sum(forward_batch.seq_lens).item()
+            total_num_tokens = forward_batch.seq_lens_sum
             attn_logits = torch.empty(
                 (self.num_head, total_num_tokens),
                 dtype=self.reduce_dtype,
@@ -61,8 +64,7 @@ class TritonAttnBackend(AttentionBackend):
             max_extend_len = None
         else:
             start_loc = attn_logits = max_seq_len = None
-            prefix_lens = forward_batch.extend_prefix_lens
-            max_extend_len = torch.max(forward_batch.seq_lens - prefix_lens).item()
+            max_extend_len = torch.max(forward_batch.extend_seq_lens).item()
 
         self.forward_metadata = start_loc, attn_logits, max_seq_len, max_extend_len
 

sglang/srt/layers/custom_op_util.py

@@ -0,0 +1,26 @@
+"""
+Copyright 2023-2024 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from vllm.model_executor.custom_op import CustomOp
+
+
+def register_custom_op(op_name):
+    def decorator(cls):
+        if hasattr(CustomOp, "register"):
+            return CustomOp.register(op_name)(cls)
+        else:
+            return cls
+
+    return decorator

sglang/srt/layers/fused_moe/fused_moe.py

@@ -250,9 +250,12 @@ def invoke_fused_moe_kernel(
     assert topk_weights.stride(1) == 1
     assert sorted_token_ids.stride(0) == 1
 
+    padded_size = padding_size
     if not use_fp8:
         assert A_scale is None
         assert B_scale is None
+        # MOE_PADDING FP8 only
+        padded_size = 0
     else:
         A, A_scale = ops.scaled_fp8_quant(A, A_scale)
         assert B_scale is not None
@@ -262,7 +265,7 @@ def invoke_fused_moe_kernel(
         * triton.cdiv(B.shape[1], META["BLOCK_SIZE_N"]),
     )
 
-    K = B.shape[2] - padding_size
+    K = B.shape[2] - padded_size
     if K % config["BLOCK_SIZE_K"] == 0:
         even_ks = True
     else:
@@ -279,7 +282,7 @@ def invoke_fused_moe_kernel(
         expert_ids,
         num_tokens_post_padded,
         B.shape[1],
-        B.shape[2] - padding_size,
+        B.shape[2] - padded_size,
         sorted_token_ids.shape[0],
         topk_ids.numel(),
         A.stride(0),
@@ -480,8 +483,12 @@ def fused_experts(
     a1_scale: Optional[torch.Tensor] = None,
     a2_scale: Optional[torch.Tensor] = None,
 ):
+    padded_size = padding_size
+    if not use_fp8:
+        # MOE_PADDING FP8 only
+        padded_size = 0
     # Check constraints.
-    assert hidden_states.shape[1] == w1.shape[2] - padding_size, "Hidden size mismatch"
+    assert hidden_states.shape[1] == w1.shape[2] - padded_size, "Hidden size mismatch"
     assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
     assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
     assert w1.is_contiguous(), "Expert weights1 must be contiguous"
@@ -498,7 +505,7 @@ def fused_experts(
     get_config_func = functools.partial(
         try_get_optimal_moe_config,
         w1.shape,
-        (w2.shape[0], w2.shape[1], w2.shape[2] - padding_size),
+        (w2.shape[0], w2.shape[1], w2.shape[2] - padded_size),
         topk_ids.shape[1],
         "float8" if use_fp8 else None,
         override_config=override_config,

sglang/srt/layers/layernorm.py

@@ -33,9 +33,12 @@ if is_flashinfer_available():
 
 from vllm.model_executor.custom_op import CustomOp
 
+from sglang.srt.layers.custom_op_util import register_custom_op
+
 logger = logging.getLogger(__name__)
 
 
+@register_custom_op("sglang_rmsnorm")
 class RMSNorm(CustomOp):
     def __init__(
         self,
@@ -78,6 +81,7 @@ class RMSNorm(CustomOp):
             return x, residual
 
 
+@register_custom_op("sglang_gemma_rmsnorm")
 class GemmaRMSNorm(CustomOp):
     def __init__(
         self,

sglang/srt/layers/logits_processor.py

@@ -62,21 +62,21 @@ class LogitsMetadata:
 
     @classmethod
     def from_forward_batch(cls, forward_batch: ForwardBatch):
+        extend_logprob_pruned_lens_cpu = None
+
         if forward_batch.return_logprob:
             return_top_logprob = any(x > 0 for x in forward_batch.top_logprobs_nums)
+            if forward_batch.forward_mode.is_extend():
+                extend_logprob_pruned_lens_cpu = [
+                    extend_len - start_len
+                    for extend_len, start_len in zip(
+                        forward_batch.extend_seq_lens_cpu,
+                        forward_batch.extend_logprob_start_lens_cpu,
+                    )
+                ]
         else:
             return_top_logprob = False
 
-        if forward_batch.forward_mode.is_extend():
-            extend_logprob_pruned_lens_cpu = [
-                extend_len - start_len
-                for extend_len, start_len in zip(
-                    forward_batch.extend_seq_lens,
-                    forward_batch.extend_logprob_start_lens_cpu,
-                )
-            ]
-        else:
-            extend_logprob_pruned_lens_cpu = None
         return cls(
             forward_mode=forward_batch.forward_mode,
             top_logprobs_nums=forward_batch.top_logprobs_nums,

sglang/srt/layers/sampler.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from typing import Union
 
 import torch
@@ -8,7 +7,7 @@ from torch import nn
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
-from sglang.srt.utils import is_flashinfer_available
+from sglang.srt.utils import crash_on_warnings, is_flashinfer_available
 
 if is_flashinfer_available():
     from flashinfer.sampling import (
@@ -19,17 +18,13 @@ if is_flashinfer_available():
     )
 
 
-# Crash on warning if we are running CI tests
-crash_on_warning = os.getenv("SGLANG_IS_IN_CI", "false") == "true"
-
-
 logger = logging.getLogger(__name__)
 
 
 class Sampler(nn.Module):
     def __init__(self):
         super().__init__()
-        self.use_nan_detectioin = not global_server_args_dict["disable_nan_detection"]
+        self.use_nan_detectioin = global_server_args_dict["enable_nan_detection"]
 
     def forward(
         self,
@@ -46,7 +41,8 @@ class Sampler(nn.Module):
             logits = torch.where(
                 torch.isnan(logits), torch.full_like(logits, -1e5), logits
             )
-            exit(1) if crash_on_warning else None
+            if crash_on_warnings():
+                raise ValueError("Detected errors during sampling! NaN in the logits.")
 
         if sampling_info.is_all_greedy:
             # Use torch.argmax if all requests use greedy sampling

sglang/srt/layers/torchao_utils.py

@@ -62,6 +62,8 @@ def torchao_quantize_param_data(param: torch.Tensor, torchao_config: str):
                 granularity=GRANULARITY_MAP[granularity]
             ),
         )
+    else:
+        raise ValueError(f"Unexpected config: {torchao_config}")
 
     return dummy_linear.weight