[MLX] Fix step-bounded profiling for bench tools on Apple Silicon (#29217)

python/sglang/benchmark/offline_throughput.py

@@ -284,7 +284,15 @@ def throughput_test_once(
         dir = os.getenv("SGLANG_TORCH_PROFILER_DIR")
         if not profile_steps:
             known_files = set(os.listdir(dir))
+        # With --profile-steps the scheduler auto-stops mid-run after N steps, so
+        # a second stop here raises "not in progress"; a run shorter than N steps
+        # never hit the target and still needs this explicit stop. Either way we
+        # must stop before monitor_trace_file, which loops forever waiting for a
+        # trace that would otherwise never be finalized.
+        try:
             backend.stop_profile()
+        except RuntimeError:
+            pass
         monitor_trace_file(known_files, dir)
 
     if backend_name == "runtime":

python/sglang/srt/hardware_backend/mlx/scheduler_mixin.py

@@ -86,6 +86,16 @@ class SchedulerMlxOverlapMixin:
     """Mixin that adds MLX overlap scheduling to :class:`Scheduler`."""
 
     def _finalize_mlx_pending_job(self: Scheduler, pending: MlxPendingJob):
+        # Account for this completed forward step. The standard scheduler does
+        # this inside run_batch(), but the MLX overlap loop bypasses run_batch,
+        # so without this forward_ct never advances on MLX. That stalls the
+        # watchdog liveness counter and, more importantly, breaks step-bounded
+        # profiling: _profile_batch_predicate auto-starts/stops based on
+        # forward_ct, so `--profile-steps` (and the server /start_profile
+        # num_steps path) only takes effect once the counter moves here.
+        self.forward_ct += 1
+        self.profiler_manager._profile_batch_predicate(pending.schedule_batch)
+
         result = self.tp_worker.finalize_mlx_result(
             pending.prefills,
             pending.extends,

test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py

@@ -0,0 +1,78 @@
+"""Unit tests for the MLX overlap scheduler mixin (hardware_backend/mlx/scheduler_mixin.py).
+
+Covers:
+  - _finalize_mlx_pending_job advances forward_ct once per completed step
+  - _finalize_mlx_pending_job calls the profiler batch predicate with the
+    finalized batch, so step-bounded profiling (``--profile-steps`` /
+    ``/start_profile`` num_steps) can auto-stop on the MLX overlap loop, which
+    bypasses the standard Scheduler.run_batch().
+
+Skips on non-Apple-Silicon platforms and when ``mlx`` is missing (importing
+scheduler_mixin requires ``mlx.core``).
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import platform
+import unittest
+from unittest.mock import MagicMock
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+_IS_APPLE_SILICON = platform.system() == "Darwin" and platform.machine() == "arm64"
+_HAS_MLX = importlib.util.find_spec("mlx") is not None
+_SKIP_REASON = "requires Apple Silicon and mlx"
+
+
+@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
+class TestFinalizeMlxPendingJob(unittest.TestCase):
+    """forward_ct accounting + profiler predicate wiring in the overlap loop."""
+
+    def _make_scheduler(self):
+        scheduler = MagicMock()
+        scheduler.forward_ct = 0
+        result = MagicMock()
+        result.next_token_ids = None
+        scheduler.tp_worker.finalize_mlx_result.return_value = result
+        return scheduler
+
+    def test_finalize_advances_forward_ct_and_runs_predicate(self):
+        from sglang.srt.hardware_backend.mlx.scheduler_mixin import (
+            SchedulerMlxOverlapMixin,
+        )
+
+        scheduler = self._make_scheduler()
+        pending = MagicMock()
+
+        SchedulerMlxOverlapMixin._finalize_mlx_pending_job(scheduler, pending)
+
+        # Standard run_batch() advances forward_ct and runs the profiler
+        # predicate; the MLX overlap loop must do the same here.
+        self.assertEqual(scheduler.forward_ct, 1)
+        scheduler.profiler_manager._profile_batch_predicate.assert_called_once_with(
+            pending.schedule_batch
+        )
+        # The rest of finalization still runs.
+        scheduler.process_batch_result.assert_called_once()
+
+    def test_forward_ct_advances_once_per_step(self):
+        from sglang.srt.hardware_backend.mlx.scheduler_mixin import (
+            SchedulerMlxOverlapMixin,
+        )
+
+        scheduler = self._make_scheduler()
+
+        for expected in (1, 2, 3):
+            SchedulerMlxOverlapMixin._finalize_mlx_pending_job(scheduler, MagicMock())
+            self.assertEqual(scheduler.forward_ct, expected)
+
+        self.assertEqual(
+            scheduler.profiler_manager._profile_batch_predicate.call_count, 3
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()