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 comentarios · En GitHub

BenchmarksSetup & installNVIDIA / CUDAModels & quantsDocumentationWindowsLinux

Descripción

### 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>

En el sitio

Enlaces a install, modelos, releases.