feat: serve Ministral 3; add --max-batch-tokens flag #58

Merged

rcsheets merged 3 commits from fix/ministral3-serving into main

2026-08-04 20:53:03 +00:00

rcsheets commented

2026-08-04 20:49:55 +00:00

Owner

Also fixes a group of documentation issues

rcsheets added 3 commits

2026-08-04 20:49:56 +00:00

test(cuda): drop Blackwell assumptions from CUDA test comments 98dcf4b3ba

These tests open device 0 via cuda.New(0) and gate on nothing: there is no
capability check and no skip anywhere in the CUDA test files. Calling it "the
Blackwell target" recorded which GPU happened to be device 0 on the development
box, and reads as though the tests assert something Blackwell-specific.

TestCUDAMemInfo's comment was doubly wrong: it claimed the total was checked
against "a plausible range for the Blackwell target" when the assertion is just
total >= 1 GiB. Its comment now describes what the test actually does.

Left alone: device_policy_test.go, where "Blackwell" names synthetic dInfo
fixture rows rather than the local machine, and nvfp4probe_cuda_test.go, whose
"widely supported on Blackwell" is a claim about the architecture. Non-test
mentions are all domain facts (the nvcc build target, NVIDIA's workspace
guidance, NVFP4's provenance, the below-Blackwell warning the code really does
emit) and are unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

feat(mistral): serve Ministral 3 fp8 checkpoints 3f5091011f

Ministral-3-14B-Instruct-2512 failed to load outright and would have been
subtly wrong at long context once it did.

