[cherrypick from #42349] [RL] Pass return_indexer_topk through chat completions and sessions

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

@@ -902,6 +902,7 @@ class ChatCompletionRequest(PDRoutingFields):
     return_hidden_states: Union[bool, Literal["last"]] = False
     return_routed_experts: bool = False
     routed_experts_start_len: int = 0
+    return_indexer_topk: bool = False
     return_cached_tokens_details: bool = False
     return_spec_tokens_details: bool = False
     return_prompt_token_ids: bool = False

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

@@ -1047,6 +1047,9 @@ def _validate_request(self, request: ChatCompletionRequest) -> str | None:
         if request.return_sampling_mask and not request.return_meta_info:
             return "return_sampling_mask requires return_meta_info=true."
 
+        if request.return_indexer_topk and not request.return_meta_info:
+            return "return_indexer_topk requires return_meta_info=true."
+
         if request.return_flat_raw_output_top_logprobs and not request.return_meta_info:
             return "return_flat_raw_output_top_logprobs requires return_meta_info=true."
 
@@ -1317,6 +1320,7 @@ def _convert_to_internal_request(
             return_hidden_states=request.return_hidden_states,
             return_routed_experts=request.return_routed_experts,
             routed_experts_start_len=request.routed_experts_start_len,
+            return_indexer_topk=request.return_indexer_topk,
             rid=request.rid,
             session_id=request.session_id,
             extra_key=request.extra_key,

python/sglang/srt/session/session_controller.py

@@ -317,6 +317,7 @@ def create_req(
             return_hidden_states=req.return_hidden_states,
             return_routed_experts=req.return_routed_experts,
             routed_experts_start_len=req.routed_experts_start_len,
+            return_indexer_topk=req.return_indexer_topk,
             bootstrap_host=req.bootstrap_host,
             bootstrap_port=req.bootstrap_port,
             bootstrap_room=req.bootstrap_room,

test/registered/rl/test_return_indexer_topk.py

@@ -21,7 +21,7 @@
 
 DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2"
 
-# V3.2 config — hardcoded for response decoding (mirrors test_return_routed_experts.py).
+# V3.2 config - expected response shape (mirrors test_return_routed_experts.py).
 NUM_INDEXER_LAYERS = 61
 INDEX_TOPK = 2048
 
@@ -132,10 +132,11 @@ async def _collect_async(cls):
                 for text in cls.texts
             ]
             http_results = await asyncio.gather(*tasks)
-            # Reshape raw int32 bytes into (seqlen-1, num_indexer_layers, index_topk).
+            # Reshape raw int32 bytes into (seqlen-1, num_indexer_layers, index_topk)
+            # with the response's own layer count, which _check_shape_and_range verifies.
             return [
                 extract_indexer_topk_from_meta_info(res).reshape(
-                    -1, NUM_INDEXER_LAYERS, INDEX_TOPK
+                    -1, res["meta_info"]["indexer_topk_num_layers"], INDEX_TOPK
                 )
                 for res in http_results
             ]

test/registered/unit/entrypoints/openai/test_serving_chat.py

@@ -610,12 +610,14 @@ def test_convert_to_internal_request_single(self):
 
             self.basic_req.return_sampling_mask = True
             self.basic_req.sampling_logprobs_mode = "support"
+            self.basic_req.return_indexer_topk = True
             self.basic_req.return_meta_info = True
             adapted, processed = self.chat._convert_to_internal_request(self.basic_req)
             self.assertIsInstance(adapted, GenerateReqInput)
             self.assertFalse(adapted.stream)
             self.assertTrue(adapted.return_sampling_mask)
             self.assertEqual(adapted.sampling_logprobs_mode, "support")
+            self.assertTrue(adapted.return_indexer_topk)
             self.assertEqual(adapted.session_id, "session-1")
             self.assertEqual(processed, self.basic_req)
 
@@ -759,6 +761,18 @@ def test_validate_request_rejects_sampling_mask_without_meta_info(self):
             "return_sampling_mask requires return_meta_info=true.",
         )
 
+    def test_validate_request_rejects_indexer_topk_without_meta_info(self):
+        req = ChatCompletionRequest(
+            model="x",
+            messages=[{"role": "user", "content": "Hi?"}],
+            return_indexer_topk=True,
+        )
+
+        self.assertEqual(
+            self.chat._validate_request(req),
+            "return_indexer_topk requires return_meta_info=true.",
+        )
+
     def test_validate_request_rejects_flat_output_top_logprobs_without_meta_info(
         self,
     ):

test/registered/unit/managers/test_multi_tokenizer_mixin.py

@@ -232,6 +232,17 @@ def test_batch_str_output_keeps_sampling_distribution_aligned(self):
             chunk.to_lists(support_logprobs=True), ([[7, 8]], [[-0.5, -1.0]])
         )
 
