perf: keep homogeneous batches on the vectorized sampling path

omlx/scheduler.py

@@ -1726,6 +1726,9 @@ def __init__(
         # BatchGenerator - the actual batching engine
         self.batch_generator: BatchGenerator | None = None
         self._current_sampler_params: tuple | None = None
+        # Sampling-param key of the current BatchGenerator's fallback sampler;
+        # refreshed by _create_batch_generator (see _row_sampler_for_insert).
+        self._fallback_sampler_params: tuple | None = None
         # Boundary cache snapshots for stateful non-sliceable caches (e.g., ArraysCache).
         # request_id -> {token_count -> snapshot_cache_or_None}
         # Multiple snapshots per request to support per-block ArraysCache state storage.
@@ -2550,6 +2553,40 @@ def _get_xtc_special_tokens(self) -> list[int]:
             tokens.extend(self._output_parser_factory.stop_token_ids)
         return tokens
 
+    @staticmethod
+    def _sampler_param_key(sampling_params: SamplingParams) -> tuple:
+        """Parameters that fully determine a sampler built by omlx_make_sampler.
+
+        Two requests with equal keys produce behaviorally identical samplers
+        (xtc_special_tokens is scheduler-level, so it is shared by all rows).
+        """
+        return (
+            sampling_params.temperature,
+            sampling_params.top_p,
+            sampling_params.min_p,
+            sampling_params.top_k,
+            sampling_params.xtc_probability,
+            sampling_params.xtc_threshold,
+        )
+
+    def _row_sampler_for_insert(
+        self, sampler: Any, sampling_params: SamplingParams
+    ) -> Any:
+        """Return the per-row sampler to pass to BatchGenerator.insert().
+
+        None when the row can run on the BatchGenerator's fallback sampler
+        (same sampling params as the generator was created with). Rows with a
+        None sampler let mlx-lm's GenerationBatch._step keep the vectorized
+        ``fallback_sampler(logprobs)`` path — one sampling op for the whole
+        batch — instead of the per-row slice/sample/concatenate loop that any
+        single per-row sampler forces on every decode step.
+        """
+        if self._fallback_sampler_params is None:
+            return sampler
+        if self._sampler_param_key(sampling_params) == self._fallback_sampler_params:
+            return None
+        return sampler
+
     def _create_batch_generator(
         self, sampling_params: SamplingParams
     ) -> BatchGenerator:
@@ -2608,6 +2645,11 @@ def _create_batch_generator(
             stream=self._stream,
         )
 
+        # Rows whose sampling params match this fallback sampler are inserted
+        # with a None row sampler so decode stays on the vectorized sampling
+        # path (see _row_sampler_for_insert).
+        self._fallback_sampler_params = self._sampler_param_key(sampling_params)
+
         return bg
 
     def _on_prompt_progress(self, updates: list[tuple[int, int, int]]) -> None:
@@ -4212,17 +4254,20 @@ def _insert_prefilled_request(
             mx.random.seed(request.sampling_params.seed)
 
         per_row_lps = state.per_row_lps if state.per_row_lps is not None else []
+        row_sampler = self._row_sampler_for_insert(
+            state.sampler, request.sampling_params
+        )
         uids = self.batch_generator.insert(
             [state.last_token],
             max_tokens=[request.sampling_params.max_tokens],
             caches=[state.cache] if state.cache else None,
             all_tokens=[_batch_generator_all_tokens(request)],
-            samplers=[state.sampler],
+            samplers=[row_sampler],
             logits_processors=[per_row_lps],
             state_machines=[state.sm],
         )
         if uids:
-            _register_uid_rows(self.model, uids, [state.sampler], [per_row_lps])
+            _register_uid_rows(self.model, uids, [row_sampler], [per_row_lps])
             uid = uids[0]
             self.request_id_to_uid[request.request_id] = uid
             self.uid_to_request_id[uid] = request.request_id
@@ -8669,17 +8714,20 @@ def _sparse_progress(processed: int, total: int) -> None:
             # See vllm-mlx-patched commit 8d4052b for the same root cause
             # in a sibling project, and #934 for the user-visible symptom.
             per_row_lps = list(logits_processors) if logits_processors else []
+            row_sampler = self._row_sampler_for_insert(
+                sampler, request.sampling_params
+            )
             uids = self.batch_generator.insert(
                 [tokens_to_process],
                 max_tokens=[request.sampling_params.max_tokens],
                 caches=[cache_to_use] if cache_to_use else None,
                 all_tokens=[_batch_generator_all_tokens(request)],
-                samplers=[sampler],
+                samplers=[row_sampler],
                 logits_processors=[per_row_lps],
                 state_machines=[sm],
             )
             if uids:
-                _register_uid_rows(self.model, uids, [sampler], [per_row_lps])
+                _register_uid_rows(self.model, uids, [row_sampler], [per_row_lps])
                 uid = uids[0]
                 self.request_id_to_uid[request.request_id] = uid
                 self.uid_to_request_id[uid] = request.request_id

tests/test_scheduler_logits_processors.py

@@ -909,3 +909,127 @@ def test_model_wide_paths_release_every_row(self, func_name):
             "wholesale but leaves the registry rows behind; release by model "
             "so nothing survives a reset, recovery, or shutdown. See #1823."
         )
+
+
+class TestHomogeneousRowSamplerElision:
+    """Rows whose sampling params match the BatchGenerator fallback sampler
+    are inserted with a None row sampler so ``GenerationBatch._step`` keeps
+    the vectorized ``fallback_sampler(logprobs)`` path instead of the
+    per-row slice/sample/concatenate loop (docs/perf-analysis F2 stage 1).
+    """
+
+    def _scheduler_with_fallback(self, key):
+        """Bare Scheduler carrying only the attribute under test."""
+        from omlx.scheduler import Scheduler
+
+        scheduler = Scheduler.__new__(Scheduler)
+        scheduler._fallback_sampler_params = key
+        return scheduler
+
+    def _params(self, **overrides):
+        from omlx.request import SamplingParams
+
+        return SamplingParams(**overrides)
+
+    def test_matching_params_elide_row_sampler(self):
+        from omlx.scheduler import Scheduler
+
+        params = self._params(temperature=0.7, top_p=0.9)
+        scheduler = self._scheduler_with_fallback(
+            Scheduler._sampler_param_key(params)
+        )
+        sampler = object()
+
+        assert scheduler._row_sampler_for_insert(sampler, params) is None
+
+    def test_differing_params_keep_row_sampler(self):
+        from omlx.scheduler import Scheduler
+
+        fallback = self._params(temperature=0.7, top_p=0.9)
+        scheduler = self._scheduler_with_fallback(
+            Scheduler._sampler_param_key(fallback)
+        )
+        sampler = object()
+        request_params = self._params(temperature=0.0, top_p=0.9)
+
+        assert (
+            scheduler._row_sampler_for_insert(sampler, request_params) is sampler
+        )
+
+    def test_no_fallback_key_keeps_row_sampler(self):
+        scheduler = self._scheduler_with_fallback(None)
+        sampler = object()
+
+        assert (
+            scheduler._row_sampler_for_insert(sampler, self._params()) is sampler
+        )
+
+    def test_sampler_param_key_ignores_non_sampler_fields(self):
+        """max_tokens/stop/penalties do not shape the sampler callable, so
+        they must not break homogeneity."""
+        from omlx.scheduler import Scheduler
+
+        a = self._params(max_tokens=10, stop=["x"], repetition_penalty=1.2)
+        b = self._params(max_tokens=999, stop=[], repetition_penalty=1.0)
+
+        assert Scheduler._sampler_param_key(a) == Scheduler._sampler_param_key(b)
+
+    def test_realign_restores_sampler_after_remove_then_extend(self):
+        """Mirror of test_processor_lands_on_its_own_row_after_remove_then_extend
+        for samplers: mlx-lm's ``filter`` reindexes ``samplers`` only when
+        ``any(self.samplers)`` — with all-None rows (the new steady state)
+        the stale list survives removal, and a later per-row sampler joins
+        behind its own index. The per-step realign chokepoint must rebuild
+        the row from the uid registry.
+        """
+        import omlx.scheduler as scheduler_mod
+        from omlx.scheduler import (
+            _omlx_realign_generation_batch_rows,
+            _register_uid_rows,
+            _unregister_uid_rows_for_model,
+        )
+
+        model = object()
+        try:
+            # Request A: homogeneous row (None sampler) that finishes.
+            survivor = _bare_generation_batch(uid=0, logits_processors=[[]])
+            survivor.model = model
+            survivor.samplers = [None]
+            _register_uid_rows(model, [0], [None], [[]])
+            survivor.filter([])
+
+            # Request B: heterogeneous row with its own sampler.
+            def row_sampler(x):
+                return x
+
+            joiner = _bare_generation_batch(
+                uid=1, logits_processors=[[]]
+            )
+            joiner.model = model
+            joiner.samplers = [row_sampler]
+            _register_uid_rows(model, [1], [row_sampler], [[]])
+            survivor.extend(joiner)
+
+            _omlx_realign_generation_batch_rows(survivor)
+
+            assert survivor.uids == [1]
+            assert survivor.samplers == [row_sampler]
+            assert survivor.logits_processors == [[]]
+        finally:
+            _unregister_uid_rows_for_model(model)
+
+    def test_scheduler_source_elides_row_sampler_at_insert(self):
+        """Source-level guard: both insert sites must route the row sampler
+        through _row_sampler_for_insert (and register the same value)."""
+        from pathlib import Path
+
+        scheduler_src = (
+            Path(__file__).resolve().parents[1] / "omlx" / "scheduler.py"
+        ).read_text()
+        assert scheduler_src.count("self._row_sampler_for_insert(") >= 2, (
+            "Both BatchGenerator.insert call sites must decide the row "
+            "sampler via _row_sampler_for_insert so homogeneous batches "
+            "stay on the vectorized sampling path."
+        )
+        assert scheduler_src.count("samplers=[row_sampler]") == 2
+        assert scheduler_src.count("[row_sampler], [per_row_lps]") == 2