Issues / #3664
#3664 [Engine] mtp graph capture OOM exits the engine (reported as "no error") with STRATA_QSA_SPLIT_ATTN + --batch-mtp + a fresh long prompt
open · @Adam34332 · 2 comments · View on GitHub
BenchmarksSetup & installNVIDIA / CUDAModels & quantsDocumentationWindowsLinux
Description
### Area
[Engine] (speed, kernels, models, the server and API)
### Desktop app, Terminal or source?
Source (START-HERE.bat, setup.sh or a build)
### What happened
With `STRATA_QSA_SPLIT_ATTN=1` and `--batch-mtp`, when a **fresh** long prompt (no prompt-cache reuse; past the resident KV window - we tested 155K-176K tokens) arrives while several slots are decoding, the engine exits with status 1 and this log line:
```
mtp: round capture: no error
```
The stock engine crashes in this scenario 5 out of 5 runs. Expected: the opt-in switches are documented as measured hints, so this should degrade (disable the switch for the session, print a WARNING) instead of exiting - and the error message should show the real error, not "no error".
### Smallest steps or command that reproduce it
1. Cold engine, `--batch 6 --batch-mtp`, `STRATA_QSA_SPLIT_ATTN=1` (plus any of the other opt-ins - they do not matter).
2. Start 6 concurrent short generations.
3. 3 s later, send one fresh (never seen, no cache reuse) 176K-token prompt.
4. Engine exits 1 with `mtp: round capture: no error`.
Binary search over 10 rounds: `STRATA_QSA_SPLIT_ATTN` alone fails (5/5); any combination of the other 7 switches passes; `STRATA_PROMPT_ATTN_IMMA` alone passes; zero opt-ins + `--batch-mtp` + a fresh 155K prompt passes (peer mode is not a factor).
### Reproduction status
Reproduced locally
### Strata version
engine 0.1.45, built from source (CUDA, archs 86;89) - upstream commit 477ac13c
### Graphics card, memory and system
RTX 4090 48 GB (primary, expert cache ~98% of VRAM by design) + RTX A4000 16 GB (expert-cache peer via `--peer-device`) · 128 GB RAM · Linux
### Model, size and settings
Qwen3.8-Flash-Next IQ3_S (GSQ-RCO pack) · 262144 context · `--batch 6 --batch-mtp --kv int8 --kv-resident 16384` · pictures on (on the peer card) · opt-ins: `STRATA_QSA_SPLIT_ATTN=1` plus 7 others (`STRATA_ATTN_DAILY`, `STRATA_EXPDEC`, `STRATA_B6_MMVQ_ROWS`, `STRATA_PROMPT_ATTN_IMMA`, `STRATA_HC_REQ8`, `STRATA_HC_UPMIX`, `STRATA_QSA_ROWS`)
### Logs
Stock engine (the swallowed error):
```
mtp: round capture: no error
```
After fixing the reporting, the real error:
```
mtp: round capture: out of memory (instantiate: out of memory, vram free 2 MiB)
```
## Analysis: two bugs
**Bug 1 - `finish_capture()` reports the wrong error.** `src/core/mtp.cpp`:
```cpp
if (ce != cudaSuccess || cudaGraphInstantiate(&exec, graph, 0) != cudaSuccess) {
if (graph) cudaGraphDestroy(graph);
err = std::string("mtp: ") + what + " capture: " + cudaGetErrorString(ce);
```
`ce` holds the result of `cudaStreamEndCapture` (a success), not of `cudaGraphInstantiate`. When instantiate fails with `cudaErrorMemoryAllocation`, the message therefore prints `no error`.
**Bug 2 - no graceful degradation on instantiate OOM.** At that moment the primary card's VRAM is intentionally near-full (the expert cache by design holds almost everything), and the fresh long prompt simultaneously needs KV writes, prefill buffers, and a new slot drafter's graph captures. For reference, `instantiate_evicting` in `verify.cpp` (#997) already frees-and-retries for the verify graphs; the MTP drafter path has no equivalent.
## Our patch (below, against 0.1.45 / 477ac13c)
Three small changes in `src/core/mtp.cpp` + `include/strata/core/mtp.hpp`, in the spirit of #997:
1. `finish_capture()` keeps the real `cudaGraphInstantiate` result and reports `vram free` alongside it (fixes Bug 1);
2. on instantiate OOM, the drafter frees **its own** captured graphs (they are captured again when needed) and retries once;
3. for a fresh drafter with nothing to free (our first fix attempt still failed on cold engines for exactly this reason), it waits briefly (80 ms x 20) for the adaptive tier to free VRAM before giving up.
We also raised `--vram-reserve-mib` from 713 to 1280 on this box.
**After the patch**: the killer scenario (cold engine, 6 concurrent generations, then a fresh 161K-token prompt) passes 8/8, zero engine exits since, aggregate throughput unchanged (conc6 ~= 195-200 tok/s with `--batch-mtp` kept on).
## Question: is `--peer-device` a maintained path going forward?
`docs/MULTI_GPU.md` refers to the `--peer-device` / helper expert caches as "the older helper-GPU caches" (see `docs/SECOND_GPU.md`). Is that path still maintained? It works very well for us on heterogeneous cards - the 4090 runs the model while the A4000 holds the experts the primary cannot fit (23929/24576 routed experts on-card across the two), with decode about 1.55x our single-card split setup - and we would like to keep relying on it. If it is maintained, we may report peer-specific findings separately.
<details><summary>Patch (git format, 2 files)</summary>
```diff
diff --git a/include/strata/core/mtp.hpp b/include/strata/core/mtp.hpp
index 7216f015..dc6a0b0d 100644
--- a/include/strata/core/mtp.hpp
+++ b/include/strata/core/mtp.hpp
@@ -180,6 +180,11 @@ private:
bool capture_prefill_dev(int T, std::string& err); ///< E-4: without the mapped staging (inputs copied on device)
bool capture_round(int T, bool coupled, std::string& err);
bool capture_step(int j, bool coupled, std::string& err);
+ /// Free this drafter's own captured exec graphs (they are captured again when needed).
+ void drop_graphs();
+ /// #997-style (verify.cpp's instantiate_evicting): when instantiate fails for VRAM, free this drafter's
+ /// own other graphs and retry, then wait briefly for the adaptive tier to free room, before giving up.
+ bool finish_evicting(cudaStream_t cs, bool ok, cudaGraphExec_t& exec, const char* what, std::string& err);
cudaGraphExec_t step_exec_[9] = {};
// coupled draft sampling: its own round/step graphs (the argmax ones stay as they were), the request's
// parameters and the penalty ring (mapped staging + device copies), the split scratch, token id -> subset index
diff --git a/src/core/mtp.cpp b/src/core/mtp.cpp
index 68be4420..cb1f3d51 100644
--- a/src/core/mtp.cpp
+++ b/src/core/mtp.cpp
@@ -1227,16 +1227,21 @@ bool MtpDrafter::record_rest(int step_row, cudaStream_t cs, std::string& err) {
}
namespace {
-bool finish_capture(cudaStream_t cs, bool ok, cudaGraphExec_t& exec, const char* what, std::string& err) {
+bool finish_capture(cudaStream_t cs, bool ok, cudaGraphExec_t& exec, const char* what, std::string& err,
+ cudaGraph_t* keep_on_oom = nullptr) {
cudaGraph_t graph = nullptr;
const cudaError_t ce = cudaStreamEndCapture(cs, &graph);
if (!ok) {
if (graph) cudaGraphDestroy(graph);
return false;
}
- if (ce != cudaSuccess || cudaGraphInstantiate(&exec, graph, 0) != cudaSuccess) {
+ const cudaError_t ie = cudaGraphInstantiate(&exec, graph, 0);
+ if (ce != cudaSuccess || ie != cudaSuccess) {
+ size_t mfree = 0, mtotal = 0; cudaMemGetInfo(&mfree, &mtotal);
+ if (keep_on_oom != nullptr && ie == cudaErrorMemoryAllocation) { *keep_on_oom = graph; return false; }
if (graph) cudaGraphDestroy(graph);
- err = std::string("mtp: ") + what + " capture: " + cudaGetErrorString(ce);
+ err = std::string("mtp: ") + what + " capture: " + (ce != cudaSuccess ? cudaGetErrorString(ce) : cudaGetErrorString(ie)) +
+ " (instantiate: " + cudaGetErrorString(ie) + ", vram free " + std::to_string(mfree >> 20) + " MiB)";
return false;
}
mtp_list_nodes(graph, what);
@@ -1248,6 +1253,36 @@ bool finish_capture(cudaStream_t cs, bool ok, cudaGraphExec_t& exec, const char*
}
} // namespace
+void MtpDrafter::drop_graphs() {
+ for (auto* arr : {&prefill_exec_, &prefill_dev_exec_, &round_exec_, &step_exec_, &round_exec_c_, &step_exec_c_})
+ for (auto& e : *arr) if (e) { cudaGraphExecDestroy(e); e = nullptr; }
+}
+
+bool MtpDrafter::finish_evicting(cudaStream_t cs, bool ok, cudaGraphExec_t& exec, const char* what, std::string& err) {
+ cudaGraph_t graph = nullptr;
+ if (finish_capture(cs, ok, exec, what, err, &graph)) return true;
+ if (graph == nullptr) return false; // the capture body failed or not an instantiate OOM: already reported
+ drop_graphs(); // #997-style: make VRAM room (freed graphs are captured again when needed)
+ (void) cudaGetLastError();
+ cudaError_t ie = cudaGraphInstantiate(&exec, graph, 0);
+ for (int w = 0; ie == cudaErrorMemoryAllocation && w < 20; ++w) { // a fresh drafter has no graphs to free; the
+ (void) cudaGetLastError(); // adaptive tier frees VRAM within seconds -
+ std::this_thread::sleep_for(std::chrono::milliseconds(80)); // wait for room instead of failing the admission
+ ie = cudaGraphInstantiate(&exec, graph, 0);
+ }
+ if (ie != cudaSuccess) {
+ cudaGraphDestroy(graph);
+ err = std::string("mtp: ") + what + " capture (after freeing own graphs + waiting): " + cudaGetErrorString(ie);
+ return false;
+ }
+ std::fprintf(stderr, "strata mtp: no VRAM for a %s graph: freed own graphs, retried\n", what);
+ mtp_list_nodes(graph, what);
+ cudaGraphDestroy(graph);
+ cudaGraphUpload(exec, cs);
+ cudaStreamSynchronize(cs);
+ return true;
+}
+
bool MtpDrafter::capture_prefill(int T, std::string& err) {
if (prefill_exec_[T]) return true;
using namespace strata::kernels;
@@ -1256,14 +1291,14 @@ bool MtpDrafter::capture_prefill(int T, std::string& err) {
copy_i32_from_mapped(step_, m_step_, (int64_t) T * 4, cs_);
copy_i32_from_mapped(pos_, m_pos_, (int64_t) T * g_->n_head, cs_);
const bool ok = record_forward(T, -1, cs_, err); // K/V only, rows [0, T)
- return finish_capture(cs_, ok, prefill_exec_[T], "prefill", err);
+ return finish_evicting(cs_, ok, prefill_exec_[T], "prefill", err);
}
bool MtpDrafter::capture_prefill_dev(int T, std::string& err) {
if (prefill_dev_exec_[T]) return true;
if (cudaStreamBeginCapture(cs_, cudaStreamCaptureModeThreadLocal) != cudaSuccess) { err = "mtp: begin capture"; return false; }
const bool ok = record_forward(T, -1, cs_, err); // K/V only, rows [0, T); tok_/step_/pos_ filled before launch
- return finish_capture(cs_, ok, prefill_dev_exec_[T], "prefill (device inputs)", err);
+ return finish_evicting(cs_, ok, prefill_dev_exec_[T], "prefill (device inputs)", err);
}
bool MtpDrafter::capture_round(int T, bool coupled, std::string& err) {
@@ -1311,7 +1346,7 @@ bool MtpDrafter::capture_round(int T, bool coupled, std::string& err) {
if (ok) g_mp.stamp(16, cs_);
g_mp.base = -1;
if (ok && force_on_ && !coupled) force_token(tok_, m_force_, 0, cs_); // chain_launch: step 1's input
- return finish_capture(cs_, ok, exec, coupled ? "round (coupled)" : "round", err);
+ return finish_evicting(cs_, ok, exec, coupled ? "round (coupled)" : "round", err);
}
// Chain step j (1..max_t-2): one row at the cell staged in step row `max_t + j - 1`, from the previous step's
@@ -1341,7 +1376,7 @@ bool MtpDrafter::capture_step(int j, bool coupled, std::string& err) {
if (ok) g_mp.stamp(16, cs_);
g_mp.base = -1;
if (ok && force_on_ && !coupled) force_token(tok_, m_force_, j, cs_); // chain_launch: step j+1's input
- return finish_capture(cs_, ok, exec, coupled ? "step (coupled)" : "step", err);
+ return finish_evicting(cs_, ok, exec, coupled ? "step (coupled)" : "step", err);
}
void MtpDrafter::kv_restore(int64_t upto) {
```
</details>
Related on this site
Quick links to related pages.