diff --git a/omlx/scheduler.py b/omlx/scheduler.py index 41fd09b08f..a6dcc43f82 100644 --- a/omlx/scheduler.py +++ b/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 diff --git a/tests/test_scheduler_logits_processors.py b/tests/test_scheduler_logits_processors.py index c45dea4e6d..f5580700fc 100644 --- a/tests/test_scheduler_logits_processors.py +++ b/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