+    def test_batch_str_output_keeps_indexer_topk_layer_count(self):
+        """The batch-wide layer count reaches every per-worker split."""
+        output = _make_batch_str_output()
+        output.indexer_topk = ["AAAA", "BBBB"]
+        output.indexer_topk_num_layers = 7
+
+        for i in range(2):
+            self.assertEqual(
+                _handle_output_by_index(output, i).indexer_topk_num_layers, 7
+            )
+
     def test_batch_str_output_without_weight_versions_stays_none(self):
         """An output from an older server without the field splits into None."""
         output = _make_batch_str_output()

test/registered/unit/managers/test_output_streamer_customized_info.py

@@ -534,6 +534,61 @@ def test_sends_once_interval_tokens_are_unsent(self):
 }
 
 
+class TestOutputStreamerIndexerTopk(unittest.TestCase):
+    def setUp(self):
+        streamer = "sglang.srt.managers.scheduler_components.output_streamer"
+        for target, value in (
+            (
+                f"{streamer}.get_serving",
+                SimpleNamespace(stream_interval=1, weight_version="v"),
+            ),
+            (
+                f"{streamer}.get_observability",
+                SimpleNamespace(enable_request_time_stats_logging=False),
+            ),
+            # to_payload imports the capturer lazily from its owning module.
+            (
+                "sglang.srt.state_capturer.indexer_topk.get_global_indexer_capturer",
+                SimpleNamespace(num_indexer_layers=7),
+            ),
+        ):
+            patcher = patch(target, return_value=value)
+            patcher.start()
+            self.addCleanup(patcher.stop)
+        enter_scope(self, published_topology(ranks={"dp_rank": 0}))
+
+    def test_payload_carries_the_capturer_layer_count(self):
+        """Clients reshape the flat indexer_topk bytes by this layer count, so it must
+        cross either IPC codec in the batch that carries indexer_topk."""
+
+        class Streamer(SchedulerOutputStreamer):
+            def get_cached_tokens_details(self, req):
+                return None
+
+        for codec, round_trip in _IPC_ROUND_TRIPS.items():
+            with (
+                self.subTest(codec=codec),
+                patch.object(io_struct, "_USE_PICKLE_IPC", codec == "pickle"),
+            ):
+                outputs = []
+                streamer = Streamer(
+                    send_to_detokenizer=SimpleNamespace(send_output=outputs.append),
+                    tree_cache=None,
+                    server_args=SimpleNamespace(),
+                    is_generation=True,
+                    spec_algorithm=SpeculativeAlgorithm.NONE,
+                    disaggregation_mode=DisaggregationMode.NULL,
+                    enable_hicache_storage=lambda: False,
+                )
+                req = _FakeReq("r0", array("q", [10]), finished=True)
+                req.return_indexer_topk = True
+                req.indexer_topk = torch.zeros((1, 7, 4), dtype=torch.int32)
+
+                streamer._stream_output_generation([req], False)
+
+                self.assertEqual(round_trip(outputs[0]).indexer_topk_num_layers, 7)
+
+
 class TestOutputStreamerSamplingMasks(unittest.TestCase):
     def test_rows_stream_once_and_expand_to_response_lists(self):
         """Queued rows cross either IPC codec once, in batch order, as per-token lists;

test/registered/unit/managers/test_tokenizer_manager_rid_cleanup.py

@@ -451,6 +451,25 @@ def test_abort_output_ids_match_the_streaming_mode(self):
                 self.assertEqual(out["meta_info"]["completion_tokens"], 3)
 
 
+class TestIndexerTopkMetaInfo(CustomTestCase):
+    def test_indexer_topk_arrives_with_its_layer_count(self):
+        """meta_info["indexer_topk"] is flat int32 bytes; clients cannot reshape it
+        unless the layer count arrives in the same chunk."""
+        tm = _make_tokenizer_manager(self)
+        rid = "indexer_topk_layers_rid"
+        state = _make_req_state(rid)
+        tm.rid_to_state[rid] = state
+        batch_output = _make_batch_str_output(rid)
+        batch_output.indexer_topk = ["AAAA"]
+        batch_output.indexer_topk_num_layers = 7
+
+        asyncio.run(tm._handle_batch_output(batch_output))
+
+        meta_info = state.out_list[-1]["meta_info"]
+        self.assertEqual(meta_info["indexer_topk"], "AAAA")
+        self.assertEqual(meta_info["indexer_topk_num_layers"], 7)
+
+
 class TestRidToStateCleanupOnBatchOutput(CustomTestCase):
     """Test that _handle_batch_output removes rid from rid_to_state on completion."""
 

test/registered/unit/mem_cache/test_session_token_share_unit.py

@@ -44,6 +44,7 @@ def _recv(rid, input_ids, max_new_tokens=8):
         return_hidden_states=False,
         return_routed_experts=False,
         routed_experts_start_len=0,
+        return_indexer_topk=False,
         bootstrap_host=None,
         bootstrap_port=None,
         bootstrap_room=None,