diff --git a/mlx_lm/chat.py b/mlx_lm/chat.py index 328ee36d4..ba47e9a3b 100644 --- a/mlx_lm/chat.py +++ b/mlx_lm/chat.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import argparse +import json import mlx.core as mx @@ -80,6 +81,34 @@ def setup_arg_parser(): default=None, help="System prompt to be used for the chat template", ) + parser.add_argument( + "--chat-template-config", + help="Additional JSON config for apply_chat_template, e.g. '{\"enable_thinking\": false}'", + default=None, + ) + parser.add_argument( + "--draft-type", + choices=["none", "ngram-simple", "ngram-mod"], + default="none", + help="Draft strategy for speculative decoding.", + ) + parser.add_argument( + "--num-draft-tokens", + type=int, + default=3, + help="Number of draft tokens to propose.", + ) + parser.add_argument( + "--ngram-size", + type=int, + default=None, + help="N-gram window size. Defaults to 3 for ngram-simple and 16 for ngram-mod.", + ) + parser.add_argument( + "--disable-adaptive-gate", + action="store_true", + help="Disable the adaptive speculative decoding gate.", + ) parser.add_argument( "--pipeline", action="store_true", @@ -118,6 +147,10 @@ def main(): with ChatUI(args, rank=rank) as ui: prompt_cache = make_prompt_cache(model, args.max_kv_size) + template_kwargs = json.loads(args.chat_template_config or "{}") + messages = [] + if args.system_prompt is not None: + messages.append({"role": "system", "content": args.system_prompt}) while True: query = ui.prompt() if query == "q": @@ -125,19 +158,29 @@ def main(): break if query == "r": prompt_cache = make_prompt_cache(model, args.max_kv_size) + messages = [] + if args.system_prompt is not None: + messages.append({"role": "system", "content": args.system_prompt}) ui.say_reset() continue if query == "h": ui.say_help() continue - messages = [] - if args.system_prompt is not None: - messages.append({"role": "system", "content": args.system_prompt}) messages.append({"role": "user", "content": query}) prompt = tokenizer.apply_chat_template( messages, add_generation_prompt=True, + **template_kwargs, ) + generate_kwargs = {} + if args.draft_type != "none": + generate_kwargs["draft_type"] = args.draft_type + generate_kwargs["num_draft_tokens"] = args.num_draft_tokens + generate_kwargs["disable_adaptive_gate"] = args.disable_adaptive_gate + if args.ngram_size is not None: + generate_kwargs["ngram_size"] = args.ngram_size + response_text = [] + accepted = 0 last_response = None for response in stream_generate( model, @@ -154,10 +197,25 @@ def main(): ), ), prompt_cache=prompt_cache, + **generate_kwargs, ): ui.stream_token(response.text) + response_text.append(response.text) + accepted += 1 if response.from_draft else 0 last_response = response ui.end_turn(last_response) + if last_response is not None and rank == 0: + generated = last_response.generation_tokens + acceptance = 100 * accepted / generated if generated else 0.0 + print( + "[stats] " + f"prompt={last_response.prompt_tokens} tok " + f"generated={generated} tok " + f"tok/s={last_response.generation_tps:.2f} " + f"accepted={accepted}/{generated} ({acceptance:.1f}%) " + f"peak={last_response.peak_memory:.2f} GB" + ) + messages.append({"role": "assistant", "content": "".join(response_text)}) if __name__ == "__main__": diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py index 38792a160..bfc1ff258 100644 --- a/mlx_lm/generate.py +++ b/mlx_lm/generate.py @@ -1,11 +1,14 @@ # Copyright © 2023-2024 Apple Inc. import argparse +import array import contextlib import copy import functools import json +import logging import sys +import threading import time from collections import deque from dataclasses import dataclass @@ -16,6 +19,7 @@ Generator, List, Optional, + Protocol, Sequence, Tuple, Union, @@ -55,6 +59,43 @@ DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit" DEFAULT_QUANTIZED_KV_START = 5000 +# Draft strategy defaults for prompt lookup decoding (PLD). +DEFAULT_DRAFT_TYPE = "none" +DRAFT_TYPES = ("none", "ngram-simple", "ngram-mod") +NGRAM_SIMPLE_DEFAULT_SIZE = 3 +NGRAM_GATE_DEFAULT_SIZE = 3 +# Minimum fraction of n-gram windows in the prompt that must recur for +# speculation to be worthwhile. Empirically derived placeholder; tune later. +NGRAM_GATE_THRESHOLD = 0.02 + +# ngram-mod (rolling-hash drafter, port of llama.cpp common_ngram_mod). +# Recommended ngram size is >= 16; collisions dominate at small n. +NGRAM_MOD_DEFAULT_SIZE = 16 +NGRAM_MOD_DEFAULT_TABLE_SIZE = 4 * 1024 * 1024 # 4M entries (~16 MB) +NGRAM_MOD_HASH_MULTIPLIER = 6364136223846793005 # Knuth LCG multiplier +NGRAM_MOD_HASH_MASK = 0xFFFFFFFFFFFFFFFF +NGRAM_MOD_EMPTY = -1 +NGRAM_MOD_CHUNK_GAP = 32 # ingest only when i_last + 32 < cur_len +NGRAM_MOD_OCCUPANCY_RESET = 0.25 +NGRAM_MOD_LOW_ACCEPT_THRESHOLD = 0.5 +NGRAM_MOD_LOW_ACCEPT_STREAK = 3 + + +def _non_trimmable_prompt_cache_types( + model: nn.Module, prompt_cache: Optional[Any] = None +) -> set[str]: + if prompt_cache is None: + model_cache = cache.make_prompt_cache(model) + else: + model_cache = prompt_cache[: len(model.layers)] + if cache.can_trim_prompt_cache(model_cache): + return set() + return {type(c).__name__ for c in model_cache if not c.is_trimmable()} + + +def _clone_prompt_cache(prompt_cache: List[Any]) -> List[Any]: + return copy.deepcopy(prompt_cache) + def str2bool(string): return string.lower() not in ["false", "f"] @@ -219,6 +260,34 @@ def setup_arg_parser(): help="Number of tokens to draft when using speculative decoding.", default=3, ) + parser.add_argument( + "--draft-type", + type=str, + choices=list(DRAFT_TYPES), + default=DEFAULT_DRAFT_TYPE, + help=( + "Draft source for speculative decoding. 'none' disables " + "draft-based speculation (a draft model may still be used via " + "--draft-model). 'ngram-simple' uses prompt-lookup decoding." + ), + ) + parser.add_argument( + "--disable-adaptive-gate", + action="store_true", + help=( + "Disable the adaptive 3-gram repetition gate that skips " + "speculative decoding on prompts unlikely to benefit." + ), + ) + parser.add_argument( + "--ngram-size", + type=int, + default=None, + help=( + "N-gram window size for the ngram drafter. If unset, defaults " + "to 3 for ngram-simple and 16 for ngram-mod." + ), + ) return parser @@ -470,10 +539,430 @@ def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None): n += 1 +def ngram_repeat_score( + tokens: Sequence[int], n: int = NGRAM_GATE_DEFAULT_SIZE +) -> float: + """Fraction of size-``n`` windows in ``tokens`` that recur at least once. + + Used by the adaptive gate to decide if a prompt is repetitive enough for + n-gram speculative decoding to be worthwhile. Returns 0.0 when ``tokens`` + is too short to contain a single window. + """ + if len(tokens) < n + 1: + return 0.0 + seen = set() + repeats = 0 + total = 0 + for i in range(len(tokens) - n + 1): + gram = tuple(tokens[i : i + n]) + total += 1 + if gram in seen: + repeats += 1 + else: + seen.add(gram) + return repeats / total if total else 0.0 + + +@dataclass +class _DraftContext: + """Helpers shared between speculative_generate_step and a DraftStrategy. + + ``prev_tokens_ref`` is a single-element list used as a mutable holder for + the running token tensor that feeds ``logits_processors``. Strategies that + sample with logits processors must read and update ``prev_tokens_ref[0]`` + so the verifier path stays in sync. + """ + + sampler: Callable[[mx.array], mx.array] + logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] + quantize_cache_fn: Callable[[List], None] + prefill_step_size: int + prev_tokens_ref: List[Optional[mx.array]] + + +class DraftStrategy(Protocol): + """Pluggable draft-token source for :func:`speculative_generate_step`. + + Implementations cover both neural draft models and token-only drafters + (e.g. prompt lookup decoding). The verifier path in + ``speculative_generate_step`` is shared across strategies. + """ + + def make_cache(self, model_layers: int, prompt_cache: Optional[List]) -> List: ... + + def prefill(self, y: mx.array, ctx: _DraftContext) -> mx.array: ... + + def propose(self, y: mx.array, n: int, ctx: _DraftContext) -> mx.array: ... + + def rewind(self, num_draft: int, num_accept: int) -> None: ... + + def observe(self, tokens: Sequence[int]) -> None: + """Ingest tokens the strategy may want to learn from (e.g. prompt + tokens before generation, or accepted tokens during generation). + Token-only drafters use this to update their lookup memory; neural + drafters can no-op. Called by ``stream_generate`` before the first + ``propose`` and continuously via :meth:`accept`.""" + ... + + def accept(self, tokens: List[int]) -> None: ... + + def cache_must_be_trimmable(self) -> bool: ... + + def advances_draft_input(self) -> bool: + """Whether the strategy benefits from re-feeding the last accepted + draft token on the next call to :meth:`propose` (true for neural + drafters that must run a forward pass over it). PLD-style drafters + return False.""" + ... + + +class ModelDraftStrategy: + """Draft tokens produced by a small neural draft model.""" + + def __init__(self, draft_model: nn.Module): + self.draft_model = draft_model + self._cache: List = [] + + def make_cache(self, model_layers: int, prompt_cache: Optional[List]) -> List: + if prompt_cache is None: + self._cache = cache.make_prompt_cache(self.draft_model) + else: + self._cache = prompt_cache[model_layers:] + return self._cache + + def prefill(self, y: mx.array, ctx: _DraftContext) -> mx.array: + while y.size > 1: + n_to_process = min(ctx.prefill_step_size, y.size - 1) + self.draft_model(y[:n_to_process][None], cache=self._cache) + ctx.quantize_cache_fn(self._cache) + mx.eval([c.state for c in self._cache]) + y = y[n_to_process:] + mx.clear_cache() + return y + + def propose(self, y: mx.array, n: int, ctx: _DraftContext) -> mx.array: + if n == 0: + return mx.array([], mx.uint32) + # Snapshot prev_tokens so logits processors see the in-flight draft + # tokens during sampling, then restore it before returning so the + # verifier's _step sees the pre-draft state. + prev_before = ctx.prev_tokens_ref[0] + ys = [] + for _ in range(n): + with mx.stream(generation_stream): + logits = self.draft_model(y[None], cache=self._cache) + logits = logits[:, -1:, :] + ctx.quantize_cache_fn(self._cache) + if ctx.logits_processors: + prev = ctx.prev_tokens_ref[0] + prev = mx.concatenate([prev, y]) if prev is not None else y + ctx.prev_tokens_ref[0] = prev + y, _ = _process_and_sample_logits( + prev, logits[:, 0, :], ctx.sampler, ctx.logits_processors + ) + else: + y, _ = _process_and_sample_logits( + None, logits.squeeze(0), ctx.sampler, None + ) + mx.async_eval(y) + ys.append(y) + # Restore prev_tokens so the verifier's _step re-appends from the + # pre-draft state. This decouples prev_tokens accounting from + # draft_strategy so PLD-style drafters don't need to manage it. + ctx.prev_tokens_ref[0] = prev_before + return mx.concatenate(ys) + + def rewind(self, num_draft: int, num_accept: int) -> None: + cache.trim_prompt_cache(self._cache, max(num_draft - num_accept - 1, 0)) + + def observe(self, tokens: Sequence[int]) -> None: # noqa: D401 - protocol noop + # Neural draft model has nothing to learn from raw tokens; its + # state is the KV cache, updated by forward passes. + return None + + def accept(self, tokens: List[int]) -> None: # noqa: D401 - protocol noop + return None + + def cache_must_be_trimmable(self) -> bool: + return True + + def advances_draft_input(self) -> bool: + return True + + +class NgramSimpleStrategy: + """Prompt-lookup decoding: draft tokens are copied from token history. + + Searches the running token history for the last occurrence of the most + recent ``ngram_size`` tokens and returns the up-to-``n`` tokens that + followed that match. No neural forward pass, no KV cache. + """ + + def __init__(self, ngram_size: int = NGRAM_SIMPLE_DEFAULT_SIZE): + if ngram_size < 1: + raise ValueError("ngram_size must be >= 1") + self.ngram_size = ngram_size + self._history: List[int] = [] + + def observe(self, tokens: Sequence[int]) -> None: + self._history.extend(int(t) for t in tokens) + + def make_cache(self, model_layers: int, prompt_cache: Optional[List]) -> List: + return [] + + def prefill(self, y: mx.array, ctx: _DraftContext) -> mx.array: + return y + + def propose(self, y: mx.array, n: int, ctx: _DraftContext) -> mx.array: + if n == 0: + return mx.array([], mx.uint32) + history = self._history + size = self.ngram_size + if len(history) < size: + return mx.array([], mx.uint32) + pattern = history[-size:] + # Search backwards (excluding the trailing pattern itself) for the + # most recent occurrence of `pattern`. + for i in range(len(history) - size - 1, -1, -1): + if history[i : i + size] == pattern: + start = i + size + draft = history[start : start + n] + if not draft: + return mx.array([], mx.uint32) + return mx.array(draft, mx.uint32) + return mx.array([], mx.uint32) + + def rewind(self, num_draft: int, num_accept: int) -> None: + return None + + def accept(self, tokens: List[int]) -> None: + self._history.extend(int(t) for t in tokens) + + def cache_must_be_trimmable(self) -> bool: + return False + + def advances_draft_input(self) -> bool: + return False + + +class NgramModTable: + """Process-global rolling-hash n-gram table (port of llama.cpp common_ngram_mod). + + Stores ``hash(ngram[0..n-1]) -> next_token`` in a fixed-size flat array of + int32. No tag verification: collisions silently overwrite, and bad lookups + just waste a verifier round (the verifier catches them). + + Thread-safety: ``add`` and ``reset`` take a lock. ``get`` reads without a + lock — stale reads are harmless because the verifier rejects bad drafts. + """ + + def __init__( + self, + n: int = NGRAM_MOD_DEFAULT_SIZE, + size: int = NGRAM_MOD_DEFAULT_TABLE_SIZE, + ): + if n < 1: + raise ValueError("n must be >= 1") + if size < 1: + raise ValueError("size must be >= 1") + self._n = n + # Signed int32 storage (token ids fit in 31 bits; -1 is the EMPTY + # sentinel). ``array.array`` keeps this compact in memory and faster + # than a Python list for index/assign. + self._entries = array.array("i", [NGRAM_MOD_EMPTY] * size) + self._used = 0 + self._lock = threading.Lock() + + @property + def n(self) -> int: + return self._n + + @property + def size(self) -> int: + return len(self._entries) + + @property + def used(self) -> int: + return self._used + + @property + def occupancy(self) -> float: + return self._used / self.size if self.size else 0.0 + + def _idx(self, tokens: Sequence[int]) -> int: + """Hash ``tokens[0..n-1]`` to a table slot. Matches llama.cpp idx().""" + res = 0 + for i in range(self._n): + res = (res * NGRAM_MOD_HASH_MULTIPLIER + tokens[i]) & NGRAM_MOD_HASH_MASK + return res % len(self._entries) + + def add(self, tokens: Sequence[int]) -> None: + """Insert ``tokens[0..n-1] -> tokens[n]``. + + ``tokens`` must have at least ``n + 1`` elements. + """ + i = self._idx(tokens) + next_tok = int(tokens[self._n]) + with self._lock: + if self._entries[i] == NGRAM_MOD_EMPTY: + self._used += 1 + self._entries[i] = next_tok + + def get(self, tokens: Sequence[int]) -> int: + """Return the next-token stored at hash(tokens[0..n-1]) or EMPTY.""" + i = self._idx(tokens) + return self._entries[i] + + def reset(self) -> None: + with self._lock: + for i in range(len(self._entries)): + self._entries[i] = NGRAM_MOD_EMPTY + self._used = 0 + + +class NgramModStrategy: + """Per-request strategy backed by a (potentially shared) :class:`NgramModTable`. + + Maintains request-local state — token history, ingest cursor, and + acceptance-streak counters — while delegating the actual lookup memory to + the shared table. Multiple concurrent requests can hold their own + ``NgramModStrategy`` pointing at the same ``NgramModTable``. + """ + + def __init__( + self, + table: NgramModTable, + ngram_size: Optional[int] = None, + n_min: int = 0, + ): + if ngram_size is not None and ngram_size != table.n: + raise ValueError( + f"ngram_size={ngram_size} does not match table.n={table.n}; " + "rebuild the table or pass a matching strategy ngram_size." + ) + self.table = table + self.ngram_size = table.n + self.n_min = max(0, int(n_min)) + self._history: List[int] = [] + self._i_last = 0 + self._n_draft_last = 0 + self._n_low = 0 + + def make_cache(self, model_layers: int, prompt_cache: Optional[List]) -> List: + return [] + + def prefill(self, y: mx.array, ctx: _DraftContext) -> mx.array: + return y + + def observe(self, tokens: Sequence[int]) -> None: + """Ingest prompt tokens. Mirrors common_speculative_state_ngram_mod::begin.""" + self._history.extend(int(t) for t in tokens) + self._i_last = 0 + self._n_draft_last = 0 + n = self.ngram_size + if len(self._history) < n + 1: + return + # Insert every ngram [i, i+n+1) for i in [0, len - n). + for i in range(len(self._history) - n): + self.table.add(self._history[i : i + n + 1]) + self._i_last = len(self._history) - n + # Begin-time occupancy reset (matches llama.cpp's f_thold = 0.25). + if self.table.occupancy > NGRAM_MOD_OCCUPANCY_RESET: + logging.info( + "ngram-mod occupancy %.3f exceeds %.2f — resetting shared table", + self.table.occupancy, + NGRAM_MOD_OCCUPANCY_RESET, + ) + self.table.reset() + + def propose(self, y: mx.array, n: int, ctx: _DraftContext) -> mx.array: + """Chunked ingest of new history, then chained hash lookups.""" + self._n_draft_last = 0 + if n == 0: + return mx.array([], mx.uint32) + size = self.ngram_size + cur_len = len(self._history) + if cur_len < size: + return mx.array([], mx.uint32) + + # Chunked insert of newly-accepted tokens (mirrors llama.cpp's `+32` + # gate). Insert ngrams [i_last, cur_len - size). + if self._i_last + NGRAM_MOD_CHUNK_GAP < cur_len: + for i in range(self._i_last, cur_len - size): + self.table.add(self._history[i : i + size + 1]) + self._i_last = cur_len - size + + # Form the initial lookup pattern: last (size) tokens of history. + # The C++ code uses prompt_tgt[cur_len - n + 1 ..] plus id_last, which + # is exactly the last ``size`` tokens of the running history since + # accept() appended id_last for us. + pattern: List[int] = list(self._history[cur_len - size :]) + drafted: List[int] = [] + for i in range(n): + tok = self.table.get(pattern) + if tok == NGRAM_MOD_EMPTY: + if i < self.n_min: + return mx.array([], mx.uint32) + break + drafted.append(tok) + # Slide window: drop oldest, append the just-looked-up token. + pattern.pop(0) + pattern.append(tok) + if not drafted: + return mx.array([], mx.uint32) + self._n_draft_last = len(drafted) + return mx.array(drafted, mx.uint32) + + def rewind(self, num_draft: int, num_accept: int) -> None: + return None + + def accept(self, tokens: List[int]) -> None: + """Append accepted tokens to history. Track acceptance fraction and + trigger a table reset on sustained low acceptance.""" + self._history.extend(int(t) for t in tokens) + if self._n_draft_last > 0: + # tokens may include the verifier's bonus correction; only the + # first n_draft_last positions correspond to drafted tokens that + # the verifier was asked to confirm. + n_accepted_drafts = min(self._n_draft_last, max(0, len(tokens) - 1)) + f_acc = n_accepted_drafts / self._n_draft_last + if f_acc < NGRAM_MOD_LOW_ACCEPT_THRESHOLD: + self._n_low += 1 + if self._n_low >= NGRAM_MOD_LOW_ACCEPT_STREAK: + logging.info( + "ngram-mod low-acceptance streak (%d) — resetting " + "shared table", + self._n_low, + ) + self.table.reset() + self._n_low = 0 + else: + self._n_low = 0 + + def cache_must_be_trimmable(self) -> bool: + return False + + def advances_draft_input(self) -> bool: + return False + + +def _process_and_sample_logits(tokens, logits, sampler, logits_processors): + """Module-level twin of the inner ``_process_and_sample`` helper. + + Kept here so :class:`ModelDraftStrategy` can sample without rebinding the + closures used inside :func:`speculative_generate_step`. + """ + if logits_processors: + for processor in logits_processors: + logits = processor(tokens, logits) + logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) + return sampler(logprobs), logprobs + + def speculative_generate_step( prompt: mx.array, model: nn.Module, - draft_model: nn.Module, + draft_strategy: DraftStrategy, *, num_draft_tokens: int = 2, max_tokens: int = 256, @@ -491,7 +980,9 @@ def speculative_generate_step( 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. + draft_strategy (DraftStrategy): The draft-token source. Use + :class:`ModelDraftStrategy` for neural draft models or + :class:`NgramSimpleStrategy` for prompt-lookup 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 @@ -516,21 +1007,28 @@ def speculative_generate_step( """ y = prompt.astype(mx.uint32) - prev_tokens = None + # Mutable holder so the draft strategy and the verifier _step can share + # the running token tensor used by logits processors. + prev_tokens_ref: List[Optional[mx.array]] = [None] - # Create the KV cache for generation + # Create the KV cache for generation. The strategy owns the draft-side + # cache (which may be empty for token-only drafters such as PLD). 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) :] + draft_strategy.make_cache(len(model.layers), prompt_cache) + requires_trimmable_cache = draft_strategy.cache_must_be_trimmable() + use_cache_snapshot = False if not cache.can_trim_prompt_cache(model_cache): types = {type(c).__name__ for c in model_cache if not c.is_trimmable()} - raise ValueError( - f"Speculative decoding requires a trimmable prompt cache " f"(got {types})." - ) + if requires_trimmable_cache: + raise ValueError( + "Speculative decoding requires a trimmable prompt cache " + f"(got {types})." + ) + use_cache_snapshot = True sampler = sampler or (lambda x: mx.argmax(x, axis=-1)) @@ -541,6 +1039,14 @@ def speculative_generate_step( kv_bits=kv_bits, ) + ctx = _DraftContext( + sampler=sampler, + logits_processors=logits_processors, + quantize_cache_fn=quantize_cache_fn, + prefill_step_size=prefill_step_size, + prev_tokens_ref=prev_tokens_ref, + ) + def _process_and_sample(tokens, logits): if logits_processors: for processor in logits_processors: @@ -557,17 +1063,14 @@ def _step(model, cache, y, n_predict=1): 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.concatenate([prev_tokens, y]) - if prev_tokens is not None - else y - ) - y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :]) + prev = prev_tokens_ref[0] + prev = mx.concatenate([prev, y]) if prev is not None else y + prev_tokens_ref[0] = prev + y, logprobs = _process_and_sample(prev, logits[:, i, :]) out_y.append(y) out_logprobs.append(logprobs) return mx.concatenate(out_y, axis=0), mx.concatenate( @@ -588,20 +1091,22 @@ def _prefill(model, cache, 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)) + draft_strategy.rewind(num_draft, num_accept) - 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) + def _restore_cache(snapshot): + model_cache[:] = _clone_prompt_cache(snapshot) + + def _replay_cache_inputs(tokens: mx.array): + if tokens.size == 0: + return + with mx.stream(generation_stream): + model(tokens[None], cache=model_cache) + quantize_cache_fn(model_cache) + mx.eval([c.state for c in model_cache]) + mx.clear_cache() with mx.stream(generation_stream): - draft_y = _prefill(draft_model, draft_cache, y) + draft_y = draft_strategy.prefill(y, ctx) y = _prefill(model, model_cache, y) ntoks = 0 @@ -610,48 +1115,86 @@ def _draft_generate(y, num_draft): n = 0 try: while True: + cache_snapshot = _clone_prompt_cache(model_cache) if use_cache_snapshot else None + step_y = y 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] + draft_tokens = draft_strategy.propose(draft_y, num_draft, ctx) + # The strategy may return fewer than ``num_draft`` tokens (e.g. + # PLD with no matching n-gram). Re-sync num_draft so the verifier + # only processes what the strategy actually produced. + num_draft = int(draft_tokens.size) + # Note: the strategy is responsible for restoring prev_tokens to + # its pre-propose state if it touched it (see ModelDraftStrategy). 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 + accepted: List[int] = [] while n < num_draft: tn, dtn, lpn = tokens[n], draft_tokens[n], logprobs[n] if tn != dtn: break n += 1 ntoks += 1 + accepted.append(tn) yield tn, lpn, True if ntoks == max_tokens: break if ntoks < max_tokens: ntoks += 1 + accepted.append(tokens[n]) yield tokens[n], logprobs[n], False + # Let the strategy observe the accepted tokens (e.g. PLD updates + # its token history). Includes the verifier's bonus/correction + # token at position ``n``. + if accepted: + draft_strategy.accept(accepted) + 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: + # If we accepted all the draft tokens, neural drafters need the + # last accepted draft re-fed so the draft model can process it. + # Token-only drafters (PLD) skip this — they don't need it. + extends_draft = ( + num_draft > 0 + and n == num_draft + and draft_strategy.advances_draft_input() + ) + if extends_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) + if prev_tokens_ref[0] is not None: + # Base trim removes the rejected suffix that the verifier + # appended to prev_tokens but never yielded. Neural drafters + # trim one extra when fully accepted to balance the re-fed + # last draft token above. + trim_amount = num_draft - n + if extends_draft: + trim_amount += 1 + if trim_amount > 0: + prev = prev_tokens_ref[0] + prev_tokens_ref[0] = prev[:-trim_amount] + if use_cache_snapshot: + _restore_cache(cache_snapshot) + replay_tokens = mx.concatenate([step_y, mx.array(draft_tokens[:n], mx.uint32)]) + _replay_cache_inputs(replay_tokens) + draft_strategy.rewind(num_draft, n) + else: + _rewind_cache(num_draft, n) finally: - _rewind_cache(num_draft, n) + if use_cache_snapshot: + draft_strategy.rewind(num_draft, n) + else: + _rewind_cache(num_draft, n) def stream_generate( @@ -660,6 +1203,11 @@ def stream_generate( prompt: Union[str, mx.array, List[int]], max_tokens: int = 256, draft_model: Optional[nn.Module] = None, + draft_type: Optional[str] = None, + draft_strategy: Optional[DraftStrategy] = None, + disable_adaptive_gate: bool = False, + ngram_size: Optional[int] = None, + ngram_mod_table: Optional[NgramModTable] = None, **kwargs, ) -> Generator[GenerationResponse, None, None]: """ @@ -672,9 +1220,26 @@ def stream_generate( integer tokens. max_tokens (int): The maximum number of tokens to generate. Default: ``256``. - draft_model (Optional[nn.Module]): An optional draft model. If provided - then speculative decoding is used. The draft model must use the same - tokenizer as the main model. Default: ``None``. + draft_model (Optional[nn.Module]): An optional neural draft model. + If provided, speculative decoding with a draft model is used. The + draft model must share the main tokenizer. Default: ``None``. + draft_type (Optional[str]): One of ``"none"``, ``"ngram-simple"``, + ``"ngram-mod"``. Selects a token-only draft source (prompt-lookup + decoding). Ignored when ``draft_model`` or ``draft_strategy`` is + set. Default: ``None``. + draft_strategy (Optional[DraftStrategy]): A pre-built strategy + instance, used as-is. Takes precedence over ``draft_type``. The + server passes shared-state strategies (e.g. ``ngram-mod``) here. + Default: ``None``. + disable_adaptive_gate (bool): Disable the 3-gram repetition gate + that suppresses ngram speculation on non-repetitive prompts. + Default: ``False``. + ngram_size (Optional[int]): N-gram window size. If ``None``, uses + per-strategy defaults (3 for ngram-simple, 16 for ngram-mod). + ngram_mod_table (Optional[NgramModTable]): A pre-built shared + ngram-mod table to use for ``draft_type="ngram-mod"``. Servers + should construct one at startup and pass it here. When ``None``, + a fresh per-call table is constructed (CLI/one-shot usage). kwargs: The remaining options get passed to :func:`generate_step`. See :func:`generate_step` for more details. @@ -698,7 +1263,72 @@ def stream_generate( kwargs["max_tokens"] = max_tokens - if draft_model is None: + # Resolve the draft strategy. Precedence (highest first): + # 1. ``draft_model`` — wrap as ModelDraftStrategy + # 2. ``draft_strategy`` — caller-supplied instance (used as-is) + # 3. ``draft_type`` — construct a fresh ngram strategy + # The adaptive gate only applies when we're constructing an ngram + # strategy ourselves; a caller passing a pre-built strategy is assumed + # to know what they're doing. + if draft_model is not None: + draft_strategy = ModelDraftStrategy(draft_model) + elif draft_strategy is not None: + pass # use caller-supplied instance as-is + elif draft_type and draft_type != "none": + if draft_type not in DRAFT_TYPES: + raise ValueError( + f"Unknown draft_type {draft_type!r}; expected one of {DRAFT_TYPES}." + ) + use_speculation = True + if not disable_adaptive_gate: + score = ngram_repeat_score(prompt.tolist()) + if score < NGRAM_GATE_THRESHOLD: + use_speculation = False + logging.info( + "Adaptive ngram gate disabled %s speculation " + "(repeat score %.4f < %.4f).", + draft_type, + score, + NGRAM_GATE_THRESHOLD, + ) + if use_speculation: + if draft_type == "ngram-simple": + size = ( + ngram_size if ngram_size is not None else NGRAM_SIMPLE_DEFAULT_SIZE + ) + strategy = NgramSimpleStrategy(ngram_size=size) + strategy.observe(prompt.tolist()) + draft_strategy = strategy + elif draft_type == "ngram-mod": + size = ngram_size if ngram_size is not None else NGRAM_MOD_DEFAULT_SIZE + if size < NGRAM_MOD_DEFAULT_SIZE: + logging.warning( + "ngram-mod with n=%d is below recommended %d; " + "draft quality is likely to suffer (see " + "llama.cpp PR 19164).", + size, + NGRAM_MOD_DEFAULT_SIZE, + ) + table = ngram_mod_table or NgramModTable(n=size) + strategy = NgramModStrategy(table) + strategy.observe(prompt.tolist()) + draft_strategy = strategy + + if ( + draft_strategy is not None + and draft_strategy.cache_must_be_trimmable() + and ( + incompatible_cache_types := _non_trimmable_prompt_cache_types( + model, kwargs.get("prompt_cache") + ) + ) + ): + raise ValueError( + "Speculative decoding requires a trimmable prompt cache " + f"(got {incompatible_cache_types})." + ) + + if draft_strategy is None: kwargs.pop("num_draft_tokens", None) token_generator = generate_step(prompt, model, **kwargs) # from_draft always false for non-speculative generation @@ -709,7 +1339,7 @@ def stream_generate( kwargs.pop("max_kv_size", None) kwargs.pop("prompt_progress_callback", None) token_generator = speculative_generate_step( - prompt, model, draft_model, **kwargs + prompt, model, draft_strategy, **kwargs ) with wired_limit(model, [generation_stream]): tic = time.perf_counter() @@ -2104,7 +2734,10 @@ def main(): kv_group_size=args.kv_group_size, quantized_kv_start=args.quantized_kv_start, draft_model=draft_model, + draft_type=args.draft_type, num_draft_tokens=args.num_draft_tokens, + disable_adaptive_gate=args.disable_adaptive_gate, + ngram_size=args.ngram_size, ) if not args.verbose: print(response) diff --git a/mlx_lm/server.py b/mlx_lm/server.py index c000f4c10..95b68ef0f 100644 --- a/mlx_lm/server.py +++ b/mlx_lm/server.py @@ -15,7 +15,7 @@ from pathlib import Path from queue import Empty as QueueEmpty from queue import Queue -from threading import Thread +from threading import Lock, Thread from typing import ( Any, Callable, @@ -34,7 +34,12 @@ from ._version import __version__ from .generate import ( + DEFAULT_DRAFT_TYPE, + DRAFT_TYPES, + NGRAM_MOD_DEFAULT_SIZE, + NGRAM_MOD_DEFAULT_TABLE_SIZE, BatchGenerator, + NgramModTable, SequenceStateMachine, stream_generate, ) @@ -192,6 +197,9 @@ class GenerationArguments: top_logprobs: int seed: Optional[int] chat_template_kwargs: Optional[Dict[str, Any]] + draft_type: Optional[str] = None + disable_adaptive_gate: bool = False + ngram_size: Optional[int] = None @dataclass @@ -321,6 +329,32 @@ def __init__(self, cli_args: argparse.Namespace): if cli_args.chat_template: self._tokenizer_config["chat_template"] = cli_args.chat_template + # Shared ngram-mod hash table (allocated lazily on first request that + # uses draft_type=ngram-mod). One per server process — survives across + # requests, reset only by the strategy on quality collapse. + self._ngram_mod_table: Optional[NgramModTable] = None + self._ngram_mod_lock = Lock() + + def get_ngram_mod_table(self, n: int) -> NgramModTable: + """Return the process-shared ngram-mod table, allocating it on first + use. If a table exists with a different ``n``, it is rebuilt — the + hash slots are not portable across different ngram sizes.""" + with self._ngram_mod_lock: + if self._ngram_mod_table is None or self._ngram_mod_table.n != n: + size = getattr( + self.cli_args, + "ngram_mod_table_size", + NGRAM_MOD_DEFAULT_TABLE_SIZE, + ) + logging.info( + "Allocating ngram-mod shared table: n=%d size=%d (~%.1f MB)", + n, + size, + (size * 4) / (1024 * 1024), + ) + self._ngram_mod_table = NgramModTable(n=n, size=size) + return self._ngram_mod_table + def _load(self, model_path, adapter_path=None, draft_model_path=None): if self.is_distributed and ( adapter_path is not None or draft_model_path is not None @@ -972,6 +1006,17 @@ def progress(tokens_processed, tokens_total): if self.model_provider.draft_model is not None: cache += make_prompt_cache(self.model_provider.draft_model) + # For ngram-mod, fetch the process-shared hash table. The table + # is sized by ngram_size (per-request override or CLI default). + ngram_mod_table = None + if args.draft_type == "ngram-mod": + n = ( + args.ngram_size + if args.ngram_size is not None + else NGRAM_MOD_DEFAULT_SIZE + ) + ngram_mod_table = self.model_provider.get_ngram_mod_table(n) + # Process the prompt and generate tokens for gen in stream_generate( model=model, @@ -983,6 +1028,10 @@ def progress(tokens_processed, tokens_total): prompt_cache=cache, draft_model=draft_model, num_draft_tokens=args.num_draft_tokens, + draft_type=args.draft_type, + disable_adaptive_gate=args.disable_adaptive_gate, + ngram_size=args.ngram_size, + ngram_mod_table=ngram_mod_table, prompt_progress_callback=progress, prefill_step_size=self.cli_args.prefill_step_size, ): @@ -1165,6 +1214,18 @@ def do_POST(self): self.num_draft_tokens = self.body.get( "num_draft_tokens", self.response_generator.cli_args.num_draft_tokens ) + self.draft_type = self.body.get( + "draft_type", + getattr(self.response_generator.cli_args, "draft_type", DEFAULT_DRAFT_TYPE), + ) + self.disable_adaptive_gate = self.body.get( + "disable_adaptive_gate", + getattr(self.response_generator.cli_args, "disable_adaptive_gate", False), + ) + self.ngram_size = self.body.get( + "ngram_size", + getattr(self.response_generator.cli_args, "ngram_size", None), + ) self.adapter = self.body.get("adapters", None) self.max_tokens = self.body.get("max_completion_tokens", None) if self.max_tokens is None: @@ -1247,6 +1308,12 @@ def validate_model_parameters(self): self._validate("xtc_threshold", float, min_val=0, max_val=1) self._validate("requested_model", str) self._validate("adapter", str, optional=True) + if self.draft_type is not None and self.draft_type not in DRAFT_TYPES: + raise ValueError( + f"draft_type must be one of {DRAFT_TYPES}, got {self.draft_type!r}" + ) + self._validate("disable_adaptive_gate", bool) + self._validate("ngram_size", int, min_val=1, optional=True) self._validate("seed", int, optional=True) self._validate("logit_bias", dict, optional=True) @@ -1403,6 +1470,9 @@ def handle_completion(self, request: CompletionRequest, stop_words: List[str]): top_logprobs=self.top_logprobs, seed=self.seed, chat_template_kwargs=self.chat_template_kwargs, + draft_type=self.draft_type, + disable_adaptive_gate=self.disable_adaptive_gate, + ngram_size=self.ngram_size, ) # Keep connection allive during long prompt processing (and also log @@ -1790,6 +1860,45 @@ def main(): help="Number of tokens to draft when using speculative decoding.", default=3, ) + parser.add_argument( + "--draft-type", + type=str, + choices=list(DRAFT_TYPES), + default=DEFAULT_DRAFT_TYPE, + help=( + "Default draft source for speculative decoding when no neural " + "draft model is configured. Per-request overrides via " + "'draft_type' in the request body." + ), + ) + parser.add_argument( + "--disable-adaptive-gate", + action="store_true", + help=( + "Disable the 3-gram repetition gate that suppresses ngram " + "speculative decoding on non-repetitive prompts." + ), + ) + parser.add_argument( + "--ngram-size", + type=int, + default=None, + help=( + "Default n-gram window size for the ngram drafter. If unset, " + "defaults to 3 for ngram-simple and 16 for ngram-mod. " + "Per-request overrides via 'ngram_size' in the request body." + ), + ) + parser.add_argument( + "--ngram-mod-table-size", + type=int, + default=NGRAM_MOD_DEFAULT_TABLE_SIZE, + help=( + "Number of entries in the shared ngram-mod hash table. " + "Each entry is 4 bytes. Default: 4194304 (~16 MB). " + "Allocated once at server startup, shared across requests." + ), + ) parser.add_argument( "--trust-remote-code", action="store_true", diff --git a/tests/test_chat.py b/tests/test_chat.py index b4ddcf3ae..438fb6e4d 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -1,3 +1,4 @@ +import copy import argparse import unittest from unittest.mock import MagicMock, patch @@ -74,7 +75,13 @@ def test_system_prompt_in_messages( # Mock the model and tokenizer mock_model = MagicMock() mock_tokenizer = MagicMock() - mock_tokenizer.apply_chat_template.return_value = "processed_prompt" + captured_messages = [] + + def apply_chat_template(messages, *args, **kwargs): + captured_messages.append(copy.deepcopy(messages)) + return "processed_prompt" + + mock_tokenizer.apply_chat_template.side_effect = apply_chat_template mock_load.return_value = (mock_model, mock_tokenizer) # Mock prompt cache @@ -104,9 +111,7 @@ def test_system_prompt_in_messages( # Verify that apply_chat_template was called with system prompt mock_tokenizer.apply_chat_template.assert_called() - call_args = mock_tokenizer.apply_chat_template.call_args[0][ - 0 - ] # First positional arg (messages) + call_args = captured_messages[-1] # Check that the messages contain both system and user messages self.assertEqual(len(call_args), 2) @@ -133,7 +138,13 @@ def test_no_system_prompt_in_messages( # Mock the model and tokenizer mock_model = MagicMock() mock_tokenizer = MagicMock() - mock_tokenizer.apply_chat_template.return_value = "processed_prompt" + captured_messages = [] + + def apply_chat_template(messages, *args, **kwargs): + captured_messages.append(copy.deepcopy(messages)) + return "processed_prompt" + + mock_tokenizer.apply_chat_template.side_effect = apply_chat_template mock_load.return_value = (mock_model, mock_tokenizer) # Mock prompt cache @@ -161,9 +172,7 @@ def test_no_system_prompt_in_messages( # Verify that apply_chat_template was called without system prompt mock_tokenizer.apply_chat_template.assert_called() - call_args = mock_tokenizer.apply_chat_template.call_args[0][ - 0 - ] # First positional arg (messages) + call_args = captured_messages[-1] # Check that the messages contain only user message self.assertEqual(len(call_args), 1) diff --git a/tests/test_generate.py b/tests/test_generate.py index 67fc939fe..7c9a4da6b 100644 --- a/tests/test_generate.py +++ b/tests/test_generate.py @@ -3,6 +3,7 @@ import random import unittest from typing import List +from unittest.mock import MagicMock, patch import mlx.core as mx @@ -15,8 +16,9 @@ generate_step, stream_generate, ) -from mlx_lm.models.cache import KVCache, RotatingKVCache +from mlx_lm.models.cache import ArraysCache, KVCache, RotatingKVCache from mlx_lm.sample_utils import make_logits_processors, make_sampler +from mlx_lm.tokenizer_utils import TokenizerWrapper from mlx_lm.utils import load @@ -179,6 +181,207 @@ def progress_callback(processed: int, total: int) -> None: num_embeddings / prefill_step_size < num_prompt_processing_callbacks ) + +class TestNgramSpeculationBehavior(unittest.TestCase): + + def test_stream_generate_ngram_supports_arrays_cache(self): + class CycleModel: + def __init__(self): + self.layers = [object()] + + def __call__(self, y, cache): + history = cache[0][0] + if history is None: + cache[0][0] = y + else: + cache[0][0] = mx.concatenate([history, y], axis=1) + + logits = mx.full((1, y.shape[1], 8), -1e9, dtype=mx.float32) + for i, token in enumerate(y[0].tolist()): + logits[0, i, (token % 3) + 1] = 0.0 + return logits + + class DummyTokenizer: + eos_token_id = 99 + eos_token_ids = {99} + bos_token = None + chat_template = None + clean_up_tokenization_spaces = False + + def decode(self, tokens): + return "".join(str(t) for t in tokens) + + def get_vocab(self): + return {} + + class DummyDetokenizer: + def __init__(self, tokenizer): + self.text = "" + self.tokens = [] + self.offset = 0 + + def reset(self): + self.text = "" + self.tokens = [] + self.offset = 0 + + def add_token(self, token): + self.tokens.append(token) + self.text += str(token) + + def finalize(self): + return None + + @property + def last_segment(self): + segment = self.text[self.offset :] + self.offset = len(self.text) + return segment + + model = CycleModel() + tokenizer = TokenizerWrapper( + DummyTokenizer(), detokenizer_class=DummyDetokenizer + ) + prompt = mx.array([1, 2, 3, 1, 2, 3], dtype=mx.uint32) + + baseline = [ + resp.token + for resp in stream_generate( + model=model, + tokenizer=tokenizer, + prompt=prompt, + max_tokens=5, + prompt_cache=[ArraysCache(size=1)], + ) + ] + speculative = list( + stream_generate( + model=model, + tokenizer=tokenizer, + prompt=prompt, + max_tokens=5, + draft_type="ngram-simple", + ngram_size=1, + disable_adaptive_gate=True, + prompt_cache=[ArraysCache(size=1)], + ) + ) + + self.assertEqual([resp.token for resp in speculative], baseline) + self.assertTrue(any(resp.from_draft for resp in speculative)) + + def test_stream_generate_rejects_untrimmable_cache_for_trim_required_strategy(self): + class UntrimmableCache: + def is_trimmable(self): + return False + + model = MagicMock() + model.layers = [object(), object()] + tokenizer = MagicMock() + tokenizer.eos_token_id = 0 + tokenizer.eos_token_ids = {0} + tokenizer.bos_token = None + tokenizer.chat_template = None + tokenizer.encode.return_value = [1, 2, 3, 1, 2, 3] + tokenizer.decode.return_value = "" + draft_strategy = MagicMock() + draft_strategy.cache_must_be_trimmable.return_value = True + + with self.assertRaisesRegex(ValueError, "UntrimmableCache"): + list( + stream_generate( + model=model, + tokenizer=tokenizer, + prompt="repeat repeat repeat repeat", + max_tokens=1, + draft_strategy=draft_strategy, + prompt_cache=[KVCache(), UntrimmableCache()], + ) + ) + + @patch("mlx_lm.generate.logging.info") + @patch("mlx_lm.generate.generate_step") + def test_stream_generate_logs_when_adaptive_gate_skips_ngram( + self, mock_generate_step, mock_info + ): + model = MagicMock() + model.layers = [object()] + tokenizer = MagicMock() + tokenizer.eos_token_id = 0 + tokenizer.eos_token_ids = {99} + tokenizer.bos_token = None + tokenizer.chat_template = None + tokenizer.encode.return_value = [1, 2, 3, 4] + tokenizer.decode.return_value = "x" + mock_generate_step.return_value = iter( + [(7, mx.array([0.0], dtype=mx.float32))] + ) + + with patch("mlx_lm.generate.wired_limit", return_value=unittest.mock.MagicMock()): + responses = list( + stream_generate( + model=model, + tokenizer=tokenizer, + prompt="one two three four", + max_tokens=1, + draft_type="ngram-simple", + ) + ) + + self.assertEqual(len(responses), 1) + mock_info.assert_called_once() + self.assertIn("Adaptive ngram gate disabled", mock_info.call_args.args[0]) + + +class TestGenerateCli(unittest.TestCase): + + @classmethod + def setUpClass(cls): + cls.HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit" + cls.model, cls.tokenizer = load(cls.HF_MODEL_PATH) + cls.model.set_dtype(mx.float32) + + @patch("builtins.print") + @patch("mlx_lm.generate.generate") + @patch("mlx_lm.generate.make_sampler") + @patch("mlx_lm.generate.load") + def test_main_forwards_ngram_flags( + self, mock_load, mock_make_sampler, mock_generate, mock_print + ): + from mlx_lm.generate import main + + mock_model = MagicMock() + mock_tokenizer = MagicMock() + mock_tokenizer.has_chat_template = False + mock_tokenizer.encode.side_effect = lambda text, **_: [1, 2, 3] + mock_tokenizer.eos_token_ids = [0] + mock_load.return_value = (mock_model, mock_tokenizer) + mock_make_sampler.return_value = MagicMock() + mock_generate.return_value = "ok" + + with patch( + "sys.argv", + [ + "generate.py", + "--prompt", + "hello", + "--draft-type", + "ngram-simple", + "--ngram-size", + "5", + "--disable-adaptive-gate", + "--num-draft-tokens", + "4", + ], + ): + main() + + _, kwargs = mock_generate.call_args + self.assertEqual(kwargs["draft_type"], "ngram-simple") + self.assertEqual(kwargs["ngram_size"], 5) + self.assertTrue(kwargs["disable_adaptive_gate"]) + self.assertEqual(kwargs["num_draft_tokens"], 4) + def test_batch_matches_single(self): prompts = [