Runtime MTP fallback for tensor parallelism (try MTP, recover if it crashes)#6324
Conversation
Follow-up to #6040 (Studio tensor-parallel support). MTP-draft speculative decoding plus --split-mode tensor crashes the CUDA flash-attn kernel at decode time. The startup /health probe only checks that llama-server comes up, so the existing MTP-drop fallback (keyed on startup health) never fires and the server dies on the first generation instead. Gate MTP off when a tensor attempt actually engages: this runs before the VRAM planner (so no drafter memory is reserved) and before the speculative flag build (so no --model-draft / --spec-type is emitted). Ngram modes use no draft model and are kept, and mtp+ngram degrades to ngram rather than off. The layer-split fallback re-runs with tensor_parallel False and restores MTP. The reason is surfaced as spec_fallback_reason "tensor_parallel" so the settings sheet explains why MTP is off instead of prompting a llama.cpp update. Verified on unsloth/gemma-4-26B-A4B-it-GGUF:UD-Q4_K_XL across 4x B200: the load now emits --split-mode tensor with no MTP flags and generation completes without the prior decode crash.
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request disables MTP speculative decoding when tensor parallelism is active to prevent crashes in the CUDA flash-attention kernel, updates the frontend to display a warning banner explaining this fallback, and adds a backend test to verify the logic. The review feedback identifies a critical bug where disabling MTP overwrites self._requested_spec_mode, leading to expensive, spurious model reloads and hiding the frontend warning banner. To resolve this, it is recommended to restore self._requested_spec_mode to _mtp_canonical when MTP is disabled under tensor parallel, and to add a corresponding assertion in the test suite.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # is off so the UI explains it (not a stale-binary prompt). | ||
| self._spec_fallback_reason = "tensor_parallel" | ||
|
|
||
| # Apply custom chat template override if provided. |
There was a problem hiding this comment.
�1�c Critical Bug: Spurious Model Reloads & Hidden UI Warning Banner
When _tensor_disabled_mtp is True, speculative_type is modified to "off" or "ngram" before calling _build_speculative_flags. Consequently, _build_speculative_flags overwrites self._requested_spec_mode to "off" or "ngram".
This introduces two major issues:
- Spurious Model Reloads: Any subsequent check or duplicate load request via
_already_in_target_statewill compare the user's original requested mode (e.g.,"mtp") againstself._requested_spec_mode(now"off"or"ngram"). Since they do not match, it returnsFalseand triggers an expensive model reload on every request. - Hidden UI Warning Banner: The frontend settings sheet only displays the warning banner explaining why MTP is disabled under tensor parallel if
speculativeTypeis"mtp"or"mtp+ngram". Sinceself._requested_spec_modeis overwritten to"off"or"ngram", the banner is never shown.
Overwriting self._requested_spec_mode back to _mtp_canonical when _tensor_disabled_mtp is True resolves both issues perfectly.
| # is off so the UI explains it (not a stale-binary prompt). | |
| self._spec_fallback_reason = "tensor_parallel" | |
| # Apply custom chat template override if provided. | |
| if _tensor_disabled_mtp: | |
| # _build_speculative_flags cleared the reason; record why MTP | |
| # is off so the UI explains it (not a stale-binary prompt). | |
| self._spec_fallback_reason = "tensor_parallel" | |
| self._requested_spec_mode = _mtp_canonical |
| assert "_mtp_will_engage = False" in body | ||
| assert 'self._spec_fallback_reason = "tensor_parallel"' in src |
There was a problem hiding this comment.
Add an assertion to verify that self._requested_spec_mode is correctly restored to _mtp_canonical when MTP is disabled under tensor parallel. This ensures that _already_in_target_state works correctly and prevents future regressions.
| assert "_mtp_will_engage = False" in body | |
| assert 'self._spec_fallback_reason = "tensor_parallel"' in src | |
| assert "_mtp_will_engage = False" in body | |
| assert 'self._spec_fallback_reason = "tensor_parallel"' in src | |
| assert 'self._requested_spec_mode = _mtp_canonical' in src |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 370e3c1699
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| speculative_type = ( | ||
| "ngram" if _mtp_canonical == "mtp+ngram" else "off" | ||
| ) |
There was a problem hiding this comment.
In the tensor-parallel + MTP path this overwrites the user's requested mode before _build_speculative_flags, so self._requested_spec_mode is recorded as off (or ngram) instead of the original auto/mtp/mtp+ngram. The load/status responses echo that requested mode, and the settings sheet only shows the new tensor_parallel fallback banner when speculativeType is still auto, mtp, or mtp+ngram; after a successful load the frontend therefore hides the explanation and loses the user's MTP preference, so turning Tensor Parallelism off later keeps sending off/ngram rather than restoring MTP.
Useful? React with 👍 / 👎.
| # (ngram needs no draft model, so it stays); the layer-split | ||
| # fallback re-runs with tensor_parallel False and restores MTP. | ||
| # See llama.cpp common_speculative_impl_draft_mtp. | ||
| if tensor_parallel and _mtp_will_engage: |
There was a problem hiding this comment.
Because this gate lives inside the GPU-selection try and after calls like _get_gguf_size_bytes() / _get_gpu_free_memory(), any probe/planning exception skips it; the existing except then continues with tensor_parallel still true, and the later command still appends --split-mode tensor plus the original MTP spec flags. On hosts where the probe fails but llama-server --fit can still start, this leaves the exact tensor+MTP decode-time crash path enabled instead of applying the new fallback.
Useful? React with 👍 / 👎.
| # (ngram needs no draft model, so it stays); the layer-split | ||
| # fallback re-runs with tensor_parallel False and restores MTP. | ||
| # See llama.cpp common_speculative_impl_draft_mtp. | ||
| if tensor_parallel and _mtp_will_engage: |
There was a problem hiding this comment.
This condition relies on _mtp_will_engage, but that flag is deliberately false when llama_extra_args contains --spec-type / --spec-default; those extras are then appended after Studio's flags, so a tensor load with llama_extra_args=["--spec-type", "draft-mtp"] still launches --split-mode tensor plus MTP and can hit the same decode-time crash this change is trying to avoid. Tensor mode already strips conflicting extra cache/split flags for safety, so the MTP spec extras need the same treatment when tensor parallelism is active.
Useful? React with 👍 / 👎.
The assertion pinned the multi-line `speculative_type = (` form, but ruff collapses it onto one line, so match `speculative_type =` instead.
…g MTP MTP-draft speculative decoding under --split-mode tensor usually works, but can crash llama-server's CUDA flash-attn kernel at decode time (the prompt-cache checkpoint-restore path). The earlier fix statically disabled MTP whenever tensor parallelism was on, which is not future-proof and gives up the MTP speedup even though it normally works. Replace the static ban with a try/recover, mirroring the existing load-time MTP-drop fallback: - Load-time decode probe: after the server passes /health under tensor + MTP, run one tiny /completion to exercise the draft path. A failure flips the load unhealthy so the existing fallback respawns with --spec-default. Catches a hard incompatibility that crashes on the first decode. - Generation-time recovery: snapshot the load kwargs after a healthy load, and if llama-server exits mid-generation while MTP + tensor parallelism were active, quietly reload the same model with speculative decoding off (one single-flight background reload) and surface spec_fallback_reason=runtime_error. Catches the rare mid-generation crash the probe and load-time fallback miss. No persistent ban: a later fresh load re-tries MTP, so this self-heals if a future llama.cpp supports the combo. Verified on gemma-4-26B-A4B + 4x B200: MTP runs normally, and killing llama-server mid-generation reloads it without MTP and serves the next request cleanly.
danielhanchen
changed the title
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 81cd30e9bd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| url = f"http://127.0.0.1:{self._port}/completion" | ||
| payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False} | ||
| try: | ||
| resp = httpx.post(url, json = payload, timeout = timeout) |
There was a problem hiding this comment.
When UNSLOTH_DIRECT_STREAM=1 is enabled, this same load path starts llama-server with --api-key, and the other internal llama.cpp requests in this backend attach Authorization: Bearer <self._api_key>. This new probe is unauthenticated, so tensor-parallel MTP loads in direct-stream mode will see the probe rejected (for example 401), mark healthy = False, and incorrectly fall back to no MTP even when the server is otherwise healthy.
Useful? React with 👍 / 👎.
| snapshot["speculative_type"] = "off" | ||
| self.load_model(**snapshot) |
There was a problem hiding this comment.
If the user unloads/cancels after _maybe_recover_from_mtp_crash has scheduled the background thread but while _recover is waiting for proc.poll() (up to five seconds), the captured snapshot is still replayed here; load_model() then clears _cancel_event and reloads the model the user just unloaded. Re-check the cancel flag or that the original load snapshot is still current immediately before calling load_model.
Useful? React with 👍 / 👎.
| except Exception as e: | ||
| # 200 headers already flushed; errors must go in the SSE body. | ||
| logger.error("openai passthrough stream error: %s", e) | ||
| get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) |
There was a problem hiding this comment.
For OpenAI-compatible streaming passthrough, the usual mid-stream upstream death is surfaced as httpx.RemoteProtocolError, ReadError, or CloseError, but those are caught by the preceding typed handler and re-raised when the request was not cancelled. Because the recovery call is only in the later generic except, tensor-parallel MTP crashes on this path do not schedule the no-MTP reload and subsequent requests keep the same crashing configuration.
Useful? React with 👍 / 👎.
| return | ||
| # Server died mid-generation? Quietly reload without MTP if this was | ||
| # a tensor-parallel + MTP crash; re-raise unchanged for this request. | ||
| self._maybe_recover_from_mtp_crash(e) |
There was a problem hiding this comment.
In the standard GGUF streaming path, a server that has already exited before the next request connects is caught by the except httpx.ConnectError branch immediately above this added recovery call, which re-raises RuntimeError without invoking _maybe_recover_from_mtp_crash; the route-level streaming catch only emits an SSE error. For tensor-parallel MTP crashes that leave the server down before reconnect, no background no-MTP reload is scheduled, so subsequent streaming requests keep hitting the dead/crashing configuration.
Useful? React with 👍 / 👎.
| except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: | ||
| if not cancel_event.is_set(): | ||
| logger.error("anthropic_messages passthrough stream error: %s", e) | ||
| get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) |
There was a problem hiding this comment.
This recovery call only runs for the typed mid-body stream errors; if the tensor-parallel MTP crash happens before the Anthropic stream receives response headers (or the first-token wait raises a ReadTimeout), control falls into the generic except Exception immediately below, which only emits an error event and never schedules the no-MTP reload. In that pre-header crash case, subsequent Anthropic requests keep using the same dead/crashing configuration.
Useful? React with 👍 / 👎.
- Authenticate the decode probe: direct-stream mode runs llama-server with --api-key, so the unauthenticated /completion probe got a 401 and falsely dropped MTP. Attach the same bearer auth the other internal requests use. - Re-check the cancel flag inside the recovery thread after the death poll, so an /unload that races the reload can't resurrect the dropped model. - Schedule the no-MTP recovery on the connection-error paths it was missing: generate_chat_completion's ConnectError branch, the OpenAI passthrough typed (RemoteProtocolError/ReadError/CloseError) stream catch, and the Anthropic passthrough generic stream catch. Previously a server that died before reconnect, or a typed mid-stream error, skipped the reload.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 48ce0139c9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| snapshot["speculative_type"] = "off" | ||
| self.load_model(**snapshot) |
There was a problem hiding this comment.
If a tensor+MTP crash schedules this background recovery and the user starts loading a different model during the 5-second poll, _cancel_event is clear and this thread still replays the stale snapshot; load_model(**snapshot) will serialize behind the user's load, kill that newly loaded server, and restore the old crashed model with spec off. Re-check that the captured process/snapshot is still the active load immediately before replaying it, not just that unload/cancel was not requested.
Useful? React with 👍 / 👎.
| # Watcher closed resp on cancel. Emit nothing extra; the client | ||
| # initiated the cancel or already disconnected. | ||
| if not cancel_event.is_set(): | ||
| get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) |
There was a problem hiding this comment.
This recovery call only covers errors after the OpenAI passthrough stream has received 200 headers. If llama-server has already exited before _send_stream_with_preheader_cancel returns, the separate except httpx.RequestError branch above raises a 502 without calling _maybe_recover_from_mtp_crash, so tensor-parallel MTP crashes on that path leave the backend on the same dead/crashing configuration for subsequent OpenAI streaming requests.
Useful? React with 👍 / 👎.
| logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") | ||
| # If llama-server died mid-stream from a tensor-parallel + MTP | ||
| # crash, quietly reload without MTP so the next request works. | ||
| get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) |
There was a problem hiding this comment.
This wires the runtime reload for OpenAI tool streaming, but the Anthropic Messages server-tool path calls the same generate_chat_completion_with_tools generator and then catches its exceptions in _anthropic_tool_stream / _collect_anthropic_events without invoking _maybe_recover_from_mtp_crash. When a local GGUF Anthropic tool request hits the tensor-parallel MTP crash, it only emits/maps an error and leaves the next request on the same dead/crashing server; add the same recovery there or move it into the tool generator itself.
Useful? React with 👍 / 👎.
| and self._tensor_parallel | ||
| and _spec_requested_mtp | ||
| and not self._cancel_event.is_set() | ||
| and not self._probe_mtp_decode() |
There was a problem hiding this comment.
When an unload/cancel arrives after this condition has checked _cancel_event but while _probe_mtp_decode() is blocked in its 60-second httpx.post, unload_model() cannot kill the subprocess because load_model() still holds _lock. In the hung-first-decode case this leaves the UI waiting for the probe timeout before the cancelled load unwinds; make the probe cancellable or run it where unload can interrupt the process.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3bd2c7ba54
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if self._cancel_event.is_set(): | ||
| logger.info("MTP-crash reload skipped: load was cancelled/unloaded.") | ||
| return | ||
| snapshot["speculative_type"] = "off" |
There was a problem hiding this comment.
Fresh evidence for the requested-mode regression is that this runtime recovery mutates the saved load request to off before replaying it. In the tensor-parallel MTP crash case, load_model() commits _requested_spec_mode from that mutated value, and /status / /load responses report speculative_type='off' even though the user's request was auto/mtp/mtp+ngram; the settings flow then treats MTP as user-disabled, so turning Tensor Parallelism off later won't re-try MTP despite this fallback being intended to be non-persistent. Keep the original requested mode separate from the launch-time no-MTP override.
Useful? React with 👍 / 👎.
| return False | ||
| return True | ||
|
|
||
| def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: |
There was a problem hiding this comment.
This recovery helper only works on generation paths that invoke it, but the direct llama-server proxy endpoints still bypass it: checked openai_completions in routes/inference.py (its stream/non-stream catches at 5203-5214 and 5240-5248) and _responses_stream (6140-6157 and 6314-6331), and those paths emit an error without calling _maybe_recover_from_mtp_crash. When a tensor-parallel MTP crash is triggered through /v1/completions or streaming /v1/responses, the server remains on the same dead/crashing configuration instead of scheduling the no-MTP reload.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 22743ca4d9
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| "MTP speculative decoding crashed on the first decode " | ||
| "under tensor parallelism; retrying without it." | ||
| ) | ||
| healthy = False |
There was a problem hiding this comment.
When the requested mode is mtp+ngram, this probe-failure path falls through to the generic MTP fallback below, which replaces the entire spec block with --spec-default. In the tensor-parallel decode-crash case only the MTP drafter is known-bad, while the requested ngram half is still safe and useful; as written, users who explicitly selected mtp+ngram lose all speculative decoding instead of falling back to ngram-only.
Useful? React with 👍 / 👎.
| except Exception as e: | ||
| if not cancel_event.is_set(): | ||
| logger.error("anthropic_messages passthrough stream error: %s", e) | ||
| get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) |
There was a problem hiding this comment.
This recovery is only installed in the Anthropic streaming passthrough path, but the stream=false client-tool path still calls _anthropic_passthrough_non_streaming, which posts to the same local llama-server endpoint without catching transport failures or invoking _maybe_recover_from_mtp_crash. When an Anthropic Messages request with client tools and stream=false triggers the tensor-parallel MTP crash, that request errors and the backend remains on the same dead/crashing server for subsequent calls.
Useful? React with 👍 / 👎.
…chdog The runtime MTP-crash recovery only fired from request handlers that observed the failure, so the direct llama-server proxy endpoints (/v1/completions, /v1/responses, the OpenAI/Anthropic passthrough transports) -- and a crash with no request in flight -- could leave a dead server. Add a single background watchdog, armed only on a healthy MTP + tensor-parallel load, that polls the subprocess and routes an unexpected death into the existing single-flight no-MTP reload. It is stopped inside _kill_process (the one deliberate-termination chokepoint) so a planned reload/unload is never mistaken for a crash, and re-checks the stop flag after a detected exit to close the kill-vs-poll race. The reload turns MTP off, so the replacement server arms no watchdog and the fallback cannot loop; a later fresh load still re-tries MTP.
# Conflicts: # studio/backend/routes/inference.py
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
|
/gemini review |
|
@codex review |
There was a problem hiding this comment.
Code Review
This pull request implements a robust recovery mechanism for mid-generation llama-server crashes when running Multi-Token Prediction (MTP) speculative decoding under tensor parallelism. It introduces a startup decode probe, a background watchdog thread to monitor process liveness, and automatic quiet reloads without MTP upon detecting a crash. Feedback from the review highlights a potential race condition where a background recovery could overwrite a newer model load, and suggests storing a shallow copy of extra_args in the load snapshot to prevent mutation issues.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| snapshot["speculative_type"] = "off" | ||
| self.load_model(**snapshot) |
There was a problem hiding this comment.
There is a potential race condition where a background recovery reload (_recover) could overwrite a newer model load that was initiated or completed by the user while the recovery was waiting for the lock. Checking if self._process has changed (i.e., self._process is not proc) before calling self.load_model prevents this from happening.
| snapshot["speculative_type"] = "off" | |
| self.load_model(**snapshot) | |
| if self._process is not proc: | |
| logger.info("MTP-crash reload skipped: a newer load is in progress or completed.") | |
| return | |
| snapshot["speculative_type"] = "off" | |
| self.load_model(**snapshot) |
| "n_threads": n_threads, | ||
| "n_gpu_layers": n_gpu_layers, | ||
| "n_parallel": n_parallel, | ||
| "extra_args": extra_args, |
There was a problem hiding this comment.
The extra_args list is stored by reference in _pending_load_kwargs. If the caller mutates the extra_args list after calling load_model, it could inadvertently mutate the snapshot stored in self._last_load_kwargs. It is safer to store a shallow copy of the list.
| "extra_args": extra_args, | |
| "extra_args": list(extra_args) if extra_args is not None else None, |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4632cdb483
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if self._cancel_event.is_set(): | ||
| logger.info("MTP-crash reload skipped: load was cancelled/unloaded.") | ||
| return | ||
| snapshot["speculative_type"] = "off" | ||
| self.load_model(**snapshot) |
There was a problem hiding this comment.
When a request-side recovery is queued and, during the 5s poll, the user starts another load/unload, that path kills this proc and then load_model clears _cancel_event for the new load; this check can therefore pass and the stale snapshot below is replayed, killing the newly selected model. Verify the watched process is still current (or use a generation token) before calling load_model.
Useful? React with 👍 / 👎.
| snapshot["speculative_type"] = "off" | ||
| self.load_model(**snapshot) |
There was a problem hiding this comment.
When this runtime fallback reloads with speculative_type="off", load_model resets _requested_spec_mode to off; after a page/status refresh the UI reads that value and the fallback warning is only shown for auto/mtp/mtp+ngram, so users who had Auto/MTP selected can see plain Off with no runtime-error explanation and future Apply calls won't retry MTP. Keep the reload command MTP-free, but restore the original requested mode from the snapshot, matching the startup MTP fallback behavior.
Useful? React with 👍 / 👎.
…requested mode Address review findings on the runtime MTP-crash recovery: - Stale-load race: the recovery thread snapshotted the crashed load, waited up to 5s for the process to confirm dead, then only checked the cancel flag before replaying load_model. A concurrent user load clears that flag, so the stale snapshot could reload the old model over the user's new one. Make the load lock re-entrant and run the staleness check (cancel + same process + unchanged snapshot) under it, atomically with the reload. - Pass-through MTP: MTP can also be requested via a user --spec-type in extra_args or LLAMA_ARG_SPEC_TYPE, where Studio emits no spec flags and _speculative_type stays unset, so the probe/watchdog/recovery never engaged. Track _mtp_runtime_fallback_active from the actual launched config and gate on it; on the no-MTP reload, append a last-wins --spec-default so the replay drops MTP regardless of source (and the load-time fallback does the same). - Requested mode: the off-reload reset _requested_spec_mode to off, so after a status refresh the UI showed a bare Off with the runtime-error note suppressed and would not retry MTP. Restore the original requested mode after the reload, matching the startup MTP fallback. - Snapshot the extra_args list by value so a caller mutating it cannot corrupt the recovery snapshot. Tests: test_tensor_parallel.py + test_llama_server_args.py green (303 passed).
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
Tighten the docstrings and inline comments added for the runtime MTP recovery (watchdog, probe, reload, gating) to succinct one/two-line forms; no code change (verified comment-only).
2 participants
Follow-up to #6040 (Studio tensor-parallel support).
Problem
MTP-draft speculative decoding combined with
--split-mode tensorcrashes the CUDA flash-attn kernel at decode time (ggml-cuda/fattn.cu, incommon_speculative_impl_draft_mtp). The server loads fine, then dies on the first generation.The existing MTP-drop fallback is keyed on startup health (the
/healthprobe only checks thatllama-servercomes up), so it never fires for a decode-time crash. For example, loadinggemma-4-26B-A4B-itwith the toggle on resolves to a separate MTP drafter, passes/health, and then crashes on the first request withpeer closed connection without sending complete message body.Fix
Gate MTP off when a tensor attempt actually engages. The gate sits after the
< 2 GPUdowngrade (so it only triggers when--split-mode tensoris really emitted) and runs:--model-draft/--spec-typeis emitted.Ngram modes use no draft model and are kept;
mtp+ngramdegrades tongramrather thanoff. The existing layer-split fallback re-runs withtensor_parallelFalse and restores MTP, so layer mode is unaffected.The reason is surfaced as
spec_fallback_reason = "tensor_parallel"and the settings sheet now explains why MTP is off instead of prompting a llama.cpp update.Testing
test_mtp_is_disabled_under_tensor_parallel; fulltest_tensor_parallel.py+test_llama_server_args.pysuite green (257 passed).unsloth/gemma-4-26B-A4B-it-GGUF:UD-Q4_K_XLacross 4x B200: the load now emits--split-mode tensorwith no MTP flags, the status reportsspeculative_type=off/spec_fallback_reason=tensor_parallel, and generation completes without the prior decode crash.