Loading: the checkpoint spells its per-tensor fp8 scale weight_scale_inv
(DeepSeek's name) as a 0-d BF16 scalar, not weight_scale, so every projection
errored with "has no scale tensor". Despite the _inv, it multiplies exactly as
weight_scale does -- verified numerically, fp8_absmax * scale lands at 0.05-0.16
against ~1e5 if divided. loader.quantized now accepts either spelling and
normalizes a 0-d shape to [1]. The .activation_scale shipped alongside stays
unused: this is the weight-only path.

llama_4 scaling: the checkpoint sets llama_4_scaling_beta, which only mistral4
implemented, so the dense path silently ignored it. It is 1 + beta*log(1 +
floor(pos/original_max_position_embeddings)), applied to the query only, after
RoPE, in every layer -- confirmed against transformers 5.14.1's
Ministral3Attention.get_llama_4_attn_scale. Being exactly 1 below the original
context, it is invisible under 16384 tokens and degrades quietly above it. The
formula moves to config.RopeParameters.Llama4AttnScale, shared with mistral4
rather than duplicated; Mistral.llama4Scales returns nil for a short batch so
the multiply is skipped outright.

Pinned by a tiny-llama4 fixture with HF goldens, covering prefill and decode.
It sets original_max_position_embeddings to 4 and an exaggerated beta: at the
real 0.1 the effect on a random tiny model is ~4e-5, under the goldens' 1e-4
tolerance, so the test would have passed against a gllm ignoring llama_4
entirely. Verified by mutation -- disabling the scaling, dropping the floor(),
or scaling the key as well each fail it.

The real checkpoint now loads in 38s and serves coherent text, including a
30072-token prompt whose planted fact it retrieves (positions well past 16384,
so llama_4 scaling is live across most of the prompt).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

feat(serve): add --max-batch-tokens

ci / test_and_build (pull_request) Successful in 2m41s

Details

5ff360c688

MaxTokensPerBatch was already an engine.Options field, already a flag on
gllm plan, and already auto-capped when the activation scratch would not fit --
serve just never exposed it. So `gllm plan --max-batch-tokens 32768` would model
a configuration serve could not be told to run.

It is also the ceiling on a single prompt, since a sequence larger than one
batch is rejected rather than chunked (scheduler.go's chunked-prefill TODO), so
the default 8192 caps prompts well below what a 262144-context model accepts.
Raising the default instead would be wrong: scratch is linear in the cap
(~322 KiB/token for this model) and comes straight out of the KV cache, so it
is a per-deployment trade, not a better constant. On a 96 GiB card 8192 costs
2.8 GiB of scratch for 250,384 KV tokens; 32768 costs 10.1 GiB for 226,416.

Chunked prefill remains the real fix for long prompts -- it would turn this
into a pure throughput/memory knob rather than a limit on prompt length.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

pr-reviewer-bot commented

2026-08-04 20:49:58 +00:00

Collaborator

Automated review by pr-reviewer v0.44.2 | Safety Check | Claude Sonnet 5 | tracking id r-725075-4c5604
This is an AI-generated review and may contain mistakes.

Status: ✅ Completed

⚠️ The primary review (Mistral Small) failed. The review below is from Claude Sonnet 5, which reviewed the same diff.

✅ Verdict: LGTM — The changes are additive and gated behind nil/zero checks preserving prior behavior; no secrets, breaking changes, or obvious bugs found.

Went through the diff carefully given its size.

Core logic change: llama4Scale in mistral4 was replaced with a call to config.RopeParameters.Llama4AttnScale, which is documented as nil-safe (rp == nil check). mistral4.go now stores ropeParams directly (possibly nil) instead of copying out llama4Beta/llama4OrigCtx with zero defaults — since Llama4AttnScale is a method with a nil-receiver guard, this is safe and equivalent behavior for models without rope_parameters.
The FP8 scale-loading change adds a fallback list of suffixes (.weight_scale, .weight_scale_inv) and handles 0-d scalar shapes by coercing to [1]. This looks like a reasonable, backward-compatible extension (old suffix still tried first).
llama4Scales batch logic: uses max position in the batch to decide whether to skip scaling; correctly documented as monotonic in position, so this early-exit is safe.
New --max-batch-tokens flag simply wires an int into an existing struct field, default 0 meaning "use built-in default" — no red flags.
Large forward_vectors_llama4.json and .safetensors fixture: these are generated test goldens/fixtures accompanied by a generator script (generate_llama4.py), consistent with the existing testdata pattern in this repo (similar fixtures exist elsewhere). Not vendored/foreign code, just a large test golden — flagged by the automated diff stats but appears legitimate given context and thorough docstring explaining generation parameters.
Comment-only cleanups (removing "Blackwell target" references) are harmless.
No secrets, credentials, or suspicious debug code spotted.

Nothing here looks like an accidental commit, security issue, or logic bug that would break existing behavior; the new code paths are only additive (new flag, new optional scaling, new fallback for FP8 scale names) and gated behind nil/zero checks that preserve old behavior for models without llama_4 scaling.

*Automated review by [pr-reviewer](https://git.brooktrails.org/brooktrails/pr-reviewer) v0.44.2 | Safety Check | Claude Sonnet 5 | tracking id `r-725075-4c5604`* *This is an AI-generated review and may contain mistakes.* **Status:** ✅ Completed --- > ⚠️ The primary review (**Mistral Small**) failed. The review below is from **Claude Sonnet 5**, which reviewed the same diff. **✅ Verdict: LGTM** — The changes are additive and gated behind nil/zero checks preserving prior behavior; no secrets, breaking changes, or obvious bugs found. Went through the diff carefully given its size. - Core logic change: `llama4Scale` in mistral4 was replaced with a call to `config.RopeParameters.Llama4AttnScale`, which is documented as nil-safe (`rp == nil` check). `mistral4.go` now stores `ropeParams` directly (possibly nil) instead of copying out `llama4Beta`/`llama4OrigCtx` with zero defaults — since `Llama4AttnScale` is a method with a nil-receiver guard, this is safe and equivalent behavior for models without rope_parameters. - The FP8 scale-loading change adds a fallback list of suffixes (`.weight_scale`, `.weight_scale_inv`) and handles 0-d scalar shapes by coercing to `[1]`. This looks like a reasonable, backward-compatible extension (old suffix still tried first). - `llama4Scales` batch logic: uses max position in the batch to decide whether to skip scaling; correctly documented as monotonic in position, so this early-exit is safe. - New `--max-batch-tokens` flag simply wires an int into an existing struct field, default 0 meaning "use built-in default" — no red flags. - Large `forward_vectors_llama4.json` and `.safetensors` fixture: these are generated test goldens/fixtures accompanied by a generator script (`generate_llama4.py`), consistent with the existing testdata pattern in this repo (similar fixtures exist elsewhere). Not vendored/foreign code, just a large test golden — flagged by the automated diff stats but appears legitimate given context and thorough docstring explaining generation parameters. - Comment-only cleanups (removing "Blackwell target" references) are harmless. - No secrets, credentials, or suspicious debug code spotted. Nothing here looks like an accidental commit, security issue, or logic bug that would break existing behavior; the new code paths are only additive (new flag, new optional scaling, new fallback for FP8 scale names) and gated behind nil/zero checks that preserve old behavior for models without llama_4 scaling.

rcsheets merged commit 92dfbacb83 into main

2026-08-04 20:53:03 +00:00

rcsheets deleted branch fix/ministral3-serving

2026-08-04 20:53:04 +00:00

rcsheets referenced this pull request from a commit

2026-08-04 20:53:05 +00:00

Merge pull request 'feat: serve Ministral 3; add --max-batch-tokens flag' (#58) from fix/ministral3-serving into main