diff --git a/CHANGELOG.md b/CHANGELOG.md index c0da328..a8640fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,12 +6,28 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +## [0.3.2] - 2026-08-13 + ### Changed - `init`, `sync-models`, and `pull` now pair a sibling drafter (gemma4 assistant, DSpark, DFlash) into the model it serves as that model's `draft_gguf`, which turns speculative decoding on. `sync-models` also adds the key to a model the config already carries. Before, you added it by hand. +- Muse-Glimmer DFlash drafts the block its checkpoint was trained for, which + is 16 tokens per round on the current one, instead of 2. On a machine with + no NAX tile a verify step costs about the same for 16 rows as for 8, so a + full block accepts more tokens for the same forward. + +### Fixed + +- `--draft-block-size N` could only lower the draft depth. Each drafter froze + its trained depth to the depth it loaded with, so a deeper request did + nothing at all, and said nothing. A drafter now carries its trained depth + apart from the depth it drafts by default, a request up to that depth is + honored, and a deeper one clamps and warns. `gmlx run`, `gmlx chat` and + `gmlx serve` (`--draft-block-size`, `GMLX_DRAFT_BLOCK_SIZE`) all resolve the + depth the same way now. ## [0.3.1] - 2026-08-12 diff --git a/docs/cli.md b/docs/cli.md index 5bfcad4..0964259 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -142,7 +142,7 @@ with a warning); `--no-speculative`/`--no-mtp` forces it off. | `--speculative` / `--mtp` | Force MTP speculative decoding on. Native-head models (qwen3.5/3.6 `nextn`) need no companion; gemma4 and muse-glimmer need `--draft-gguf`. Native heads are auto-enabled without this. Use it to force the path when a sampler flag would otherwise defer. | | `--no-speculative` / `--no-mtp` | Disable MTP. Overrides the native-head auto-enable and config `speculative: true`. | | `--draft-gguf PATH` | Separate assistant-drafter GGUF (gemma4 two-GGUF MTP shape, a muse-glimmer DFlash drafter, or a deepseek4 DSpark/MTP sidecar - gmlx `deepseek4-dspark`, llama.cpp `dflash`, or legacy `deepseek4_mtp_support`); implies `--speculative` (same as `serve`). A sidecar in the target's directory is autodetected without the flag. | -| `--draft-block-size N` | Override the MTP draft block size. | +| `--draft-block-size N` | MTP draft tokens per round. Raises or lowers the drafter's own default, up to the deepest block it can produce. | Speculative generation takes only `--temp`/`--top-p`/`--top-k`/`--min-p` plus a baked `--system-prompt`; mlx-vlm's verify walk has no stop/bias/penalty/KV @@ -542,7 +542,7 @@ gmlx serve Qwen3.6-27B-Q4_K_S.gguf --speculative | `--hf-source REPO` | Processor/config override for a single VLM model (rarely needed). | | `--speculative` | Serve a single positional model with MTP (native-head qwen3.5/3.6; gemma4 also needs `--draft-gguf`). | | `--draft-gguf PATH` | Companion drafter GGUF for assistant-shape MTP (gemma4); implies `--speculative`. | -| `--draft-block-size N` | MTP draft tokens per round (analogous to llama-server `--spec-draft-n-max`). Default: the drafter's own block size; muse-glimmer defaults to 2 drafts and caps N at the loaded block (raise it with `GMLX_MUSE_DFLASH_BLOCK` at load). Also via `GMLX_DRAFT_BLOCK_SIZE`. | +| `--draft-block-size N` | MTP draft tokens per round (analogous to llama-server `--spec-draft-n-max`). Raises or lowers the drafter's own default, up to the deepest block it can produce (a deeper request clamps, with a warning). Also via `GMLX_DRAFT_BLOCK_SIZE`. | | `--adapter PATH` | GGUF LoRA adapter applied live over a single positional model at load (text only, not `--mmproj`/`--speculative`). In config mode set `adapter:` per model instead. | | `--stream-cpu` | Run a single positional model entirely on the CPU device: the over-RAM MoE path, same semantics as [`run --stream-cpu`](#loading). In config mode set `stream: cpu` per model instead; see [server-config.md](server-config.md#models). | | `--stream-experts` | Routed-expert stacks stream from disk while the every-token layers and KV cache stay on GPU; the decode feeder (default) serves decode from a wired expert arena and makes this the faster placement once warm. Config mode: `stream: experts`. Mutually exclusive with `--stream-cpu`. | diff --git a/gmlx/cli.py b/gmlx/cli.py index d02d693..1d26b20 100644 --- a/gmlx/cli.py +++ b/gmlx/cli.py @@ -241,7 +241,8 @@ def add_speculative_args(ap: argparse.ArgumentParser) -> None: type=int, default=None, metavar="N", - help="Override the MTP draft block size.", + help="MTP draft tokens per round. Raises or lowers the drafter's own " + "default, up to the deepest block it can produce.", ) ap.add_argument( "--stochastic-mtp", diff --git a/gmlx/deepseek_v4_dspark.py b/gmlx/deepseek_v4_dspark.py index db37f12..bfea595 100644 --- a/gmlx/deepseek_v4_dspark.py +++ b/gmlx/deepseek_v4_dspark.py @@ -57,6 +57,7 @@ import mlx.nn as nn from . import deepseek_v4_model as v4 +from .drafter_protocol import native_block_size from .deepseek_v4_hyper_connection import HyperHead @@ -71,7 +72,8 @@ def _env_float(name: str, default: float) -> float: class DeepseekV4DSparkConfig: """``text`` is the target's ModelArgs with ``compress_ratios`` POST-init extended by one 0 per stage. ``block_size`` is the engine block TOTAL - (drafts + bonus), at most ``draft_len + 1``.""" + (drafts + bonus), at most ``draft_len + 1``. ``native_block_size`` is the + deepest block the stages can produce, also at most ``draft_len + 1``.""" text: Any n_stages: int = 3 @@ -80,6 +82,7 @@ class DeepseekV4DSparkConfig: target_layer_ids: tuple = (40, 41, 42) markov_rank: int = 256 block_size: int = 6 + native_block_size: int | None = None class DSparkLocalAttention(v4.LocalAttention): @@ -180,12 +183,14 @@ def __init__(self, config: DeepseekV4DSparkConfig): f"post-init extended with {n_stages} zeros (got " f"{len(args.compress_ratios)} entries, {n} trunk layers)" ) - if config.block_size > config.draft_len + 1: + self._native_block_size = ( + native_block_size(config) or int(config.block_size)) + if max(int(config.block_size), self._native_block_size) > config.draft_len + 1: raise ValueError( - f"block_size {config.block_size} exceeds draft_len + 1 " + f"block_size {config.block_size} / native_block_size " + f"{self._native_block_size} exceeds draft_len + 1 " f"({config.draft_len + 1})" ) - self._native_block_size = int(config.block_size) self._draft_len = int(config.draft_len) self._noise_token_id = int(config.noise_token_id) self._sliding_window = int(args.sliding_window) diff --git a/gmlx/deepseek_v4_mtp.py b/gmlx/deepseek_v4_mtp.py index d7f3820..e4afcf5 100644 --- a/gmlx/deepseek_v4_mtp.py +++ b/gmlx/deepseek_v4_mtp.py @@ -51,6 +51,7 @@ from mlx_lm.models.base import create_attention_mask from . import deepseek_v4_model as v4 +from .drafter_protocol import native_block_size from .deepseek_v4_cache import ensure_rollback_attached, set_undo_armed from .deepseek_v4_hyper_connection import HyperHead @@ -266,7 +267,8 @@ def __init__(self, config: DeepseekV4MTPConfig): "post-init extended with the MTP layer's ratio 0 " f"(got {len(args.compress_ratios)} entries for layer index {n})" ) - self._native_block_size = int(config.block_size) + self._native_block_size = ( + native_block_size(config) or int(config.block_size)) self._sliding_window = int(args.sliding_window) # Engine capture seams keep only this many trailing prompt hiddens: # the head attends through a sliding window with relative RoPE, so diff --git a/gmlx/drafter_protocol.py b/gmlx/drafter_protocol.py index 357a55d..afb908b 100644 --- a/gmlx/drafter_protocol.py +++ b/gmlx/drafter_protocol.py @@ -68,7 +68,13 @@ class BatchDrafterProtocol(Protocol): draft_eval_state -- default None -> sampler_state_attrs sampler_state_attrs -- default ("_seed_token",) draft_lens -- hasattr guard - config.runtime_block_size -- default None + config.runtime_block_size -- default None: the depth drafted per + round when the caller asks for none, + overriding config.block_size + config.native_block_size -- default None (unbounded): the deepest + block this drafter can produce + correctly, from the checkpoint or the + architecture cap_at_configured_depth -- default False _native_block_size -- default configured_block_total mtp_width_cap -- default 0 (uncapped): speculate @@ -116,7 +122,31 @@ def bind(self, target_model: Any) -> "BatchDrafterProtocol": ... -# 2. DrafterAdapter -- pure delegator, per-instance +# 2. Block-depth accessors -- the two quantities config carries + +def native_block_size(config: Any) -> int | None: + """The deepest block the drafter can produce, or None when unbounded. + + A bound comes from the checkpoint (a trained diffusion block) or from the + architecture (one draft row per stage). Families without one leave the + field unset, and the caller falls back to their configured depth. + """ + declared = getattr(config, "native_block_size", None) + if declared is None: + return None + declared = int(declared) + return declared if declared > 0 else None + + +def default_block_size(config: Any) -> int: + """The block total to draft when the caller requests no depth.""" + runtime = getattr(config, "runtime_block_size", None) + if runtime is None: + return int(config.block_size) + return max(1, int(runtime)) + + +# 3. DrafterAdapter -- pure delegator, per-instance def _check_accepts_left_padding(inner: Any) -> bool: """True if inner.reset accepts a left_padding kwarg. @@ -182,7 +212,7 @@ def __getattr__(self, name: str) -> Any: return getattr(inner, name) -# 3. validate_drafter -- load-time crash prevention +# 4. validate_drafter -- load-time crash prevention def validate_drafter(drafter: Any) -> None: """Check required drafter members exist at load time. @@ -206,9 +236,20 @@ def validate_drafter(drafter: Any) -> None: errors.append("missing config") else: try: - int(drafter.config.block_size) + block = int(drafter.config.block_size) except (AttributeError, TypeError, ValueError): errors.append("config.block_size must be int-castable") + else: + try: + ceiling = native_block_size(drafter.config) + except (TypeError, ValueError): + errors.append("config.native_block_size must be int-castable") + else: + if ceiling is not None and ceiling < block: + errors.append( + f"config.native_block_size {ceiling} is below " + f"config.block_size {block}" + ) if not hasattr(drafter, "accept_lens"): errors.append("missing accept_lens (must be a list)") diff --git a/gmlx/generation.py b/gmlx/generation.py index 9ed14d4..cb41bbf 100644 --- a/gmlx/generation.py +++ b/gmlx/generation.py @@ -818,6 +818,7 @@ def _generate_speculative( ) from mlx_lm.sample_utils import make_sampler + from .spec_helpers import _resolve_block_total from mlx_vlm.generate.ar import generate_step if ( @@ -859,7 +860,7 @@ def _generate_speculative( if temp == 0.0 else make_sampler(temp=temp, top_p=top_p, top_k=top_k, min_p=min_p) ) - block = draft_block_size or int(getattr(drafter.config, "block_size", 3)) + block = _resolve_block_total(drafter, draft_block_size) eos_ids = set(getattr(tokenizer, "eos_token_ids", None) or [tokenizer.eos_token_id]) detok = tokenizer.detokenizer @@ -966,6 +967,7 @@ def generate_speculative_owned( generate_step. This is the bench_tg_depth default (matches serve); GMLX_OWNED_ROUND=0 opts back to mlx-vlm's generate_speculative.""" from mlx_lm.sample_utils import make_sampler + from .spec_helpers import _resolve_block_total from mlx_vlm.models import cache as _cache from .speculative import annotate_sampling_params, stream_speculative @@ -1001,7 +1003,7 @@ def generate_speculative_owned( ) annotate_sampling_params( sampler, temp=temp, top_p=top_p, top_k=top_k, min_p=min_p) - block = draft_block_size or int(getattr(drafter.config, "block_size", 3)) + block = _resolve_block_total(drafter, draft_block_size) eos_ids = set(getattr(tokenizer, "eos_token_ids", None) or [tokenizer.eos_token_id]) lm = model.language_model if hasattr(model, "language_model") else model @@ -1149,6 +1151,7 @@ def _stream_generate_speculative_owned( surface as the stock round, but drives ``stream_speculative`` (which does its own chunked prefill through the persistent ``prompt_cache``).""" from mlx_lm.sample_utils import make_sampler + from .spec_helpers import _resolve_block_total from .speculative import annotate_sampling_params, stream_speculative @@ -1168,7 +1171,7 @@ def _stream_generate_speculative_owned( ) annotate_sampling_params( sampler, temp=temp, top_p=top_p, top_k=top_k, min_p=min_p) - block = draft_block_size or int(getattr(drafter.config, "block_size", 2)) + block = _resolve_block_total(drafter, draft_block_size) eos_ids = set(getattr(tokenizer, "eos_token_ids", None) or [tokenizer.eos_token_id]) detok = tokenizer.detokenizer @@ -1276,6 +1279,7 @@ def _stream_generate_speculative( return from mlx_lm.sample_utils import make_sampler + from .spec_helpers import _resolve_block_total from mlx_vlm.generate.ar import generate_step if isinstance(prompt, str): @@ -1295,7 +1299,7 @@ def _stream_generate_speculative( if temp == 0.0 else make_sampler(temp=temp, top_p=top_p, top_k=top_k, min_p=min_p) ) - block = draft_block_size or int(getattr(drafter.config, "block_size", 3)) + block = _resolve_block_total(drafter, draft_block_size) eos_ids = set(getattr(tokenizer, "eos_token_ids", None) or [tokenizer.eos_token_id]) detok = tokenizer.detokenizer diff --git a/gmlx/hy_v3_mtp.py b/gmlx/hy_v3_mtp.py index 6326014..2991d97 100644 --- a/gmlx/hy_v3_mtp.py +++ b/gmlx/hy_v3_mtp.py @@ -43,6 +43,7 @@ from mlx_lm.models.base import create_attention_mask from . import hy_v3_model as hy +from .drafter_protocol import native_block_size from .mtp_drafter import QwenMTPDrafter @@ -157,7 +158,8 @@ class HyV3MTPDrafter(QwenMTPDrafter): def __init__(self, config: HyV3MTPConfig): nn.Module.__init__(self) self.config = config - self._native_block_size = int(config.block_size) + self._native_block_size = ( + native_block_size(config) or int(config.block_size)) args = config.text_config hidden_size = args.hidden_size diff --git a/gmlx/mtp_drafter.py b/gmlx/mtp_drafter.py index ec39c8e..addfd29 100644 --- a/gmlx/mtp_drafter.py +++ b/gmlx/mtp_drafter.py @@ -48,6 +48,7 @@ from mlx_vlm.models.cache import BatchKVCache, KVCache from . import prefill_decay +from .drafter_protocol import native_block_size from .envflags import env_bool, env_int from mlx_vlm.models.qwen3_5.language import Qwen3_5DecoderLayer from mlx_vlm.models.qwen3_5_moe.language import Qwen3_5MoeDecoderLayer @@ -86,7 +87,8 @@ class QwenMTPDrafter(nn.Module): def __init__(self, config): super().__init__() self.config = config - self._native_block_size = int(config.block_size) + self._native_block_size = ( + native_block_size(config) or int(config.block_size)) text_config = config.text_config if text_config is None: raise ValueError("MTP drafter config.text_config must be set") diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index 2e34280..9b864c5 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -84,12 +84,22 @@ # losing regime. Uncapped is earned by measurement, not inherited by default. _MTP_WIDTH_CAP_FALLBACK = 2 -# Drafted depth per DFlash round. Verify cost on the 30B target rises ~36% -# from block 3 to 4 (the kquant small-M kernels hold near-flat only through -# M=3 at these projection shapes), which outweighs block 4's extra accepted -# tokens; llama.cpp defaults to 4 (n_max=3). GMLX_MUSE_DFLASH_BLOCK -# overrides, up to the GGUF's dflash.block_size. -_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT = 3 +# Drafted depth per DFlash round: the GGUF's trained block, capped at 16. +# Verify cost on the 30B target is flat from 8 to 16 rows (the kquant split-K +# tile holds 16 rows in one MMA row-tile). Row 17 starts a second row-tile and +# costs ~55% more. +_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT = 16 + + +def _drafter_block_depths(native_total, preferred_total=None) -> tuple[int, int]: + """Return (deepest block the drafter can produce, depth drafted per round). + + The runtime depth is the family's preferred depth, bounded by the ceiling. + --draft-block-size moves it at run time. + """ + native_total = int(native_total) + preferred = min(int(preferred_total or native_total), native_total) + return native_total, max(2, min(preferred, native_total)) def _stamp_mtp_width_cap(drafter, model_type: str, *, target=None, @@ -782,10 +792,8 @@ def _load_muse_glimmer_dflash_drafter( "sliding_attention" if bool(t) else "full_attention" for t in pattern ] or ["full_attention"] * n_layers window = int(meta.get("dflash.attention.sliding_window") or 0) or None - native_total = int(block_size) - default_total = min(_MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT, native_total) - block_total = max( - 2, min(env_int("GMLX_MUSE_DFLASH_BLOCK", default_total), native_total)) + native_total, block_total = _drafter_block_depths( + block_size, _MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT) config = MuseGlimmerDFlashConfig( hidden_size=int(target_config_dict["hidden_size"]), @@ -801,6 +809,7 @@ def _load_muse_glimmer_dflash_drafter( rope_theta=float(meta["dflash.rope.freq_base"]), tie_word_embeddings=False, block_size=block_total, + native_block_size=native_total, mask_token_id=int(mask_token_id), target_layer_ids=list(layer_ids), num_target_layers=n_target_layers, @@ -1050,8 +1059,7 @@ def _load_deepseek4_dspark_drafter( f"strictly increasing and < {n}" ) args.compress_ratios = list(args.compress_ratios) + [0] * n_stages - native_total = draft_len + 1 - block_total = max(2, min(env_int("GMLX_DSPARK_BLOCK", native_total), native_total)) + native_total, block_total = _drafter_block_depths(draft_len + 1) drafter = DeepseekV4DSparkDrafter( DeepseekV4DSparkConfig( text=args, @@ -1061,6 +1069,7 @@ def _load_deepseek4_dspark_drafter( target_layer_ids=layer_ids, markov_rank=int(_dspark_meta(meta, "markov_rank", 256)), block_size=block_total, + native_block_size=native_total, ) ) log( diff --git a/gmlx/muse_glimmer_dflash.py b/gmlx/muse_glimmer_dflash.py index c506771..3fa3d1f 100644 --- a/gmlx/muse_glimmer_dflash.py +++ b/gmlx/muse_glimmer_dflash.py @@ -41,14 +41,18 @@ from mlx_vlm.speculative.drafters.qwen3_dflash.dflash import DFlashDraftModel from . import muse_glimmer_model as mg +from .drafter_protocol import native_block_size @dataclass class MuseGlimmerDFlashConfig(DFlashConfig): """``DFlashConfig`` plus the Glimmer logit tail. The drafter borrows the - target's LM head, so it must reproduce the target's scale and softcap.""" + target's LM head, so it must reproduce the target's scale and softcap. + ``native_block_size`` is the checkpoint's trained diffusion block, the + deepest block this drafter can produce.""" output_multiplier: float = 1.0 + native_block_size: int | None = None class MuseGlimmerDFlashDrafter(DFlashDraftModel): @@ -65,7 +69,8 @@ class MuseGlimmerDFlashDrafter(DFlashDraftModel): def __init__(self, config: MuseGlimmerDFlashConfig): super().__init__(config) - self._native_block_size = int(config.block_size) + self._native_block_size = ( + native_block_size(config) or int(config.block_size)) self._hidden = int(config.hidden_size) self._n_targets = len(config.target_layer_ids) # Only the trailing window of the prompt capture is usable; the engine diff --git a/gmlx/server.py b/gmlx/server.py index aab9a42..10a0342 100644 --- a/gmlx/server.py +++ b/gmlx/server.py @@ -734,8 +734,9 @@ def _add_serve_args(ap: argparse.ArgumentParser) -> None: "streamed installs; GMLX_GPU_KEEPWARM=0 disables.") ap.add_argument("--draft-block-size", type=int, default=None, metavar="N", help="MTP draft tokens per round (analogous to llama-server " - "--spec-draft-n-max). Default: the drafter's own block " - "size. Also via GMLX_DRAFT_BLOCK_SIZE.") + "--spec-draft-n-max). Raises or lowers the drafter's own " + "default, up to the deepest block it can produce. Also " + "via GMLX_DRAFT_BLOCK_SIZE.") ap.add_argument("--chat-template", default=None, metavar="STR|PATH", help="Inline Jinja template, or a path to a .jinja/.txt file, " "replacing a single positional model's GGUF template " diff --git a/gmlx/server_bridge_vlm.py b/gmlx/server_bridge_vlm.py index ff110c2..d692073 100644 --- a/gmlx/server_bridge_vlm.py +++ b/gmlx/server_bridge_vlm.py @@ -45,6 +45,7 @@ from __future__ import annotations import json +import logging import os import sys from contextvars import ContextVar @@ -53,8 +54,11 @@ from mlx_vlm.models.text_only import Model as TextOnlyModel from mlx_vlm.utils import StoppingCriteria +from .drafter_protocol import native_block_size from .loader import load_model +_log = logging.getLogger(__name__) + def _is_gguf(path) -> bool: return isinstance(path, str) and path.endswith(".gguf") @@ -729,10 +733,10 @@ def load_drafter(path_or_repo, kind=None, **kwargs): def _apply_draft_block_size_override(result) -> None: """Honor GMLX_DRAFT_BLOCK_SIZE (serve --draft-block-size): set the loaded - drafter's config block size so the engine drafts N tokens/round. _dflash_block_total - reads config.block_size when no explicit override is passed, so this covers both - native-head (nextn) and two-GGUF assistant drafters. Best-effort; a frozen config - or unset env is a no-op.""" + drafter's runtime block size so the engine drafts N tokens/round, clamped to + the deepest block the drafter can produce. This covers both native-head + (nextn) and two-GGUF assistant drafters. Best-effort; a frozen config or an + unset env is a no-op.""" raw = os.environ.get("GMLX_DRAFT_BLOCK_SIZE", "").strip() if not raw: return @@ -746,10 +750,17 @@ def _apply_draft_block_size_override(result) -> None: cfg = getattr(drafter, "config", None) if cfg is None: return + ceiling = native_block_size(cfg) + if ceiling is not None and n > ceiling: + _log.warning( + "--draft-block-size %d is deeper than this drafter can produce " + "(%d); drafting %d token(s)/round", n, ceiling, ceiling - 1) + n = ceiling try: - cfg.block_size = n if hasattr(cfg, "runtime_block_size"): cfg.runtime_block_size = n + else: + cfg.block_size = n except Exception: pass # frozen/odd config object -> keep the drafter's own default diff --git a/gmlx/spec_helpers.py b/gmlx/spec_helpers.py index 00161df..f99db96 100644 --- a/gmlx/spec_helpers.py +++ b/gmlx/spec_helpers.py @@ -9,6 +9,7 @@ Logic is a faithful copy (acceptance must stay token-identical to the validated path); keep it in sync when mlx-vlm's round changes in a way we want to track. """ +import logging from dataclasses import dataclass from typing import Any from collections.abc import Callable @@ -19,6 +20,9 @@ from mlx_vlm.models import cache from .cache_compat import cache_types +from .drafter_protocol import default_block_size, native_block_size + +_log = logging.getLogger(__name__) # Shared generation stream (the drafter model pins no stream of its own, so the # round and verify forward run here; cross-stream deps are handled by MLX events). @@ -137,15 +141,38 @@ def _record_speculative_round( draft_model.draft_lens.append(int(draft_count)) -def _dflash_block_total(draft_model: nn.Module, draft_block_size: int | None) -> int: - if draft_block_size is not None: - return int(draft_block_size) +def _drafter_block_ceiling(draft_model: nn.Module) -> int | None: + if not getattr(draft_model, "cap_at_configured_depth", False): + return None + declared = native_block_size(draft_model.config) + if declared is not None: + return declared + return int(getattr(draft_model, "_native_block_size", + draft_model.config.block_size)) - configured = int(draft_model.config.block_size) - runtime = getattr(draft_model.config, "runtime_block_size", None) - if runtime is None: - return configured - return min(configured, max(1, int(runtime))) + +def _warn_block_over_depth(draft_model: nn.Module, requested: int, + ceiling: int) -> None: + if getattr(draft_model, "_depth_cap_warned", False): + return + draft_model._depth_cap_warned = True + _log.warning( + "draft block %d is deeper than this drafter can produce (%d); " + "drafting %d token(s)/round", + requested, ceiling, ceiling - 1) + + +def _resolve_block_total(draft_model: nn.Module, draft_block_size: int | None) -> int: + """The block total for a run: the caller's depth, else the drafter's own + default, clamped to the deepest block the drafter can produce.""" + requested = (int(draft_block_size) if draft_block_size is not None + else default_block_size(draft_model.config)) + ceiling = _drafter_block_ceiling(draft_model) + if ceiling is None or requested <= ceiling: + return requested + if draft_block_size is not None: + _warn_block_over_depth(draft_model, requested, ceiling) + return ceiling def _effective_mtp_block_size( @@ -189,7 +216,9 @@ def _mtp_next_block_size( ) -> int: budget = min(requested_block_total, remaining_budget) if getattr(draft_model, "cap_at_configured_depth", False): - native = getattr(draft_model, "_native_block_size", configured_block_total) + native = _drafter_block_ceiling(draft_model) or configured_block_total + if requested_block_total > native: + _warn_block_over_depth(draft_model, requested_block_total, native) return min(budget, native) if getattr(draft_model, "prefer_requested_block_size", False): return budget diff --git a/gmlx/speculative.py b/gmlx/speculative.py index 65b385f..171928b 100644 --- a/gmlx/speculative.py +++ b/gmlx/speculative.py @@ -38,7 +38,7 @@ from .spec_helpers import ( _SpeculativeSamplerRNG, _buffer_mtp_target_cache, - _dflash_block_total, + _resolve_block_total, _mtp_cache_offset_max, _mtp_draft_hidden, _mtp_draft_position, @@ -685,7 +685,7 @@ def _owned_decode_rounds( # moot -- disable it to skip the per-round save/restore. greedy_draft = greedy or _FORCE_GREEDY_DRAFT - block_total = _dflash_block_total(drafter, draft_block_size) + block_total = _resolve_block_total(drafter, draft_block_size) configured_block_total = int(getattr(drafter.config, "block_size", block_total)) drafter.reset(model) @@ -1647,7 +1647,7 @@ def _owned_decode_rounds_batch( gen_rows: list[list[int]] = [[int(t)] for t in b] retired = [False] * B_orig - block_total = _dflash_block_total(drafter, draft_block_size) + block_total = _resolve_block_total(drafter, draft_block_size) configured_block_total = int( getattr(drafter.config, "block_size", block_total)) diff --git a/pyproject.toml b/pyproject.toml index 95e5dd6..3f2188d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "gmlx" -version = "0.3.1" +version = "0.3.2" description = "A local inference platform for Apple Silicon: run, chat with, serve, and fine-tune the GGUF ecosystem's quantized models natively on MLX, straight off the file." readme = "README.md" requires-python = ">=3.11" # 3.10 EOLs 2026-10; mlx-kquant's lower floor is a library floor @@ -31,7 +31,7 @@ classifiers = [ # mlx / mlx-lm / mlx-vlm move in lockstep with the exact mlx-vlm pin; bump them # together per docs/upstream-upgrades.md. dependencies = [ - "mlx-kquant>=0.3.12,<0.4", + "mlx-kquant>=0.3.13,<0.4", "mlx-lm>=0.31", # Direct imports of both throughout (mx.*, np.*); declared explicitly # rather than inherited through mlx-lm. diff --git a/tests/test_draft_block_depth.py b/tests/test_draft_block_depth.py new file mode 100644 index 0000000..1b9c984 --- /dev/null +++ b/tests/test_draft_block_depth.py @@ -0,0 +1,168 @@ +"""Draft-depth resolution across the drafter families. + +Two quantities travel on the drafter config: ``block_size`` (the depth drafted +per round when nobody asks for one) and ``native_block_size`` (the deepest block +the drafter can produce). ``--draft-block-size`` moves the first one in either +direction and can never pass the second. + +Pure Python -- the stubs carry only the flags the resolver reads. +""" + +from __future__ import annotations + +import logging + +import pytest + +from gmlx.drafter_protocol import default_block_size, native_block_size +from gmlx.mtp_load import _drafter_block_depths +from gmlx.spec_helpers import _mtp_next_block_size, _resolve_block_total + + +class _Cfg: + def __init__(self, block_size, native=None, runtime=None): + self.block_size = block_size + if native is not None: + self.native_block_size = native + if runtime is not None: + self.runtime_block_size = runtime + + +class _Drafter: + def __init__(self, cfg, *, cap=True): + self.config = cfg + self.cap_at_configured_depth = cap + self.prefer_requested_block_size = not cap + self.accept_lens = [] + + +def _muse(**kw): + """Declares a ceiling: muse-glimmer dflash, dspark.""" + return _Drafter(_Cfg(kw.pop("block_size", 16), native=kw.pop("native", 16), **kw)) + + +def _ds4_mtp(**kw): + """Caps at its load-time depth but declares no ceiling: ds4 mtp, hy3.""" + return _Drafter(_Cfg(kw.pop("block_size", 3), **kw)) + + +def _qwen(**kw): + """Honors any requested depth: qwen native head.""" + return _Drafter(_Cfg(kw.pop("block_size", 3), **kw), cap=False) + + +# --- the two config quantities --------------------------------------------- + +class TestAccessors: + + def test_ceiling_is_none_when_undeclared(self): + assert native_block_size(_Cfg(3)) is None + + def test_ceiling_is_none_when_zero(self): + assert native_block_size(_Cfg(3, native=0)) is None + + def test_ceiling_reads_the_declared_depth(self): + assert native_block_size(_Cfg(4, native=16)) == 16 + + def test_default_is_the_block_size(self): + assert default_block_size(_Cfg(4, native=16)) == 4 + + def test_runtime_overrides_the_default_upward(self): + assert default_block_size(_Cfg(4, native=16, runtime=12)) == 12 + + def test_runtime_overrides_the_default_downward(self): + assert default_block_size(_Cfg(16, native=16, runtime=4)) == 4 + + +# --- what a run resolves to ------------------------------------------------- + +class TestResolveBlockTotal: + + @pytest.mark.parametrize("build", [_muse, _ds4_mtp, _qwen]) + def test_no_request_uses_the_configured_default(self, build): + drafter = build(block_size=3) + assert _resolve_block_total(drafter, None) == 3 + + @pytest.mark.parametrize("build", [_muse, _ds4_mtp, _qwen]) + def test_no_request_prefers_the_runtime_default(self, build): + drafter = build(block_size=3, runtime=2) + assert _resolve_block_total(drafter, None) == 2 + + def test_request_below_the_ceiling_is_honored(self): + assert _resolve_block_total(_muse(block_size=4, native=16), 12) == 12 + + def test_request_at_the_ceiling_is_honored(self): + assert _resolve_block_total(_muse(block_size=4, native=16), 16) == 16 + + def test_request_past_the_ceiling_clamps(self): + assert _resolve_block_total(_muse(block_size=4, native=16), 24) == 16 + + def test_request_below_the_default_lowers_it(self): + assert _resolve_block_total(_muse(block_size=16, native=16), 4) == 4 + + def test_undeclared_ceiling_still_caps_at_the_load_time_depth(self): + assert _resolve_block_total(_ds4_mtp(block_size=3), 8) == 3 + + def test_uncapped_family_honors_any_depth(self): + assert _resolve_block_total(_qwen(block_size=3), 8) == 8 + + def test_runtime_default_does_not_cap_an_explicit_request(self): + drafter = _muse(block_size=16, native=16, runtime=4) + assert _resolve_block_total(drafter, None) == 4 + assert _resolve_block_total(drafter, 16) == 16 + + +# --- per-round sizing ------------------------------------------------------- + +class TestNextBlockSize: + + def test_round_honors_a_depth_within_the_ceiling(self): + assert _mtp_next_block_size(_muse(native=16), 16, 4, 128) == 16 + + def test_round_clamps_to_the_ceiling(self): + assert _mtp_next_block_size(_muse(native=16), 24, 16, 128) == 16 + + def test_round_clamps_to_the_remaining_budget(self): + assert _mtp_next_block_size(_muse(native=16), 16, 16, 5) == 5 + + def test_uncapped_family_keeps_the_requested_depth(self): + assert _mtp_next_block_size(_qwen(), 8, 3, 128) == 8 + + +class TestDepthWarning: + + def test_a_deeper_request_warns_once(self, caplog): + drafter = _muse(block_size=4, native=16) + with caplog.at_level(logging.WARNING, logger="gmlx.spec_helpers"): + assert _resolve_block_total(drafter, 24) == 16 + assert _mtp_next_block_size(drafter, 24, 16, 128) == 16 + warned = [r for r in caplog.records if "deeper than" in r.message] + assert len(warned) == 1 + + def test_a_request_within_the_ceiling_is_silent(self, caplog): + with caplog.at_level(logging.WARNING, logger="gmlx.spec_helpers"): + assert _resolve_block_total(_muse(block_size=4, native=16), 16) == 16 + assert not [r for r in caplog.records if "deeper than" in r.message] + + def test_the_configured_default_is_silent(self, caplog): + drafter = _muse(block_size=16, native=16, runtime=32) + with caplog.at_level(logging.WARNING, logger="gmlx.spec_helpers"): + assert _resolve_block_total(drafter, None) == 16 + assert not [r for r in caplog.records if "deeper than" in r.message] + + +# --- what the loaders stamp ------------------------------------------------- + +class TestLoaderDepths: + + def test_the_family_default_bounds_the_runtime_depth(self): + assert _drafter_block_depths(32, 16) == (32, 16) + + def test_the_ceiling_bounds_the_family_default(self): + assert _drafter_block_depths(8, 16) == (8, 8) + + def test_no_family_default_drafts_the_full_block(self): + assert _drafter_block_depths(6) == (6, 6) + + def test_the_runtime_depth_keeps_one_draft(self): + assert _drafter_block_depths(32, 1) == (32, 2) diff --git a/tests/test_drafter_protocol.py b/tests/test_drafter_protocol.py index 7a032f5..cf45cce 100644 --- a/tests/test_drafter_protocol.py +++ b/tests/test_drafter_protocol.py @@ -11,6 +11,7 @@ from gmlx.drafter_protocol import ( DrafterAdapter, _check_accepts_left_padding, + native_block_size, validate_drafter, ) @@ -176,6 +177,25 @@ def test_multiple_errors_all_reported(self): assert "missing accept_lens" in msg assert "draft_block" in msg + def test_declared_ceiling_passes(self): + d = _FullDrafter() + d.config.native_block_size = 16 + validate_drafter(d) + assert native_block_size(d.config) == 16 + + def test_ceiling_below_the_configured_depth_fails(self): + d = _FullDrafter() + d.config.block_size = 8 + d.config.native_block_size = 4 + with pytest.raises(RuntimeError, match="native_block_size 4 is below"): + validate_drafter(d) + + def test_bad_ceiling(self): + d = _FullDrafter() + d.config.native_block_size = "deep" + with pytest.raises(RuntimeError, match="native_block_size must be int"): + validate_drafter(d) + def test_error_names_adapter(self): d = _NoLeftPaddingDrafter() del d.config diff --git a/tests/test_muse_glimmer_mtp.py b/tests/test_muse_glimmer_mtp.py index 6bf15a1..900ceff 100644 --- a/tests/test_muse_glimmer_mtp.py +++ b/tests/test_muse_glimmer_mtp.py @@ -196,7 +196,7 @@ def test_verify_walk_is_token_identical_to_greedy(armed): # --- the drafter side of the same seam --------------------------------------- -def _build_drafter(cfg, n_layers=2): +def _build_drafter(cfg, n_layers=2, block_size=BLOCK, native_block_size=None): from gmlx.muse_glimmer_dflash import ( MuseGlimmerDFlashConfig, MuseGlimmerDFlashDrafter, @@ -214,7 +214,8 @@ def _build_drafter(cfg, n_layers=2): max_position_embeddings=1024, rope_theta=10000.0, tie_word_embeddings=False, - block_size=BLOCK, + block_size=block_size, + native_block_size=native_block_size, mask_token_id=7, target_layer_ids=list(CAPTURE), num_target_layers=cfg["num_hidden_layers"], @@ -272,3 +273,82 @@ def test_drafter_satisfies_the_protocol(): validate_drafter(drafter) assert drafter.uses_shared_kv is False assert drafter.requires_owned_engine is True + + +def test_a_shallow_default_still_drafts_the_trained_block_on_request(): + """The loader defaults the runtime depth below the checkpoint's block. A + deeper ``--draft-block-size`` must reach the drafter and produce that many + drafts, which is what the depth ceiling exists for.""" + from gmlx.spec_helpers import _resolve_block_total + + lm, cfg = _build() + drafter = _build_drafter(cfg, block_size=2, native_block_size=BLOCK) + mx.eval(drafter.parameters()) + drafter.reset(lm) + + assert _resolve_block_total(drafter, None) == 2 + block_total = _resolve_block_total(drafter, BLOCK) + assert block_total == BLOCK + + drafts = drafter.draft_block(3, None, None, block_total, None, greedy=True) + mx.eval(drafts) + assert drafts.shape == (1, BLOCK - 1) + + +def _load_drafter_config(cfg, monkeypatch, block_size): + """Drive the GGUF loader far enough to capture the config it builds.""" + from gmlx import muse_glimmer_dflash as mgd + from gmlx import mtp_load + + captured = {} + + class _Stop(Exception): + pass + + def _capture(config): + captured["config"] = config + raise _Stop + + monkeypatch.setattr(mgd, "MuseGlimmerDFlashDrafter", _capture) + meta = { + "dflash.target_layers": [1, 3], + "dflash.block_size": block_size, + "tokenizer.ggml.mask_token_id": 7, + "dflash.feed_forward_length": 64, + "dflash.attention.head_count": 4, + "dflash.attention.head_count_kv": 2, + "dflash.attention.key_length": 16, + "dflash.attention.layer_norm_rms_epsilon": 1e-6, + "dflash.rope.freq_base": 10000.0, + "dflash.attention.sliding_window": 512, + } + with pytest.raises(_Stop): + mtp_load._load_muse_glimmer_dflash_drafter( + "draft.gguf", None, cfg, + arrays={"blk.0.attn_q.weight": None, "blk.1.attn_q.weight": None}, + kquant_meta={}, meta=meta, log=lambda *a, **k: None) + return captured["config"] + + +def test_the_loader_stamps_the_gguf_block_as_the_drafter_ceiling(monkeypatch): + """The runtime depth is a tuned default; the checkpoint's trained block is + the ceiling a deeper ``--draft-block-size`` is allowed to reach.""" + from gmlx.mtp_load import _MUSE_GLIMMER_DFLASH_BLOCK_DEFAULT as TUNED + + _, cfg = _build() + deep = _load_drafter_config(cfg, monkeypatch, TUNED + 16) + assert deep.native_block_size == TUNED + 16 + assert deep.block_size == TUNED + + shallow = _load_drafter_config(cfg, monkeypatch, 3) + assert shallow.native_block_size == 3 + assert shallow.block_size == 3 + + +def test_an_undeclared_ceiling_keeps_the_configured_depth(): + from gmlx.spec_helpers import _resolve_block_total + + _, cfg = _build() + drafter = _build_drafter(cfg, block_size=BLOCK) + assert drafter._native_block_size == BLOCK + assert _resolve_block_total(drafter, BLOCK + 8) == BLOCK