fix(mtp): --draft-block-size was silently ignored above the default depth

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
 

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`. |

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",

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)

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

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)")

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

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

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")

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(

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