diff --git a/CHANGELOG.md b/CHANGELOG.md index b0ebb7a..c1a21dd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,6 +25,22 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). The server always decoded up to 32 requests together, which slows every stream past the width where total throughput stops growing. +### Fixed +- Sampled speculative verify and the server's unfiltered sampler computed + logprob math at the activation dtype; float16 rows reached categorical + unwidened. Both now widen to float32 first (greedy paths unchanged). +- MiniMax MSA indexer projections stay F32 under float16 activations; the + bfloat16 exact-bit narrowing does not exist at float16. +- Non-streamed replies lost the leading indentation of their first content + line: the reasoning splitters (ATEM/harmony and the think-tag path all + served models use) stripped all content whitespace. Verbatim code answers + came back misindented on line one. Content now keeps first-line indent + and still drops leading blank lines and trailing whitespace. +- gemma-4 applied the final logit softcap at the activation dtype, rounding + every logit by up to ~0.12 nats at bfloat16 and flipping near-tie top-1 + picks. The softcap now computes in float32 and emits float32 logits, like + muse-glimmer. GMLX_G4_SOFTCAP_F32=0 restores the old behavior. + ## [0.3.2] - 2026-08-13 ### Changed diff --git a/gmlx/loader.py b/gmlx/loader.py index 1a7a060..71ee2af 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -40,6 +40,7 @@ from .sparse_sdpa import install_sparse_sdpa from .arrays_cache_fix import install_arrays_cache_fix from .gemma4_sync import install_gemma4_nosync +from .softcap_f32 import install_gemma4_softcap_f32 from .quantized_sdpa_fix import install_quantized_sdpa_mask_fix from .rope_batch_fix import install_rope_batch_fix from .rotating_cache_fix import install_rotating_cache_fix @@ -2930,6 +2931,8 @@ def _install_and_load( log("[install] gemma-4 host-sync-free masks/rope offsets active") if install_gemma4_batched_sdpa() and _gemma4_target(model): log("[install] gemma-4 hd512 batched-decode row route active") + if install_gemma4_softcap_f32() and _gemma4_target(model): + log("[install] gemma-4 float32 logit softcap active") if install_cascade_sdpa() and install_cascade_stamp(): log("[install] shared-prefix cascade decode route active") if install_sparse_sdpa(): @@ -3391,6 +3394,8 @@ def load_model( _log("[install] gemma-4 host-sync-free masks/rope offsets active") if install_gemma4_batched_sdpa() and _gemma4_target(model): _log("[install] gemma-4 hd512 batched-decode row route active") + if install_gemma4_softcap_f32() and _gemma4_target(model): + _log("[install] gemma-4 float32 logit softcap active") if install_cascade_sdpa() and install_cascade_stamp(): _log("[install] shared-prefix cascade decode route active") if install_sparse_sdpa(): diff --git a/gmlx/minimax_m3_model.py b/gmlx/minimax_m3_model.py index 8cc8f42..c678fc4 100644 --- a/gmlx/minimax_m3_model.py +++ b/gmlx/minimax_m3_model.py @@ -637,12 +637,13 @@ def rename(k): and v.dtype == mx.float32 ): # Some MSA GGUFs store the indexer projections as F32. The - # source checkpoint is BF16, so narrowing to the activation - # dtype halves the memory and keeps the indexer matmul at the - # graph's width; at bfloat16 it recovers the original bits - # exactly. Norms stay F32 (tiny, and the norm math runs in F32 - # regardless). - v = v.astype(activation_dtype()) + # source checkpoint is BF16, so narrowing to bfloat16 recovers + # the original bits exactly and halves the memory. float16 + # cannot represent those bits; keep F32 there (MLX promotes + # the indexer matmul). Norms stay F32 (tiny, and the norm math + # runs in F32 regardless). + if activation_dtype() != mx.float16: + v = v.astype(activation_dtype()) renamed[rename(k)] = v weights = renamed diff --git a/gmlx/reasoning.py b/gmlx/reasoning.py index 044aa3b..515aa8a 100644 --- a/gmlx/reasoning.py +++ b/gmlx/reasoning.py @@ -275,10 +275,20 @@ def split_harmony_reply(text: str, *, spans = filt.feed(text) spans += filt.flush() reasoning = "".join(t for t, m in spans if m == _REASON).strip() - content = "".join(t for t, m in spans if m == _ANSWER).strip() + content = trim_content_ws("".join(t for t, m in spans if m == _ANSWER)) return (reasoning or None, content) +def trim_content_ws(text: str) -> str: + """Trim reply content: drop leading blank lines and trailing whitespace, + keep the first content line's indent (a full strip eats verbatim-code + indentation).""" + lines = (text or "").split("\n") + while lines and not lines[0].strip(): + lines.pop(0) + return "\n".join(lines).rstrip() + + _DIM = "\x1b[2m" _RESET = "\x1b[0m" _CLEAR_EOL = "\x1b[K" diff --git a/gmlx/server_patches/chat_behavior.py b/gmlx/server_patches/chat_behavior.py index 12a8dc6..19fec76 100644 --- a/gmlx/server_patches/chat_behavior.py +++ b/gmlx/server_patches/chat_behavior.py @@ -259,7 +259,38 @@ def _make_thinking_budget_criteria(self, args, input_ids): def _strip_xtml_sections(content: str) -> str: for m in _XTML_SECTION_MARKERS: content = content.replace(m, "") - return content.strip() + return _trim_content_ws(content) + + +def _trim_content_ws(text: str) -> str: + from ..reasoning import trim_content_ws + return trim_content_ws(text) + + +def _split_thinking_keep_indent(rs, cls, text, start_token, end_token): + """rs._split_thinking with content whitespace preserved. Same marker + walk as stock; only the content-side .strip() calls are narrowed + (upstream body seam-fingerprinted via app._split_thinking_text).""" + if not text: + return None, text + for start_marker, end_marker in cls._build_open_close_markers( + start_token, end_token): + start = text.find(start_marker) + end = text.find(end_marker, start if start >= 0 else 0) + if 0 <= start < end: + reasoning = text[start + len(start_marker):end].strip() + content = rs._strip_content_markers( + text[:start] + text[end + len(end_marker):]) + return reasoning or None, _trim_content_ws(content) + if end_marker in text: + reasoning, content = text.split(end_marker, 1) + reasoning = rs._clean_reasoning(reasoning, start_marker) + return (reasoning or None, + _trim_content_ws(rs._strip_content_markers(content))) + if start_marker in text: + reasoning = rs._clean_reasoning(text, start_marker) + return reasoning or None, "" + return None, _trim_content_ws(rs._strip_content_markers(text)) _LAST_RENDERED_PROMPT: contextvars.ContextVar = contextvars.ContextVar( "kq_last_rendered_prompt", default=None) @@ -370,8 +401,10 @@ def _split_thinking_text(text, thinking_start_token=None, # so this gate cannot misfire on it.) from ..reasoning import split_harmony_reply return split_harmony_reply(text, start_in_header=in_header) - reasoning, content = split( - text, thinking_start_token, thinking_end_token) + # Stock walk with content whitespace preserved: the stock + # splitter's .strip() eats first-line code indent. + reasoning, content = _split_thinking_keep_indent( + rs, cls, text, thinking_start_token, thinking_end_token) if reasoning is None and content and retire_key.truncated_thinking( text, cls._build_open_close_markers( thinking_start_token, thinking_end_token), @@ -383,6 +416,9 @@ def _split_thinking_text(text, thinking_start_token=None, _split_thinking_text.__dict__[_STREAM_SEED_FLAG] = True app._split_thinking_text = _split_thinking_text + # The /v1/responses output builder calls the responses_state module + # global directly; give it the same splitter. + rs._split_thinking = _split_thinking_text # ignore-eos: forced-length decode (server-level) diff --git a/gmlx/server_patches/sampling.py b/gmlx/server_patches/sampling.py index b683e7c..be4c59f 100644 --- a/gmlx/server_patches/sampling.py +++ b/gmlx/server_patches/sampling.py @@ -283,7 +283,9 @@ def __call__(self, logprobs): # them so the returned ids keep the caller's shape (matches top_p_sampling). import mlx.core as mx if not self._has_filter: - return mx.random.categorical(logprobs * (1.0 / self.temperature), axis=-1) + # f32: float16 logprobs would otherwise reach categorical unwidened + lp = logprobs.astype(mx.float32) + return mx.random.categorical(lp * (1.0 / self.temperature), axis=-1) lead = logprobs.shape[:-1] lp2 = logprobs.reshape(-1, logprobs.shape[-1]) masked, part, order = self._filtered(lp2) @@ -300,7 +302,7 @@ def _cat(row, key): return mx.random.categorical(row, key=key) if not self._has_filter: - scaled = logprobs * (1.0 / self.temperature) + scaled = logprobs.astype(mx.float32) * (1.0 / self.temperature) return mx.vmap(_cat, in_axes=(0, 0))(scaled, keys) # only this tiny categorical-over-k is per-row keyed; the filter is batched. masked, part, order = self._filtered(logprobs) diff --git a/gmlx/softcap_f32.py b/gmlx/softcap_f32.py new file mode 100644 index 0000000..9ca48a0 --- /dev/null +++ b/gmlx/softcap_f32.py @@ -0,0 +1,36 @@ +"""Float32 final-logit softcap for gemma-4. + +Upstream computes tanh(x/cap)*cap at the activation dtype. Each 16-bit +logit then rounds by ~cap*eps: ~0.12 nats at bfloat16 with cap 30, enough +to flip near-tie top-1 picks. Compute in float32 and emit float32 logits, +matching the muse-glimmer softcap path. Covers the mlx-lm gemma4_text and +mlx-vlm gemma4 seams; install-once, GMLX_G4_SOFTCAP_F32=0 disables. +""" +from __future__ import annotations + +import mlx.core as mx + +from .envflags import env_bool + +_installed = False + + +def _logit_softcap_f32(softcap, x): + return mx.tanh(x.astype(mx.float32) / softcap) * softcap + + +def install_gemma4_softcap_f32() -> bool: + global _installed + if not env_bool("GMLX_G4_SOFTCAP_F32", True): + return False + if _installed: + return True + from mlx_lm.models import gemma4_text as _lm + _lm.logit_softcap = _logit_softcap_f32 + try: + from mlx_vlm.models.gemma4 import language as _vlm + _vlm.logit_softcap = _logit_softcap_f32 + except ImportError: + pass + _installed = True + return True diff --git a/gmlx/speculative.py b/gmlx/speculative.py index 37be22b..7088023 100644 --- a/gmlx/speculative.py +++ b/gmlx/speculative.py @@ -580,6 +580,7 @@ def _coupled_walk(lm, verify, draft_tokens: mx.array, sampler, budget: int, pq_arr = _pq_graph(logits[:n_draft], pq) else: _pq_stats["misaligned"] += 1 + logits = logits.astype(mx.float32) # 16-bit logprob math skews sampled verify logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) if sampler is None: target = mx.argmax(logprobs, axis=-1) # [n_pos] @@ -649,6 +650,7 @@ def _coupled_walk_batch( target = verify.target_tokens # [B, n_pos] else: logits = lm.speculative_logits_from_hidden(verify.hidden) # [B, n_pos, V] + logits = logits.astype(mx.float32) # 16-bit logprob math skews sampled verify logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) if sampler is None: target = mx.argmax(logprobs, axis=-1) # [B, n_pos] @@ -1817,6 +1819,7 @@ def _gated_step(inputs): logits = getattr(out, "logits", out)[:, -1, :] if greedy: return mx.argmax(logits, axis=-1).astype(token_dtype) + logits = logits.astype(mx.float32) # 16-bit logprob math skews sampled verify logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) return sampler(logprobs).reshape(-1).astype(token_dtype) diff --git a/tests/test_long_context.py b/tests/test_long_context.py index ef1356d..1a854ad 100644 --- a/tests/test_long_context.py +++ b/tests/test_long_context.py @@ -20,7 +20,9 @@ real attention bug diverges in the first few characters. Both are ``integration`` + ``slow`` and skip unless the env points at real -models (see ``conftest``). Select one arch with ``-k qwen2`` and/or shrink the +models (see ``conftest``). Both sweep the 16-bit activation dtypes (float16 is +the pre-Apple9 auto choice; its narrower range must hold at depth too). Select +one arch with ``-k qwen2``, one dtype with ``-k float16``, and/or shrink the length with ``KQUANT_LONGCTX_TOKENS=4096`` for a quick run; the defaults sweep every arch whose GGUF is present at >=16k and can take minutes on large models. """ @@ -48,6 +50,10 @@ TARGET = int(os.environ.get("KQUANT_LONGCTX_TOKENS", "16384")) +# Both 16-bit widths. float16 is the auto choice on pre-Apple9 GPUs and has +# a narrower exponent range, so depth signatures must hold under both. +ACTIVATION_DTYPES = ("bfloat16", "float16") + # Arches whose mlx-lm model has a known upper-context limitation vs the GGUF # reference, so token-for-token parity is only expected up to this length: # gemma2 - mlx-lm's gemma2 has no sliding-window attention; output is exact @@ -206,8 +212,10 @@ def _llama_complete(binary, model_path, prompt_text, n, ctx): # tests +@pytest.mark.parametrize("dtype", ACTIVATION_DTYPES) @pytest.mark.parametrize("arch", CANDIDATE_ARCHES) -def test_long_decode_integrity(arch, gguf_index): +def test_long_decode_integrity(arch, dtype, gguf_index, monkeypatch): + monkeypatch.setenv("GMLX_ACTIVATION_DTYPE", dtype) path = _require(gguf_index, arch) model, config, tok = _load(path) vocab = int(config["vocab_size"]) @@ -251,10 +259,12 @@ def _suppress_eos(_tokens, logits): f"- degeneration") +@pytest.mark.parametrize("dtype", ACTIVATION_DTYPES) @pytest.mark.parametrize("arch", CANDIDATE_ARCHES) -def test_long_prefill_parity(arch, gguf_index, llamacpp_bin): +def test_long_prefill_parity(arch, dtype, gguf_index, llamacpp_bin, monkeypatch): from mlx_lm.generate import stream_generate + monkeypatch.setenv("GMLX_ACTIVATION_DTYPE", dtype) if arch in PARITY_SKIP: pytest.skip(f"{arch}: {PARITY_SKIP[arch]}") path = _require(gguf_index, arch) diff --git a/tests/test_reasoning.py b/tests/test_reasoning.py index 5a0397e..533fe6d 100644 --- a/tests/test_reasoning.py +++ b/tests/test_reasoning.py @@ -539,6 +539,20 @@ def test_split_harmony_reply_shapes(): assert split_harmony_reply(plain) == (None, "Six.") +def test_split_harmony_reply_keeps_first_line_code_indent(): + from gmlx.reasoning import split_harmony_reply + full = ('<|channel|>analysis<|message|>Recall the body.' + "<|end|><|start|>assistant<|channel|>final<|message|>" + '\n """\n Docstring.\n') + r, c = split_harmony_reply(full) + assert r == "Recall the body." + assert c == ' """\n Docstring.' + # whitespace-only content collapses to empty + ws = ('<|channel|>analysis<|message|>x' + "<|end|><|start|>assistant<|channel|>final<|message|>\n \n") + assert split_harmony_reply(ws) == ("x", "") + + # --- muse-glimmer ATEM channel ------------------------------------------------ # # Routing is on the message HEADER, not on a "to=self" marker: any marker whose diff --git a/tests/test_server_patches.py b/tests/test_server_patches.py index 0d6834d..b0f80b1 100644 --- a/tests/test_server_patches.py +++ b/tests/test_server_patches.py @@ -330,6 +330,39 @@ def test_nonstream_split_strips_xtml_section_markers(): sp_chat._LAST_RENDERED_PROMPT.reset(tok) +def test_nonstream_split_keeps_first_line_code_indent(): + """The stock splitter .strip()s content, which deletes the first line's + leading indent from verbatim-code replies. The seeded splitter trims + newlines only, for every marker branch and the no-marker fallthrough.""" + app_mod = importlib.import_module("mlx_vlm.server.app") + rs = importlib.import_module("mlx_vlm.server.responses_state") + sp.install_stream_thinking_seed() + split = app_mod._split_thinking_text + tok = sp_chat._LAST_RENDERED_PROMPT.set(None) + try: + # open+close pair + r, c = split("plan\n\n if x:\n y()") + assert r == "plan" + assert c == " if x:\n y()" + # close-only (prompt-opened block) + r, c = split("plan\n\n\n\tdoc = doc || document;") + assert r == "plan" + assert c == "\tdoc = doc || document;" + # no marker at all + assert split(" indented") == (None, " indented") + # whitespace-only content collapses to empty + assert split("plan\n \n") == ("plan", "") + # XTML section strip keeps the indent too + text = ("plan<|close|>think<|sep|><|open|>response<|sep|>" + " return 4;<|close|>response<|sep|><|end_of_msg|>") + r, c = split(text, "<|open|>think<|sep|>", "<|close|>think<|sep|>") + assert c == " return 4;" + # the /v1/responses module global got the same splitter + assert rs._split_thinking is split + finally: + sp_chat._LAST_RENDERED_PROMPT.reset(tok) + + def test_install_thinking_budget_fix_applies_and_idempotent(): # Fail-loud guard: asserts the seam bound to the REAL mlx-vlm symbol, so a # rename of ResponseGenerator._make_thinking_budget_criteria turns into a CI