diff --git a/driver/cuda/src/entry.cpp b/driver/cuda/src/entry.cpp index 36deacea8..fdd59e890 100644 --- a/driver/cuda/src/entry.cpp +++ b/driver/cuda/src/entry.cpp @@ -626,7 +626,8 @@ int run_impl(int argc, // concurrent speculation. const int q35_planned_slots = std::max(1, mem_plan.state_slots); qwen3_5_runtime_rs_slots = q35_planned_slots; - const int q35_alloc_slots = qwen3_5_runtime_rs_slots; + const int q35_alloc_slots = + qwen3_5_runtime_rs_slots + (graph_pad_slot >= 0 ? 1 : 0); qwen3_5_state_cache = pie_cuda_driver::RecurrentStateCache::allocate( qwen3_5_layer_is_linear, conv_dim, cfg_q.linear_conv_kernel_dim, local_linear_value_heads, diff --git a/driver/cuda/src/executor/executor.cpp b/driver/cuda/src/executor/executor.cpp index 3efe0af91..c4b375f9c 100644 --- a/driver/cuda/src/executor/executor.cpp +++ b/driver/cuda/src/executor/executor.cpp @@ -2130,7 +2130,7 @@ inline void write_probes(pie_driver::PieForwardResponseView& out, } } // namespace -void handle_fire_batch( +bool handle_fire_batch( std::uint32_t req_id, const pie_driver::PieForwardRequestView& view, pie_driver::PieForwardResponseView& out_resp, @@ -2276,6 +2276,7 @@ void handle_fire_batch( const int N = static_cast(tok_view.size()); const int num_sampling = static_cast(sidx_view.size()); + const bool has_sampling = num_sampling > 0; dbg_R = R; dbg_N = N; // Qwen3-VL: assemble the per-token [N,3] M-RoPE positions. Text rows @@ -2315,13 +2316,12 @@ void handle_fire_batch( // Empty batch — emit a zero-request response view. std::vector empty(std::max(R, 0)); executor.response_builder.build(empty, out_resp); - return; + return true; } if (N > max_workspace_tokens) { std::cerr << "[pie-driver-cuda] batch tokens=" << N << " exceeds workspace=" << max_workspace_tokens << "\n"; - out_resp = pie_driver::PieForwardResponseView{}; - return; + return false; } // Compute max KV length across requests for shmem sizing. @@ -2438,6 +2438,16 @@ void handle_fire_batch( slot_ids_h[r] = executor.graph_pad_slot; is_fresh_h[r] = 0u; } + const int max_rs_slots = executor.rs_cache->max_slots(); + for (int r = 0; r < slot_count; ++r) { + if (slot_ids_h[r] < 0 || slot_ids_h[r] >= max_rs_slots) { + throw std::runtime_error( + "rs_cache slot id " + std::to_string(slot_ids_h[r]) + + " for request " + std::to_string(r) + + " is outside allocated slot range [0, " + + std::to_string(max_rs_slots) + ")"); + } + } pi.slot_ids.copy_from_host(std::span(slot_ids_h)); pi.is_fresh.copy_from_host(std::span(is_fresh_h)); } @@ -2520,16 +2530,14 @@ void handle_fire_batch( << logit_rows_required << " logit rows, exceeding workspace capacity " << tensor_rows(ws.logits) << "\n"; - out_resp = pie_driver::PieForwardResponseView{}; - return; + return false; } if (prob_rows_required > tensor_rows(ws.probs)) { std::cerr << "[pie-driver-cuda] fire_batch needs " << prob_rows_required << " probability rows, exceeding workspace capacity " << tensor_rows(ws.probs) << "\n"; - out_resp = pie_driver::PieForwardResponseView{}; - return; + return false; } const SamplingPlan sample_plan{ @@ -2677,7 +2685,7 @@ void handle_fire_batch( // pass) produces no logits and no sampled tokens — skip every argmax / // sampling launch (they would read the unwritten, undersized // `ws.logits` over all N rows and fault). - if (num_sampling == 0) { + if (!has_sampling) { // nothing to sample } else if (tp_greedy_argmax) { kernels::launch_select_global_argmax_pairs( @@ -2733,19 +2741,25 @@ void handle_fire_batch( // `h_per_*`, `h_sample_idx`) are still in scope here for the // response builder. - // Only copy the first N entries — `pi.sampled` is sized for - // max_workspace_tokens, but only [0, N) are valid this fire. - // Async on the same stream the sampler ran on so it slots into - // the stream's FIFO; we sync immediately after because the - // response payload depends on these tokens. (Future work moves - // the sync past the host-side response-prep so the host - // and GPU can overlap.) - std::int32_t* sampled_host = - sampled_pinned_buf(static_cast(N)); - CUDA_CHECK(cudaMemcpyAsync(sampled_host, pi.sampled.data(), - sizeof(std::int32_t) * N, - cudaMemcpyDeviceToHost, - cublas.stream())); + // Only sampled fires produce valid `pi.sampled` rows. Prefill-only + // fires (for example Context::flush) still need the stream sync below + // to observe KV writes, but must not copy sampled rows that no kernel + // wrote. + std::int32_t* sampled_host = nullptr; + if (has_sampling) { + // Copy the first N entries — `pi.sampled` is sized for + // max_workspace_tokens, but only [0, N) are valid this fire. + // Async on the same stream the sampler ran on so it slots into + // the stream's FIFO; we sync immediately after because the + // response payload depends on these tokens. (Future work moves + // the sync past the host-side response-prep so the host + // and GPU can overlap.) + sampled_host = sampled_pinned_buf(static_cast(N)); + CUDA_CHECK(cudaMemcpyAsync(sampled_host, pi.sampled.data(), + sizeof(std::int32_t) * N, + cudaMemcpyDeviceToHost, + cublas.stream())); + } const auto t_kernel_launch_end = clock::now(); verify_timer.finish(cublas.stream()); CUDA_CHECK(cudaStreamSynchronize(cublas.stream())); @@ -2758,7 +2772,7 @@ void handle_fire_batch( // and benchmark traffic keep the GPU sampler result. const std::int32_t* all_sampled = sampled_host; std::vector sampled_override; - if (!logit_masks_view.empty()) { + if (has_sampling && !logit_masks_view.empty()) { sampled_override.assign(sampled_host, sampled_host + N); apply_logit_mask_overrides( ws, sampled_override, logit_masks_view, logit_mask_indptr_view, @@ -2794,7 +2808,7 @@ void handle_fire_batch( << " sampled=" << num_sampling << " max_kv=" << max_kv_len << "\n"; } - return; + return true; } // Flat-path arrays: token sampler is the only slot type allowed @@ -3186,7 +3200,7 @@ void handle_fire_batch( << " sampled=" << num_sampling << " max_kv=" << max_kv_len << "\n"; } - return; + return true; } executor.response_builder.build_token_only( @@ -3202,13 +3216,12 @@ void handle_fire_batch( << " sampled=" << num_sampling << " max_kv=" << max_kv_len << "\n"; } - return; + return true; } catch (const std::exception& e) { std::cerr << "[pie-driver-cuda] fire_batch failed for req_id=" << req_id << ": " << e.what() << "\n"; - out_resp = pie_driver::PieForwardResponseView{}; - return; + return false; } } diff --git a/driver/cuda/src/executor/executor.hpp b/driver/cuda/src/executor/executor.hpp index faa424ba2..1ac8844c9 100644 --- a/driver/cuda/src/executor/executor.hpp +++ b/driver/cuda/src/executor/executor.hpp @@ -349,7 +349,7 @@ struct Executor { // cumulative fire counter used as PRNG offset and logging-cadence gate. // The caller's inproc transport hands `out_resp` to `send_response` // immediately after this returns. -void handle_fire_batch( +bool handle_fire_batch( std::uint32_t req_id, const pie_driver::PieForwardRequestView& view, pie_driver::PieForwardResponseView& out_resp, diff --git a/driver/cuda/src/kernels/gated_delta_net.cu b/driver/cuda/src/kernels/gated_delta_net.cu index 45584f946..4ae17ab2b 100644 --- a/driver/cuda/src/kernels/gated_delta_net.cu +++ b/driver/cuda/src/kernels/gated_delta_net.cu @@ -104,15 +104,14 @@ bool qwen_gdn_fused_step_enabled() { // FLA chunked-prefill default), so strictly less quantization than // the legacy per-element-round kernel. // -// Default ON: +32% end-to-end throughput on Qwen/Qwen3.5-4B -// (6924 -> 9166 tok/s). The launcher only routes here for the GQA -// bf16 V-last decode shape (V_d==K_d==128, !k_last); everything else -// falls back to the legacy kernel. Set PIE_QWEN35_GDN_SMEM_STEP=0 to -// force the fallback. +// Opt-in: this path is faster on saturated decode microbenches, but it +// has not yet had the same forked/batched-context parity coverage as the +// legacy BF16 GQA kernel. Keep the production default on the covered path; +// set PIE_QWEN35_GDN_SMEM_STEP=1 to benchmark the SMEM variant. bool qwen_gdn_smem_step_enabled() { static const bool enabled = [] { const char* v = std::getenv("PIE_QWEN35_GDN_SMEM_STEP"); - if (v == nullptr || v[0] == '\0') return true; + if (v == nullptr || v[0] == '\0') return false; return v[0] != '0'; }(); return enabled; diff --git a/driver/cuda/src/model/llama_like_model.cpp b/driver/cuda/src/model/llama_like_model.cpp index 53998b161..1ed552d85 100644 --- a/driver/cuda/src/model/llama_like_model.cpp +++ b/driver/cuda/src/model/llama_like_model.cpp @@ -48,8 +48,11 @@ void LlamaLikeModel::body(Qwen3Workspace& ws, AttentionWorkspace& attn_ws, ops::CublasHandle& cublas, const ForwardFn::ForwardInputs& in) { + LlamaLikeForwardCfg fwd = fwd_cfg_; + fwd.emit_logits = in.emit_logits; + llama_like_forward_paged( - weights_, hf_config_, fwd_cfg_, plan_, + weights_, hf_config_, fwd, plan_, ws, kv, attn_ws, cublas, in.token_ids, in.positions, in.qo_indptr_d, in.kv_page_indices_d, in.kv_page_indptr_d, diff --git a/driver/cuda/src/model/qwen3_5_forward.cpp b/driver/cuda/src/model/qwen3_5_forward.cpp index c4fc881f9..681ec69c5 100644 --- a/driver/cuda/src/model/qwen3_5_forward.cpp +++ b/driver/cuda/src/model/qwen3_5_forward.cpp @@ -1367,12 +1367,11 @@ void qwen3_5_forward_paged( } // 3. Final norm. - // State-only forward (num_logit_rows < 0): the speculative repair re-runs - // the backbone purely to advance the rs_cache state (written in the layer - // loop above). Its logits/hidden are discarded, so skip final_norm, - // lm_head, and the final hidden copy entirely — that dense lm_head over all - // N repair tokens is the single largest piece of wasted work. - if (num_logit_rows < 0 || commit_advance) { + // State-only/no-sampler forward: the caller only needs KV/recurrent-state + // side effects, so skip final_norm, lm_head, and the final hidden copy. + // This covers speculative repair (num_logit_rows < 0) and ordinary + // Context::flush/KV-fill fires (emit_logits=false). + if (!fwd_cfg.emit_logits || num_logit_rows < 0 || commit_advance) { profile.end(stream); maybe_print_forward_profile(profile); return; diff --git a/driver/cuda/src/model/qwen3_5_forward.hpp b/driver/cuda/src/model/qwen3_5_forward.hpp index e6b3b1b67..4995b044d 100644 --- a/driver/cuda/src/model/qwen3_5_forward.hpp +++ b/driver/cuda/src/model/qwen3_5_forward.hpp @@ -87,6 +87,12 @@ struct Qwen3_5ForwardCfg { // cache lookup should stay pinned to the verified source prefix while // draft positions advance for RoPE and history masking. bool mtp_global_cache_uses_prefix_position = false; + + // False for KV-fill/flush fires that have no sampler rows. Those fires + // only need cache/state side effects; materializing dense logits is both + // wasted work and can violate the executor's no-sampling workspace + // contract. + bool emit_logits = true; }; // Persistent decode-plan cache. Owned by main.cpp's serving setup so diff --git a/driver/cuda/src/model/qwen3_5_model.cpp b/driver/cuda/src/model/qwen3_5_model.cpp index 44608d065..9eb8415f6 100644 --- a/driver/cuda/src/model/qwen3_5_model.cpp +++ b/driver/cuda/src/model/qwen3_5_model.cpp @@ -47,8 +47,11 @@ void Qwen35Model::body(Qwen3Workspace& ws, AttentionWorkspace& attn_ws, ops::CublasHandle& cublas, const ForwardFn::ForwardInputs& in) { + Qwen3_5ForwardCfg fwd = fwd_cfg_; + fwd.emit_logits = in.emit_logits; + qwen3_5_forward_paged( - weights_, hf_config_, fwd_cfg_, plan_state_, + weights_, hf_config_, fwd, plan_state_, ws, la_ws_, kv, state_cache_, attn_ws, cublas, in.token_ids, in.positions, diff --git a/driver/cuda/src/model/qwen3_5_moe_forward.cpp b/driver/cuda/src/model/qwen3_5_moe_forward.cpp index e398eb02c..c30784870 100644 --- a/driver/cuda/src/model/qwen3_5_moe_forward.cpp +++ b/driver/cuda/src/model/qwen3_5_moe_forward.cpp @@ -1774,10 +1774,8 @@ void qwen3_5_moe_forward_paged( } } - // State-only forward (num_logit_rows < 0) or recurrent-only commit-advance: - // logits/hidden are discarded, so skip final_norm + lm_head + final copy. - // See qwen3_5_forward. - if (num_logit_rows < 0 || commit_advance) { + // State-only/no-sampler forward: see qwen3_5_forward. + if (!fwd_cfg.emit_logits || num_logit_rows < 0 || commit_advance) { profile.end(stream); maybe_print_profile(profile); return; diff --git a/driver/cuda/src/model/qwen3_5_moe_model.cpp b/driver/cuda/src/model/qwen3_5_moe_model.cpp index f25efe89a..04ec10060 100644 --- a/driver/cuda/src/model/qwen3_5_moe_model.cpp +++ b/driver/cuda/src/model/qwen3_5_moe_model.cpp @@ -49,8 +49,11 @@ void Qwen35MoeModel::body(Qwen3Workspace& ws, AttentionWorkspace& attn_ws, ops::CublasHandle& cublas, const ForwardFn::ForwardInputs& in) { + Qwen3_5ForwardCfg fwd = fwd_cfg_; + fwd.emit_logits = in.emit_logits; + qwen3_5_moe_forward_paged( - weights_, hf_config_, fwd_cfg_, plan_state_, + weights_, hf_config_, fwd, plan_state_, ws, la_ws_, moe_ws_, kv, state_cache_, attn_ws, cublas, in.token_ids, in.positions, diff --git a/driver/cuda/src/service/inproc_service.cpp b/driver/cuda/src/service/inproc_service.cpp index eb358fb6b..9b21718a5 100644 --- a/driver/cuda/src/service/inproc_service.cpp +++ b/driver/cuda/src/service/inproc_service.cpp @@ -36,9 +36,27 @@ void InProcService::serve_forever(pie_driver::InProcServer& server) { switch (req.method) { case pie_driver::PIE_METHOD_FORWARD: { ++handled_; - handle_fire_batch( + const bool ok = handle_fire_batch( req_id, req.forward, out.forward, executor_, handled_); - out.status = 0; + if (ok) { + const auto qo = req.forward.qo_indptr.as(); + const auto expected = qo.empty() + ? 0u + : static_cast(qo.size() - 1); + if (out.forward.num_requests == expected) { + out.status = 0; + } else { + std::cerr << "[pie-driver-cuda] fire_batch response count " + << "mismatch for req_id=" << req_id + << ": expected " << expected + << ", got " << out.forward.num_requests << "\n"; + out.method = pie_driver::PIE_METHOD_HEALTH; + out.status = -1; + } + } else { + out.method = pie_driver::PIE_METHOD_HEALTH; + out.status = -1; + } break; } case pie_driver::PIE_METHOD_COPY_D2H: { diff --git a/driver/cuda/tests/test_executor_prefill_only.py b/driver/cuda/tests/test_executor_prefill_only.py new file mode 100644 index 000000000..b7ee6dc4d --- /dev/null +++ b/driver/cuda/tests/test_executor_prefill_only.py @@ -0,0 +1,91 @@ +from pathlib import Path + + +CUDA_SRC = Path(__file__).resolve().parents[1] / "src" + + +def test_prefill_only_fire_batch_does_not_copy_sampled_tokens(): + source = (CUDA_SRC / "executor" / "executor.cpp").read_text() + + guard = "const bool has_sampling = num_sampling > 0;" + assert guard in source + + copy = "cudaMemcpyAsync(sampled_host, pi.sampled.data()" + copy_index = source.index(copy) + guard_index = source.rfind("if (has_sampling)", 0, copy_index) + sync_index = source.rfind("cudaStreamSynchronize(cublas.stream())", 0, copy_index) + + assert guard_index > sync_index + + +def test_llama_like_model_honors_per_fire_emit_logits_flag(): + source = (CUDA_SRC / "model" / "llama_like_model.cpp").read_text() + + assert "LlamaLikeForwardCfg fwd = fwd_cfg_;" in source + assert "fwd.emit_logits = in.emit_logits;" in source + assert "weights_, hf_config_, fwd, plan_," in source + assert "weights_, hf_config_, fwd_cfg_, plan_," not in source + + +def test_forward_failures_do_not_return_zero_request_forward_response(): + executor_source = (CUDA_SRC / "executor" / "executor.cpp").read_text() + service_source = (CUDA_SRC / "service" / "inproc_service.cpp").read_text() + + signature = "bool handle_fire_batch(" + assert signature in executor_source + assert signature in (CUDA_SRC / "executor" / "executor.hpp").read_text() + + assert "return false;" in executor_source + assert "out_resp = pie_driver::PieForwardResponseView{};" not in executor_source + + assert "const bool ok = handle_fire_batch(" in service_source + assert "out.forward.num_requests == expected" in service_source + assert "out.method = pie_driver::PIE_METHOD_HEALTH;" in service_source + assert "out.status = -1;" in service_source + + +def test_qwen35_allocates_graph_pad_recurrent_slot(): + source = (CUDA_SRC / "entry.cpp").read_text() + + assert "qwen3_5_runtime_rs_slots = q35_planned_slots;" in source + assert ( + "const int q35_alloc_slots =\n" + " qwen3_5_runtime_rs_slots + (graph_pad_slot >= 0 ? 1 : 0);" + ) in source + assert "static_cast(qwen3_5_runtime_rs_slots)" in source + + +def test_rs_cache_slot_ids_are_validated_before_cuda_upload(): + source = (CUDA_SRC / "executor" / "executor.cpp").read_text() + + validation = "const int max_rs_slots = executor.rs_cache->max_slots();" + copy = "pi.slot_ids.copy_from_host(std::span(slot_ids_h));" + assert validation in source + assert "rs_cache slot id " in source + assert source.index(validation) < source.index(copy) + + +def test_qwen35_smem_gqa_decode_kernel_is_opt_in(): + source = (CUDA_SRC / "kernels" / "gated_delta_net.cu").read_text() + + marker = "bool qwen_gdn_smem_step_enabled()" + start = source.index(marker) + end = source.index("}();", start) + body = source[start:end] + + assert "PIE_QWEN35_GDN_SMEM_STEP" in body + assert "return false;" in body + + +def test_qwen35_models_honor_per_fire_emit_logits_flag(): + header = (CUDA_SRC / "model" / "qwen3_5_forward.hpp").read_text() + qwen_model = (CUDA_SRC / "model" / "qwen3_5_model.cpp").read_text() + moe_model = (CUDA_SRC / "model" / "qwen3_5_moe_model.cpp").read_text() + qwen_forward = (CUDA_SRC / "model" / "qwen3_5_forward.cpp").read_text() + moe_forward = (CUDA_SRC / "model" / "qwen3_5_moe_forward.cpp").read_text() + + assert "bool emit_logits = true;" in header + assert "fwd.emit_logits = in.emit_logits;" in qwen_model + assert "fwd.emit_logits = in.emit_logits;" in moe_model + assert "if (!fwd_cfg.emit_logits || num_logit_rows < 0 || commit_advance)" in qwen_forward + assert "if (!fwd_cfg.emit_logits || num_logit_rows < 0 || commit_advance)" in moe_forward diff --git a/runtime/src/api/inference.rs b/runtime/src/api/inference.rs index 5eace9125..431912a81 100644 --- a/runtime/src/api/inference.rs +++ b/runtime/src/api/inference.rs @@ -186,6 +186,7 @@ pub struct FutureOutput { spec_positions_for_fill: Vec, adapter_id: Option, adapter_seed: Option, + error: Option, } impl FutureOutput { @@ -262,17 +263,17 @@ impl FutureOutput { self.done = true; } - fn finish_empty(&mut self) { + fn finish_error(&mut self, error: String) { self.release_pin(); - self.result = Some(pie::core::inference::Output { - slots: Vec::new(), - spec_tokens: Vec::new(), - spec_positions: Vec::new(), - }); + self.error = Some(error); self.done = true; } } +fn future_output_stderr_message(error: &str) -> String { + format!("Pie async forward failed: {error}") +} + fn empty_forward_request() -> pie_bridge::ForwardRequest { pie_bridge::ForwardRequest { adapter_bindings: vec![pie_bridge::AdapterBinding { @@ -306,10 +307,10 @@ impl Pollable for FutureOutput { Ok(Ok(resp)) => self.finish_ok(resp), Ok(Err(e)) => { tracing::warn!("future output failed: {e:#}"); - self.finish_empty(); + self.finish_error(e.to_string()); } Err(_) => { - self.finish_empty(); + self.finish_error("future output channel closed".to_string()); } } } else { @@ -593,8 +594,8 @@ impl pie::core::inference::HostForwardPass for InstanceState { let img = self.ctx().table.get(&image)?; ( img.span.token_count, - img.span.grid, // merged LLM grid (for M-RoPE positions) - img.patch_grid, // pre-merge patch grid (for the driver encoder) + img.span.grid, // merged LLM grid (for M-RoPE positions) + img.patch_grid, // pre-merge patch grid (for the driver encoder) img.uses_mrope, img.pixels.clone(), img.positions.clone(), @@ -671,7 +672,8 @@ impl pie::core::inference::HostForwardPass for InstanceState { // Log-mel features (f32 → little-endian bytes). req.audio_features .extend_from_slice(bytemuck::cast_slice(&mel)); - req.audio_feature_indptr.push(req.audio_features.len() as u32); + req.audio_feature_indptr + .push(req.audio_features.len() as u32); Ok(()) } @@ -890,7 +892,10 @@ impl pie::core::inference::HostForwardPass for InstanceState { let (system_spec_supported, system_spec_enabled) = crate::model::get_model(model_id) .map(|m| { - (m.system_speculation_supported(), m.enable_system_speculation()) + ( + m.system_speculation_supported(), + m.enable_system_speculation(), + ) }) .unwrap_or((false, false)); if !system_spec_supported { @@ -986,6 +991,7 @@ impl pie::core::inference::HostForwardPass for InstanceState { spec_positions_for_fill, adapter_id, adapter_seed, + error: None, }; let postprocess_start = profiling.then(Instant::now); let pushed = self.ctx().table.push(future_output)?; @@ -1024,16 +1030,20 @@ impl pie::core::inference::HostFutureOutput for InstanceState { Ok(Err(e)) => { result.rx = None; tracing::warn!("future output failed: {e:#}"); - result.finish_empty(); + result.finish_error(e.to_string()); } Err(TryRecvError::Closed) => { result.rx = None; - result.finish_empty(); + result.finish_error("future output channel closed".to_string()); } Err(TryRecvError::Empty) => {} } } } + if let Some(error) = result.error.take() { + crate::process::stderr(self.id(), future_output_stderr_message(&error)); + return Err(anyhow::anyhow!(error)); + } if result.done { Ok(take(&mut result.result)) } else { @@ -1198,3 +1208,46 @@ impl pie::core::inference::HostMatcher for InstanceState { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn future_output_records_forward_errors() { + let (tx, rx) = oneshot::channel(); + tx.send(Err(anyhow::anyhow!("driver exploded"))).unwrap(); + + let mut future = FutureOutput { + result: None, + rx: Some(rx), + samplers: Vec::new(), + done: false, + model_id: 0, + context_id: None, + was_pinned: false, + fill_tokens: Vec::new(), + fill_positions: Vec::new(), + fill_masks: Vec::new(), + spec_tokens_for_fill: Vec::new(), + spec_positions_for_fill: Vec::new(), + adapter_id: None, + adapter_seed: None, + error: None, + }; + + future.ready().await; + + assert!(future.result.is_none()); + assert_eq!(future.error.as_deref(), Some("driver exploded")); + } + + #[test] + fn future_output_stderr_preserves_driver_error() { + let message = + future_output_stderr_message("cuda illegal memory access at executor.cpp:2758"); + + assert!(message.contains("Pie async forward failed")); + assert!(message.contains("cuda illegal memory access at executor.cpp:2758")); + } +} diff --git a/runtime/src/inference/scheduler/chunked.rs b/runtime/src/inference/scheduler/chunked.rs index db3fb8e54..c2799eb41 100644 --- a/runtime/src/inference/scheduler/chunked.rs +++ b/runtime/src/inference/scheduler/chunked.rs @@ -10,11 +10,11 @@ use std::collections::BTreeMap; use anyhow::Result; use tokio::sync::{mpsc, oneshot}; -use crate::context::pagestore::{PhysicalPageId, compute_last_page_len}; +use crate::context::pagestore::{compute_last_page_len, PhysicalPageId}; use crate::driver::SchedulerLimits; use crate::inference::ForwardOutput; -use super::{Completion, PendingRequest, RequestCapacityUsage, request_capacity_usage}; +use super::{request_capacity_usage, Completion, PendingRequest, RequestCapacityUsage}; pub(super) struct ChunkContinuation { original_request: pie_bridge::ForwardRequest, @@ -56,15 +56,16 @@ impl PendingRequest { } let usage = request_capacity_usage(&self, page_size); - if usage.forward_tokens <= limits.max_forward_tokens { - return Ok(self); - } + let chunk_size = match chunk_size_for_limits(&usage, limits) { + Ok(Some(chunk_size)) => chunk_size, + Ok(None) => return Ok(self), + Err(msg) => return Err((self, msg)), + }; - if let Err(msg) = validate_chunkable_request(&self.request, limits.max_forward_tokens) { + if let Err(msg) = validate_chunkable_request(&self.request, chunk_size) { return Err((self, msg)); } - let chunk_size = limits.max_forward_tokens; let chunk_end = chunk_size.min(self.request.token_ids.len()); let sampler_slots_by_chunk = chunk_sampler_slots_by_chunk(&self.request, chunk_size); if let Err(msg) = validate_chunk_capacity( @@ -658,11 +659,44 @@ fn response_nested_slot_count( Ok(count) } +fn chunk_size_for_limits( + usage: &RequestCapacityUsage, + limits: SchedulerLimits, +) -> std::result::Result, String> { + let mut chunk_size = limits.max_forward_tokens; + let mut needs_chunking = usage.forward_tokens > limits.max_forward_tokens; + + if usage.sampler_rows > 0 { + if usage.has_dense_logit_requirement && usage.forward_tokens > limits.max_logit_rows { + if limits.max_logit_rows == 0 { + return Err("driver max logit rows is zero".to_string()); + } + chunk_size = chunk_size.min(limits.max_logit_rows); + needs_chunking = true; + } + if usage.has_prob_sampling && usage.forward_tokens > limits.max_prob_rows { + if limits.max_prob_rows == 0 { + return Err("driver max probability rows is zero".to_string()); + } + chunk_size = chunk_size.min(limits.max_prob_rows); + needs_chunking = true; + } + } + + if !needs_chunking { + return Ok(None); + } + if chunk_size == 0 { + return Err("driver max forward tokens is zero".to_string()); + } + Ok(Some(chunk_size)) +} + fn validate_chunkable_request( req: &pie_bridge::ForwardRequest, - max_forward_tokens: usize, + chunk_size: usize, ) -> std::result::Result<(), String> { - if max_forward_tokens == 0 { + if chunk_size == 0 { return Err("driver max forward tokens is zero".to_string()); } if !req.spec_token_ids.is_empty() { @@ -670,13 +704,13 @@ fn validate_chunkable_request( "forward request has {} input tokens and {} speculative tokens, exceeding driver limit {}; chunked prefill does not yet support speculative drafts", req.token_ids.len(), req.spec_token_ids.len(), - max_forward_tokens + chunk_size )); } validate_chunk_request_shape(req)?; let total = req.token_ids.len(); - if total <= max_forward_tokens { + if total <= chunk_size { return Ok(()); } @@ -769,6 +803,20 @@ fn chunk_limit_error(usage: RequestCapacityUsage, limits: SchedulerLimits) -> Op )); } + if usage.logit_rows > limits.max_logit_rows { + return Some(format!( + "forward request chunk needs {} logit rows, exceeding driver limit {}", + usage.logit_rows, limits.max_logit_rows + )); + } + + if usage.prob_rows > limits.max_prob_rows { + return Some(format!( + "forward request chunk needs {} probability rows, exceeding driver limit {}", + usage.prob_rows, limits.max_prob_rows + )); + } + if usage.sampler_rows > limits.max_sampler_rows { return Some(format!( "forward request chunk has {} sampler rows, exceeding driver limit {}", @@ -1019,8 +1067,9 @@ fn chunk_capacity_usage( .copied() .unwrap_or(false); } + let carries_sampler = !sampler_slots.is_empty(); let has_dense_logit_requirement = original.has_user_mask - || !original.logit_masks.is_empty() + || (carries_sampler && !original.logit_masks.is_empty()) || has_output_spec || !all_samplers_token; let user_custom_mask_bytes = if original.has_user_mask && chunk_len > 1 { @@ -1034,16 +1083,14 @@ fn chunk_capacity_usage( Ok(RequestCapacityUsage { forward_tokens: chunk_len, page_refs: chunk_pages, - logit_rows: if has_dense_logit_requirement { - sampler_slots.len() - } else { + logit_rows: if sampler_slots.is_empty() { 0 - }, - prob_rows: if has_prob_sampling { - sampler_slots.len() + } else if has_dense_logit_requirement { + chunk_len } else { - 0 + sampler_slots.len() }, + prob_rows: if has_prob_sampling { chunk_len } else { 0 }, sampler_rows: sampler_slots.len(), logprob_labels, user_custom_mask_bytes, @@ -1278,7 +1325,9 @@ mod tests { fn chunk_sampler_slots(req: &PendingRequest) -> &[usize] { match &req.completion { Completion::Chunk { sampler_slots, .. } => sampler_slots, - Completion::Direct(_) | Completion::Chain { .. } => panic!("expected chunk continuation"), + Completion::Direct(_) | Completion::Chain { .. } => { + panic!("expected chunk continuation") + } } } @@ -1305,14 +1354,78 @@ mod tests { assert_eq!(cont.physical_page_ids, vec![100, 101, 102]); assert_eq!(cont.final_last_page_len, 2); } - Completion::Direct(_) | Completion::Chain { .. } => panic!("expected chunk continuation"), + Completion::Direct(_) | Completion::Chain { .. } => { + panic!("expected chunk continuation") + } } } + #[test] + fn dense_logit_prefill_chunks_by_logit_row_limit() { + let mut pending = positioned_pending(10, 4); + pending.request.sampling_indices = vec![9]; + pending.request.sampling_indptr = vec![0, 1]; + pending.request.samplers = vec![pie_bridge::Sampler::TopK { + temperature: 0.0, + k: 1, + }]; + pending.request.sampler_indptr = vec![0, 1]; + pending.request.logit_masks = vec![pie_bridge::Brle::from_vec(vec![1, u32::MAX])]; + pending.request.logit_mask_indptr = vec![0, 1]; + + let mut capped = limits(8, 100, 100); + capped.max_logit_rows = 4; + + let chunked = match pending.maybe_start_chunking(capped, 4) { + Ok(p) => p, + Err((_, msg)) => panic!("{msg}"), + }; + + assert_eq!(chunked.request.token_ids, vec![0, 1, 2, 3]); + assert!(chunked.request.sampling_indices.is_empty()); + assert!(chunked.request.samplers.is_empty()); + assert!(chunked.request.logit_masks.is_empty()); + assert_eq!(chunked.request.logit_mask_indptr, vec![0, 0]); + assert!(chunk_sampler_slots(&chunked).is_empty()); + + let (submit_tx, submit_rx) = crossbeam::channel::unbounded(); + let weak_submit_tx = submit_tx.clone(); + chunked.send_result( + Ok(pie_bridge::ForwardResponse { + num_requests: 1, + tokens_indptr: vec![0, 0], + ..Default::default() + }), + Some(&weak_submit_tx), + 4, + ); + + let second = submit_rx.try_recv().expect("second chunk"); + assert_eq!(second.request.token_ids, vec![4, 5, 6, 7]); + assert!(second.request.sampling_indices.is_empty()); + second.send_result( + Ok(pie_bridge::ForwardResponse { + num_requests: 1, + tokens_indptr: vec![0, 0], + ..Default::default() + }), + Some(&weak_submit_tx), + 4, + ); + + let final_chunk = submit_rx.try_recv().expect("final chunk"); + assert_eq!(final_chunk.request.token_ids, vec![8, 9]); + assert_eq!(final_chunk.request.sampling_indices, vec![1]); + assert_eq!(final_chunk.request.samplers.len(), 1); + assert_eq!(final_chunk.request.logit_masks.len(), 1); + assert_eq!(final_chunk.request.logit_mask_indptr, vec![0, 1]); + assert_eq!(chunk_sampler_slots(&final_chunk), &[0]); + } + #[test] fn oversized_prefill_with_chain_completion_degrades_to_chunk() { - use std::sync::Arc; use std::sync::atomic::AtomicBool; + use std::sync::Arc; use dashmap::DashMap; @@ -1368,7 +1481,9 @@ mod tests { assert_eq!(cont.physical_page_ids, vec![100, 101, 102]); assert_eq!(cont.final_last_page_len, 2); } - Completion::Direct(_) | Completion::Chain { .. } => panic!("expected chunk continuation"), + Completion::Direct(_) | Completion::Chain { .. } => { + panic!("expected chunk continuation") + } } } @@ -1383,7 +1498,9 @@ mod tests { Completion::Chunk { continuation: cont, .. } => cont, - Completion::Direct(_) | Completion::Chain { .. } => panic!("expected chunk continuation"), + Completion::Direct(_) | Completion::Chain { .. } => { + panic!("expected chunk continuation") + } }; let next = match cont.into_next_pending(4) { @@ -1398,7 +1515,9 @@ mod tests { Completion::Chunk { continuation: cont, .. } => assert_eq!(cont.chunk_end, 8), - Completion::Direct(_) | Completion::Chain { .. } => panic!("expected chunk continuation"), + Completion::Direct(_) | Completion::Chain { .. } => { + panic!("expected chunk continuation") + } } } @@ -1790,7 +1909,9 @@ mod tests { offset += chunk_len; match current.completion { - Completion::Direct(_) | Completion::Chain { .. } => panic!("expected chunk continuation"), + Completion::Direct(_) | Completion::Chain { .. } => { + panic!("expected chunk continuation") + } Completion::Chunk { continuation: cont, .. } => {