wehub-resource-sync
181 lines
6.4 KiB
Python
181 lines
6.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""omlx sampling utilities — mx.compile-free re-implementation of mlx-lm samplers.
|
|
|
|
mlx-lm 0.31.x decorates ``categorical_sampling`` and the apply_* helpers with
|
|
``@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)``. In
|
|
the omlx server environment the decorator stops advancing the RNG state after
|
|
the first call: all subsequent samples reuse the same state, so identical
|
|
prompts produce character-identical output even at temperature > 1. Direct
|
|
calls to the underlying primitives advance the state correctly.
|
|
|
|
This module mirrors the mlx-lm implementation but drops the ``mx.compile``
|
|
wrappers, keeping behavior identical otherwise. ``make_sampler`` matches
|
|
``mlx_lm.sample_utils.make_sampler`` so it can replace the import in scheduler
|
|
without further changes.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from typing import Callable, List
|
|
|
|
import mlx.core as mx
|
|
|
|
|
|
def apply_top_p(logprobs: mx.array, top_p: float) -> mx.array:
|
|
"""Top-p (nucleus) filtering — keep the smallest set of tokens whose
|
|
cumulative probability mass is at least ``top_p``."""
|
|
probs = mx.exp(logprobs)
|
|
sorted_indices = mx.argsort(logprobs, axis=-1)
|
|
sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
|
|
|
|
cumulative_probs = mx.cumsum(sorted_probs, axis=-1)
|
|
|
|
inverse_indices = mx.put_along_axis(
|
|
mx.zeros_like(sorted_indices),
|
|
sorted_indices,
|
|
mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
|
|
axis=-1,
|
|
)
|
|
cumulative_probs = mx.take_along_axis(cumulative_probs, inverse_indices, axis=-1)
|
|
|
|
return mx.where(
|
|
cumulative_probs > 1 - top_p,
|
|
logprobs,
|
|
-float("inf"),
|
|
)
|
|
|
|
|
|
def apply_min_p(
|
|
logprobs: mx.array,
|
|
min_p: float,
|
|
min_tokens_to_keep: int = 1,
|
|
) -> mx.array:
|
|
"""Min-p filtering — drop tokens with probability below
|
|
``max(p) * min_p``, while always keeping ``min_tokens_to_keep`` tokens."""
|
|
if not (0 <= min_p <= 1.0):
|
|
raise ValueError(
|
|
f"`min_p` has to be a float in the [0, 1] interval, but is {min_p}"
|
|
)
|
|
if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
|
|
raise ValueError(
|
|
f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"
|
|
)
|
|
|
|
top_logprobs = mx.max(logprobs, axis=-1, keepdims=True)
|
|
scaled_min_p = top_logprobs + math.log(min_p)
|
|
tokens_to_remove = logprobs < scaled_min_p
|
|
|
|
if min_tokens_to_keep > 1:
|
|
top_indices = mx.argpartition(logprobs, kth=-min_tokens_to_keep, axis=-1)
|
|
top_indices = top_indices[..., -min_tokens_to_keep:]
|
|
tokens_to_remove = mx.put_along_axis(
|
|
tokens_to_remove,
|
|
top_indices,
|
|
False,
|
|
axis=-1,
|
|
)
|
|
|
|
return mx.where(tokens_to_remove, -float("inf"), logprobs)
|
|
|
|
|
|
def apply_top_k(logprobs: mx.array, top_k: int) -> mx.array:
|
|
"""Top-k filtering — keep only the ``top_k`` highest-probability tokens."""
|
|
vocab_size = logprobs.shape[-1]
|
|
if not isinstance(top_k, int) or not (0 < top_k < vocab_size):
|
|
raise ValueError(
|
|
f"`top_k` has to be an integer in the (0, {vocab_size}] interval,"
|
|
f" but is {top_k}."
|
|
)
|
|
mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]
|
|
masked_logprobs = mx.put_along_axis(
|
|
logprobs, mask_idx, mx.array(-float("inf"), logprobs.dtype), axis=-1
|
|
)
|
|
return masked_logprobs
|
|
|
|
|
|
def apply_xtc(
|
|
logits: mx.array,
|
|
xtc_probability: float,
|
|
xtc_threshold: float,
|
|
xtc_special_tokens: List[int],
|
|
) -> mx.array:
|
|
"""XTC sampling — with ``xtc_probability``, mask out all but the lowest
|
|
above-threshold token to encourage diversity."""
|
|
if not (0 <= xtc_threshold <= 0.5):
|
|
raise ValueError(
|
|
f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}"
|
|
)
|
|
if not (0 <= xtc_probability <= 1.0):
|
|
raise ValueError(
|
|
f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}"
|
|
)
|
|
|
|
probs = mx.softmax(logits, -1)
|
|
mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
|
|
if xtc_special_tokens:
|
|
mask[..., xtc_special_tokens] = False
|
|
|
|
return mx.where(
|
|
mx.random.uniform(0, 1) > xtc_probability,
|
|
logits,
|
|
mx.where(mask, -mx.inf, logits),
|
|
)
|
|
|
|
|
|
def categorical_sampling(logits: mx.array, temp: float) -> mx.array:
|
|
"""Sample a token id from the categorical distribution defined by
|
|
``logits / temp``. RNG state is advanced through ``mx.random.categorical``."""
|
|
return mx.random.categorical(logits * (1 / temp))
|
|
|
|
|
|
def make_sampler(
|
|
temp: float = 0.0,
|
|
top_p: float = 0.0,
|
|
min_p: float = 0.0,
|
|
min_tokens_to_keep: int = 1,
|
|
top_k: int = 0,
|
|
xtc_probability: float = 0.0,
|
|
xtc_threshold: float = 0.0,
|
|
xtc_special_tokens: List[int] = [],
|
|
) -> Callable[[mx.array], mx.array]:
|
|
"""Build a sampler callable matching ``mlx_lm.sample_utils.make_sampler``.
|
|
|
|
Returns ``argmax`` when ``temp == 0``; otherwise composes optional
|
|
top-p / min-p / xtc / top-k filters and finishes with categorical sampling.
|
|
"""
|
|
if temp == 0:
|
|
sampler = lambda x: mx.argmax(x, axis=-1)
|
|
else:
|
|
sampling_methods = []
|
|
if top_p > 0 and top_p < 1.0:
|
|
sampling_methods.append(lambda x: apply_top_p(x, top_p))
|
|
if min_p != 0.0:
|
|
sampling_methods.append(
|
|
lambda x: apply_min_p(x, min_p, min_tokens_to_keep)
|
|
)
|
|
if xtc_probability > 0.0:
|
|
sampling_methods.append(
|
|
lambda x: apply_xtc(
|
|
x, xtc_probability, xtc_threshold, xtc_special_tokens
|
|
)
|
|
)
|
|
if top_k > 0:
|
|
sampling_methods.append(lambda x: apply_top_k(x, top_k))
|
|
|
|
def sampler(logprobs: mx.array) -> mx.array:
|
|
for method in sampling_methods:
|
|
logprobs = method(logprobs)
|
|
return categorical_sampling(logprobs, temp)
|
|
|
|
# Expose sampling params on the returned callable so downstream code
|
|
# (e.g. MTP acceptance check) can rebuild the filtered distribution
|
|
# without re-plumbing the params through the BatchGenerator contract.
|
|
# Lambda functions accept attribute assignment in CPython.
|
|
sampler.temp = temp
|
|
sampler.top_p = top_p
|
|
sampler.min_p = min_p
|
|
sampler.top_k = top_k
|
|
sampler.min_tokens_to_keep = min_tokens_to_keep
|
|
return sampler
|