From 348e6fd29bd6b4e175f6a17079ac35cb18ac5fc4 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 11 Jul 2026 12:22:13 -0700 Subject: [PATCH] [Fix] Guard kernel OOB accesses and harden runtime edge cases (#30847) --- python/sglang/benchmark/one_batch_server.py | 12 ++++++--- .../csrc/deepseek_v4/main_norm_rope.cuh | 6 +++++ .../ops/kvcache/trtllm_mha_page_table.py | 6 ++++- python/sglang/srt/eplb/expert_location.py | 23 +++++++++++----- .../srt/layers/vocab_parallel_embedding.py | 2 ++ .../srt/model_executor/forward_batch_info.py | 27 ++++++++++++------- .../sglang/srt/model_executor/model_runner.py | 11 ++++---- .../srt/utils/hf_transformers/common.py | 10 +++++++ 8 files changed, 73 insertions(+), 24 deletions(-) diff --git a/python/sglang/benchmark/one_batch_server.py b/python/sglang/benchmark/one_batch_server.py index ca7f8131f..b2fdd4dfd 100644 --- a/python/sglang/benchmark/one_batch_server.py +++ b/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) diff --git a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh b/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh index edcb6b0a9..848cb3646 100644 --- a/python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh +++ b/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; diff --git a/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py b/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py index fdf41a811..67d23a576 100644 --- a/python/sglang/kernels/ops/kvcache/trtllm_mha_page_table.py +++ b/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, diff --git a/python/sglang/srt/eplb/expert_location.py b/python/sglang/srt/eplb/expert_location.py index 6f21898fe..78b38e7f4 100644 --- a/python/sglang/srt/eplb/expert_location.py +++ b/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 ] diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 8f553060e..259254195 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/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, ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 6f97b707b..8a3a7b02c 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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"): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 43509007f..4574da6c1 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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": diff --git a/python/sglang/srt/utils/hf_transformers/common.py b/python/sglang/srt/utils/hf_transformers/common.py index e33ac276d..cfca6af56 100644 --- a/python/sglang/srt/utils/hf_transformers/common.py +++ b/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)