[Session R3] Add routed_experts_start_len for absolute routing slice control (#24851)

Co-authored-by: Byron Hsu <byron@periodiclabs.ai>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: zyzshishui <zyzshishui@gmail.com>
Co-authored-by: Yuzhen Zhou <82826991+zyzshishui@users.noreply.github.com>

test/registered/rl/test_return_routed_experts.py

@@ -5,6 +5,8 @@ import unittest
 from typing import List
 
 import aiohttp
+import numpy as np
+import requests
 import torch
 from torch.nn.utils.rnn import pad_sequence
 
@@ -15,6 +17,7 @@ from sglang.srt.state_capturer.routed_experts import (
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import (
+    DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
     CustomTestCase,
@@ -32,6 +35,9 @@ SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered"
 SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json"
 logger = logging.getLogger(__name__)
 
+_QWEN3_30B_A3B_NUM_LAYERS = 48
+_QWEN3_30B_A3B_TOPK = 8
+
 
 class TestReturnRoutedExperts(CustomTestCase):
     """End-to-end check that --enable-return-routed-experts stays correct
@@ -263,5 +269,195 @@ def compare_baseline_w_reference(baseline, reference):
     return num_total_mismatches
 
 
+class TestRoutedExpertsStartLen(CustomTestCase):
+    """Verify the `routed_experts_start_len` parameter:
+
+    - default (0) returns the full sequence
+    - explicit start_len crops the response and the cropped tail matches
+      the corresponding tail of the full response
+    """
+
+    MAX_NEW_TOKENS = 8
+
+    @classmethod
+    def setUpClass(cls):
+        cls.process = popen_launch_server(
+            DEFAULT_ENABLE_ROUTED_EXPERTS_MODEL_NAME_FOR_TEST,
+            DEFAULT_URL_FOR_TEST,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=[
+                "--enable-return-routed-experts",
+                "--enable-deterministic-inference",
+                "--tp",
+                2,
+            ],
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def _send(self, payload: dict) -> dict:
+        resp = requests.post(
+            f"{DEFAULT_URL_FOR_TEST}/generate", json=payload, timeout=120
+        )
+        return resp
+
+    def _build_payload(self, **extra) -> dict:
+        payload = {
+            "text": "User: Tell me a fact about cats.\nAssistant:",
+            "sampling_params": {
+                "temperature": 0,
+                "max_new_tokens": self.MAX_NEW_TOKENS,
+                "ignore_eos": True,
+            },
+            "return_routed_experts": True,
+        }
+        payload.update(extra)
+        return payload
+
+    def _routed_experts(self, resp_json: dict):
+        return extract_routed_experts_from_meta_info(resp_json).reshape(
+            -1, _QWEN3_30B_A3B_NUM_LAYERS, _QWEN3_30B_A3B_TOPK
+        )
+
+    def _seqlen(self, resp_json: dict) -> int:
+        meta = resp_json["meta_info"]
+        return meta["prompt_tokens"] + meta["completion_tokens"]
+
+    def test_start_len_zero_is_default(self):
+        """Omitting the field must match `routed_experts_start_len=0`,
+        which returns the full sequence (start_len=0)."""
+        resp_default = self._send(self._build_payload()).json()
+        resp_zero = self._send(self._build_payload(routed_experts_start_len=0)).json()
+
+        rows_default = self._routed_experts(resp_default)
+        rows_zero = self._routed_experts(resp_zero)
+
+        seqlen_default = self._seqlen(resp_default)
+        seqlen_zero = self._seqlen(resp_zero)
+        self.assertEqual(seqlen_default, seqlen_zero)
+        self.assertEqual(rows_default.shape[0], seqlen_default - 1)
+        self.assertEqual(rows_zero.shape[0], seqlen_zero - 1)
+        self.assertTrue(
+            np.array_equal(rows_default, rows_zero),
+            "default and explicit 0 must produce identical routed experts",
+        )
+
+    def test_start_len_controls_row_count(self):
+        """`routed_experts_start_len=N` must return `seqlen - 1 - N` rows
+        and the returned tail must match the corresponding tail of the
+        full sequence (start_len omitted)."""
+        full_resp = self._send(self._build_payload()).json()
+        full_rows = self._routed_experts(full_resp)
+        seqlen = self._seqlen(full_resp)
+        self.assertEqual(full_rows.shape[0], seqlen - 1)
+
+        start_len = max(1, full_resp["meta_info"]["prompt_tokens"] // 2)
+
+        cropped_resp = self._send(
+            self._build_payload(routed_experts_start_len=start_len)
+        ).json()
+        cropped_rows = self._routed_experts(cropped_resp)
+        cropped_seqlen = self._seqlen(cropped_resp)
+
+        self.assertEqual(seqlen, cropped_seqlen)
+        expected_rows = seqlen - 1 - start_len
+        self.assertEqual(
+            cropped_rows.shape[0],
+            expected_rows,
+            f"expected {expected_rows} rows, got {cropped_rows.shape[0]}",
+        )
+        self.assertTrue(
+            np.array_equal(full_rows[start_len:], cropped_rows),
+            "cropped routed experts must match the tail of the full sequence",
+        )
+
+    def test_start_len_exceeds_prompt_tokens_aborts(self):
+        """`routed_experts_start_len > prompt_tokens` must abort the request:
+        the caller cannot meaningfully reference positions that don't exist
+        in the prompt yet."""
+        baseline = self._send(self._build_payload()).json()
+        prompt_tokens = baseline["meta_info"]["prompt_tokens"]
+
+        ok = self._send(self._build_payload(routed_experts_start_len=prompt_tokens))
+        self.assertEqual(
+            ok.status_code,
+            200,
+            f"start_len=={prompt_tokens} should pass, got {ok.text}",
+        )
+
+        too_big = self._send(
+            self._build_payload(routed_experts_start_len=prompt_tokens + 1)
+        )
+        self._assert_aborted(too_big, "is higher than the number of input tokens")
+
+    def test_start_len_with_cache_hit(self):
+        """`start_len` must allow the radix prefix to extend past it. The
+        first request seeds the cache; the second sends the same prompt
+        with `start_len` somewhere inside the prompt. We verify:
+
+          - meta_info.cached_tokens > start_len (would be impossible if a
+            cap forced the prefix match to <= start_len),
+          - the response row count still equals `seqlen - 1 - start_len`.
+        """
+        cache_salt = "cache-hit-test"
+        first = self._send(self._build_payload(extra_key=cache_salt)).json()
+        self.assertEqual(
+            first["meta_info"].get("cached_tokens", 0),
+            0,
+            "first request must be a cold miss",
+        )
+
+        prompt_tokens = first["meta_info"]["prompt_tokens"]
+        start_len = max(1, prompt_tokens // 2)
+        second = self._send(
+            self._build_payload(
+                extra_key=cache_salt,
+                routed_experts_start_len=start_len,
+            )
+        ).json()
+
+        cached = second["meta_info"].get("cached_tokens", 0)
+        self.assertGreater(
+            cached,
+            start_len,
+            f"expected radix prefix past start_len={start_len}, "
+            f"got cached_tokens={cached} (cap not removed?)",
+        )
+
+        rows = self._routed_experts(second)
+        expected = self._seqlen(second) - 1 - start_len
+        self.assertEqual(
+            rows.shape[0],
+            expected,
+            f"expected {expected} rows, got {rows.shape[0]}",
+        )
+
+    def _assert_aborted(self, resp, expected_substring: str):
+        """Assert a request was aborted with `expected_substring` in the
+        error message."""
+        if resp.status_code == 200:
+            body = resp.json()
+            meta = body.get("meta_info", {})
+            finish_reason = meta.get("finish_reason") or {}
+            message = (
+                str(finish_reason.get("message", ""))
+                + " "
+                + str(body.get("text", ""))
+                + " "
+                + str(body.get("error", ""))
+            )
+            self.assertIn(
+                expected_substring,
+                message,
+                f"expected abort with '{expected_substring}', got body={body}",
+            )
+        else:
+            self.assertGreaterEqual(resp.status_code, 400)
+            self.assertIn(expected_substring, resp.text)
+
+
 if __name__ == "__main__":
     unittest.main()