[Fix] Guard kernel OOB accesses and harden runtime edge cases (#30847)

python/sglang/benchmark/one_batch_server.py

@@ -785,9 +785,15 @@ def run_one_case(
         response.raise_for_status()
         server_info = response.json()
         internal_states = server_info.get("internal_states", [])
-        internal_state = internal_states[0] if internal_states else {}
-        last_gen_throughput = internal_state.get("last_gen_throughput", None) or -1
-        acc_length = internal_state.get("avg_spec_accept_length", None) or -1
+        acc_length = -1
+        last_gen_throughput = -1
+        for internal_state in internal_states:
+            val_acc = internal_state.get("avg_spec_accept_length")
+            if val_acc is not None:
+                acc_length = val_acc
+            val_thr = internal_state.get("last_gen_throughput")
+            if val_thr is not None:
+                last_gen_throughput = val_thr
 
     # Calculate cache hit rate from before/after metrics delta
     metrics_after = get_cache_tokens_from_metrics(url)

python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh

@@ -317,6 +317,12 @@ K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeF
     }
   }
 
+  // A negative out_loc marks a slot with no KV write target (e.g. the -1
+  // sentinel from the full->SWA translation for out-of-window tokens or
+  // padded rows); skip the row instead of writing out of bounds. Checked
+  // here, not at the load, so the out_loc prefetch overlaps the norm above.
+  if (out_loc < 0) return;
+
   const int32_t page = out_loc >> kPageBits;
   const int32_t offset = out_loc & ((1 << kPageBits) - 1);
   const auto page_ptr = params.kvcache + page * kPageBytes;

python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py

