Keep non-MLA draft K caches quantized

When an MLA target uses a separate non-MLA drafter, such as the supported Gemma drafter path, mla is derived only from the target's _kv_lora_rank, but this list also rewrites the drafter's K flags. The K/V equality restriction is evaluated against the drafter's own model, so changing its quantized K cache to f16 is unnecessary and can substantially increase memory use during the flash-attention fallback, potentially turning the recovery into an OOM. Determine MLA status separately for the draft model or exclude draft K flags from the target-MLA rewrite.

Useful? React with 👍 / 👎.

Correct, and it is worse than stated in the other direction. Fixed in 9328343.

The target and the drafter are separate models with separate contexts and llama.cpp applies the restriction per context, so gating the draft flags on the target's answer was wrong both ways: an MLA target with a non-MLA drafter needlessly doubles the drafter's K cache, which is the OOM the size argument exists to avoid, and a non-MLA target with an MLA drafter leaves the draft K quantized against an f16 draft V and aborts the draft context outright.

Each side is now gated on its own metadata. Studio already reads the drafter's GGUF via _draft_backend_for for KV sizing, so the drafter's own signal was available. When the drafter cannot be read the draft signal is None and falls back to the target, which is the safe direction: an unnecessary reset costs memory, a missing one aborts.

Include DeepSeek4 in the symmetry signal

For a DeepSeek4 GGUF without attention.kv_lora_rank, this passes mla=False even though the upstream restriction quoted in this change explicitly applies when model->arch == LLM_ARCH_DEEPSEEK4 independently of is_mla(). A quantized-K/V launch that reaches the flash-attention fallback therefore still becomes K=q8_0/V=f16 and aborts instead of recovering. Include self._architecture == "deepseek4" when deriving this signal at every call site.

Useful? React with 👍 / 👎.

Confirmed against the upstream source and fixed in 9328343.

The arch term is not redundant with is_mla(). is_mla() is n_embd_head_k_mla_impl and n_embd_head_v_mla_impl both being non-zero (llama-hparams.cpp), and only deepseek2, deepseek32, glm-dsa, kimi-k3 and kimi-linear set those. DeepSeek4 has its own KV cache in llama-kv-cache-dsv4.cpp and sets neither, which is exactly why upstream needed the explicit arch clause.

The converter side confirms the kv_lora_rank probe cannot see it: DeepseekV4Model subclasses TextModel rather than DeepseekV2Model, and add_kv_lora_rank is only called in DeepseekV2Model. DeepSeek4 writes q_lora_rank and attention_output_lora_rank and no kv_lora_rank at all.

_requires_symmetric_kv now mirrors the upstream condition. The sweep carries a deepseek4 model and scores llama.cpp's ground truth separately from what Studio infers, and DeepSeek4 picks up the same 8064 fixed launches as the is_mla() models.

Preserve MLA symmetry for environment-only cache types

When an MLA launch inherits equal quantized K/V types through LLAMA_ARG_CACHE_TYPE_K and LLAMA_ARG_CACHE_TYPE_V, this call can only rewrite argv; the recovery block immediately afterward drops the quantized V environment variable while preserving K. The retry consequently runs with K=q8_0 and the default V=f16 and hits the same MLA mismatch abort this change is intended to prevent. On MLA retries, clear the corresponding quantized K environment variable whenever the V variable is cleared, including the draft pair and the flagless-build path.

Useful? React with 👍 / 👎.

Correct. This was already fixed in 191c0c4, which landed after the commit this review is anchored to (2b395dc).

_drop_env_quantized_v_cache now takes the K vars down with the V ones, and the argv reset no longer makes the K rewrite conditional on having just reset V there, which covers the argv-K-only shape: -ctk q8_0 on the command line with LLAMA_ARG_CACHE_TYPE_V in the environment, where nothing quantized appears on argv for V so a V-triggered reset never fires. As of 9328343 the draft var is gated on the drafter's own model rather than the target's.

Found independently by the same sweep before the review landed, which is a good sign for both.

Derive symmetry from the drafter that actually launches

