[misc] Fix tool-call index, graph padded-row count, and prefill-graph input_embeds refresh (#39574)

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,

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:

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]

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

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
 

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,

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,
             )

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."""