@@ -44,6 +44,7 @@ def create_trtllm_mha_kv_indices_triton(
     full_to_swa_ptr,  # full->SWA token-slot lookup table, or dummy when not SWA
     page_table_ptr,  # [bs, num_pages] int32 block ids (output)
     swa_page_table_ptr,  # [bs, num_pages] int32 SWA block ids (output), or dummy
+    full_to_swa_numel,
     req_to_token_stride: tl.constexpr,
     page_table_stride: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
@@ -85,7 +86,8 @@ def create_trtllm_mha_kv_indices_triton(
     out_off = pid_req * page_table_stride + page_idx
     tl.store(page_table_ptr + out_off, (slot // PAGE_SIZE).to(tl.int32), mask=mask)
     if HAS_SWA:
-        swa_slot = tl.load(full_to_swa_ptr + slot.to(tl.int64), mask=mask)
+        swa_index = tl.minimum(tl.maximum(slot, 0), full_to_swa_numel - 1)
+        swa_slot = tl.load(full_to_swa_ptr + swa_index.to(tl.int64), mask=mask)
         tl.store(
             swa_page_table_ptr + out_off,
             (swa_slot // PAGE_SIZE).to(tl.int32),
@@ -118,6 +120,7 @@ def build_trtllm_mha_page_table(
         _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0
     ), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
     bs, num_pages = page_table.shape
+    full_to_swa_numel = full_to_swa.numel() if has_swa else 0
     create_trtllm_mha_kv_indices_triton[
         (bs, get_num_mha_kv_index_blocks(num_pages, page_size))
     ](
@@ -127,6 +130,7 @@ def build_trtllm_mha_page_table(
         full_to_swa,
         page_table,
         swa_page_table,
+        full_to_swa_numel,
         req_to_token.stride(0),
         page_table.stride(0),
         PAGE_SIZE=page_size,

python/sglang/srt/eplb/expert_location.py

@@ -289,18 +289,29 @@ class ExpertLocationMetadata:
         require_global_experts: bool = False,
     ) -> List[int]:
         # Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario
+        cpu_map = self.logical_to_all_physical_map_cpu
+        # Draft workers can query MoE layers whose layer_id lies beyond the
+        # target-sized expert map; fall back to the identity mapping (no EPLB
+        # rebalancing for those layers) instead of indexing out of range.
+        if layer_id >= cpu_map.shape[0]:
+            if require_global_experts:
+                num_physical_experts = cpu_map.shape[-1]
+                return list(
+                    range(
+                        logical_expert_id,
+                        num_physical_experts,
+                        self.num_logical_experts,
+                    )
+                )
+            return [logical_expert_id]
         if require_global_experts:
-            num_physical_experts = self.logical_to_all_physical_map_cpu[layer_id].shape[
-                -1
-            ]
+            num_physical_experts = cpu_map[layer_id].shape[-1]
             return list(
                 range(logical_expert_id, num_physical_experts, self.num_logical_experts)
             )
         return [
             physical_expert_id
-            for physical_expert_id in self.logical_to_all_physical_map_cpu[
-                layer_id, logical_expert_id
-            ].tolist()
+            for physical_expert_id in cpu_map[layer_id, logical_expert_id].tolist()
             if physical_expert_id != -1
         ]
 

python/sglang/srt/layers/vocab_parallel_embedding.py

@@ -571,6 +571,7 @@ class ParallelLMHead(VocabParallelEmbedding):
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        enable_tp: bool = True,
         use_attn_tp_group: bool = False,
         use_presharded_weights: bool = False,
     ):
@@ -582,6 +583,7 @@ class ParallelLMHead(VocabParallelEmbedding):
             padding_size=padding_size,
             quant_config=quant_config,
             prefix=prefix,
+            enable_tp=enable_tp,
             use_attn_tp_group=use_attn_tp_group,
             use_presharded_weights=use_presharded_weights,
         )

python/sglang/srt/model_executor/forward_batch_info.py

@@ -1302,7 +1302,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
         # padding
         self.input_ids = self._pad_tensor_to_size(self.input_ids, num_tokens)
         self.req_pool_indices = self._pad_tensor_to_size(self.req_pool_indices, bs)
-        self.lora_ids.extend((bs - len(self.lora_ids)) * [None])
+        if self.lora_ids is not None:
+            self.lora_ids.extend((bs - len(self.lora_ids)) * [None])
 
         seq_len_fill_value = (
             model_runner.attn_backend.get_cuda_graph_seq_len_fill_value()
@@ -1417,22 +1418,30 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
                 self.req_pool_indices = self.req_pool_indices[:bs]
                 if self.seq_lens_cpu is not None:
                     self.seq_lens_cpu = self.seq_lens_cpu[:bs]
-                logits_output.next_token_logits = logits_output.next_token_logits[
-                    :num_tokens
-                ]
+                if logits_output.next_token_logits is not None:
+                    logits_output.next_token_logits = logits_output.next_token_logits[
+                        :num_tokens
+                    ]
                 logits_output.hidden_states = logits_output.hidden_states[:num_tokens]
             elif self.forward_mode.is_target_verify():  # verify
                 num_tokens = bs * self.spec_info.draft_token_num
-                logits_output.next_token_logits = logits_output.next_token_logits[
-                    :num_tokens
-                ]
+                if logits_output.next_token_logits is not None:
+                    logits_output.next_token_logits = logits_output.next_token_logits[
+                        :num_tokens
+                    ]
                 logits_output.hidden_states = logits_output.hidden_states[:num_tokens]
             elif self.forward_mode.is_draft_extend_v2():  # draft extend_v2
                 bs = bs * self.spec_info.num_tokens_per_req
-                logits_output.next_token_logits = logits_output.next_token_logits[:bs]
+                if logits_output.next_token_logits is not None:
+                    logits_output.next_token_logits = logits_output.next_token_logits[
+                        :bs
+                    ]
                 logits_output.hidden_states = logits_output.hidden_states[:bs]
             elif self.forward_mode.is_extend() or self.forward_mode.is_idle():
-                logits_output.next_token_logits = logits_output.next_token_logits[:bs]
+                if logits_output.next_token_logits is not None:
+                    logits_output.next_token_logits = logits_output.next_token_logits[
+                        :bs
+                    ]
                 logits_output.hidden_states = logits_output.hidden_states[:bs]
 
             if hasattr(self, "hidden_states_backup"):

python/sglang/srt/model_executor/model_runner.py

@@ -748,14 +748,15 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # the first forward (`set_mla_kv_buffer` -> `self.kv_buffer[layer_id - self.start_layer]`).
         _nnpl = self.model_config.num_nextn_predict_layers
         model_has_mtp_layers = _nnpl is not None and _nnpl > 0
-        model_num_layers = (
-            self.model_config.num_nextn_predict_layers
-            if self.is_draft_worker and model_has_mtp_layers
-            else max(
+        if self.is_draft_worker and model_has_mtp_layers:
+            model_num_layers = getattr(
+                self.model, "num_stages", self.model_config.num_nextn_predict_layers
+            )
+        else:
+            model_num_layers = max(
                 self.model_config.num_hidden_layers,
                 self.model_config.num_attention_layers,
             )
-        )
         if self.model_config.hf_config.architectures[0] == "MiMoV2MTP":
             model_num_layers = 1
         elif self.model_config.hf_config.architectures[0] == "Step3p5MTP":

python/sglang/srt/utils/hf_transformers/common.py

@@ -142,6 +142,16 @@ try:
 except ImportError:
     pass
 
+try:
+    from transformers import Gemma4Config as _HFGemma4Config
+
+    class _Gemma4UnifiedConfigAlias(_HFGemma4Config):
+        model_type = "gemma4_unified"
+
+    _CONFIG_REGISTRY["gemma4_unified"] = _Gemma4UnifiedConfigAlias
+except ImportError:
+    pass
+
 for name, cls in _CONFIG_REGISTRY.items():
     try:
         AutoConfig.register(name, cls)