nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc8__cp310-cp310-macosx_14_0_universal2.whl

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3vl_moe/llm_common/generate.py

@@ -1,701 +0,0 @@
-import contextlib
-import functools
-from dataclasses import dataclass
-from typing import (
-    Any,
-    Callable,
-    Generator,
-    List,
-    Optional,
-    Tuple,
-    Union,
-)
-
-import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_reduce
-from transformers import PreTrainedTokenizer
-
-from .cache import (
-    QuantizedKVCache,
-    load_prompt_cache,
-)
-from . import cache
-from .sample_utils import make_sampler
-from .tokenizer_utils import TokenizerWrapper
-
-DEFAULT_PROMPT = "hello"
-DEFAULT_MAX_TOKENS = 100
-DEFAULT_TEMP = 0.0
-DEFAULT_TOP_P = 1.0
-DEFAULT_MIN_P = 0.0
-DEFAULT_TOP_K = 0
-DEFAULT_XTC_PROBABILITY = 0.0
-DEFAULT_XTC_THRESHOLD = 0.0
-DEFAULT_MIN_TOKENS_TO_KEEP = 1
-DEFAULT_SEED = None
-DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
-DEFAULT_QUANTIZED_KV_START = 5000
-
-
-def str2bool(string):
-    return string.lower() not in ["false", "f"]
-
-
-# A stream on the default device just for generation
-generation_stream = mx.new_stream(mx.default_device())
-
-
-@contextlib.contextmanager
-def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
-    """
-    A context manager to temporarily change the wired limit.
-
-    Note, the wired limit should not be changed during an async eval.  If an
-    async eval could be running pass in the streams to synchronize with prior
-    to exiting the context manager.
-    """
-    model_bytes = tree_reduce(
-        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
-    )
-    max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
-    if model_bytes > 0.9 * max_rec_size:
-        model_mb = model_bytes // 2**20
-        max_rec_mb = max_rec_size // 2**20
-        print(
-            f"[WARNING] Generating with a model that requires {model_mb} MB "
-            f"which is close to the maximum recommended size of {max_rec_mb} "
-            "MB. This can be slow. See the documentation for possible work-arounds: "
-            "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
-        )
-    old_limit = mx.set_wired_limit(max_rec_size)
-    try:
-        yield None
-    finally:
-        if streams is not None:
-            for s in streams:
-                mx.synchronize(s)
-        else:
-            mx.synchronize()
-        mx.set_wired_limit(old_limit)
-
-
-@dataclass
-class GenerationResponse:
-    """
-    The output of :func:`stream_generate`.
-
-    Args:
-        text (str): The next segment of decoded text. This can be an empty string.
-        token (int): The next token.
-        from_draft (bool): Whether the token was generated by the draft model.
-        logprobs (mx.array): A vector of log probabilities.
-        prompt_tokens (int): The number of tokens in the prompt.
-        prompt_tps (float): The prompt processing tokens-per-second.
-        generation_tokens (int): The number of generated tokens.
-        generation_tps (float): The tokens-per-second for generation.
-        peak_memory (float): The peak memory used so far in GB.
-        finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
-    """
-
-    text: str
-    token: int
-    logprobs: mx.array
-    from_draft: bool
-    prompt_tokens: int
-    prompt_tps: float
-    generation_tokens: int
-    generation_tps: float
-    peak_memory: float
-    finish_reason: Optional[str] = None
-
-
-def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
-    if (
-        kv_bits is not None
-        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
-        and prompt_cache[0].offset > quantized_kv_start
-    ):
-        for i in range(len(prompt_cache)):
-            if isinstance(prompt_cache[i], cache.KVCache):
-                prompt_cache[i] = prompt_cache[i].to_quantized(
-                    group_size=kv_group_size, bits=kv_bits
-                )
-
-
-def generate_step(
-    prompt: mx.array,
-    model: nn.Module,
-    *,
-    max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
-    max_kv_size: Optional[int] = None,
-    prompt_cache: Optional[Any] = None,
-    prefill_step_size: int = 2048,
-    kv_bits: Optional[int] = None,
-    kv_group_size: int = 64,
-    quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[int, int]] = None,
-    input_embeddings: Optional[mx.array] = None,
-) -> Generator[Tuple[mx.array, mx.array], None, None]:
-    """
-    A generator producing token ids based on the given prompt from the model.
-
-    Args:
-        prompt (mx.array): The input prompt.
-        model (nn.Module): The model to use for generation.
-        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
-          generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
-          token from a vector of log probabilities. Default: ``None``.
-        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed
-          logits. Default: ``None``.
-        max_kv_size (int, optional): Maximum size of the key-value cache. Old
-          entries (except the first 4 tokens) will be overwritten.
-        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
-          provided, the cache will be updated in place.
-        prefill_step_size (int): Step size for processing the prompt.
-        kv_bits (int, optional): Number of bits to use for KV cache quantization.
-          None implies no cache quantization. Default: ``None``.
-        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
-        quantized_kv_start (int): Step to begin using a quantized KV cache.
-           when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[int, int]): A call-back which takes the
-           prompt tokens processed so far and the total number of prompt tokens.
-        input_embeddings (mx.array, optional): Input embeddings to use in place of
-           prompt tokens. Default: ``None``.
-
-    Yields:
-        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
-    """
-    assert (prompt is not None) ^ (
-        input_embeddings is not None
-    ), "Exactly one of prompt or input_embeddings must be provided, not both"
-
-    tokens = None
-
-    # Create the KV cache for generation
-    if prompt_cache is None:
-        prompt_cache = cache.make_prompt_cache(
-            model,
-            max_kv_size=max_kv_size,
-        )
-
-    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
-
-    quantize_cache_fn = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=quantized_kv_start,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
-
-    def _model_call(y):
-        if y.ndim == 3:
-            return model(None, cache=prompt_cache, input_embeddings=y)
-        else:
-            return model(y, cache=prompt_cache)
-
-    def _step(y):
-        nonlocal tokens
-
-        with mx.stream(generation_stream):
-            logits = _model_call(y[None])
-
-            logits = logits[:, -1, :]
-
-            if logits_processors and input_embeddings is None:
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
-                for processor in logits_processors:
-                    logits = processor(tokens, logits)
-
-            quantize_cache_fn(prompt_cache)
-
-            logprobs = logits - mx.logsumexp(logits, keepdims=True)
-            y = sampler(logprobs)
-            return y, logprobs.squeeze(0)
-
-    using_embeddings = input_embeddings is not None
-
-    y = input_embeddings if using_embeddings else prompt
-    with mx.stream(generation_stream):
-        total_prompt_tokens = y.shape[0]
-        prompt_processed_tokens = 0
-        while y.shape[0] > prefill_step_size:
-            _model_call(y[:prefill_step_size][None])
-            quantize_cache_fn(prompt_cache)
-            mx.eval([c.state for c in prompt_cache])
-            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            y = y[prefill_step_size:]
-            mx.clear_cache()
-
-        y, logprobs = _step(y)
-
-    mx.async_eval(y, logprobs)
-    n = 0
-    while True:
-        if n != max_tokens:
-            next_y, next_logprobs = _step(y)
-            mx.async_eval(next_y, next_logprobs)
-        if n == 0:
-            mx.eval(y)
-            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
-        if n == max_tokens:
-            break
-        yield y.item(), logprobs
-        if n % 256 == 0:
-            mx.clear_cache()
-        y, logprobs = next_y, next_logprobs
-        n += 1
-
-
-def nexa_generate_step(
-    model: nn.Module,
-    *,  # enforces explicit parameter naming
-    prompt: Optional[mx.array] = None,
-    max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
-    max_kv_size: Optional[int] = None,
-    prompt_cache: Optional[Any] = None,
-    prefill_step_size: int = 2048,
-    kv_bits: Optional[int] = None,
-    kv_group_size: int = 64,
-    quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[int, int]] = None,
-    input_embeddings: Optional[mx.array] = None,
-    visual_pos_masks: Optional[mx.array] = None,
-    deepstack_visual_embeds: Optional[List[mx.array]] = None,
-    cos: Optional[mx.array] = None,
-    sin: Optional[mx.array] = None,
-    rope_deltas: Optional[mx.array] = None,
-) -> Generator[Tuple[mx.array, mx.array], None, None]:
-    """
-    A generator producing token ids based on the given prompt from the model.
-
-    Args:
-        prompt (mx.array): The input prompt.
-        model (nn.Module): The model to use for generation.
-        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
-          generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
-          token from a vector of log probabilities. Default: ``None``.
-        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed
-          logits. Default: ``None``.
-        max_kv_size (int, optional): Maximum size of the key-value cache. Old
-          entries (except the first 4 tokens) will be overwritten.
-        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
-          provided, the cache will be updated in place.
-        prefill_step_size (int): Step size for processing the prompt.
-        kv_bits (int, optional): Number of bits to use for KV cache quantization.
-          None implies no cache quantization. Default: ``None``.
-        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
-        quantized_kv_start (int): Step to begin using a quantized KV cache.
-           when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[int, int]): A call-back which takes the
-           prompt tokens processed so far and the total number of prompt tokens.
-        input_embeddings (mx.array, optional): Input embeddings to use in place of
-           prompt tokens. Default: ``None``.
-
-    Yields:
-        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
-    """
-    assert (prompt is not None) ^ (
-        input_embeddings is not None
-    ), "Exactly one of prompt or input_embeddings must be provided, not both"
-
-    tokens = None
-
-    # Create the KV cache for generation
-    if prompt_cache is None:
-        prompt_cache = cache.make_prompt_cache(
-            model,
-            max_kv_size=max_kv_size,
-        )
-
-    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
-
-    quantize_cache_fn = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=quantized_kv_start,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
-
-    def _model_call(y):
-        if y.ndim == 4:
-            y = y[0]
-            return model(None, cache=prompt_cache, inputs_embeds=y, cos=cos, sin=sin, visual_pos_masks=visual_pos_masks, deepstack_visual_embeds=deepstack_visual_embeds, rope_deltas=rope_deltas)
-        elif y.ndim == 3:
-            return model(None, cache=prompt_cache, inputs_embeds=y, cos=cos, sin=sin, visual_pos_masks=visual_pos_masks, deepstack_visual_embeds=deepstack_visual_embeds, rope_deltas=rope_deltas)
-        else:
-            return model(y, cache=prompt_cache, rope_deltas=rope_deltas)
-
-    def _step(y):
-        nonlocal tokens
-
-        with mx.stream(generation_stream):
-            logits = _model_call(y[None])
-
-            logits = logits[:, -1, :]
-
-            if logits_processors and input_embeddings is None:
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
-                for processor in logits_processors:
-                    logits = processor(tokens, logits)
-
-            quantize_cache_fn(prompt_cache)
-
-            logprobs = logits - mx.logsumexp(logits, keepdims=True)
-            y = sampler(logprobs)
-            return y, logprobs.squeeze(0)
-
-    using_embeddings = input_embeddings is not None
-
-    y = input_embeddings if using_embeddings else prompt
-    with mx.stream(generation_stream):
-        total_prompt_tokens = y.shape[0]
-        prompt_processed_tokens = 0
-        while y.shape[0] > prefill_step_size:
-            _model_call(y[:prefill_step_size][None])
-            quantize_cache_fn(prompt_cache)
-            mx.eval([c.state for c in prompt_cache])
-            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            y = y[prefill_step_size:]
-            mx.clear_cache()
-
-        y, logprobs = _step(y)
-
-    mx.async_eval(y, logprobs)
-    n = 0
-    while True:
-        if n != max_tokens:
-            next_y, next_logprobs = _step(y)
-            mx.async_eval(next_y, next_logprobs)
-        if n == 0:
-            mx.eval(y)
-            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
-        if n == max_tokens:
-            break
-        yield y.item(), logprobs
-        if n % 256 == 0:
-            mx.clear_cache()
-        y, logprobs = next_y, next_logprobs
-        n += 1
-
-
-
-## Explicit parameter naming means we need to specify the parameter names. 
-def nexa_multimodal_generate_step(
-    model: nn.Module,
-    *,  # enforces explicit parameter naming
-    prompt: Optional[mx.array] = None,
-    max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
-    max_kv_size: Optional[int] = None,
-    prompt_cache: Optional[Any] = None,
-    prefill_step_size: int = 2048,
-    kv_bits: Optional[int] = None,
-    kv_group_size: int = 64,
-    quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[int, int]] = None,
-    input_embeddings: Optional[mx.array] = None,
-    cos: Optional[mx.array] = None,
-    sin: Optional[mx.array] = None,
-    rope_deltas: Optional[mx.array] = None,
-) -> Generator[Tuple[mx.array, mx.array], None, None]:
-    """
-    A generator producing token ids based on the given prompt from the model.
-
-    Args:
-        prompt (mx.array): The input prompt.
-        model (nn.Module): The model to use for generation.
-        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
-          generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
-          token from a vector of log probabilities. Default: ``None``.
-        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed
-          logits. Default: ``None``.
-        max_kv_size (int, optional): Maximum size of the key-value cache. Old
-          entries (except the first 4 tokens) will be overwritten.
-        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
-          provided, the cache will be updated in place.
-        prefill_step_size (int): Step size for processing the prompt.
-        kv_bits (int, optional): Number of bits to use for KV cache quantization.
-          None implies no cache quantization. Default: ``None``.
-        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
-        quantized_kv_start (int): Step to begin using a quantized KV cache.
-           when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[int, int]): A call-back which takes the
-           prompt tokens processed so far and the total number of prompt tokens.
-        input_embeddings (mx.array, optional): Input embeddings to use in place of
-           prompt tokens. Default: ``None``.
-
-    Yields:
-        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
-    """
-    assert (prompt is not None) ^ (
-        input_embeddings is not None
-    ), "Exactly one of prompt or input_embeddings must be provided, not both"
-
-    tokens = None
-
-    # Create the KV cache for generation
-    if prompt_cache is None:
-        prompt_cache = cache.make_prompt_cache(
-            model,
-            max_kv_size=max_kv_size,
-        )
-
-    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
-
-    quantize_cache_fn = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=quantized_kv_start,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
-
-    def _model_call(y):
-        if y.ndim == 3:
-            return model(None, cache=prompt_cache, input_embeddings=y, cos=cos, sin=sin, rope_deltas=rope_deltas)
-        else:
-            return model(y, cache=prompt_cache, rope_deltas=rope_deltas)
-
-    def _step(y):
-        nonlocal tokens
-
-        with mx.stream(generation_stream):
-            logits = _model_call(y[None])
-
-            logits = logits[:, -1, :]
-
-            if logits_processors and input_embeddings is None:
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
-                for processor in logits_processors:
-                    logits = processor(tokens, logits)
-
-            quantize_cache_fn(prompt_cache)
-
-            logprobs = logits - mx.logsumexp(logits, keepdims=True)
-            y = sampler(logprobs)
-            return y, logprobs.squeeze(0)
-
-    using_embeddings = input_embeddings is not None
-
-    y = input_embeddings if using_embeddings else prompt
-    with mx.stream(generation_stream):
-        total_prompt_tokens = y.shape[0]
-        prompt_processed_tokens = 0
-        while y.shape[0] > prefill_step_size:
-            _model_call(y[:prefill_step_size][None])
-            quantize_cache_fn(prompt_cache)
-            mx.eval([c.state for c in prompt_cache])
-            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            y = y[prefill_step_size:]
-            mx.clear_cache()
-
-        y, logprobs = _step(y)
-
-    mx.async_eval(y, logprobs)
-    n = 0
-    while True:
-        if n != max_tokens:
-            next_y, next_logprobs = _step(y)
-            mx.async_eval(next_y, next_logprobs)
-        if n == 0:
-            mx.eval(y)
-            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
-        if n == max_tokens:
-            break
-        yield y.item(), logprobs
-        if n % 256 == 0:
-            mx.clear_cache()
-        y, logprobs = next_y, next_logprobs
-        n += 1
-
-
-
-
-
-def speculative_generate_step(
-    prompt: mx.array,
-    model: nn.Module,
-    draft_model: nn.Module,
-    *,
-    num_draft_tokens=2,
-    max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
-    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
-    prompt_cache: Optional[Any] = None,
-    prefill_step_size: int = 512,
-    kv_bits: Optional[int] = None,
-    kv_group_size: int = 64,
-    quantized_kv_start: int = 0,
-) -> Generator[Tuple[mx.array, mx.array, bool], None, None]:
-    """
-    A generator producing token ids based on the given prompt from the model.
-
-    Args:
-        prompt (mx.array): The input prompt.
-        model (nn.Module): The model to use for generation.
-        draft_model (nn.Module): The draft model for speculative decoding.
-        num_draft_tokens (int, optional): The number of draft tokens for
-          speculative decoding. Default: ``2``.
-        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
-          generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
-          token from a vector of log probabilities. Default: ``None``.
-        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
-          A list of functions that take tokens and logits and return the processed
-          logits. Default: ``None``.
-        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
-          provided, the cache will be updated in place. The cache must be trimmable.
-        prefill_step_size (int): Step size for processing the prompt.
-        kv_bits (int, optional): Number of bits to use for KV cache quantization.
-          None implies no cache quantization. Default: ``None``.
-        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
-        quantized_kv_start (int): Step to begin using a quantized KV cache.
-           when ``kv_bits`` is non-None. Default: ``0``.
-
-    Yields:
-        Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
-          and a bool indicating if the token was generated by the draft model
-    """
-
-    y = prompt.astype(mx.uint32)
-    prev_tokens = None
-
-    # Create the KV cache for generation
-    if prompt_cache is None:
-        model_cache = cache.make_prompt_cache(model)
-        draft_cache = cache.make_prompt_cache(draft_model)
-    else:
-        model_cache = prompt_cache[: len(model.layers)]
-        draft_cache = prompt_cache[len(model.layers) :]
-
-    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
-
-    quantize_cache_fn = functools.partial(
-        maybe_quantize_kv_cache,
-        quantized_kv_start=quantized_kv_start,
-        kv_group_size=kv_group_size,
-        kv_bits=kv_bits,
-    )
-
-    def _process_and_sample(tokens, logits):
-        if logits_processors:
-            for processor in logits_processors:
-                logits = processor(tokens, logits)
-
-        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-        y = sampler(logprobs)
-        return y, logprobs
-
-    def _step(model, cache, y, n_predict=1):
-        with mx.stream(generation_stream):
-            logits = model(y[None], cache=cache)
-            logits = logits[:, -n_predict:, :]
-
-            quantize_cache_fn(cache)
-            if logits_processors:
-                nonlocal prev_tokens
-                out_y, out_logprobs = [], []
-                if n_predict > 1:
-                    y = y[: -(n_predict - 1)]
-                for i in range(n_predict):
-                    prev_tokens = mx.concat([prev_tokens, y]) if prev_tokens is not None else y
-                    y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
-                    out_y.append(y)
-                    out_logprobs.append(logprobs)
-                return mx.concatenate(out_y, axis=0), mx.concatenate(out_logprobs, axis=0)
-            else:
-                return _process_and_sample(None, logits.squeeze(0))
-
-    def _prefill(model, cache, y):
-        while y.size > prefill_step_size:
-            model(y[:prefill_step_size][None], cache=cache)
-            quantize_cache_fn(cache)
-            mx.eval([c.state for c in cache])
-            y = y[prefill_step_size:]
-            mx.clear_cache()
-        return y
-
-    def _rewind_cache(num_draft, num_accept):
-        cache.trim_prompt_cache(model_cache, num_draft - num_accept)
-        cache.trim_prompt_cache(draft_cache, max(num_draft - num_accept - 1, 0))
-
-    def _draft_generate(y, num_draft):
-        if num_draft == 0:
-            return mx.array([], mx.uint32)
-        ys = []
-        for _ in range(num_draft):
-            y, _ = _step(draft_model, draft_cache, y)
-            mx.async_eval(y)
-            ys.append(y)
-        return mx.concatenate(ys)
-
-    with mx.stream(generation_stream):
-        draft_y = _prefill(draft_model, draft_cache, y)
-        y = _prefill(model, model_cache, y)
-
-    ntoks = 0
-    # Set these so the finally block doesn't raise
-    num_draft = 0
-    n = 0
-    try:
-        while True:
-            num_draft = min(max_tokens - ntoks, num_draft_tokens)
-            draft_tokens = _draft_generate(draft_y, num_draft)
-            if prev_tokens is not None:
-                prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1]
-            y = mx.concatenate([y, draft_tokens])
-            tokens, logprobs = _step(model, model_cache, y, num_draft + 1)
-            mx.eval(tokens, draft_tokens)
-            draft_tokens = draft_tokens.tolist()
-            tokens = tokens.tolist()
-            n = 0
-            while n < num_draft:
-                tn, dtn, lpn = tokens[n], draft_tokens[n], logprobs[n]
-                if tn != dtn:
-                    break
-                n += 1
-                ntoks += 1
-                yield tn, lpn, True
-                if ntoks == max_tokens:
-                    break
-            if ntoks < max_tokens:
-                ntoks += 1
-                yield tokens[n], logprobs[n], False
-
-            if ntoks == max_tokens:
-                break
-
-            y = mx.array([tokens[n]], mx.uint32)
-            draft_y = y
-
-            # If we accepted all the draft tokens, include the last
-            # draft token in the next draft step since it hasn't been
-            # processed yet by the draft model
-            if n == num_draft:
-                draft_y = mx.concatenate([mx.array(draft_tokens[-1:], mx.uint32), draft_y])
-
-            if prev_tokens is not None:
-                prev_tokens = prev_tokens[: -max(num_draft - n, 1)]
-            _rewind_cache(num_draft, n)
-    finally:
-        _rewind_cache(num_draft, n)