diff --git a/CMakeLists.txt b/CMakeLists.txt index e5e8430b..e584d2e7 100644 --- a/CMakeLists.txt +++ b/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 diff --git a/src/exec/executor.h b/src/exec/executor.h index d04b1334..0447d440 100644 --- a/src/exec/executor.h +++ b/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(). diff --git a/src/exec/executor_forward_moe.cu b/src/exec/executor_forward_moe.cu index 65e19a94..834c7875 100644 --- a/src/exec/executor_forward_moe.cu +++ b/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) diff --git a/src/runtime/cuda_graph.cu b/src/runtime/cuda_graph.cu index 88cf7dfb..c35ee398 100644 --- a/src/runtime/cuda_graph.cu +++ b/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 diff --git a/src/runtime/cuda_graph.h b/src/runtime/cuda_graph.h index 0cfada79..eae93ae9 100644 --- a/src/runtime/cuda_graph.h +++ b/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; diff --git a/src/runtime/engine.cpp b/src/runtime/engine.cpp index e38758db..dbf4f7e0 100644 --- a/src/runtime/engine.cpp +++ b/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) { diff --git a/src/runtime/engine_scheduler.cpp b/src/runtime/engine_scheduler.cpp index 2347c3ee..a94e944d 100644 --- a/src/runtime/engine_scheduler.cpp +++ b/src/runtime/engine_scheduler.cpp @@ -813,8 +813,12 @@ void Engine::step_prefill_one(std::shared_ptr& 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(block_table.size()); if (chunk_len != last_prefill_chunk_len_ || block_count != last_prefill_block_count_) { diff --git a/src/runtime/vram_budget.cpp b/src/runtime/vram_budget.cpp index 18c602fd..b139223c 100644 --- a/src/runtime/vram_budget.cpp +++ b/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(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(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; + } + } } } diff --git a/tests/test_capture_abort.cu b/tests/test_capture_abort.cu new file mode 100644 index 00000000..80ca9d2a --- /dev/null +++ b/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 +#include + +#include + +#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); +} diff --git a/tests/test_vram_budget_reserve.cpp b/tests/test_vram_budget_reserve.cpp new file mode 100644 index 00000000..94c632ab --- /dev/null +++ b/tests/test_vram_budget_reserve.cpp @@ -0,0 +1,116 @@ +// KV-pool sizing vs weight-cache demand (Ornith-35B Q4_K_M regression). +// +// Q4_K sources are not nvfp4_beneficial, so the heuristic weight-cache +// estimate in compute_vram_budget is ~0 while the StoragePlanner (source- +// aware) routes them to the FP16 cache. Before the fix the KV backstop +// sized the pool to available-minus-heuristic — i.e. it ate all post-weight +// VRAM, phases 1/3 built 0 cache tensors, and decode fell back to +// on-the-fly dequant (11 tok/s instead of a cached decode path). +// +// The budget must reserve the planner projection (capped so one full +// max_seq_len sequence always fits) before the KV pool fills the rest. + +#include + +#include "model/model.h" +#include "model/model_config.h" +#include "runtime/engine.h" +#include "runtime/storage_planner.h" +#include "runtime/vram_budget.h" + +#include "test_cuda_skip.h" + +using namespace imp; + +namespace { + +Tensor make_weight(TensorKind kind, int64_t rows, int64_t cols, QType qt, uintptr_t sentinel) { + Tensor t; + t.data = reinterpret_cast(sentinel); + t.qtype = qt; + t.ndim = 2; + t.shape[0] = rows; + t.shape[1] = cols; + t.on_device = false; + t.kind = kind; + return t; +} + +// 32 layers; attention weights of `attn_qt`, FFN weights of `ffn_qt` +// (Model is non-copyable — populate the caller's instance). +void fill_model(Model& m, QType attn_qt, QType ffn_qt) { + m.config_.n_layers = 32; + m.config_.n_kv_heads = 8; + for (int i = 0; i < 32; ++i) { + TransformerLayer L; + uintptr_t base = static_cast(i) * 100 + 1; + L.wq = make_weight(TensorKind::WQ, 4096, 4096, attn_qt, base + 0); + L.w_gate = make_weight(TensorKind::W_GATE, 11008, 4096, ffn_qt, base + 1); + L.w_up = make_weight(TensorKind::W_UP, 11008, 4096, ffn_qt, base + 2); + L.w_down = make_weight(TensorKind::W_DOWN, 4096, 11008, ffn_qt, base + 3); + m.layers_.push_back(std::move(L)); + } +} + +} // namespace + +TEST(VramBudgetReserve, PlannerDemandKeepsKvPoolFromEatingWeightCacheRoom) { + SKIP_IF_NO_CUDA(); // compute_vram_budget queries total VRAM + + // Attention Q6_K (nvfp4-beneficial, small heuristic footprint), FFN Q4_K + // (heuristic-invisible, planner routes FP16) — the Ornith-35B shape. + Model m; + fill_model(m, QType::Q6_K, QType::Q4_K); + + EngineConfig config; + config.max_seq_len = 32768; + config.max_batch_size = 8; + config.use_nvfp4_decode = 2; // mode 2: reserve floor is the flat 512 MiB + config.use_cuda_graphs = false; + config.kv_cache_dtype = QType::F16; + + const int n_kv_layers = 32; + const int head_dim = 128; + const size_t GiB = 1024ull * 1024 * 1024; + const size_t free_vram = 8 * GiB; + + VRAMBudget budget = compute_vram_budget(m, config, n_kv_layers, head_dim, free_vram); + + // per-block cost as compute_vram_budget derives it (F16, block size 16): + const size_t per_block = 16ull * 8 * 128 * 2 /*K+V*/ * 2 /*bytes*/ * n_kv_layers; + const size_t kv_bytes = static_cast(budget.kv_max_blocks) * per_block; + const int blocks_per_seq = 32768 / 16; + + // The FFN FP16 demand here is ~8.6 GiB (> available), so the reservation + // is capped at available minus one full sequence — the KV pool must come + // out at roughly one max_seq_len sequence, not fill all of free_vram. + EXPECT_GE(budget.kv_max_blocks, blocks_per_seq) + << "one full max_seq_len sequence must always fit in the KV pool"; + EXPECT_LE(kv_bytes, free_vram - 3 * GiB) + << "KV pool ate the weight-cache room again (heuristic-only backstop)"; +} + +TEST(VramBudgetReserve, BeneficialSourcesKeepFullKvPool) { + SKIP_IF_NO_CUDA(); + + // All-Q6_K model: heuristic covers the demand, planner projection stays + // in the same ballpark — the reservation must NOT fire and shrink KV. + Model m; + fill_model(m, QType::Q6_K, QType::Q6_K); + + EngineConfig config; + config.max_seq_len = 32768; + config.max_batch_size = 8; + config.use_nvfp4_decode = 2; + config.use_cuda_graphs = false; + config.kv_cache_dtype = QType::F16; + + const size_t GiB = 1024ull * 1024 * 1024; + VRAMBudget with_planner = compute_vram_budget(m, config, 32, 128, 8 * GiB); + + // Q6_K routes to the NVFP4 tier: heuristic ≈ planner, so KV keeps most of + // the post-reserve VRAM (well above the one-sequence floor). + const size_t per_block = 16ull * 8 * 128 * 2 * 2 * 32; + const size_t kv_bytes = static_cast(with_planner.kv_max_blocks) * per_block; + EXPECT_GE(kv_bytes, 4 * GiB) << "reservation must not fire for heuristic-covered sources"; +} diff --git a/tools/imp-server/batching_engine.cpp b/tools/imp-server/batching_engine.cpp index 6bb8965c..6636b3a0 100644 --- a/tools/imp-server/batching_engine.cpp +++ b/tools/imp-server/batching_engine.cpp @@ -40,6 +40,7 @@ void BatchingEngine::start(ImpContext ctx) { } ctx_ = ctx; stop_requested_.store(false); + faulted_.store(false); pause_requested_.store(false); paused_.store(false); running_.store(true); @@ -223,6 +224,7 @@ void BatchingEngine::worker_loop() { IMP_LOG_ERROR("BatchingEngine: CUDA context poisoned (%s / %s) — stopping the " "worker; the server must be restarted.", cudaGetErrorString(sync_err), cudaGetErrorString(sticky)); + faulted_.store(true, std::memory_order_relaxed); stop_requested_.store(true, std::memory_order_relaxed); } } @@ -251,6 +253,7 @@ void BatchingEngine::worker_loop() { IMP_LOG_ERROR("BatchingEngine: CUDA context poisoned (%s / %s) — stopping the " "worker; the server must be restarted.", cudaGetErrorString(sync_err), cudaGetErrorString(sticky)); + faulted_.store(true, std::memory_order_relaxed); stop_requested_.store(true, std::memory_order_relaxed); } } diff --git a/tools/imp-server/batching_engine.h b/tools/imp-server/batching_engine.h index d20836a6..a7cd7bab 100644 --- a/tools/imp-server/batching_engine.h +++ b/tools/imp-server/batching_engine.h @@ -109,6 +109,12 @@ class BatchingEngine { bool is_running() const { return running_.load(std::memory_order_relaxed); } + // True after the worker declared the CUDA context poisoned and stopped + // (#874). /health reports unhealthy so an orchestrator can restart the + // process — before this the server answered "ok" while every request + // failed with internal_error. + bool faulted() const { return faulted_.load(std::memory_order_relaxed); } + private: void worker_loop(); @@ -117,6 +123,7 @@ class BatchingEngine { std::thread worker_thread_; std::atomic running_{false}; std::atomic stop_requested_{false}; + std::atomic faulted_{false}; // Graceful pause handshake (see pause()/resume()). pause_requested_ tells // the worker to drain in-flight work and park; paused_ reports that it has. diff --git a/tools/imp-server/handlers.cpp b/tools/imp-server/handlers.cpp index 15352717..05c228f1 100644 --- a/tools/imp-server/handlers.cpp +++ b/tools/imp-server/handlers.cpp @@ -44,13 +44,20 @@ int64_t unix_timestamp() { void handle_health(const httplib::Request& /*req*/, httplib::Response& res, ServerState& state) { bool loaded; int queue = 0; + bool faulted = false; { std::lock_guard lock(state.mtx); loaded = state.model_loaded(); - if (state.batching) + if (state.batching) { queue = state.batching->queue_depth(); + faulted = state.batching->faulted(); + } } - json body = {{"status", "ok"}, {"model_loaded", loaded}, {"queue_depth", queue}}; + json body = {{"status", faulted ? "unhealthy" : "ok"}, + {"model_loaded", loaded}, + {"queue_depth", queue}}; + if (faulted) + res.status = 503; // let orchestrators restart a wedged server (#874) res.set_content(dump_safe(body), "application/json"); }