fix(mtp): serialize multi-ubatch decode execution

[taylorsatula]

RCA: Whole-Host Lock During MTP Prefill

Model: Qwopus3.6-27B-Fusion-Q5_K_M on dual RTX 3090 (tensor split)
(tested with multiple models and they all do it. This is not a bug in a particular model's mtp head as detailed below.)

Context: llama.cpp b9745, speculative decoding via MTP

Impact: Full host lock requiring power cycle during long-prompt prefill with MTP active.

I have read and agree with the contributing guideline = Yes
AI usage disclosure = Yes

Summary

Huge prefills (100k+ tokens) reliably locked the system within a short period when MTP speculative decoding was enabled. MTP-off completed identical requests fine (~295 s).

Root cause: MTP draft catch-up queued two GPU graphs against the same mutable KV cache without waiting for either to finish. Under prolonged tensor-split prefill, the overlapping work hit a dependency error inside the CUDA/meta backend and at least one GPU entered a non-completing kernel.

Note: The issue inside the backend is still unknown. I spent probably 6 hours and a dozen of hard-wedges solving the sync issue. I'm spent. Additionally, I didn't want to mix a speculative CUDA fix into a concrete MTP repair.

Fix: serialize MTP ubatches so each finishes before the next starts.

---

How it was finally pinpointed and reproduced reliably

Added explicit llama_synchronize(ctx_dft) after the catch-up call. The host locked at that synchronize confirming the failing work lives inside the draft context, not the target.

Both draft ubatches had already logged success (CTX-COMPUTE-EXIT status=0, AR-EXIT ok=1) because the API accepts queued work without waiting for GPU completion.

To confirm serialization alone was the cure, I added per-ubatch synchronization inside llama_context::decode() no eval callback, normal graph UID reuse. I then ran three consecutive high-prefill sweeps (100k-180k). Each completed successfully, zero failures. Great success.

Earlier tests using scheduler eval callbacks succeeded too, including a terminal-only callback creating no interior graph cut which rules out theories about K/V or attention boundaries being required.

None of these fixes were necessary when MTP was off. I can say with confidence that it is a race issue on systems with multiple GPUs and MTP.

---

Other theories disproven

Disabling CUDA graphs, switching to fallback/butterfly allreduce, disabling ASPM L1, reverting from kernel 7.0.0-29 to 7.0.0-28. None of these prevented the wedge. Memory and VRAM remained available throughout. Failure position varied from ~78k to ~143k tokens with no deterministic boundary.

These results narrow the failure to the overlap condition itself.

---

When the code was introduced

PR #22673 added llama_decode(ctx_dft, batch) for MTP catch-up on May 16, 2026 without synchronization between internally split ubatches. Has been there since MTP support was added to llama.cpp. Short-context operations fit in one ubatch so they never triggered it. This workload guarantees multiple ubatches (2048 batch vs. 1024 ubatch). This bug has been present since MTP support was added to llama.

---

Fix

Synchronize after each internal ubatch for MTP contexts:

bool mtp_multi_ubatch = false;
do {
    const auto & ubatch = mctx->get_ubatch();
    // ... process ...
    bool has_next = mctx->next();
    mtp_multi_ubatch |= has_next;
    if (cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && mtp_multi_ubatch) {
        synchronize();
    }
} while (has_next);

Only activates for MTP contexts with multiple ubatches. Target decode, non-MTP speculative decoding, and short MTP drafts pass through unchanged.

---

Regression test

The upstream-ready checkout adds test-mtp-ubatch-sync to tests/test-llama-archs.cpp.

The test builds a tiny synthetic Qwen3.5 model with one MTP layer and runs a four-token MTP batch on CPU:

• n_batch=4, n_ubatch=2: two internal ubatches. The test requires synchronization before the second ubatch starts.
• n_batch=4, n_ubatch=4: one internal ubatch. The test verifies that this path remains asynchronous.

The CPU abort callback runs while each graph is executing. synchronize() accounts the queued tokens in llama_perf_context, so the callback on the second ubatch can tell whether the first ubatch was synchronized. This tests ordering inside the actual decode loop without requiring CUDA or reproducing a host lock in CI.

Before:

MTP ubatches were not synchronized
test-mtp-ubatch-sync: Failed

After:

test-mtp-ubatch-sync: Passed (0.02 s)

Upstream issue

The existing report that most closely matches this failure is #23210: llama-server crashes on CUDA with Qwen3.6-27B. It includes Qwen3.6, CUDA, MTP, long-context prefill, multi-GPU configurations, and hard GPU/host hangs. No open PR matching the ubatch serialization fix was found.

Issue #26558 is a different MTP crash involving CUDA graph-cache growth and an invalid cuBLAS stream. Its GGML_CUDA_DISABLE_GRAPHS=1 workaround did not fix this reproducer in two controlled atts. I wanted to note it anyway for reviewers looking at the PR.

Disclosure Note

Investigative analysis was done manually and in conjunction with GPT 5.6 Sol. I sign off on the code used to patch llama. It is the proper chokepoint for the repair. I explored a bunch of different approaches and this was the right one as it does not have a performance impact in day-to-day usage and fixes the issue as close to the root (without getting into repairing the backend) as possible.

[ggml-gh-bot]

Hi @taylorsatula, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• AI-generated content: While code is allowed to be generated by AI, please write the PR description and commit messages on your own without the help of AI.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[taylorsatula]

I amended the commit message text content (still retained the structured pattern I've used forever) and pushed an update.

The code itself has not changed and I still believe the sync issue is the root cause of the issue.

Thank you for reviewing this PR and I hope it makes it into main!

[taylorsatula]

I turned the PR into a draft because I wanted to make sure the fix was durable. Since I made the initial PR I have had no hard lockups! Marking as Ready for Review.

[Anbeeld]

I think I may have a useful adjacent reproducer for the backend problem mentioned in this PR:

#27750

I've been independently investigating multi-GPU tensor-split instability on essentially the same hardware topology: 2x RTX 3090.

The main difference is that I reproduced a related failure on official upstream b10595 with MTP completely disabled.

Config was Qwen3.8-27B, Windows CUDA, tensor split 1,1, b=4096, ub=1024, mmproj enabled, MTP off. Canary and the first 24,344-token request passed, then the following image-first request reset the connection and upstream eventually aborted with a CUDA error. There was still ~8.5 GiB free on each GPU.

So I don't think this contradicts the MTP race fixed by this PR. The serialization here can still be the correct fix for the MTP-specific overlapping-ubatch problem.

My MTP-off reproduction suggests MTP overlap may be one trigger for a deeper tensor-split/backend ordering or lifetime problem rather than the only way to reach it.

Some other results from my investigation:

• disabling CUDA graphs did not fix it;
• LLAMA_GRAPH_REUSE_DISABLE=1 was also insufficient in my fork reproducer;
• ample VRAM remained available;
• failure position was not deterministic;
• the projected Q5_K operation observed near one failure works standalone on both 3090s at ~79-80 TFLOPS;
• slowing/synchronizing execution often suppresses the failure, but at a large performance cost.

Investigation checkpoint is here if useful:

https://github.com/Anbeeld/beellama.cpp/tree/v0.4.4-wddm

Would be interested to know whether this MTP-off case looks like it could be exercising the same unknown backend failure you hit underneath the MTP overlap.

[taylorsatula]

Are you on the most recent CUDA kernel (6XX)? I was having crashes with MTP off but when I updated fully it went away and never returned.

[Anbeeld]

Are you on the most recent CUDA kernel (6XX)? I was having crashes with MTP off but when I updated fully it went away and never returned.

Just updated aaand...

E CUDA error: an illegal memory access was encountered
6.31.039.965 E   current device: 0, in function ggml_cuda_ar_acquire_slot at C:/Users/anbee/projects/beellama.cpp/ggml/src/ggml-cuda/allreduce.cu:368
6.31.039.966 E   cudaEventSynchronize(p->ev_pool[i][slot].ker)

Might need to try CUDA 13.3 as well, I guess.

[47Hunter47]

Single-GPU data point: we run draft-mtp on one RTX 3090 (Qwen3.8-27B, -ub 1024, -np 1, FA on) and have not seen this on our end. If the serialization changes ubatch behavior, we can run our long-prompt prefill tests against the branch and report.

[taylorsatula]

Since yours is MTP-off, this PR's serialization won't cover it, and from your last update the driver bump didn't clear it either which is a bummer. If you can get a command line that fails on stock master (model, quant, params) plus a rough trigger rate, it's please post it as a standalone issue and link it here. I'd be happy to compare failure signatures with the MTP one.

[taylorsatula]

@am17an Hi Aman, this change touches the MTP path you originally authored. Would you please review the change and kick off the CI workflows if everything looks good.?

I am a first-time contributor to llama.cpp so the checks hang in purgatory till they're manually kicked off by a maintainer.
Thanks!

[am17an]

Please create an issue with a reproduction first. Thanks

[taylorsatula]

Issue created: #28252

Contains the full reproduction (hardware, model, server command, trigger), causal isolation controls, what was disproven, when the bug was introduced, and the proposed fix.

Happy to iterate on the issue if anything's missing.