When extra_args supplies its own --model-draft/--hf-repo-draft, that model is the drafter loaded by llama.cpp, but this lookup uses only self._mtp_draft_path, which is assigned from launch_mtp_draft_path and may be absent or point to an auto-resolved sidecar that _build_speculative_flags does not emit. Consequently, a non-MLA target with a user-supplied MLA drafter leaves the draft K cache quantized when V is reset and the recovery still aborts; the inverse combination can unnecessarily expand a non-MLA drafter's K cache. The new per-model gating is fresh evidence beyond the earlier comment, but its signal must be derived from the effective draft flag/path in the launch command rather than this stored managed path.

Useful? React with 👍 / 👎.

Correct, and the stored path was wrong for a second reason you did not name. Fixed in 0334942.

self._mtp_draft_path is assigned at the end of the launch path, well after the launch-site reset runs, so at that call site it held whatever the previous load left, or nothing at all. So it was stale there even with no extra_args involved.

The signal now comes from the command being launched rather than any stored path, via the existing _extra_args_mtp_draft_source, which already resolves the local and HF draft flags last-wins and falls back to LLAMA_ARG_SPEC_DRAFT_MODEL. _build_speculative_flags emits --model-draft into that same command, so one source covers the managed sidecar and a user-supplied drafter alike, and every call site passes the command it is about to launch or retry. An HF repo id is not a local file whose metadata can be read, so it returns None and falls back to the target, as does an unreadable or absent drafter.

I removed _kv_symmetry_signals rather than leaving it beside the new path, since a second way to derive the same signal is how the stale source got used to begin with.

Initialize the child environment before inspecting the drafter

When the capability probe authoritatively reports a build without --flash-attn, _flash_attn_known_off is true and this branch executes before env is assigned at line 15759. Because env is a local variable in load_model, evaluating this argument raises UnboundLocalError, so every load through the flagless-build compatibility path fails before spawning llama-server. Construct the child environment before this fixup or pass an already initialized environment view.

Useful? React with 👍 / 👎.

Confirmed and fixed in 7e23c97. This one was mine, introduced in 0334942.

Verified with an AST walk over load_model rather than by eye: exactly one Load of env sits before the first Store, at the line this comment is on. The apparent earlier binding a few lines up is a keyword argument (env = _mem_env inside a call), not an assignment, so there is no earlier binding to save it. On a build whose probe authoritatively reports no --flash-attn the branch does execute, so the flagless-build path would raise before llama-server was ever spawned.

The lookup no longer takes env there. os.environ is the correct source at that point anyway, since Studio never rewrites LLAMA_ARG_SPEC_DRAFT_MODEL and the child env is a copy of it plus loader paths.

Worth recording why the tests missed it: the block-extraction harness the flagless tests use seeds env into its exec scope, so it passed while the real function raised. I replaced that with a static check on load_model itself and confirmed it fails on the broken code and passes on the fix, rather than assuming a new test works.

Reset draft K conservatively when drafter metadata is unknown

When a non-MLA target launches an MLA drafter through --hf-repo-draft, or a local drafter whose metadata cannot be read, _draft_kv_symmetry returns None and this fallback converts it to the target's False. A flash-attention retry then resets the quantized draft V cache to f16 while leaving draft K quantized, causing the drafter context to abort on the K/V mismatch. Fresh evidence after the earlier thread is that the current implementation explicitly returns None for every remote drafter and still maps that unknown result to the target signal; unknown drafters should take the conservative symmetric path instead.

Useful? React with 👍 / 👎.

Agreed, fixed in 7e23c97. The implementation contradicted its own stated rationale: the docstring argued that an unnecessary reset only costs memory while a missing one aborts, and then mapped the unknown case to the target's possibly-False answer.

A named drafter whose metadata cannot be read now resolves to True, covering both the Hugging Face repo id and the unreadable local GGUF. None is returned only when no drafter is named at all, where there are no draft flags to reset and the value cannot matter.

This is consistent with the earlier thread rather than a reversal of it: that one was about a KNOWN non-MLA drafter, where the size argument applies and the draft K is left quantized. Only the unknown case takes the conservative path.