[feat] Init true on policy with qwen_dense

python/sglang/srt/debug_utils/dumper.py

@@ -442,11 +442,12 @@ def _dump_inner(
                 meta_only_fields={**(value_meta_only_fields or {}), **recompute_meta},
             )
 
-        if (
-            enable_curr_grad
-            and isinstance(value, torch.Tensor)
-            and (g := value.grad) is not None
-        ):
+        if enable_curr_grad and isinstance(value, torch.Tensor):
+            g = value.grad if value.grad is not None else getattr(value, "main_grad", None)
+        else:
+            g = None
+
+        if g is not None:
             self._dump_single(
                 tag=grad_tag,
                 tags={**tags, "name": f"grad__{name}"},

python/sglang/srt/disaggregation/encode_receiver.py

@@ -945,6 +945,7 @@ def create_req(self, recv_req: TokenizedGenerateReqInput):
             require_reasoning=recv_req.require_reasoning,
             return_hidden_states=recv_req.return_hidden_states,
             return_routed_experts=recv_req.return_routed_experts,
+            routed_experts_start_len=recv_req.routed_experts_start_len,
             eos_token_ids=self.scheduler.model_config.hf_eos_token_id,
             bootstrap_host=recv_req.bootstrap_host,
             bootstrap_port=recv_req.bootstrap_port,

python/sglang/srt/distributed/communication_op.py

@@ -5,6 +5,9 @@
 import torch
 import torch.distributed
 
+from sglang.srt.tp_invariant_ops import tree_all_reduce_sum
+from sglang.srt.true_on_policy import should_use_tp_invariant_tree_all_reduce
+
 from .parallel_state import (
     get_attn_tp_group,
     get_moe_ep_group,
@@ -15,6 +18,8 @@
 
 def tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
     """All-reduce the input tensor across model parallel group."""
+    if should_use_tp_invariant_tree_all_reduce():
+        return tree_all_reduce_sum(input_, device_group=get_tp_group().device_group)
     return get_tp_group().all_reduce(input_)
 
 
@@ -57,6 +62,10 @@ def broadcast_tensor_dict(
 
 def attention_tensor_model_parallel_all_reduce(input_: torch.Tensor) -> torch.Tensor:
     """All-reduce the input tensor across attention parallel group."""
+    if should_use_tp_invariant_tree_all_reduce():
+        return tree_all_reduce_sum(
+            input_, device_group=get_attn_tp_group().device_group
+        )
     return get_attn_tp_group().all_reduce(input_)
 
 

python/sglang/srt/entrypoints/engine.py

@@ -295,6 +295,7 @@ def generate(
         custom_logit_processor: Optional[Union[List[str], str]] = None,
         return_hidden_states: bool = False,
         return_routed_experts: bool = False,
+        routed_experts_start_len: int = 0,
         stream: bool = False,
         bootstrap_host: Optional[Union[List[str], str]] = None,
         bootstrap_port: Optional[Union[List[int], int]] = None,
@@ -331,6 +332,7 @@ def generate(
             custom_logit_processor=custom_logit_processor,
             return_hidden_states=return_hidden_states,
             return_routed_experts=return_routed_experts,
+            routed_experts_start_len=routed_experts_start_len,
             stream=stream,
             bootstrap_host=bootstrap_host,
             bootstrap_port=bootstrap_port,
@@ -385,6 +387,7 @@ async def async_generate(
         custom_logit_processor: Optional[Union[List[str], str]] = None,
         return_hidden_states: bool = False,
         return_routed_experts: bool = False,
+        routed_experts_start_len: int = 0,
         stream: bool = False,
         bootstrap_host: Optional[Union[List[str], str]] = None,
         bootstrap_port: Optional[Union[List[int], int]] = None,
@@ -420,6 +423,7 @@ async def async_generate(
             lora_path=lora_path,
             return_hidden_states=return_hidden_states,
             return_routed_experts=return_routed_experts,
+            routed_experts_start_len=routed_experts_start_len,
             stream=stream,
             custom_logit_processor=custom_logit_processor,
             bootstrap_host=bootstrap_host,

python/sglang/srt/entrypoints/openai/protocol.py

@@ -273,6 +273,7 @@ class CompletionRequest(BaseModel):
     user: Optional[str] = None
     return_hidden_states: bool = False
     return_routed_experts: bool = False
+    routed_experts_start_len: int = 0
     return_cached_tokens_details: bool = False
 
     # Extra parameters for SRT backend only and will be ignored by OpenAI models.
@@ -588,6 +589,7 @@ class ChatCompletionRequest(BaseModel):
     parallel_tool_calls: bool = True
     return_hidden_states: bool = False
     return_routed_experts: bool = False
+    routed_experts_start_len: int = 0
     return_cached_tokens_details: bool = False
     return_prompt_token_ids: bool = False
     return_meta_info: bool = False

python/sglang/srt/entrypoints/openai/serving_chat.py

@@ -317,6 +317,7 @@ def _convert_to_internal_request(
             disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
             return_hidden_states=request.return_hidden_states,
             return_routed_experts=request.return_routed_experts,
+            routed_experts_start_len=request.routed_experts_start_len,
             rid=request.rid,
             extra_key=self._compute_extra_key(request),
             require_reasoning=self._get_reasoning_from_request(request),
@@ -1137,11 +1138,7 @@ def _process_tool_call_id(
         history_tool_calls_cnt: int,
     ) -> str:
         """Process for generating a new and unique `tool_call_id`"""
-        if self.tool_call_parser != "kimi_k2":
-            # A simple uuid is sufficient for all models except for Kimi-K2.
-            tool_call_id = f"call_{uuid.uuid4().hex[:24]}"
-            return tool_call_id
-        else:
+        if self.tool_call_parser == "kimi_k2":
             # Align with Kimi-K2 format: functions.{name}:{index}
             # Kimi-K2 allows multiple tool_calls in one message; SGLang sets call_item.tool_index to the *local* position inside that message.
             # Therefore, the index must be corrected by using `history_tool_calls_cnt + call_item.tool_index` to ensure globally unique and properly ordered.
@@ -1150,6 +1147,17 @@ def _process_tool_call_id(
                 f"Process tool call idx, parser: {self.tool_call_parser}, tool_call_id: {tool_call_id}, history_cnt: {history_tool_calls_cnt}"
             )
             return tool_call_id
+        if self.tool_call_parser == "kimi_k2_raw_id":
+            # RL training needs the model-emitted tool_call_id round-tripped verbatim,
+            # so we skip the history-based renumbering above and return whatever the
+            # detector captured. Fall back to the canonical Kimi-K2 reconstruction
+            # (without history offset) if for any reason the detector did not record
+            # a raw id — the raw id field is best-effort but the format is stable.
+            if call_item.tool_call_id:
+                return call_item.tool_call_id
+            return f"functions.{call_item.name}:{call_item.tool_index}"
+        # A simple uuid is sufficient for all other models.
+        return f"call_{uuid.uuid4().hex[:24]}"
 
     def _process_tool_calls(
         self,

python/sglang/srt/entrypoints/openai/serving_completions.py

@@ -122,6 +122,7 @@ def _convert_to_internal_request(
             disagg_prefill_dp_rank=request.disagg_prefill_dp_rank,
             return_hidden_states=request.return_hidden_states,
             return_routed_experts=request.return_routed_experts,
+            routed_experts_start_len=request.routed_experts_start_len,
             rid=request.rid,
             extra_key=self._compute_extra_key(request),
             priority=request.priority,

python/sglang/srt/function_call/core_types.py

@@ -10,6 +10,12 @@ class ToolCallItem(BaseModel):
     tool_index: int
     name: Optional[str] = None
     parameters: str  # JSON string
+    # The tool_call_id string emitted by the model, captured verbatim.
+    # Only populated by detectors whose downstream consumers need the exact
+    # model-emitted id (e.g. RL training trajectories). Existing detectors
+    # leave this as None and the serving layer falls back to its usual id
+    # generation strategy.
+    tool_call_id: Optional[str] = None
 
 
 class StreamingParseResult(BaseModel):

python/sglang/srt/function_call/function_call_parser.py

@@ -20,7 +20,10 @@
 from sglang.srt.function_call.gpt_oss_detector import GptOssDetector
 from sglang.srt.function_call.hermes_detector import HermesDetector
 from sglang.srt.function_call.internlm_detector import InternlmDetector
-from sglang.srt.function_call.kimik2_detector import KimiK2Detector
+from sglang.srt.function_call.kimik2_detector import (
+    KimiK2Detector,
+    KimiK2RawIdDetector,
+)
 from sglang.srt.function_call.lfm2_detector import Lfm2Detector
 from sglang.srt.function_call.llama32_detector import Llama32Detector
 from sglang.srt.function_call.mimo_detector import MiMoDetector
@@ -54,6 +57,7 @@ class FunctionCallParser:
         "glm47": Glm47MoeDetector,
         "gpt-oss": GptOssDetector,
         "kimi_k2": KimiK2Detector,
+        "kimi_k2_raw_id": KimiK2RawIdDetector,
         "lfm2": Lfm2Detector,
         "llama3": Llama32Detector,
         "mimo": MiMoDetector,

python/sglang/srt/function_call/kimik2_detector.py

@@ -113,6 +113,7 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult
                         tool_index=function_idx,
                         name=function_name,
                         parameters=function_args,
+                        tool_call_id=function_id,
                     )
                 )
 
@@ -177,6 +178,7 @@ def parse_streaming_increment(
                             tool_index=self.current_tool_id,
                             name=function_name,
                             parameters="",
+                            tool_call_id=function_id,
                         )
                     )
                     self.current_tool_name_sent = True
@@ -253,3 +255,23 @@ def get_info(name: str) -> StructureInfo:
             )
 
         return get_info
+
+
+class KimiK2RawIdDetector(KimiK2Detector):
+    """
+    Variant of KimiK2Detector that preserves the model-emitted tool_call_id verbatim.
+
+    The default kimi_k2 path renumbers ids via `history_tool_calls_cnt + tool_index`
+    in the serving layer so that multi-turn conversations get globally unique,
+    monotonically increasing ids (see PR #10600). That is the right behavior for
+    chat use cases.
+
+    RL training has the opposite requirement: the trajectory must round-trip the
+    exact tool_call_id the model produced (e.g. `functions.foo:5`), so that the
+    follow-up tool result turn references the same id the policy emitted. This
+    subclass exists purely as a marker so the serving layer can branch on the
+    parser name and use `ToolCallItem.tool_call_id` directly. Parsing logic
+    is identical to KimiK2Detector.
+    """
+
+    pass

python/sglang/srt/layers/activation.py

@@ -30,7 +30,7 @@
 from sglang.srt.environ import envs
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.utils import MultiPlatformOp
-from sglang.srt.server_args import get_global_server_args
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import (
     cpu_has_amx_support,
     is_cpu,
@@ -63,7 +63,7 @@
 class SiluAndMul(MultiPlatformOp):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        if get_global_server_args().rl_on_policy_target is not None:
+        if is_true_on_policy_enabled():
             self._forward_method = self.forward_native
 
     def forward_native(self, x: torch.Tensor) -> torch.Tensor:

python/sglang/srt/layers/attention/vision.py

@@ -74,6 +74,7 @@ def flash_attn_func(*args, ver: int = 3, **kwargs):
 from sglang.srt.layers.quantization import QuantizationConfig
 from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb
 from sglang.srt.server_args import get_global_server_args
+from sglang.srt.true_on_policy import is_true_on_policy_enabled
 from sglang.srt.utils import add_prefix, get_bool_env_var
 
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
@@ -865,7 +866,7 @@ def _init_qk_norm(
                 weight_dtype=torch.float32,
                 cast_x_before_out_mul=True,
             )
-            if get_global_server_args().rl_on_policy_target is not None
+            if is_true_on_policy_enabled()
             else {}
         )
         q_norm = RMSNorm(
@@ -988,10 +989,7 @@ def forward(
         if x.dim() == 2:
             x = x.unsqueeze(0)
         assert x.dim() == 3, x.shape
-        if (
-            get_global_server_args().rl_on_policy_target is not None
-            and position_embeddings is not None
-        ):
+        if is_true_on_policy_enabled() and position_embeddings is not None:
             assert isinstance(position_embeddings, tuple), (
                 "expected position_embeddings to be a tuple of two tensors,\n"
                 f"but got {type(position_embeddings)}, change if needed"

python/sglang/srt/layers/communicator.py

@@ -60,6 +60,10 @@
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.true_on_policy import (
+    should_disable_mlp_allreduce_fusion_for_on_policy,
+    should_disable_reduce_scatter_for_on_policy,
+)
 from sglang.srt.utils import (
     get_bool_env_var,
     is_cuda,
@@ -599,6 +603,9 @@ def postprocess_layer(
         )
 
     def should_use_reduce_scatter(self, forward_batch: ForwardBatch):
+        if should_disable_reduce_scatter_for_on_policy():
+            return False
+
         if not self.allow_reduce_scatter:
             return False
         if (
@@ -617,6 +624,9 @@ def should_use_reduce_scatter(self, forward_batch: ForwardBatch):
     def should_fuse_mlp_allreduce_with_next_layer(
         self, forward_batch: ForwardBatch
     ) -> bool:
+        if should_disable_mlp_allreduce_fusion_for_on_policy():
+            return False
+
         if (
             is_dp_attention_enabled()
             and self._speculative_algo is not None

python/sglang/srt/layers/layernorm.py

@@ -27,6 +27,10 @@
 from sglang.srt.environ import envs
 from sglang.srt.layers.utils import MultiPlatformOp
 from sglang.srt.server_args import get_global_server_args
+from sglang.srt.true_on_policy import (
+    get_on_policy_rms_norm_kwargs,
+    is_true_on_policy_enabled,
+)
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -154,13 +158,35 @@ def __init__(
         cast_x_before_out_mul: bool = False,
         fp32_residual: bool = True,
         has_weight: bool = True,
+        weight_dtype: Optional[torch.dtype] = None,
+        override_orig_dtype: Optional[torch.dtype] = None,
+        true_on_policy_weight_dtype: Optional[torch.dtype] = None,
+        true_on_policy_override_orig_dtype: Optional[torch.dtype] = None,
+        true_on_policy_fp32_residual: bool = False,
     ) -> None:
         super().__init__()
+        true_on_policy_kwargs = get_on_policy_rms_norm_kwargs(
+            weight_dtype=true_on_policy_weight_dtype,
+            override_orig_dtype=true_on_policy_override_orig_dtype,
+            fp32_residual=true_on_policy_fp32_residual,
+        )
+        if not cast_x_before_out_mul:
+            cast_x_before_out_mul = true_on_policy_kwargs.get(
+                "cast_x_before_out_mul", cast_x_before_out_mul
+            )
+        fp32_residual = true_on_policy_kwargs.get("fp32_residual", fp32_residual)
+        if weight_dtype is None:
+            weight_dtype = true_on_policy_kwargs.get("weight_dtype", weight_dtype)
+        if override_orig_dtype is None:
+            override_orig_dtype = true_on_policy_kwargs.get(
+                "override_orig_dtype", override_orig_dtype
+            )
         self.has_weight = has_weight
         self.cast_x_before_out_mul = cast_x_before_out_mul
         self.fp32_residual = fp32_residual
+        self.override_orig_dtype = override_orig_dtype
         if self.has_weight:
-            self.weight = nn.Parameter(torch.ones(hidden_size))
+            self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype))
         else:
             self.weight = torch.ones(hidden_size)
         self.variance_epsilon = eps
@@ -181,11 +207,14 @@ def forward_cuda(
             return x
         if self.variance_size_override is not None:
             return self.forward_native(x, residual, post_residual_addition)
+        if (
+            self.weight.dtype != x.dtype
+            or self.cast_x_before_out_mul
+            or self.override_orig_dtype is not None
+        ):
+            return self.forward_native(x, residual, post_residual_addition)
         if is_batch_invariant_mode_enabled():
-            if (
-                residual is not None
-                or get_global_server_args().rl_on_policy_target == "fsdp"
-            ):
+            if residual is not None or is_true_on_policy_enabled():
                 return self.forward_native(x, residual, post_residual_addition)
             return rms_norm_batch_invariant(
                 x,
@@ -275,7 +304,7 @@ def forward_native(
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if not x.is_contiguous():
             x = x.contiguous()
-        orig_dtype = x.dtype
+        orig_dtype = self.override_orig_dtype or x.dtype
 
         if residual is not None and not self.fp32_residual:
             x = x + residual