Release mm features on session close and support multiple /rerun-ut specs (#21501)

python/sglang/srt/managers/schedule_batch.py

@@ -366,6 +366,11 @@ class MultimodalInputs:
     mrope_position_delta: Optional[torch.Tensor] = None
     mrope_position_delta_repeated_cache: Optional[torch.Tensor] = None
 
+    def release_features(self):
+        """Release feature tensors to free GPU memory."""
+        for item in self.mm_items:
+            item.feature = None
+
     @staticmethod
     def from_dict(obj: dict):
         # Check if MM splitting is enabled

python/sglang/srt/managers/scheduler.py

@@ -1674,8 +1674,7 @@ class Scheduler(
             if req.session:
                 continue
             # For non-session requests, clear features and mm_inputs
-            for item in mm_inputs.mm_items:
+            mm_inputs.release_features()
-                item.feature = None
             req.multimodal_inputs = None
 
     def handle_generate_request(

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -438,11 +438,7 @@ class SchedulerOutputProcessorMixin:
             if req.finished():
                 # delete feature to save memory
                 if req.multimodal_inputs is not None and req.session is None:
-                    for mm_item in req.multimodal_inputs.mm_items:
+                    req.multimodal_inputs.release_features()
-                        pixel_values = mm_item.feature
-                        if isinstance(pixel_values, torch.Tensor):
-                            mm_item.feature = None
-                            del pixel_values
                 self.maybe_collect_routed_experts(req)
 
                 if self.server_args.disaggregation_decode_enable_offload_kvcache:

python/sglang/srt/managers/session_controller.py

@@ -169,7 +169,14 @@ class Session:
                 if req.mm_inputs:
                     for item in req.mm_inputs.get("mm_items", []):
                         if item.offsets:
-                            item.offsets = [(s - 1, e - 1) for s, e in item.offsets]
+                            if any(s == 0 for s, _ in item.offsets):
+                                logging.warning(
+                                    "mm_item offset starts at 0 (BOS position), "
+                                    "clamping to 0 after BOS strip"
+                                )
+                            item.offsets = [
+                                (max(0, s - 1), max(0, e - 1)) for s, e in item.offsets
+                            ]
 
             input_ids = (
                 last_req.origin_input_ids
@@ -284,6 +291,18 @@ class SessionController:
             req = next(iter(session.req_nodes.values())).req
             if not req.finished():
                 req.session = None
+
+        # Release multimodal features held by session requests.
+        # Session reqs skip the normal mm cleanup path (scheduler and
+        # output_processor) so features stay alive until the session closes.
+        seen_mm = set()
+        for node in session.req_nodes.values():
+            mm = node.req.multimodal_inputs
+            if mm is not None and id(mm) not in seen_mm:
+                seen_mm.add(id(mm))
+                mm.release_features()
+            node.req.multimodal_inputs = None
+
         if isinstance(self.tree_cache, SessionAwareCache):
             self.tree_cache.release_session(session_id)
         del self.sessions[session_id]

scripts/ci/utils/slash_command_handler.py

@@ -515,47 +515,12 @@ def detect_cuda_suite(file_path_from_test):
     return suite, runner, use_deepep, None
 
 
-def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
+def _resolve_and_dispatch_ut(gh_repo, pr, test_spec, token):
     """
-    Handles the /rerun-ut <file>::<TestClass.test_method> command.
+    Resolve a single test spec and dispatch a workflow run.
-    Dispatches a lightweight workflow to run a single test on the correct CUDA runner.
+
+    Returns a dict with keys: spec, success, test_command, runner_label, run_url, error.
     """
-    # SECURITY: For fork PRs, only allow /rerun-ut if the commenter has write+ permission.
-    # This command checks out and executes code from the PR branch on self-hosted GPU
-    # runners, so we must ensure the commenter is a trusted collaborator.
-    is_fork = pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
-    if is_fork:
-        commenter = comment.user.login
-        perm = gh_repo.get_collaborator_permission(commenter)
-        if perm not in ("admin", "write"):
-            print(f"Permission denied: /rerun-ut on fork PR by {commenter}.")
-            comment.create_reaction("confused")
-            pr.create_issue_comment(
-                "❌ `/rerun-ut` is not available for fork PRs unless the commenter "
-                "has write permission on the repo.\n\n"
-                "Please ask a maintainer to run this command, or use the normal CI flow."
-            )
-            return False
-        print(f"Fork PR, but commenter {commenter} has write+ permission. Proceeding.")
-
-    if not (
-        user_perms.get("can_rerun_ut", False)
-        or user_perms.get("can_rerun_stage", False)
-    ):
-        print("Permission denied: neither can_rerun_ut nor can_rerun_stage is true.")
-        return False
-
-    if not test_spec:
-        comment.create_reaction("confused")
-        pr.create_issue_comment(
-            "❌ Please specify a test: `/rerun-ut <file>::<TestClass.test_method>`\n\n"
-            "Examples:\n"
-            "- `/rerun-ut test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode`\n"
-            "- `/rerun-ut registered/core/test_srt_endpoint.py::TestSRTEndpoint`\n"
-            "- `/rerun-ut test_srt_endpoint.py`"
-        )
-        return False
-
     # Parse spec: split on :: to get file path and optional test selector
     if "::" in test_spec:
         file_part, test_selector = test_spec.split("::", 1)
@@ -570,16 +535,12 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
     # Resolve file path
     resolved_path, err = resolve_test_file(file_part)
     if err:
-        comment.create_reaction("confused")
+        return {"spec": test_spec, "success": False, "error": err}
-        pr.create_issue_comment(f"❌ {err}")
-        return False
 
     # Detect suite and runner
     suite, runner_label, use_deepep, err = detect_cuda_suite(resolved_path)
     if err:
-        comment.create_reaction("confused")
+        return {"spec": test_spec, "success": False, "error": err}
-        pr.create_issue_comment(f"❌ {err}")
-        return False
 
     # Build test_command: file path (+ optional test selector as unittest arg)
     test_command = resolved_path
@@ -601,13 +562,15 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
                 break
 
         if not target_workflow:
-            print(f"Error: {workflow_name} workflow not found")
+            return {
-            return False
+                "spec": test_spec,
+                "success": False,
+                "error": f"{workflow_name} workflow not found",
+            }
 
         is_fork = (
             pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
         )
-        print(f"PR is from fork: {is_fork}")
 
         pr_head_sha = None
         if is_fork:
@@ -641,13 +604,14 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
         success = dispatch_resp.status_code in (200, 204)
         if not success:
             print(f"Dispatch failed: {dispatch_resp.status_code} {dispatch_resp.text}")
+            return {
+                "spec": test_spec,
+                "success": False,
+                "error": f"Dispatch failed: {dispatch_resp.status_code}",
+            }
 
-        if success:
         print(f"Successfully triggered rerun-ut: {test_command}")
-            comment.create_reaction("+1")
 
-            # Include test_command in expected title to distinguish
-            # concurrent /rerun-ut dispatches (run-name includes test_command)
         run_url = find_workflow_run_url(
             gh_repo,
             target_workflow.id,
@@ -659,32 +623,95 @@ def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_spec, token):
             max_wait=30,
             test_command=test_command,
         )
-            if run_url:
+        return {
+            "spec": test_spec,
+            "success": True,
+            "test_command": test_command,
+            "runner_label": runner_label,
+            "run_url": run_url,
+        }
+
+    except Exception as e:
+        print(f"Error triggering rerun-ut for {test_spec}: {e}")
+        return {"spec": test_spec, "success": False, "error": str(e)}
+
+
+def handle_rerun_ut(gh_repo, pr, comment, user_perms, test_specs, token):
+    """
+    Handles the /rerun-ut command. Accepts a list of test specs and dispatches
+    a workflow run for each, posting a single consolidated comment.
+    """
+    # SECURITY: For fork PRs, only allow /rerun-ut if the commenter has write+ permission.
+    # This command checks out and executes code from the PR branch on self-hosted GPU
+    # runners, so we must ensure the commenter is a trusted collaborator.
+    is_fork = pr.head.repo is None or pr.head.repo.owner.login != gh_repo.owner.login
+    if is_fork:
+        commenter = comment.user.login
+        perm = gh_repo.get_collaborator_permission(commenter)
+        if perm not in ("admin", "write"):
+            print(f"Permission denied: /rerun-ut on fork PR by {commenter}.")
+            comment.create_reaction("confused")
             pr.create_issue_comment(
-                    f"✅ Triggered `/rerun-ut` on `{runner_label}` runner:"
+                "❌ `/rerun-ut` is not available for fork PRs unless the commenter "
-                    f" [View workflow run]({run_url})\n"
+                "has write permission on the repo.\n\n"
-                    f"```\ncd test/ && python3 {test_command}\n```"
+                "Please ask a maintainer to run this command, or use the normal CI flow."
+            )
+            return False
+        print(f"Fork PR, but commenter {commenter} has write+ permission. Proceeding.")
+
+    if not (
+        user_perms.get("can_rerun_ut", False)
+        or user_perms.get("can_rerun_stage", False)
+    ):
+        print("Permission denied: neither can_rerun_ut nor can_rerun_stage is true.")
+        return False
+
+    if not test_specs:
+        comment.create_reaction("confused")
+        pr.create_issue_comment(
+            "❌ Please specify a test: `/rerun-ut <file>::<TestClass.test_method>`\n\n"
+            "Examples:\n"
+            "- `/rerun-ut test/registered/core/test_srt_endpoint.py::TestSRTEndpoint.test_simple_decode`\n"
+            "- `/rerun-ut registered/core/test_srt_endpoint.py::TestSRTEndpoint`\n"
+            "- `/rerun-ut test_srt_endpoint.py`\n"
+            "- `/rerun-ut test_a.py test_b.py test_c.py` (multiple tests)"
+        )
+        return False
+
+    results = []
+    for spec in test_specs:
+        results.append(_resolve_and_dispatch_ut(gh_repo, pr, spec, token))
+
+    # Build consolidated comment
+    successes = [r for r in results if r["success"]]
+    failures = [r for r in results if not r["success"]]
+
+    lines = []
+    for r in successes:
+        if r.get("run_url"):
+            lines.append(
+                f"✅ `{r['runner_label']}`: [View workflow run]({r['run_url']})\n"
+                f"```\ncd test/ && python3 {r['test_command']}\n```"
             )
         else:
-                pr.create_issue_comment(
+            lines.append(
-                    f"✅ Triggered `/rerun-ut` on `{runner_label}` runner:\n"
+                f"✅ `{r['runner_label']}`:\n"
-                    f"```\ncd test/ && python3 {test_command}\n```\n"
+                f"```\ncd test/ && python3 {r['test_command']}\n```\n"
                 f"⚠️ Could not retrieve workflow run URL. "
                 f"Check the [Actions tab](https://github.com/{gh_repo.full_name}/actions) for progress."
             )
-            return True
+    for r in failures:
-        else:
+        lines.append(f"❌ `{r['spec']}`: {r['error']}")
-            print("Failed to trigger workflow_dispatch")
-            return False
 
-    except Exception as e:
+    body = "\n\n".join(lines)
-        print(f"Error triggering rerun-ut: {e}")
+
+    if successes:
+        comment.create_reaction("+1")
+    if failures and not successes:
         comment.create_reaction("confused")
-        pr.create_issue_comment(
+
-            f"❌ Failed to trigger rerun-ut: {str(e)}\n\n"
+    pr.create_issue_comment(body)
-            f"Please check the logs or contact maintainers."
+    return len(successes) > 0
-        )
-        return False
 
 
 def main():
@@ -769,9 +796,8 @@ def main():
         handle_rerun_stage(repo, pr, comment, user_perms, stage_name, token)
 
     elif first_line.startswith("/rerun-ut"):
-        parts = first_line.split(maxsplit=1)
+        test_specs = first_line.split()[1:]
-        test_spec = parts[1].strip() if len(parts) > 1 else None
+        handle_rerun_ut(repo, pr, comment, user_perms, test_specs or None, token)
-        handle_rerun_ut(repo, pr, comment, user_perms, test_spec, token)
 
     else:
         print(f"Unknown or ignored command: {first_line}")