Studio: Add Tensor-Parallel llama.cpp support#6040
Conversation
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.mcszy.xyz/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for a Tensor Parallelism toggle (--split-mode tensor) for GGUF models on multi-GPU setups, threading the option from the frontend UI to the backend. The review feedback suggests enhancing robustness by filtering out low-VRAM GPUs (less than the 5 GiB reserve) during tensor-parallel planning to avoid OOM crashes, and refactoring a fragile test that asserts on raw source code strings to use proper mocking instead.
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.
I am having trouble creating individual review comments. Click here to see my feedback.
studio/backend/core/inference/llama_cpp.py (3079-3115)
On multi-GPU setups, attempting to run tensor parallelism when one or more GPUs have less free VRAM than the required compute-graph buffer reserve (_TENSOR_PARALLEL_BUFFER_RESERVE_MIB, i.e., 5 GiB) will result in an immediate out-of-memory (OOM) error or crash on those low-VRAM devices.
We can prevent this by filtering the list of GPUs to only those that have enough free VRAM to accommodate the compute-graph buffer. If fewer than 2 GPUs meet this requirement, we should disable tensor_parallel early and fall back to the standard layer-split planning path. This avoids unnecessary process launches and crashes, providing a much smoother fallback experience.
usable_gpus = []
if tensor_parallel:
reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
usable_gpus = [g for g in gpus if g[1] >= reserve_mib]
if len(usable_gpus) < 2:
logger.info(
"Tensor parallelism requested but only %d GPU(s) have "
"enough free VRAM (>= %d MiB) for the compute buffer; "
"ignoring (needs >= 2).",
len(usable_gpus),
reserve_mib,
)
tensor_parallel = False
if tensor_parallel and gpus:
# Tensor-parallel allocation: use all GPUs, weight the split
# by (free - buffer), and cap context to the pooled VRAM
# after weights + per-device compute-graph buffers. See
# _plan_tensor_parallel for the policy + rationale.
target_ctx = (
effective_ctx
if explicit_ctx
else (self._context_length or effective_ctx)
)
(
effective_ctx,
max_available_ctx,
gpu_indices,
tp_tensor_split,
) = self._plan_tensor_parallel(
usable_gpus,
model_size,
target_ctx,
cache_type_kv = cache_type_kv,
n_parallel = n_parallel,
mtp_engaged = _mtp_will_engage,
)
use_fit = Falsestudio/backend/core/inference/llama_cpp.py (2760-2775)
To ensure _plan_tensor_parallel is fully defensive and self-contained, we should filter the input gpus list to only include devices that have enough free VRAM to hold the compute-graph buffer. This prevents planning tensor splits on low-VRAM GPUs that are guaranteed to OOM.
reserve_mib = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
usable_gpus = [g for g in gpus if g[1] >= reserve_mib]
gpu_indices = sorted(idx for idx, _ in usable_gpus)
if len(gpu_indices) < 2:
# Tensor parallelism is meaningless on <2 GPUs (the caller drops the
# toggle before this); be defensive and never emit a split here.
return (
target_ctx if target_ctx > 0 else 4096,
target_ctx if target_ctx > 0 else 4096,
gpu_indices,
None,
)
free_by_idx = {idx: free for idx, free in usable_gpus}
pool_mib = sum(free_by_idx.values())
kv_budget_b = (
(pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size
)studio/backend/tests/test_tensor_parallel.py (408-430)
Asserting on the raw source code string of routes/inference.py using string manipulation (find, rfind) is highly fragile and makes the test suite extremely brittle to minor formatting changes, refactoring, or variable renaming.
Instead, this retry behavior should be tested by mocking llama_backend.load_model to raise an exception on the first call and verify that the route catches it and retries with tensor_parallel=False.
References
- To test a function's
try...exceptblock, mock a dependency called within thetryblock to raise an exception. Do not mock the entire function and set itsside_effectto an exception, as this will not execute the function's own exception handling logic.
![chatgpt-codex-connector[bot]](https://avatars.mcszy.xyz/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: 1399280483
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6e96608243
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6e96608243
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6bf2ab4172
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2d8cb50c91
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
# Conflicts: # studio/backend/core/inference/llama_cpp.py # studio/backend/core/inference/llama_server_args.py # studio/backend/routes/inference.py
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cb0c85fe4a
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
…n rollback, and don't raise an explicit small context
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9c03e99138
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
# Conflicts: # studio/backend/core/inference/llama_cpp.py # studio/backend/core/inference/llama_server_args.py # studio/backend/routes/inference.py
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 384e167666
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| len(tp_gpus), | ||
| len(gpus), | ||
| ) | ||
| tensor_parallel = False |
There was a problem hiding this comment.
On hosts with fewer than two usable GPUs, a tensor_parallel: true load is documented as a no-op and this branch downgrades it to layer split, but the backend then reports/stores tensor_parallel == false. A client that repeats the same load request with tensor_parallel: true will miss both already-loaded guards and restart the same healthy server every time, even though the requested option cannot take effect on that hardware; preserve enough requested/effective state for this no-op downgrade to compare as satisfied.
Useful? React with 👍 / 👎.
# Conflicts: # studio/backend/core/inference/llama_cpp.py # studio/backend/routes/inference.py # studio/backend/tests/test_llama_server_args.py # studio/frontend/src/features/chat/chat-settings-sheet.tsx # studio/frontend/src/features/chat/types/api.ts
…udio-tensor-parallelism
for more information, see https://pre-commit.ci
… override the tensor planner An inherited or stale --tensor-split in llama_extra_args was appended after Studio's computed --tensor-split and won last in llama.cpp, re-introducing the asymmetric-GPU OOM tensor mode is meant to prevent. Group -ts/--tensor-split into the split-mode shadow set so it is stripped on inherit and on the layer fallback; parse_split_mode_override still keys on the mode value only.
Tensor mode aborts on a quantized KV cache, so a user with q8_0/q4_1 etc. who enabled Tensor Parallelism silently fell back to layer split. Clear the cache type (and strip inherited/explicit --cache-type) for the tensor attempt only; the layer fallback re-runs with tensor off and keeps the user's choice. Also report max_available_ctx from the native context, not an explicit small -c, so the context slider no longer warns too early in tensor mode.
When a same-model load omitted llama_extra_args, the tensor comparison resolved the raw (None) request and treated an inherited --split-mode tensor server as a mismatch, forcing a needless reload. Compare using the stored extras stripped the same way the reload strips them.
The generalized compare path loaded each GGUF without tensor_parallel, so compare ran layer split even with the toggle on and left the settings sheet stale. Send the toggle and hydrate the loaded state from the response, matching the main chat and recipe load paths.
|
Thanks for this, the tensor parallel support is really nice. I pushed a few small follow-ups onto the branch to close some edge cases we hit while testing:
All existing backend tests pass (253, including the new cases for the split-mode/tensor-split strip and the native max-context reporting). I did not run the frontend typecheck here. Thanks again for the contribution. |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4671df30af
ℹ️ 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 ( | ||
| tensor_parallel | ||
| and cache_type_kv | ||
| and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES | ||
| ): |
There was a problem hiding this comment.
When tensor mode is requested with mixed cache extras such as ['--cache-type-k', 'q8_0', '--cache-type-v', 'f16'], parse_cache_override resolves only the last cache flag (f16), so this guard does not strip the extra args. The command later emits Studio's --cache-type-k/--cache-type-v pair and then appends user extras last, leaving --cache-type-k q8_0 active; per the new comment here, quantized KV aborts tensor loads, so this avoidably falls back to layer split even though the tensor attempt was supposed to drop quantized cache overrides. Check both key and value overrides (or strip cache extras whenever any quantized cache flag is present) before the tensor attempt.
Useful? React with 👍 / 👎.
The headless one-liner could only reach tensor mode by passing --split-mode tensor as a raw llama.cpp extra. Add a first-class --tensor-parallel/ --no-tensor-parallel option that sets the tensor_parallel field on the /api/inference/load payload, forwarded through the studio-venv re-exec like the other polarity flags. Matches the web UI toggle and the API field.
|
Added a first-class It sets the |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9ba90b9e70
ℹ️ 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 ( | ||
| tensor_parallel | ||
| and cache_type_kv | ||
| and cache_type_kv.strip().lower() not in self._TENSOR_PARALLEL_KV_TYPES | ||
| ): |
There was a problem hiding this comment.
When an API caller requests tensor mode with a first-class quantized KV setting (for example tensor_parallel: true, cache_type_kv: "q8_0"), this block clears the cache type for the actual tensor launch and stores/reports self._cache_type_kv as None. A second identical load that races behind the first, or any client that repeats the same explicit request, then misses both already-loaded guards because they compare the raw requested q8_0 against the stored None, so Studio kills and restarts a healthy tensor server even though the request resolves to the same effective command. Normalize quantized KV to None in the target-state comparisons when tensor mode is effective, or reject the combination instead of silently changing it here.
Useful? React with 👍 / 👎.
Resolve conflicts by keeping both intents: - llama_server_args.py + its tests: split-mode/tensor-parallel helpers alongside main's new mmproj-disable helpers. - use-chat-model-runtime.ts: adopt main's applyActiveModelStatusToStore refactor and port the tensor_parallel status hydration into that helper.
|
Merged latest
Backend tests pass locally (256, including the new mmproj cases). The PR shows mergeable again. |
for more information, see https://pre-commit.ci
…rashes) (#6324) * Disable MTP speculative decoding under tensor parallelism 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. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make tensor-parallel MTP gate test format-independent The assertion pinned the multi-line `speculative_type = (` form, but ruff collapses it onto one line, so match `speculative_type =` instead. * Recover from MTP+tensor-parallel crashes at runtime instead of banning 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. * Address review feedback on the MTP runtime fallback - 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. * Cover every request path with the MTP+tensor crash recovery via a watchdog 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. * Harden MTP+tensor crash recovery: stale-load race, pass-through MTP, 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). * Trim verbose comments in the MTP+tensor crash recovery 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). --------- Co-authored-by: danielhanchen <michaelhan2050@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
4 participants
This adds tensor parallelism support for llama.cpp through a new "Tensor Parallelism" checkbox:
When checked, it passes
--split-mode tensortollama-server.On single GPU setups, it gets ignored even if checked.
How
The PR doesn't just pass the flag when the checkbox is checked, it also integrates it with Studio's existing memory allocation framework. Specifically, a new policy was added for the tensor parallel case:
total free VRAM - weights - N*reserve, wherereserveis the per-GPU compute buffer), sized across all GPUs at once (reusing_estimate_kv_cache_bytes).--tensor-splitweighted byfree - reserveper GPU, but only on asymmetric setups where an even split would overflow the smallest GPU.--split-mode tensor), fall back to layer split automatically so the model loads regardless.Each of these changes is necessary to prevent out of memory errors and crashes. They were obtained after extensive testing rather than static code analysis. My asymmetric multi-gpu setup was valuable for developing this PR because it's the hardest type of setup to get tensor parallelism right (it OOMs by default).
About the per-GPU buffer reserve
In tensor mode, llama.cpp allocates a compute-graph buffer on every GPU (the logits buffer,
n_batch x vocab, plus activation scratch). Unlike the KV cache, it can't be derived from the context length: llama.cpp sizes it internally viagraph_reserve, it's roughly equal on each device regardless of the split, and it barely changes with context. In practice it ranged from ~2.3 GB (gemma-3-27B) to ~3.8 GB (gemma-4-31B) on a 256k-vocab model. And--fitis disabled in tensor mode, so nothing else caps the KV cache: underestimate this and the load OOMs.So instead of computing it, the PR reserves a fixed 5 GiB per GPU, above the measured worst case. It's subtracted from each GPU's free VRAM before computing
--tensor-split, and held back per device when capping context. It scales with vocab/batch, not context, so a constant works; the auto-fallback to layer split covers any underestimate.Speed measurements
RTX 6000 Ada (48 GB) + RTX 3090 (24 GB), PCIe, no NVLink, no NCCL. Decode speed in tokens/sec, measured during a 512-token generation after a ~5.9k-token prompt (i.e. at a realistic context depth, not an empty context), layer split (default) vs tensor split. This is a weak interconnect, so NVLink/NCCL setups should see more.
Dense models (tensor parallelism improves decode, more so for larger models):
MoE models (no benefit, can be much slower):
For reference, VRAM splits across both GPUs as expected. For example, gemma-3-27B uses 29.6 GB on a single GPU with layer split, and 16.1 / 15.9 GB split across both with tensor.