fix: capture-poisoned engine wedge + planner-driven KV pool sizing

CMakeLists.txt

@@ -628,6 +628,7 @@ if(IMP_BUILD_TESTS)
         tests/test_kv_cache_write.cu
         tests/test_kv_gather.cu
         tests/test_green_ctx.cu
+        tests/test_capture_abort.cu
         tests/test_prefix_cache_equiv.cpp
         tests/test_recurrent_snapshot_store.cpp
         tests/test_vram_query.cpp
@@ -653,6 +654,7 @@ if(IMP_BUILD_TESTS)
         tests/test_degeneration.cpp
         tests/test_quant_pipeline.cpp
         tests/test_weight_registry_preservation.cpp
+        tests/test_vram_budget_reserve.cpp
         tests/test_e2e_llm_compressor.cpp
         tests/test_chunked_prefill.cu
         tests/test_mtp_forward.cpp

src/exec/executor.h

@@ -506,6 +506,12 @@ class GraphExecutor {
 public:
     bool nvfp4_dequant_uncapturable() const { return nvfp4_dequant_uncapturable_; }
 
+    // True when MoE prefill will run the legacy host-args fallback (no CUTLASS
+    // NVFP4 grouped workspace — e.g. any GGUF Q*_K MoE model). That path reads
+    // routing on the host (D2H + sync) and throws under an active capture, so
+    // the scheduler must not attempt prefill-graph capture at all (#874).
+    bool moe_prefill_uncapturable() const;
+
 private:
     // Init-time weight-quantization pipeline (D2 extraction). Owns the build-only
     // StoragePlan; fills the long-lived caches below by reference in build().

src/exec/executor_forward_moe.cu

@@ -349,6 +349,22 @@ void GraphExecutor::moe_ffn_phase3_route_(int layer, cudaStream_t stream, MoeFfn
 // The path-selection ORDER + arch/config gates are mirrored as a pure function
 // `select_moe_prefill_path` in moe_prefill_decision.h, pinned by
 // test_routing_decision.cpp (R2 / P1.4).
+bool GraphExecutor::moe_prefill_uncapturable() const {
+    if (!has_moe_)
+        return false;
+    // Mirrors the function-entry gate of try_run_moe_cutlass3x_nvfp4_prefill_:
+    // without the CUTLASS 3.x grouped path + packed/scale workspace, every
+    // MoE prefill lands in run_moe_legacy_fallback_, whose host-args guard
+    // throws under capture (moe_host_args_capture_guard).
+    if (runtime_config().moe.no_cutlass3x)
+        return true;
+    if (!cutlass_grouped_3x_nvfp4_available())
+        return true;
+    if (!moe_.cutlass3x_packed || !moe_.cutlass3x_sf)
+        return true;
+    return false;
+}
+
 bool GraphExecutor::moe_cutlass3x_will_use_device_args_(int layer,
                                                         const MoeFfnContext& ctx) const {
     if (runtime_config().moe.no_cutlass3x)

src/runtime/cuda_graph.cu

@@ -334,6 +334,34 @@ bool CudaGraphCapture::end_capture_and_update() {
     return true;
 }
 
+void CudaGraphCapture::abort_capture() {
+    if (!capture_stream_)
+        return;
+    cudaGraph_t g = nullptr;
+    // EndCapture on an invalidated capture returns an error and g stays null —
+    // either way the stream leaves capture state, which is the point here.
+    (void)cudaStreamEndCapture(capture_stream_, &g);
+    graph_diag::g_phase = graph_diag::Phase::NORMAL;
+    if (g)
+        cudaGraphDestroy(g);
+    (void)cudaGetLastError();  // clear the sticky capture-invalidated error
+    capture_stream_ = nullptr;
+}
+
+void abort_stream_capture(cudaStream_t stream) {
+    if (!stream)
+        return;
+    cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone;
+    if (cudaStreamIsCapturing(stream, &st) != cudaSuccess || st == cudaStreamCaptureStatusNone)
+        return;
+    cudaGraph_t g = nullptr;
+    (void)cudaStreamEndCapture(stream, &g);
+    if (g)
+        cudaGraphDestroy(g);
+    (void)cudaGetLastError();
+    IMP_LOG_WARN("abort_stream_capture: closed a stray open capture (stream %p)", (void*)stream);
+}
+
 void CudaGraphCapture::reset() {
     bool had_exec = (graph_exec_ != nullptr);
     if (graph_exec_) {
@@ -397,7 +425,25 @@ bool CudaGraphRunner::execute(cudaStream_t stream) {
             return true;
         }
 
-        decode_fn_(stream);
+        try {
+            decode_fn_(stream);
+        } catch (const std::exception& e) {
+            // #874: the forward threw while the stream was capturing (e.g. the
+            // MoE host-args guard, or cuBLAS failing under capture). The
+            // capture MUST be closed here — unwinding past it leaves the
+            // stream capturing forever and every later op on it fails with
+            // "operation failed due to a previous error during capture".
+            IMP_LOG_ERROR(
+                "CudaGraphRunner: forward threw during capture (%s) — aborting "
+                "capture and falling back to per-step decode.",
+                e.what());
+            graph_.abort_capture();
+            capture_failed_ = true;
+            // Nothing executed during the aborted capture; run for real now.
+            decode_fn_(stream);
+            step_count_++;
+            return true;
+        }
 
         // Prefer the update path so re-captures (after invalidate_for_update)
         // reuse the existing exec via cudaGraphExecUpdate. On first capture

src/runtime/cuda_graph.h

@@ -11,6 +11,10 @@ class GraphExecutor;
 struct InferenceState;
 
 // Low-level CUDA graph capture/replay wrapper.
+// Close a stray open capture on `stream`, if any (no-op otherwise). Safety
+// net for exception paths that unwound past an active capture (#874).
+void abort_stream_capture(cudaStream_t stream);
+
 class CudaGraphCapture {
 public:
     CudaGraphCapture() = default;
@@ -41,6 +45,12 @@ class CudaGraphCapture {
     // enabling cudaGraphExecUpdate against the retained exec.
     void mark_needs_recapture() { captured_ = false; }
 
+    // Close an in-flight capture without keeping its graph (#874). Called
+    // when the captured fn throws mid-capture: the stream must be taken out
+    // of capture state or every later async op on it fails with
+    // cudaErrorStreamCaptureInvalidated until process restart.
+    void abort_capture();
+
 private:
     cudaGraph_t graph_ = nullptr;
     cudaGraphExec_t graph_exec_ = nullptr;

src/runtime/engine.cpp

@@ -386,6 +386,16 @@ void Engine::invalidate_graphs() {
     // Captured verify-chunk graphs (#847) baked the forward as well (incl.
     // any active LoRA kernels/pointers) — recapture costs two verify steps.
     free_spec_graphs_();
+
+    // #874 safety net: if an exception unwound past an active prefill-chunk
+    // capture, the prefill stream is still in capture state and every later
+    // op on it fails ("previous error during capture") — permanently wedging
+    // the server. Close any stray capture and drop the prefill runner so the
+    // next request starts from a clean stream.
+    prefill_graph_runner_.invalidate();
+    last_prefill_chunk_len_ = -1;
+    last_prefill_block_count_ = -1;
+    abort_stream_capture(prefill_stream());
 }
 
 int Engine::lora_load(const std::string& path) {

src/runtime/engine_scheduler.cpp

@@ -813,8 +813,12 @@ void Engine::step_prefill_one(std::shared_ptr<Request>& req, int effective_chunk
         // workspace, which is illegal under capture (cublasLt status 14 →
         // cascading capture failure, observed on GGUF mxfp4 GDN). Run eager.
         const bool ends_at_snapshot = (snap_end > 0 && offset + chunk_len == snap_end);
+        // moe_prefill_uncapturable: legacy host-args MoE prefill (GGUF Q*_K
+        // MoE) reads routing on the host — its capture guard throws and the
+        // aborted capture costs a wasted forward per chunk. Run eager (#874).
         const bool can_capture = prefill_graph_enabled && pf_pool_used && config_.use_cuda_graphs &&
-                                 !ends_at_snapshot && !executor_->nvfp4_dequant_uncapturable();
+                                 !ends_at_snapshot && !executor_->nvfp4_dequant_uncapturable() &&
+                                 !executor_->moe_prefill_uncapturable();
         if (can_capture) {
             const int block_count = static_cast<int>(block_table.size());
             if (chunk_len != last_prefill_chunk_len_ || block_count != last_prefill_block_count_) {

src/runtime/vram_budget.cpp

@@ -266,6 +266,45 @@ VRAMBudget compute_vram_budget(const Model& model, const EngineConfig& config, i
                              phase3_reserve / (1024.0 * 1024.0));
                 cutlass_sf_estimate = native_need;
             }
+        } else if (!plan.failed && per_block_total > 0) {
+            // GGUF sources the heuristic can't see: Q4_K/Q3_K weights are not
+            // nvfp4_beneficial (estimate 0) but the planner routes them to the
+            // FP16 cache — exactly the divergence the log above warns about.
+            // Without reserving that demand the KV backstop below eats ALL
+            // post-weight VRAM and phases 1/3 build 0 tensors: Ornith-35B
+            // Q4_K_M served 11 tok/s via on-the-fly dequant with a 159k-token
+            // KV pool nobody asked for (#874 follow-up). Same failure class as
+            // the native-NVFP4 branch above, same fix: fold the real demand
+            // into cutlass_sf_estimate so both the strategy math and the
+            // backstop see it.
+            //
+            // Gated on a 2x divergence so heuristic-covered models (Q6_K/Q8
+            // dense — the tuned bench configs) keep their KV pools unchanged.
+            size_t heuristic = nvfp4_estimate + cutlass_sf_estimate;
+            if (plan.projected_vram_bytes > 2 * heuristic) {
+                size_t want = plan.projected_vram_bytes +
+                              std::max(total_vram / 10, static_cast<size_t>(256ULL * 1024 * 1024));
+                // Never squeeze the pool below one full max_seq_len sequence
+                // (or the min-KV floor, whichever is larger) — long-context
+                // must stay servable; concurrency is bounded by scheduler
+                // admission, not by pool size.
+                int floor_tok = config.min_kv_tokens > 0
+                                    ? config.min_kv_tokens
+                                    : std::min(16384, config.max_seq_len * 4);
+                int guarantee_blocks = std::max(blocks_per_seq, (floor_tok + bs - 1) / bs);
+                size_t guarantee_bytes = static_cast<size_t>(guarantee_blocks) * per_block_total;
+                size_t cap = (available > guarantee_bytes) ? available - guarantee_bytes : 0;
+                size_t reserve = std::min(want, cap);
+                if (reserve > heuristic) {
+                    IMP_LOG_INFO(
+                        "VRAM budget: planner-driven weight-cache reserve %.1f MiB "
+                        "(plan=%.1f MiB, heuristic=%.1f MiB, kv guarantee=%d blocks)",
+                        reserve / (1024.0 * 1024.0),
+                        plan.projected_vram_bytes / (1024.0 * 1024.0),
+                        heuristic / (1024.0 * 1024.0), guarantee_blocks);
+                    cutlass_sf_estimate = reserve - nvfp4_estimate;
+                }
+            }
         }
     }
 

tests/test_capture_abort.cu

@@ -0,0 +1,71 @@
+// #874: an exception thrown by the forward fn while CudaGraphRunner is
+// capturing must not leave the stream in capture state. Before the fix,
+// decode_fn_ throwing mid-capture skipped cudaStreamEndCapture entirely;
+// every subsequent async op on the stream then failed with "operation
+// failed due to a previous error during capture" until process restart,
+// while /health kept reporting ok (observed with Ornith-1.0-35B Q4_K_M:
+// legacy MoE host-args prefill guard throws under the prefill-chunk
+// capture).
+
+#include <gtest/gtest.h>
+#include <cuda_runtime.h>
+
+#include <stdexcept>
+
+#include "runtime/cuda_graph.h"
+
+#include "test_cuda_skip.h"
+
+using namespace imp;
+
+namespace {
+__global__ void touch_kernel(int* p) { *p = 42; }
+}  // namespace
+
+TEST(CaptureAbort, ThrowDuringCaptureLeavesStreamUsable) {
+    SKIP_IF_NO_CUDA();
+    cudaStream_t stream;
+    ASSERT_EQ(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), cudaSuccess);
+    int* d = nullptr;
+    ASSERT_EQ(cudaMalloc(&d, sizeof(int)), cudaSuccess);
+
+    CudaGraphRunner runner;
+    runner.set_warmup_steps(0);  // go straight to capture on the first execute
+    int calls = 0;
+    runner.set_decode_fn([&](cudaStream_t s) {
+        calls++;
+        touch_kernel<<<1, 1, 0, s>>>(d);
+        // Mirror the MoE host-args guard: refuse to run under capture.
+        cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone;
+        cudaStreamIsCapturing(s, &st);
+        if (st != cudaStreamCaptureStatusNone)
+            throw std::runtime_error("host-args path not graph-capturable (test)");
+    });
+
+    // The throw mid-capture must be handled inside execute(): abort the
+    // capture, fall back to eager, run the fn for real.
+    EXPECT_NO_THROW(runner.execute(stream));
+
+    // Stream must not be left capturing.
+    cudaStreamCaptureStatus st = cudaStreamCaptureStatusNone;
+    ASSERT_EQ(cudaStreamIsCapturing(stream, &st), cudaSuccess);
+    EXPECT_EQ(st, cudaStreamCaptureStatusNone);
+
+    // The eager fallback must have actually executed the work.
+    EXPECT_GE(calls, 2) << "expected capture attempt + eager re-run";
+
+    // Subsequent async ops on the same stream must succeed — this is the
+    // #874 wedge: they all returned cudaErrorStreamCaptureInvalidated.
+    int h = 0;
+    EXPECT_EQ(cudaMemcpyAsync(&h, d, sizeof(int), cudaMemcpyDeviceToHost, stream), cudaSuccess);
+    EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess);
+    EXPECT_EQ(h, 42);
+
+    // Runner keeps serving eagerly (capture permanently disabled for it).
+    EXPECT_NO_THROW(runner.execute(stream));
+    EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess);
+    EXPECT_EQ(cudaGetLastError(), cudaSuccess);
+
+    cudaFree(d);
+    cudaStreamDestroy(stream);
+}