diff --git a/python/sglang/srt/entrypoints/openai/serving_chat.py b/python/sglang/srt/entrypoints/openai/serving_chat.py index 89cac804a..4bb1b1514 100644 --- a/python/sglang/srt/entrypoints/openai/serving_chat.py +++ b/python/sglang/srt/entrypoints/openai/serving_chat.py @@ -2421,14 +2421,16 @@ class OpenAIServingChat(OpenAIServingBase): return ToolCallProcessingResult(None, text, finish_reason) tool_calls = [] - for call_info in call_info_list: + for index, call_info in enumerate(call_info_list): tool_id = self._process_tool_call_id( call_info, history_tool_calls_cnt ) + # Call ordinal, as in the streaming deltas; + # tool_index is the tool's position in the request. tool_calls.append( ToolCall( id=tool_id, - index=getattr(call_info, "tool_index", None), + index=index, function=FunctionResponse( name=call_info.name, arguments=call_info.parameters, diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 2eb19936d..66de3a131 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -98,14 +98,11 @@ def _make_deferred_finalize_output( ) -> FlashInferTrtllmDeferredFinalizeOutput: """Validate and adapt FlashInfer's ``do_finalize=False`` output ABI.""" gemm2_out, expert_weights, expanded_idx_to_permuted_idx = result[:3] - # Some FlashInfer versions size this buffer from routing_logits dtype while - # writing BF16 weights into it. Reinterpret only the live BF16 prefix. - if expert_weights.dtype == torch.float32: - n, k = expert_weights.shape - expert_weights = expert_weights.view(torch.bfloat16).view(-1, k)[:n] - if expert_weights.dtype != torch.bfloat16: + # FlashInfer >= 0.6.18 types this buffer by content (flashinfer #3595): + # bf16 for packed routing, the caller's dtype for unpacked routing. + if expert_weights.dtype not in (torch.bfloat16, torch.float32): raise RuntimeError( - "FlashInfer deferred finalize must return BF16 expert weights, got " + "FlashInfer deferred finalize must return BF16 or FP32 expert weights, got " f"{expert_weights.dtype}" ) if gemm2_out.dtype != torch.bfloat16: diff --git a/python/sglang/srt/model_executor/runner/base_runner.py b/python/sglang/srt/model_executor/runner/base_runner.py index a125c0d1d..ca2afcb8a 100644 --- a/python/sglang/srt/model_executor/runner/base_runner.py +++ b/python/sglang/srt/model_executor/runner/base_runner.py @@ -38,6 +38,7 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardMode, NgramEmbeddingInfo, PPProxyTensors, + enable_num_token_non_padded, get_server_return_hidden_states_mode, ) from sglang.srt.model_executor.forward_context import ForwardContext, forward_context @@ -99,7 +100,12 @@ def _allocate_decode_buffers( out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype) positions = torch.zeros((max_num_token,), dtype=torch.int64) mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64) - num_token_non_padded = torch.zeros((1,), dtype=torch.int32) + # Refreshed at replay only under expert parallelism. + num_token_non_padded = ( + torch.zeros((1,), dtype=torch.int32) + if enable_num_token_non_padded() + else None + ) custom_mask = torch.ones( (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_req, dtype=torch.bool, @@ -481,7 +487,8 @@ class BaseRunner(ABC): positions = buffers.positions[:num_tokens] out_cache_loc = buffers.out_cache_loc[:num_tokens] mrope_positions = buffers.mrope_positions[:, :num_tokens] - buffers.num_token_non_padded[...] = num_tokens + if buffers.num_token_non_padded is not None: + buffers.num_token_non_padded[...] = num_tokens # Batch-axis buffer views. req_pool_indices = buffers.req_pool_indices[:batch_size] diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index d5160d43b..194ca5d37 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -252,10 +252,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): ) self.enable_two_batch_overlap = get_exec().overlap.enable_two_batch_overlap self.use_ngram_embedding = model_runner.ngram_embedding_manager.enabled - if self.use_ngram_embedding: - hf_config = model_runner.model_config.hf_config - self.ngram_embedding_n = hf_config.ngram_embedding_n - self.ngram_embedding_k = hf_config.ngram_embedding_k self.speculative_algorithm = get_spec().speculative_algorithm self.enable_profile_cuda_graph = get_exec().graph.enable_profile_cuda_graph @@ -879,18 +875,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): # Localize the count when this bucket is attn-TP sharded (SP on). attn_tp_sharded = self.model_runner.attn_tp_sequence_sharded(num_tokens) - buffers.num_token_non_padded[...] = num_tokens - if ( - enable_num_token_non_padded() - and not self.enable_prefill_cp - and attn_tp_sharded - ): - local = compute_local_num_token_non_padded( - global_num_token_non_padded=buffers.num_token_non_padded, - num_tokens_per_dp=num_tokens, - sharded=True, - ) - buffers.num_token_non_padded.copy_(local) + if buffers.num_token_non_padded is not None: + buffers.num_token_non_padded[...] = num_tokens + if not self.enable_prefill_cp and attn_tp_sharded: + local = compute_local_num_token_non_padded( + global_num_token_non_padded=buffers.num_token_non_padded, + num_tokens_per_dp=num_tokens, + sharded=True, + ) + buffers.num_token_non_padded.copy_(local) pp_proxy_tensors = None # pipeline parallelism diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index ab330617b..baf37a0ec 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -1808,10 +1808,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): ie = layer_kwargs.get("input_embeds") if ie is None and ie_idx is not None and len(args) > ie_idx: ie = args[ie_idx] - if ie is not None: - self.buffer_registry.get_slot("input_embeds").slice_for( - 1, static_num_tokens - )[: ie.shape[0]].copy_(ie) + if ie is None: + # Otherwise the graph replays the previous batch's embeddings. + input_ids = args[0] if args else layer_kwargs["input_ids"] + ie = self.model_runner.model.get_input_embeddings()(input_ids) + self.buffer_registry.get_slot("input_embeds").slice_for( + 1, static_num_tokens + )[: ie.shape[0]].copy_(ie) hs = self.backend.replay(shape_key, static_forward_batch, **kwargs) return _slice_output_rows(hs, raw_num_tokens) if full_path else hs diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 7c53ac779..47f7f772b 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -28,7 +28,10 @@ from typing import Dict, List, Optional, Tuple import torch from sglang.srt.environ import envs -from sglang.srt.model_executor.forward_batch_info import NgramEmbeddingInfo +from sglang.srt.model_executor.forward_batch_info import ( + NgramEmbeddingInfo, + enable_num_token_non_padded, +) from sglang.srt.model_executor.input_buffers import ForwardInputBuffers _has_foreach_copy = hasattr(torch, "_foreach_copy_") @@ -97,7 +100,7 @@ class DecodeInputBuffers(ForwardInputBuffers): out_cache_loc: torch.Tensor positions: torch.Tensor mrope_positions: torch.Tensor - num_token_non_padded: torch.Tensor + num_token_non_padded: Optional[torch.Tensor] custom_mask: torch.Tensor next_token_logits_buffer: torch.Tensor mamba_track_indices: Optional[torch.Tensor] @@ -142,7 +145,12 @@ class DecodeInputBuffers(ForwardInputBuffers): out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype) positions = torch.zeros((max_num_token,), dtype=torch.int64) mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64) - num_token_non_padded = torch.zeros((1,), dtype=torch.int32) + # Refreshed at replay only under expert parallelism. + num_token_non_padded = ( + torch.zeros((1,), dtype=torch.int32) + if enable_num_token_non_padded() + else None + ) custom_mask = torch.ones( (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_req, dtype=torch.bool, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index bae009ce8..1f12cea55 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -892,6 +892,9 @@ class DeepseekV2MoE(nn.Module): input_ids_global=input_ids_global, ) + num_token_non_padded = ( + forward_batch.num_token_non_padded if forward_batch is not None else None + ) if not self._enable_a2a_moe: if self._can_dual_stream_graph(hidden_states): fwd = get_forward() @@ -906,12 +909,14 @@ class DeepseekV2MoE(nn.Module): and self.num_fused_shared_experts == 0 and hidden_states.shape[0] > 0 and get_is_capture_mode() + and not is_in_breakable_cuda_graph() ): return self.forward_normal_dual_stream( hidden_states, gemm_output_zero_allocator, input_ids, input_ids_global=input_ids_global, + num_token_non_padded=num_token_non_padded, ) else: return self.forward_normal( @@ -920,6 +925,7 @@ class DeepseekV2MoE(nn.Module): input_ids, input_ids_global=input_ids_global, skip_shared_experts=skip_shared_experts, + num_token_non_padded=num_token_non_padded, ) else: return self.forward_deepep( @@ -932,6 +938,7 @@ class DeepseekV2MoE(nn.Module): gemm_output_zero_allocator: BumpAllocator = None, input_ids: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None, + num_token_non_padded: Optional[torch.Tensor] = None, ) -> torch.Tensor: # Note(kpham-sgl): issue order satisfies 3 constraints: # - no stream explosion: main (routed) issued before alt block -> capture reuses 1 alt stream; @@ -973,6 +980,7 @@ class DeepseekV2MoE(nn.Module): topk_output = self.topk( hidden_states, router_logits, + num_token_non_padded=num_token_non_padded, expert_location_dispatch_info=dispatch_info, **topk_kwargs, ) @@ -1046,6 +1054,7 @@ class DeepseekV2MoE(nn.Module): input_ids: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None, skip_shared_experts: bool = False, + num_token_non_padded: Optional[torch.Tensor] = None, ) -> torch.Tensor: if hasattr(self, "shared_experts") and use_intel_amx_backend( self.shared_experts.gate_up_proj @@ -1089,6 +1098,7 @@ class DeepseekV2MoE(nn.Module): topk_output = self.topk( hidden_states, router_logits, + num_token_non_padded=num_token_non_padded, expert_location_dispatch_info=dispatch_info, **topk_kwargs, ) diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 223314353..72aa8c30f 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -1988,6 +1988,42 @@ class ServingChatTestCase(unittest.TestCase): self.assertEqual(tool_calls[1].id, "functions.get_weather:2") self.assertEqual(tool_calls[1].function.name, "get_weather") + def test_non_streaming_tool_call_index_is_the_call_ordinal(self): + """Two calls to one tool are numbered 0 and 1, as in the streaming deltas, + not by the detector's tool_index (0 for both).""" + self.chat.tool_call_parser = "deepseekv4" + tools = [{"type": "function", "function": {"name": "get_weather"}}] + with patch( + "sglang.srt.entrypoints.openai.serving_chat.FunctionCallParser" + ) as ParserMock: + parser_instance = ParserMock.return_value + calls = [] + for city in ("San Francisco", "London"): + call_info = Mock() + call_info.name = "get_weather" + call_info.parameters = json.dumps({"location": city}) + call_info.tool_index = 0 + calls.append(call_info) + parser_instance.has_tool_call.return_value = True + parser_instance.parse_non_stream.return_value = ("", calls) + + tool_calls, _, finish_reason = self.chat._process_tool_calls( + text="<｜DSML｜tool_calls>...", + tools=tools, + finish_reason={"type": "stop", "matched": None}, + history_tool_calls_cnt=0, + ) + + self.assertEqual([tc.index for tc in tool_calls], [0, 1]) + self.assertEqual( + [tc.function.arguments for tc in tool_calls], + [ + json.dumps({"location": "San Francisco"}), + json.dumps({"location": "London"}), + ], + ) + self.assertEqual(finish_reason["type"], "tool_calls") + def test_required_tool_choice_skips_json_fallback_for_native_parser(self): """A structural-tag parser owns the output format, so a missing tool call must not be pushed through the json_schema array fallback."""