diff --git a/crates/apr-cli/src/commands/serve/handler_gpu_completion.rs b/crates/apr-cli/src/commands/serve/handler_gpu_completion.rs index 8e83f6b5b9..de7a1afdad 100644 --- a/crates/apr-cli/src/commands/serve/handler_gpu_completion.rs +++ b/crates/apr-cli/src/commands/serve/handler_gpu_completion.rs @@ -563,6 +563,11 @@ fn no_vocabulary(what: &str) -> CliError { /// explicit override must still beat the flag. The 2048 default now applies /// only when neither is set. /// +/// K1 (#4603): with neither set this now returns `None`, and the caller sizes the +/// KV to the model's context within free VRAM (`OwnedQuantizedModelCuda::for_serving`). +/// `serve::run` writes `REALIZR_CONTEXT_LENGTH` only for an operator-set flag, so the +/// 4096 default no longer caps a 32K model's KV at 4096. +/// /// NOTE ON SCOPE. This is NOT the same root as #2774 even though both land in /// the same allocation. #2774 is a VRAM budget computed before the weights are /// resident; this is a flag written to one name and read from another. They are @@ -571,13 +576,39 @@ fn no_vocabulary(what: &str) -> CliError { /// while `--context-length` was ignored the batched KV was half-sized, which is /// the only reason the 7B appeared to survive c=4 on a 24 GB card at all. #[cfg(any(all(feature = "inference", feature = "cuda"), test))] -fn resolve_serve_max_seq_len(explicit_override: Option<&str>, context_length: Option<&str>) -> usize { - const DEFAULT_MAX_SEQ_LEN: usize = 2048; +fn resolve_serve_max_seq_len( + explicit_override: Option<&str>, + context_length: Option<&str>, +) -> Option { explicit_override .and_then(|v| v.parse::().ok()) .or_else(|| context_length.and_then(|v| v.parse::().ok())) .filter(|&n| n > 0) - .unwrap_or(DEFAULT_MAX_SEQ_LEN) +} + +/// GH-129 + #2762 + K1 (#4603): build the CUDA model a server runs, with its KV sized +/// from `REALIZR_MAX_SEQ_LEN` / `--context-length` when one is set, else from the +/// model's context within free VRAM. Every CUDA serve path (GGUF, APR, SafeTensors) +/// takes this; the APR and SafeTensors paths used `OwnedQuantizedModelCuda::new`, a +/// fixed 2048, which no flag could move. +#[cfg(all(feature = "inference", feature = "cuda"))] +fn serving_cuda_model( + model: realizar::gguf::OwnedQuantizedModel, +) -> std::result::Result { + use realizar::gguf::OwnedQuantizedModelCuda; + match resolve_serve_max_seq_len( + std::env::var("REALIZR_MAX_SEQ_LEN").ok().as_deref(), + std::env::var("REALIZR_CONTEXT_LENGTH").ok().as_deref(), + ) { + Some(n) => { + println!(" Max sequence length: {n}"); + OwnedQuantizedModelCuda::with_max_seq_len(model, 0, n) + }, + None => { + println!(" Max sequence length: the model's context, within free VRAM (K1)"); + OwnedQuantizedModelCuda::for_serving(model, 0) + }, + } } /// Start GGUF server with CUDA acceleration (PAR-111). @@ -625,14 +656,7 @@ fn start_gguf_server_cuda( "Enabling optimized CUDA acceleration (PAR-111)...".cyan() ); - // GH-129 + #2762: resolve the KV-cache context length. - let max_seq_len = resolve_serve_max_seq_len( - std::env::var("REALIZR_MAX_SEQ_LEN").ok().as_deref(), - std::env::var("REALIZR_CONTEXT_LENGTH").ok().as_deref(), - ); - println!(" Max sequence length: {max_seq_len}"); - - match OwnedQuantizedModelCuda::with_max_seq_len(quantized_model, 0, max_seq_len) { + match serving_cuda_model(quantized_model) { Ok(mut cuda_model) => { preload_gpu_weights(&mut cuda_model); @@ -780,7 +804,7 @@ mod ctx_length_2762_tests { fn context_length_flag_reaches_the_kv_cache() { assert_eq!( resolve_serve_max_seq_len(None, Some("4096")), - 4096, + Some(4096), "--context-length is written to REALIZR_CONTEXT_LENGTH and must be \ the KV cache's max_len; ignoring it sizes the batched KV stride \ from a constant (#2762)" @@ -792,24 +816,28 @@ mod ctx_length_2762_tests { /// 7.4 GB of unified memory silently gets 4096 back. #[test] fn explicit_override_still_beats_the_flag() { - assert_eq!(resolve_serve_max_seq_len(Some("1024"), Some("4096")), 1024); + assert_eq!( + resolve_serve_max_seq_len(Some("1024"), Some("4096")), + Some(1024) + ); } - /// DISCRIMINATION: stays GREEN both before and after. With neither set the - /// historical default is unchanged, so this fix moves no default. + /// K1 (#4603): with neither set there is no fixed length. RED before K1, which + /// returned 2048 here -- and `serve::run` always wrote the 4096 default, so a + /// Qwen3-8B (40960 context) was served from a 4096-position KV. #[test] - fn default_is_unchanged_when_nothing_is_set() { - assert_eq!(resolve_serve_max_seq_len(None, None), 2048); + fn nothing_set_means_size_from_the_model() { + assert_eq!(resolve_serve_max_seq_len(None, None), None); } /// A garbled value must not resolve to 0 -- `num_kv_heads * 0 * head_dim` /// is a zero-length KV cache, which the allocator accepts and the attention - /// kernel then reads out of. + /// kernel then reads out of. It falls back to sizing from the model. #[test] - fn junk_and_zero_fall_back_to_the_default() { - assert_eq!(resolve_serve_max_seq_len(None, Some("banana")), 2048); - assert_eq!(resolve_serve_max_seq_len(None, Some("0")), 2048); - assert_eq!(resolve_serve_max_seq_len(Some("0"), None), 2048); + fn junk_and_zero_fall_back_to_the_model() { + assert_eq!(resolve_serve_max_seq_len(None, Some("banana")), None); + assert_eq!(resolve_serve_max_seq_len(None, Some("0")), None); + assert_eq!(resolve_serve_max_seq_len(Some("0"), None), None); } } diff --git a/crates/apr-cli/src/commands/serve/handlers_include_01.rs b/crates/apr-cli/src/commands/serve/handlers_include_01.rs index 1011889e6d..e943ec4301 100644 --- a/crates/apr-cli/src/commands/serve/handlers_include_01.rs +++ b/crates/apr-cli/src/commands/serve/handlers_include_01.rs @@ -67,7 +67,7 @@ fn start_apr_server_gpu( println!("{}", "Enabling fused CUDA acceleration (GH-87)...".cyan()); - let mut cuda_model = OwnedQuantizedModelCuda::new(quantized, 0) + let mut cuda_model = serving_cuda_model(quantized) .map_err(|e| CliError::InferenceFailed(format!("CUDA init failed: {e}")))?; preload_gpu_weights(&mut cuda_model); @@ -262,7 +262,7 @@ fn start_safetensors_server_gpu( println!("{}", "Enabling fused CUDA acceleration (GH-88)...".cyan()); - let mut cuda_model = OwnedQuantizedModelCuda::new(quantized, 0) + let mut cuda_model = serving_cuda_model(quantized) .map_err(|e| CliError::InferenceFailed(format!("CUDA init failed: {e}")))?; preload_gpu_weights(&mut cuda_model); diff --git a/crates/apr-cli/src/commands/serve/mod.rs b/crates/apr-cli/src/commands/serve/mod.rs index f096601761..2ac7569ae7 100644 --- a/crates/apr-cli/src/commands/serve/mod.rs +++ b/crates/apr-cli/src/commands/serve/mod.rs @@ -296,8 +296,13 @@ pub(crate) fn run(model_path: &Path, config: &ServerConfig) -> Result<()> { eprintln!("[PMAT-297] Thread pool config: {e} (may already be initialized)"); } - // GH-286: Set env vars for realizr's KV cache and FP8 control - std::env::set_var("REALIZR_CONTEXT_LENGTH", config.context_length.to_string()); + // GH-286: Set env vars for realizr's KV cache and FP8 control. + // K1 (#4603): only an operator-set `--context-length` is written. The 4096 default + // written unconditionally sized every GPU KV to 4096 whatever the model's context; + // unset, the GGUF CUDA path sizes the KV to the model's context within free VRAM. + if config.context_length != types::DEFAULT_CONTEXT_LENGTH { + std::env::set_var("REALIZR_CONTEXT_LENGTH", config.context_length.to_string()); + } if config.no_fp8_cache { std::env::set_var("REALIZR_NO_FP8_CACHE", "1"); } diff --git a/crates/aprender-serve/src/api/batch.rs b/crates/aprender-serve/src/api/batch.rs index 599e8ae821..5381670efe 100644 --- a/crates/aprender-serve/src/api/batch.rs +++ b/crates/aprender-serve/src/api/batch.rs @@ -271,7 +271,7 @@ fn try_apr_generate( .generate_with_cache(&prompt_ids, &gen_config) .map_err(|e| { api_err( - StatusCode::INTERNAL_SERVER_ERROR, + crate::api::generation_error_status(&e), format!("APR generation failed: {e}"), ) })?; @@ -321,7 +321,7 @@ fn registry_generate( let generated = model .generate(&prompt, &config) - .map_err(|e| api_err(StatusCode::INTERNAL_SERVER_ERROR, e))?; + .map_err(|e| api_err(crate::api::generation_error_status(&e), e))?; let token_ids: Vec = generated .iter() @@ -515,7 +515,7 @@ fn try_cuda_batch_generate( .generate_gpu_resident(&prompt_ids, &q_config) .map_err(|e| { api_err( - StatusCode::INTERNAL_SERVER_ERROR, + crate::api::generation_error_status(&e), format!("CUDA generation failed: {e}"), ) })?; @@ -622,7 +622,7 @@ fn try_apr_batch_generate( .generate_with_cache(&prompt_ids, &gen_config) .map_err(|e| { api_err( - StatusCode::INTERNAL_SERVER_ERROR, + crate::api::generation_error_status(&e), format!("APR generation failed: {e}"), ) })?; @@ -682,7 +682,7 @@ fn registry_batch_generate( let prompt: Vec = prompt_ids.iter().map(|&id| id as usize).collect(); let generated = model .generate(&prompt, &config) - .map_err(|e| api_err(StatusCode::INTERNAL_SERVER_ERROR, e))?; + .map_err(|e| api_err(crate::api::generation_error_status(&e), e))?; let token_ids: Vec = generated .iter() .map(|&id| { diff --git a/crates/aprender-serve/src/api/batch_processing.rs b/crates/aprender-serve/src/api/batch_processing.rs index 4f2baa933d..25b9c19f40 100644 --- a/crates/aprender-serve/src/api/batch_processing.rs +++ b/crates/aprender-serve/src/api/batch_processing.rs @@ -355,7 +355,7 @@ pub async fn gpu_batch_completions_handler( Ok(tokens) => results.push(tokens), Err(e) => { return Err(( - StatusCode::INTERNAL_SERVER_ERROR, + crate::api::generation_error_status(&e), Json(ErrorResponse { error: format!("Generation failed: {e}"), }), @@ -492,7 +492,7 @@ fn try_cuda_generate( .generate_gpu_resident(&prompt_ids, &q_config) .map_err(|e| { api_err( - StatusCode::INTERNAL_SERVER_ERROR, + crate::api::generation_error_status(&e), format!("CUDA generation failed: {e}"), ) })?; diff --git a/crates/aprender-serve/src/api/cuda_chat_backend.rs b/crates/aprender-serve/src/api/cuda_chat_backend.rs index 6f229a0de7..0af685c43a 100644 --- a/crates/aprender-serve/src/api/cuda_chat_backend.rs +++ b/crates/aprender-serve/src/api/cuda_chat_backend.rs @@ -61,7 +61,7 @@ fn try_safetensors_cuda_backend( let msg = format!("SafeTensors CUDA generation failed: {e}"); return Some( ( - axum::http::StatusCode::INTERNAL_SERVER_ERROR, + crate::api::generation_error_status(&e), axum::Json(serde_json::json!({"error": msg})), ) .into_response(), @@ -118,8 +118,12 @@ async fn try_cuda_backend( }; // GH-319: Use actual model architecture for chat template detection let arch_hint = state.model_architecture(); + // D5: refused before any path (streaming, batch scheduler, direct) takes the + // model, so the client gets a 400 rather than a 500 or a mid-stream error. + let tokenized = + tokenize_chat_prompt(&tokenizer, &request.messages, arch_hint.as_deref(), request.thinking(), state); let prompt_ids = - match tokenize_chat_prompt(&tokenizer, &request.messages, arch_hint.as_deref(), request.thinking(), state) { + match fit_serving_context(state, tokenized) { Ok(ids) => ids, Err(r) => return Some(r), }; @@ -246,7 +250,7 @@ async fn try_cuda_backend( let generate_start = std::time::Instant::now(); let generated = match dense_cuda_turn(&mut cuda_model, &prompt_ids, &q_config, |_| true) { Ok(g) => g, - Err(e) => return Some(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e)), + Err(e) => return Some(fail_response(state, crate::api::generation_error_status(&e), e)), }; let _ = timing_tx.send(phase_split(&mut cuda_model, generate_start)); let tokens: Vec = generated.iter().skip(prompt_tokens).copied().collect(); diff --git a/crates/aprender-serve/src/api/gpu_completions_handler.rs b/crates/aprender-serve/src/api/gpu_completions_handler.rs index 701b7b0d28..b96aec9b0f 100644 --- a/crates/aprender-serve/src/api/gpu_completions_handler.rs +++ b/crates/aprender-serve/src/api/gpu_completions_handler.rs @@ -62,7 +62,7 @@ fn try_gpu_completions( })?; let generated = gpu_model .generate(&prompt, &gpu_config) - .map_err(|e| rerr(state, StatusCode::INTERNAL_SERVER_ERROR, e))?; + .map_err(|e| rerr(state, crate::api::generation_error_status(&e), e))?; let token_ids: Vec = generated .iter() @@ -246,7 +246,7 @@ fn registry_completions( let generated = model .generate(&prompt, &config) - .map_err(|e| rerr(state, StatusCode::INTERNAL_SERVER_ERROR, e))?; + .map_err(|e| rerr(state, crate::api::generation_error_status(&e), e))?; let token_ids: Vec = generated .iter() .skip(prompt_tokens) @@ -383,6 +383,13 @@ async fn try_cuda_gguf_completions( return Err(rerr(state, StatusCode::BAD_REQUEST, "Prompt cannot be empty")); } let prompt_tokens = prompt_ids.len(); + // D5: a 400 before the batch scheduler sees it (its errors come back as strings). + if let Some(msg) = state + .serving_context() + .and_then(|ctx| super::serve_context_refusal(prompt_tokens, ctx)) + { + return Err(rerr(state, StatusCode::BAD_REQUEST, msg)); + } let eos = state.cached_eos_token_id.unwrap_or(151643); let q_config = QuantizedGenerateConfig { @@ -610,7 +617,7 @@ async fn completions_inner( model.generate_gpu_resident_logprobs( &prompt_ids.iter().map(|&id| id as u32).collect::>(), &config, - ).map_err(|e| rerr(&state, StatusCode::INTERNAL_SERVER_ERROR, e))? + ).map_err(|e| rerr(&state, super::generation_error_status(&e), e))? }; let prompt_len = prompt_ids.len(); let gen_tokens: Vec = result.tokens[prompt_len..].to_vec(); @@ -699,7 +706,7 @@ pub async fn logprobs_handler( model.generate_gpu_resident_logprobs( &prompt_ids.iter().map(|&x| x as u32).collect::>(), &config, - ).map_err(|e| rerr(&state, StatusCode::INTERNAL_SERVER_ERROR, e))? + ).map_err(|e| rerr(&state, crate::api::generation_error_status(&e), e))? }; let prompt_len = prompt_ids.len(); @@ -766,7 +773,7 @@ pub async fn perplexity_handler( // realizr#203: Run BOTH paths for comparison during development let ppl_sequential = model .perplexity_gpu_resident(&token_ids) - .map_err(|e| rerr(&state, StatusCode::INTERNAL_SERVER_ERROR, e))?; + .map_err(|e| rerr(&state, crate::api::generation_error_status(&e), e))?; let ppl_batched = model.perplexity_gpu_batched(&token_ids).ok(); drop(model); diff --git a/crates/aprender-serve/src/api/mod.rs b/crates/aprender-serve/src/api/mod.rs index f2747b2335..a51c76bd10 100644 --- a/crates/aprender-serve/src/api/mod.rs +++ b/crates/aprender-serve/src/api/mod.rs @@ -235,6 +235,10 @@ pub struct AppState { qwen35_session: Option>, /// GH-330: Cached EOS token ID (avoids RwLock in hot path) cached_eos_token_id: Option, + /// D5 (ruling-3715-2010): positions a dense CUDA serve turn can reach, the + /// model's context capped by the device KV cache. Cached at construction so + /// the pre-flight length check never waits on the scheduler's write lock. + cached_serving_context: Option, /// GH-152: Enable verbose request/response logging verbose: bool, /// GH-103: Enable inference tracing (propagates into QuantizedGenerateConfig.trace) @@ -377,6 +381,27 @@ pub(crate) fn generation_error_status(err: &RealizarError) -> StatusCode { } } +/// D5 (ruling-3715-2010): the refusal for a prompt that cannot fit a serve turn, +/// or `None` when it fits with room to answer. +/// +/// `serving_context` is the model's context capped by the device KV cache +/// ([`AppState::serving_context`]). A prompt at or past it used to reach the GPU +/// forward and come back as a 500 ("the device KV cache holds 4096"), which +/// tells the client the server broke. It is decided by the request alone, so it +/// is a 400 and says which limit it hit. +pub(crate) fn serve_context_refusal( + prompt_tokens: usize, + serving_context: usize, +) -> Option { + (prompt_tokens >= serving_context).then(|| { + format!( + "context exceeds {serving_context}: the prompt is {prompt_tokens} tokens and this \ + server holds {serving_context} positions per turn (the model's context, capped by \ + the GPU KV cache). It was refused whole, not truncated." + ) + }) +} + include!("mod_app_state_gpu.rs"); include!("mod_app_state_qwen35.rs"); include!("mod_create_demo.rs"); diff --git a/crates/aprender-serve/src/api/mod_app_state_gpu.rs b/crates/aprender-serve/src/api/mod_app_state_gpu.rs index d2ce9d6496..55169b80cf 100644 --- a/crates/aprender-serve/src/api/mod_app_state_gpu.rs +++ b/crates/aprender-serve/src/api/mod_app_state_gpu.rs @@ -64,6 +64,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -123,6 +124,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -190,6 +192,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -225,6 +228,12 @@ impl AppState { // contention with the batch scheduler's write lock. let arch = Some(cuda_model.model().config.architecture.clone()); let eos = cuda_model.model().config.eos_token_id; + // D5: cached for the same reason as `arch` — the pre-flight length check + // must not wait on the scheduler's write lock. + let serving_context = crate::gguf::dense_session::cap_context( + cuda_model.model().config.context_length, + cuda_model.executor().max_kv_len(), + ); let (audit_logger, audit_sink) = create_audit_state(); Ok(Self { @@ -261,6 +270,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: eos, + cached_serving_context: Some(serving_context), verbose: false, trace: false, model_source: None, @@ -282,6 +292,12 @@ impl AppState { let tokenizer = BPETokenizer::with_merges(vocab, merges, unk)?; let arch = Some(cuda_model.model().config.architecture.clone()); let eos = cuda_model.model().config.eos_token_id; + // D5: cached for the same reason as `arch` — the pre-flight length check + // must not wait on the scheduler's write lock. + let serving_context = crate::gguf::dense_session::cap_context( + cuda_model.model().config.context_length, + cuda_model.executor().max_kv_len(), + ); let (audit_logger, audit_sink) = create_audit_state(); Ok(Self { @@ -320,6 +336,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: eos, + cached_serving_context: Some(serving_context), verbose: false, trace: false, model_source: None, @@ -384,6 +401,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -617,6 +635,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: eos_id, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -679,6 +698,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -778,6 +798,13 @@ impl AppState { None } + /// D5: positions a dense CUDA serve turn can reach (model context capped + /// by the device KV cache), or `None` when no GGUF CUDA model is loaded. + #[must_use] + pub fn serving_context(&self) -> Option { + self.cached_serving_context + } + /// GH-330: Get EOS token ID from whichever model backend is loaded. /// /// **Design by Contract (Meyer 1992)**: The model config carries its own EOS diff --git a/crates/aprender-serve/src/api/mod_app_state_new.rs b/crates/aprender-serve/src/api/mod_app_state_new.rs index 0faac77db4..b009779b9b 100644 --- a/crates/aprender-serve/src/api/mod_app_state_new.rs +++ b/crates/aprender-serve/src/api/mod_app_state_new.rs @@ -44,6 +44,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -107,6 +108,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -272,6 +274,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -349,6 +352,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -407,6 +411,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -471,6 +476,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -530,6 +536,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, @@ -603,6 +610,7 @@ impl AppState { moe_no_gpu: true, qwen35_session: None, cached_eos_token_id: None, + cached_serving_context: None, verbose: false, trace: false, model_source: None, diff --git a/crates/aprender-serve/src/api/mod_app_state_qwen35.rs b/crates/aprender-serve/src/api/mod_app_state_qwen35.rs index e8d1225ada..4a92f601a2 100644 --- a/crates/aprender-serve/src/api/mod_app_state_qwen35.rs +++ b/crates/aprender-serve/src/api/mod_app_state_qwen35.rs @@ -82,6 +82,7 @@ impl AppState { session: std::sync::Mutex::new(session), })), cached_eos_token_id: eos_token_id, + cached_serving_context: None, verbose: false, trace: false, model_source: None, diff --git a/crates/aprender-serve/src/api/openai_handlers.rs b/crates/aprender-serve/src/api/openai_handlers.rs index 81eb375e1d..375e55b8e7 100644 --- a/crates/aprender-serve/src/api/openai_handlers.rs +++ b/crates/aprender-serve/src/api/openai_handlers.rs @@ -64,6 +64,25 @@ fn require_tokenizer(state: &AppState) -> Result, Response> { }) } +/// D5 (ruling-3715-2010): pass a tokenized turn through, or refuse it with a 400 +/// that names the limit when it cannot fit the serving context +/// ([`AppState::serving_context`]; `None` means no cap, so nothing is refused). +#[allow(clippy::result_large_err)] +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub(super) fn fit_serving_context( + state: &AppState, + tokenized: Result, Response>, +) -> Result, Response> { + let ids = tokenized?; + match state + .serving_context() + .and_then(|ctx| super::serve_context_refusal(ids.len(), ctx)) + { + Some(msg) => Err(fail_response(state, StatusCode::BAD_REQUEST, msg)), + None => Ok(ids), + } +} + /// Format chat messages, tokenize, validate non-empty. #[allow(clippy::result_large_err)] fn tokenize_chat_prompt( @@ -1062,7 +1081,13 @@ fn try_gpu_backend( }; let generated = match model.generate(&prompt_usize, &gpu_config) { Ok(g) => g, - Err(e) => return Some(fail_response(state, StatusCode::INTERNAL_SERVER_ERROR, e)), + Err(e) => { + return Some(fail_response( + state, + crate::api::generation_error_status(&e), + e, + )) + }, }; let token_ids: Vec = generated diff --git a/crates/aprender-serve/src/api/stream_generate.rs b/crates/aprender-serve/src/api/stream_generate.rs index eab5c4d564..d0e2a03516 100644 --- a/crates/aprender-serve/src/api/stream_generate.rs +++ b/crates/aprender-serve/src/api/stream_generate.rs @@ -40,7 +40,7 @@ fn dense_stream_tokens( let generated = model .generate(&prompt, &config) - .map_err(|e| api_err(StatusCode::INTERNAL_SERVER_ERROR, e))?; + .map_err(|e| api_err(crate::api::generation_error_status(&e), e))?; let token_ids: Vec = generated .iter() @@ -78,6 +78,13 @@ fn try_cuda_stream_tokens( let tokenizer = require_tok(state)?; let prompt_ids = tokenize_prompt(&tokenizer, &request.prompt)?; let prompt_len = prompt_ids.len(); + // D5: over the serving context is the request's fault, so a 400. + if let Some(msg) = state + .serving_context() + .and_then(|ctx| super::serve_context_refusal(prompt_len, ctx)) + { + return Err(api_err(StatusCode::BAD_REQUEST, msg)); + } let q_config = QuantizedGenerateConfig { max_tokens: request.max_tokens, @@ -103,7 +110,7 @@ fn try_cuda_stream_tokens( .generate_gpu_resident(&prompt_ids, &q_config) .map_err(|e| { api_err( - StatusCode::INTERNAL_SERVER_ERROR, + super::generation_error_status(&e), format!("CUDA generation failed: {e}"), ) })?; diff --git a/crates/aprender-serve/src/api/tests/mod.rs b/crates/aprender-serve/src/api/tests/mod.rs index a7c2d51fff..5e06d9bfe6 100644 --- a/crates/aprender-serve/src/api/tests/mod.rs +++ b/crates/aprender-serve/src/api/tests/mod.rs @@ -32,6 +32,7 @@ mod tests_09; mod tests_10; mod tests_11; mod completion_request; +mod serve_context_d5; mod completions_invalid; mod chat_completion; mod tests_15; @@ -53,6 +54,8 @@ mod embeddings_pmat803; // PMAT-803: model-backed embeddings (semantic-similarit mod sse_stream_whitespace; // Dogfood 0.63.0: SSE deltas must reassemble with whitespace intact mod serve_dead_routes_3991; // aprender#3991: every generation route GET / lists must generate; text is the completion mod native_routes_2376; // aprender#2376: native routes on a quantized server, KV-cache budget, sampling fields +#[cfg(feature = "cuda")] +mod serve_kv_limit_k1; // K1 (#4603): serve KV = model context; at-limit 200 whole, past it a 400 mod router_flags; // --no-cors / --no-metrics must change HTTP behaviour, not just the banner mod ollama_compat_http; // Dogfood 0.63.0 (#2396/#2402): /api/tags|show|version routed, stream:true is NDJSON, /realize/* stops fabricating mod embed_and_envelope_2376; // aprender#2376(1 seventh route, 7, 8) + #2396(2): embeddings on a quantized server, one error envelope, / and /ready diff --git a/crates/aprender-serve/src/api/tests/serve_context_d5.rs b/crates/aprender-serve/src/api/tests/serve_context_d5.rs new file mode 100644 index 0000000000..f2b75e1bef --- /dev/null +++ b/crates/aprender-serve/src/api/tests/serve_context_d5.rs @@ -0,0 +1,73 @@ +//! D5 (ruling-3715-2010): a serve turn past the device KV cache is a 400 that +//! names the limit ("context exceeds 4096"), never a 500. +//! +//! The observed failure (h4-depcheck-0d): a 27155-token turn on a dense CUDA +//! model reached `reserve`, failed with "the device KV cache holds 4096" and +//! came back as HTTP 500. The handlers now refuse it with +//! [`serve_context_refusal`] before any path takes the model. + +use crate::api::openai_handlers::fit_serving_context; +use crate::api::{generation_error_status, serve_context_refusal, AppState}; +use crate::error::RealizarError; +use axum::http::StatusCode; + +#[test] +fn d5_turn_past_the_device_kv_is_refused_and_names_the_limit() { + let msg = serve_context_refusal(27155, 4096).expect("27155 tokens cannot fit 4096"); + assert!(msg.starts_with("context exceeds 4096"), "{msg}"); + assert!(msg.contains("27155"), "{msg}"); + assert!(msg.contains("refused whole, not truncated"), "{msg}"); +} + +#[test] +fn d5_boundary_a_prompt_that_fills_the_window_leaves_no_room_to_answer() { + assert!(serve_context_refusal(4097, 4096).is_some()); + assert!(serve_context_refusal(4096, 4096).is_some()); + assert_eq!(serve_context_refusal(4095, 4096), None); + assert_eq!(serve_context_refusal(1, 4096), None); +} + +#[test] +fn d5_no_cuda_model_means_no_serving_cap() { + // The pre-flight reads `serving_context()`; only a GGUF CUDA state sets it, + // so the CPU and Qwen3.5 paths keep their own context rules. + let state = AppState::demo().expect("demo state"); + assert_eq!(state.serving_context(), None); +} + +#[test] +fn d5_the_session_error_behind_it_is_a_client_error() { + let err = RealizarError::ContextLimitExceeded { + provided: 27155, + maximum: 4096, + }; + assert_eq!(generation_error_status(&err), StatusCode::BAD_REQUEST); +} + +#[test] +fn d5_chat_pre_flight_passes_a_turn_through_when_there_is_no_cap() { + let state = AppState::demo().expect("demo state"); + let ids = fit_serving_context(&state, Ok(vec![7, 8, 9])).expect("no cap, no refusal"); + assert_eq!(ids, vec![7, 8, 9]); +} + +#[test] +fn d5_chat_pre_flight_refuses_a_turn_past_the_cap_with_a_400() { + let mut state = AppState::demo().expect("demo state"); + state.cached_serving_context = Some(3); + assert_eq!( + fit_serving_context(&state, Ok(vec![1, 2])).expect("2 tokens fit 3"), + vec![1, 2] + ); + let refused = fit_serving_context(&state, Ok(vec![1, 2, 3])).expect_err("3 tokens fill 3"); + assert_eq!(refused.status(), StatusCode::BAD_REQUEST); +} + +#[test] +fn d5_chat_pre_flight_keeps_a_tokenize_error() { + let mut state = AppState::demo().expect("demo state"); + state.cached_serving_context = Some(3); + let err = axum::response::IntoResponse::into_response(StatusCode::IM_A_TEAPOT); + let kept = fit_serving_context(&state, Err(err)).expect_err("the error passes through"); + assert_eq!(kept.status(), StatusCode::IM_A_TEAPOT); +} diff --git a/crates/aprender-serve/src/api/tests/serve_kv_limit_k1.rs b/crates/aprender-serve/src/api/tests/serve_kv_limit_k1.rs new file mode 100644 index 0000000000..8c7221a22b --- /dev/null +++ b/crates/aprender-serve/src/api/tests/serve_kv_limit_k1.rs @@ -0,0 +1,206 @@ +//! K1: `apr serve` sizes the device KV from the model's context, so a prompt at +//! the limit is answered whole and one past it is refused with a 400. +//! +//! Before K1 serve fixed the KV at 4096 whatever the model declared, so on a +//! model whose context is larger a prompt past 4096 tokens overflowed the cache +//! and came back as a 500. This builds an 8192-context model (above the old +//! 4096, so a return to the fixed size fails the KV assertion), serves it +//! through the real router and batch scheduler on the device, and checks three +//! things. At the limit it answers 200 with every requested token. At the limit +//! and one past it, it answers 400 and says the prompt was refused whole. +//! +//! The model is synthetic (no file), so the test runs on any CUDA runner. +//! Without a device it prints the skip line that `cuda-unit` fails on. + +use crate::api::{create_router, AppState}; +use crate::gguf::test_helpers::create_test_model_with_config; +use crate::gguf::{ + ArchConstraints, GGUFConfig, OwnedQKVWeights, OwnedQuantizedModel, OwnedQuantizedModelCuda, + OwnedQuantizedTensor, +}; +use axum::body::Body; +use axum::http::{Request, StatusCode}; +use tower::ServiceExt; + +/// Above serve's pre-K1 fixed 4096, so a fixed size cannot pass. +const CONTEXT: usize = 8192; +const VOCAB: usize = 256; +/// Generation per request. At the limit, prompt + this == `CONTEXT`. +const MAX_TOKENS: usize = 4; + +fn config() -> GGUFConfig { + GGUFConfig { + architecture: "llama".to_string(), + constraints: ArchConstraints::from_architecture("llama"), + hidden_dim: 256, + intermediate_dim: 512, + num_heads: 4, + num_kv_heads: 4, + num_layers: 1, + vocab_size: VOCAB, + rope_theta: 10000.0, + context_length: CONTEXT, + eps: 1e-5, + rope_type: 0, + explicit_head_dim: None, + query_pre_attn_scalar: None, + bos_token_id: None, + eos_token_id: None, + } +} + +/// One token per `a`: the only single-character entries are the letters, and the +/// greedy matcher has no longer match for a run of `a`. +fn vocab() -> Vec { + let mut v: Vec = vec!["".to_string()]; + v.extend(('a'..='z').map(String::from)); + v.extend((v.len()..VOCAB).map(|i| format!(""))); + v +} + +/// Q4_K weights small and centred on zero (d = 2^-10, dmin = 2^-7, pseudo-random +/// scales and quants). The shared `create_q4k_test_data` pattern is all-positive +/// at d = 1.0, which saturates a forward pass until the device parity gate +/// (rightly) refuses it; this one keeps activations near unit scale. +fn q4k(in_dim: usize, out_dim: usize, seed: u64) -> OwnedQuantizedTensor { + let blocks = in_dim.div_ceil(256); + let mut data = vec![0u8; out_dim * blocks * 144]; + let mut x = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1); + for block in data.chunks_exact_mut(144) { + block[0..2].copy_from_slice(&0x1400_u16.to_le_bytes()); + block[2..4].copy_from_slice(&0x2000_u16.to_le_bytes()); + for b in &mut block[4..] { + x = x + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + *b = (x >> 56) as u8; + } + } + OwnedQuantizedTensor { + data, + in_dim, + out_dim, + qtype: 12, + } +} + +/// A llama-shaped model with separate Q/K/V, as a llama GGUF loads: the +/// GPU-resident path serve generates on takes nothing else +/// (`supports_gpu_resident`). +fn model() -> OwnedQuantizedModel { + let c = config(); + let (h, kv, ffn) = ( + c.hidden_dim, + c.num_kv_heads * (c.hidden_dim / c.num_heads), + c.intermediate_dim, + ); + let mut model = create_test_model_with_config(&c); + for (i, layer) in model.layers.iter_mut().enumerate() { + let s = 10 * i as u64; + layer.qkv_weight = OwnedQKVWeights::Separate { + q: q4k(h, h, s + 1), + k: q4k(h, kv, s + 2), + v: q4k(h, kv, s + 3), + }; + layer.attn_output_weight = q4k(h, h, s + 4); + layer.ffn_up_weight = q4k(h, ffn, s + 5); + layer.ffn_down_weight = q4k(ffn, h, s + 6); + layer.ffn_gate_weight = Some(q4k(h, ffn, s + 7)); + // llama normalises before the FFN; the shared fixture leaves it out. + layer.ffn_norm_weight = Some(vec![1.0; h]); + } + model.lm_head_weight = q4k(h, VOCAB, 99); + model +} + +/// The router `apr serve` builds for a CUDA GGUF: the model sized by +/// `for_serving`, behind the continuous-batching scheduler. `None` (after the +/// device-skip line) without a CUDA device. +fn served() -> Option<(axum::Router, usize)> { + let cuda = match OwnedQuantizedModelCuda::for_serving(model(), 0) { + Ok(m) => m, + Err(e) => { + eprintln!("SKIP K1: CUDA model init failed: {}", e.error); + return None; + }, + }; + let kv = cuda.executor().max_kv_len(); + let state = AppState::with_cuda_model_and_vocab(cuda, vocab()).expect("state"); + let model = state.cuda_model().expect("the CUDA model").clone(); + let tx = crate::api::cuda_batch_scheduler::spawn_cuda_batch_scheduler( + model, + crate::api::cuda_batch_scheduler::CudaBatchConfig::default(), + crate::api::InFlightCounter::new(), + ); + Some((create_router(state.with_cuda_batch_tx(tx)), kv)) +} + +async fn complete(app: &axum::Router, prompt_tokens: usize) -> (StatusCode, serde_json::Value) { + let body = serde_json::json!({ + "model": "k1", + "prompt": "a".repeat(prompt_tokens), + "max_tokens": MAX_TOKENS, + "temperature": 0.0, + }); + let response = app + .clone() + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/completions") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .expect("request"), + ) + .await + .expect("the router answers"); + let status = response.status(); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .expect("body"); + let json = serde_json::from_slice(&bytes) + .unwrap_or_else(|_| serde_json::json!({ "raw": String::from_utf8_lossy(&bytes) })); + (status, json) +} + +#[tokio::test(flavor = "multi_thread")] +async fn k1_serve_kv_is_the_model_context_and_the_limit_is_a_clean_400() { + let Some((app, kv)) = served() else { + return; + }; + assert_eq!( + kv, CONTEXT, + "K1: the serving KV is the model's context, not a fixed size" + ); + + // At the limit: prompt + generation fill the context exactly. Answered + // whole: 200, every requested token, the prompt counted as sent. + let at_limit = CONTEXT - MAX_TOKENS; + let (status, json) = complete(&app, at_limit).await; + assert_eq!( + status, + StatusCode::OK, + "at the limit ({at_limit} + {MAX_TOKENS}): {json}" + ); + assert_eq!(json["usage"]["prompt_tokens"], at_limit, "{json}"); + assert_eq!( + json["usage"]["completion_tokens"], MAX_TOKENS, + "full output: {json}" + ); + assert!(json["choices"][0]["text"].is_string(), "{json}"); + eprintln!( + "K1 at-limit: 200, prompt {at_limit}, completion {MAX_TOKENS}, finish {}", + json["choices"][0]["finish_reason"] + ); + + // At and past the limit: refused before it reaches the device, as a 400 + // that says so. Not a 500, and not a truncated 200. + for prompt in [CONTEXT, CONTEXT + 1] { + let (status, json) = complete(&app, prompt).await; + assert_eq!(status, StatusCode::BAD_REQUEST, "prompt {prompt}: {json}"); + let text = json.to_string(); + assert!(text.contains("context exceeds"), "prompt {prompt}: {text}"); + assert!(text.contains("refused whole"), "prompt {prompt}: {text}"); + eprintln!("K1 over-limit: prompt {prompt} → 400 {text}"); + } +} diff --git a/crates/aprender-serve/src/gguf/cuda/mod.rs b/crates/aprender-serve/src/gguf/cuda/mod.rs index 14f99f7b43..69336149cd 100644 --- a/crates/aprender-serve/src/gguf/cuda/mod.rs +++ b/crates/aprender-serve/src/gguf/cuda/mod.rs @@ -40,6 +40,8 @@ mod forward_qwen35_cuda; /// #3714: the Qwen3-MoE decoder resident on the GPU. mod forward_qwen3_moe_resident; mod generation; +mod resident; +mod session_kv; mod speculative; mod weights; @@ -67,6 +69,28 @@ use super::utils::verbose; // IMP-800: CUDA-Accelerated Model Wrapper // ============================================================================= +/// How `build` sizes the device KV cache. +#[derive(Debug, Clone, Copy)] +enum KvLen { + /// Exactly this many positions (`new`, `with_max_seq_len`). + Fixed(usize), + /// K1 (#4603): the model's context capped by free VRAM, budgeting the server's batched-KV slots too and + /// never below serve's old fixed length (`session_kv::serving_kv_len`). + FitServing, +} + +/// #3715: device bytes load makes resident, from [`OwnedQuantizedModelCuda::resident_estimate`]. +struct ResidentEstimate { + /// f32 K+V bytes of one position over every layer. + kv_per_pos: usize, + /// Quantized weights. + weights: usize, + /// FP16/FP8 prefill weight cache (0 when this profile warms none). + cache: usize, + /// GH-178 workspace reserve plus the chunked-prefill score budget. + reserve: usize, +} + /// Error from CUDA model initialization that preserves the unconsumed model. /// /// When `OwnedQuantizedModelCuda::new()` fails, the model is returned inside this error @@ -624,7 +648,26 @@ impl OwnedQuantizedModelCuda { max_seq_len: usize, ) -> std::result::Result { let model = Self::check_not_moe(model)?; - Self::build(model, device_ordinal, max_seq_len) + Self::build(model, device_ordinal, KvLen::Fixed(max_seq_len)) + } + + /// K1 (#4603): an `apr serve` model. The device KV takes + /// the model's context when it fits, but the VRAM budget also holds + /// `session_kv::SERVE_BATCH_SLOTS` batched-KV slots of that length, because the + /// continuous-batching scheduler allocates them beside the session KV + /// (`fit_batched_kv_alloc`), and it never drops below serve's old 4096. Before K1 + /// serve used a fixed 4096 whatever the model's context, so a > 4096-token prompt + /// on a 32K model overflowed the KV and answered 500. + /// + /// # Errors + /// + /// The same CUDA / capability / quant / MoE errors as [`Self::with_max_seq_len`]. + pub fn for_serving( + model: OwnedQuantizedModel, + device_ordinal: i32, + ) -> std::result::Result { + let model = Self::check_not_moe(model)?; + Self::build(model, device_ordinal, KvLen::FitServing) } /// Build the CUDA wrapper for a MoE model whose caller runs the MoE forward @@ -636,7 +679,7 @@ impl OwnedQuantizedModelCuda { /// /// The same CUDA / capability / quant errors as [`Self::new`]. pub fn new_for_moe_forward(model: OwnedQuantizedModel, device_ordinal: i32) -> Result { - Self::build(model, device_ordinal, 2048).map_err(|e| e.error) + Self::build(model, device_ordinal, KvLen::Fixed(2048)).map_err(|e| e.error) } /// #3992: the dense CUDA forward cannot run a Mixture-of-Experts model: its @@ -678,7 +721,7 @@ impl OwnedQuantizedModelCuda { fn build( model: OwnedQuantizedModel, device_ordinal: i32, - max_seq_len: usize, + kv_len: KvLen, ) -> std::result::Result { use crate::cuda::CudaExecutor; @@ -765,6 +808,14 @@ impl OwnedQuantizedModelCuda { let num_layers = model.layers.len(); let num_kv_heads = model.config.num_kv_heads; // PAR-021 GQA support let head_dim = model.config.hidden_dim / model.config.num_heads; + let max_seq_len = match kv_len { + KvLen::Fixed(n) => n, + KvLen::FitServing => { + // C25 revert probe: the pre-K1 fixed size. DO NOT MERGE. + let _ = Self::serving_kv_len(&executor, &model, memory_info.0); + 4096 + }, + }; if let Err(error) = Self::configure_executor(&mut executor, &model, max_seq_len) { return Err(CudaInitError { @@ -801,6 +852,73 @@ impl OwnedQuantizedModelCuda { Ok(cuda_model) } + /// K1 (#4603): the device KV length for [`Self::for_serving`] -- the model's context + /// capped by free VRAM, with the server's batched slots budgeted and serve's old 4096 as the floor + /// (see `session_kv::serving_kv_len`). + fn serving_kv_len( + executor: &crate::cuda::CudaExecutor, + model: &OwnedQuantizedModel, + free_vram: usize, + ) -> usize { + let r = Self::resident_estimate(executor, model); + let len = session_kv::serving_kv_len( + model.config.context_length, + r.kv_per_pos, + free_vram, + r.weights + r.cache, + r.reserve, + ); + eprintln!( + "[K1] serve device KV: {len} positions (context {}, {} batched slots budgeted, {:.1} GB free, {:.1} GB weights + {:.1} GB prefill cache resident)", + model.config.context_length, + session_kv::SERVE_BATCH_SLOTS, + free_vram as f64 / 1e9, + r.weights as f64 / 1e9, + r.cache as f64 / 1e9, + ); + len + } + + /// #3715: what load makes resident on the device: the quantized weights, the + /// FP16 (or FP8) prefill weight cache when this profile warms one, and the + /// GH-178 reserve plus the chunked-prefill score budget. Also the f32 K+V bytes + /// of one position over every layer. + fn resident_estimate( + executor: &crate::cuda::CudaExecutor, + model: &OwnedQuantizedModel, + ) -> ResidentEstimate { + let c = &model.config; + let head_dim = c.hidden_dim / c.num_heads.max(1); + let kv_per_pos = + 2 * c.num_kv_heads * head_dim * std::mem::size_of::() * model.layers.len(); + let (weights, elems) = resident::projection_sizes( + model + .layers + .iter() + .flat_map(resident::layer_projections) + .chain(std::iter::once(&model.lm_head_weight)), + ); + let profile = &executor.gpu_profile; + let cache_bytes_per_elem = if profile.fp8_prefill { + 1 + } else if std::env::var("HGEMM_PREFILL").as_deref() == Ok("0") + || (profile.cc >= 120 && std::env::var("FORCE_FP16_CACHE").as_deref() != Ok("1")) + { + 0 + } else { + 2 + }; + let cache = cache_bytes_per_elem * elems; + let reserve = + 3_500_000_000 + crate::gguf::cuda::forward_qwen35_cuda::PREFILL_SCORES_BUDGET_BYTES; + ResidentEstimate { + kv_per_pos, + weights, + cache, + reserve, + } + } + /// GH-129: Free CPU projection weight copies after GPU preload. /// /// On Jetson (unified memory), the CPU `Vec` copies are redundant after diff --git a/crates/aprender-serve/src/gguf/cuda/resident.rs b/crates/aprender-serve/src/gguf/cuda/resident.rs new file mode 100644 index 0000000000..57dbf3ee1b --- /dev/null +++ b/crates/aprender-serve/src/gguf/cuda/resident.rs @@ -0,0 +1,108 @@ +//! #3715: what a CUDA model's load makes resident on the device, so a decision taken +//! before the upload (the QK-norm batched-prefill fit check) can see whether the +//! FP16 prefill weight cache and its workspace will fit beside the weights. + +use crate::gguf::quantized::{OwnedQKVWeights, OwnedQuantizedLayer, OwnedQuantizedTensor}; + +/// The prefill weight cache holds the seven projection matrices of every layer +/// plus the LM head (the "197 matrices" a 28-layer model logs): the quantized +/// tensors load uploads. Returns `(quantized bytes, elements)` of those tensors; +/// the cache is `elements` × its bytes per element. Read from the tensors +/// themselves, so it cannot drift from the model the way a dims formula can. +pub(crate) fn projection_sizes<'a>( + tensors: impl IntoIterator, +) -> (usize, usize) { + tensors.into_iter().fold((0, 0), |(bytes, elems), t| { + (bytes + t.data.len(), elems + t.in_dim * t.out_dim) + }) +} + +/// The projection tensors of one layer, fused QKV or separate. +pub(crate) fn layer_projections(l: &OwnedQuantizedLayer) -> Vec<&OwnedQuantizedTensor> { + let mut v = match &l.qkv_weight { + OwnedQKVWeights::Fused(t) => vec![t], + OwnedQKVWeights::Separate { q, k, v } => vec![q, k, v], + }; + v.extend([&l.attn_output_weight, &l.ffn_up_weight, &l.ffn_down_weight]); + v.extend(l.ffn_gate_weight.as_ref()); + v +} + +#[cfg(test)] +mod tests { + use super::*; + + fn t(in_dim: usize, out_dim: usize, bytes: usize) -> OwnedQuantizedTensor { + OwnedQuantizedTensor { + data: vec![0; bytes], + in_dim, + out_dim, + qtype: 12, + } + } + + /// A Qwen3 layer: separate Q/K/V, gated FFN. `bytes` goes on every tensor. + fn layer(hidden: usize, kv: usize, inter: usize, bytes: usize) -> OwnedQuantizedLayer { + OwnedQuantizedLayer { + attn_norm_weight: Vec::new(), + attn_norm_bias: None, + qkv_weight: OwnedQKVWeights::Separate { + q: t(hidden, hidden, bytes), + k: t(hidden, kv, bytes), + v: t(hidden, kv, bytes), + }, + qkv_bias: None, + attn_output_weight: t(hidden, hidden, bytes), + attn_output_bias: None, + ffn_up_weight: t(hidden, inter, bytes), + ffn_up_bias: None, + ffn_down_weight: t(inter, hidden, bytes), + ffn_down_bias: None, + ffn_gate_weight: Some(t(hidden, inter, bytes)), + ffn_gate_bias: None, + ffn_norm_weight: None, + ffn_norm_bias: None, + attn_q_norm_weight: None, + attn_k_norm_weight: None, + post_attn_norm_weight: None, + post_ffw_norm_weight: None, + } + } + + /// FP16 cache bytes of a Qwen3 model with these dims, via the tensor walk. + fn fp16_cache(layers: usize, hidden: usize, kv: usize, inter: usize) -> usize { + let l = layer(hidden, kv, inter, 0); + let head = t(hidden, 151_936, 0); + let per_layer = projection_sizes(layer_projections(&l)).1; + 2 * (layers * per_layer + projection_sizes([&head]).1) + } + + #[test] + fn fp16_cache_estimate_matches_the_logged_qwen3_1_7b_cache() { + // [PMAT-037] FP16 weight cache: 197 matrices cached (3281.5 MB), measured on lambda. + let bytes = fp16_cache(28, 2048, 1024, 6144); + let logged = 3281.5 * 1024.0 * 1024.0; + assert!( + (bytes as f64 - logged).abs() / logged < 0.02, + "{bytes} vs {logged}" + ); + } + + #[test] + fn projection_walk_counts_every_tensor_once() { + // Seven projections, distinct dims so a dropped or doubled tensor shows. + let mut l = layer(2, 3, 5, 7); + assert_eq!(layer_projections(&l).len(), 7); + let (bytes, elems) = projection_sizes(layer_projections(&l)); + // q 2x2, k 2x3, v 2x3, o 2x2, up 2x5, down 5x2, gate 2x5 + assert_eq!((bytes, elems), (7 * 7, 4 + 6 + 6 + 4 + 10 + 10 + 10)); + // A fused QKV is one tensor; an ungated FFN has no gate. + l.qkv_weight = OwnedQKVWeights::Fused(t(2, 8, 7)); + l.ffn_gate_weight = None; + assert_eq!(layer_projections(&l).len(), 4); + assert_eq!( + projection_sizes(layer_projections(&l)), + (4 * 7, 16 + 4 + 10 + 10) + ); + } +} diff --git a/crates/aprender-serve/src/gguf/cuda/session_kv.rs b/crates/aprender-serve/src/gguf/cuda/session_kv.rs new file mode 100644 index 0000000000..cd1bced264 --- /dev/null +++ b/crates/aprender-serve/src/gguf/cuda/session_kv.rs @@ -0,0 +1,164 @@ +//! #3715: the device KV length for a chat session, whose turn length is unknown at load. +//! +//! `apr chat` built its CUDA model with `OwnedQuantizedModelCuda::new`, a flat 2048 +//! positions. The dense session refuses a turn longer than the device KV, so a +//! 26k-token chat turn on Qwen3-1.7B left the GPU for the CPU and hit the 600 s cell +//! timeout. `apr run` sizes the KV to the turn (#4268); chat cannot, because it loads +//! before the first turn is read. So chat takes the model's whole context when the +//! device has room for it, after the weights, the FP16/FP8 prefill cache and a +//! reserve, and keeps the old 2048 floor when it does not. + +/// The positions a session model's device KV holds. Pure; see the module doc. +/// +/// * `context_length` — the model's trained context (the ceiling). +/// * `kv_bytes_per_position` — f32 K+V over every layer for one position. +/// * `free_vram` — free device bytes before the KV and the weights are allocated. +/// * `resident_bytes` — the weights plus the prefill weight cache that load will upload. +/// * `reserve_bytes` — workspace, prefill score buffers and runtime headroom. +pub(crate) fn session_kv_len( + context_length: usize, + kv_bytes_per_position: usize, + free_vram: usize, + resident_bytes: usize, + reserve_bytes: usize, +) -> usize { + const FLOOR: usize = 2048; + let floor = FLOOR.min(context_length.max(1)); + let room = free_vram.saturating_sub(resident_bytes.saturating_add(reserve_bytes)); + (room / kv_bytes_per_position.max(1)).clamp(floor, context_length.max(floor)) +} + +/// K1 (#4603): batched-KV slots an `apr serve` KV budget holds beside the session KV. +/// The continuous-batching scheduler allocates slots of the same length for concurrent +/// requests (`fit_batched_kv_alloc`); two keep c=2 at full context. +pub(crate) const SERVE_BATCH_SLOTS: usize = 2; + +/// K1 (#4603): serve's fixed KV length before K1. A server never gets less, so a +/// model that fills the card serves exactly what it served before. +pub(crate) const SERVE_FLOOR: usize = 4096; + +/// K1 (#4603): the device KV length for `apr serve`: [`session_kv_len`] with each +/// position also paying for [`SERVE_BATCH_SLOTS`] batched slots, floored at +/// [`SERVE_FLOOR`] (or the model's context when shorter). +pub(crate) fn serving_kv_len( + context_length: usize, + kv_bytes_per_position: usize, + free_vram: usize, + resident_bytes: usize, + reserve_bytes: usize, +) -> usize { + let fit = session_kv_len( + context_length, + kv_bytes_per_position.saturating_mul(1 + SERVE_BATCH_SLOTS), + free_vram, + resident_bytes, + reserve_bytes, + ); + fit.max(SERVE_FLOOR.min(context_length.max(1))) +} + +#[cfg(test)] +mod tests { + use super::*; + + const GB: usize = 1_000_000_000; + // f32 K+V per position: 2 * kv_heads * head_dim * 4 * layers + const QWEN3_1_7B_KV: usize = 2 * 8 * 128 * 4 * 28; + const QWEN3_8B_KV: usize = 2 * 8 * 128 * 4 * 36; + // FP16 prefill caches as logged on lambda ([PMAT-037]): 3281.5 MB for 1.7B, + // ~15.1 GB for 8B (`resident::projection_sizes` computes them from the tensors). + const QWEN3_1_7B_FP16: usize = 3_441_000_000; + const QWEN3_8B_FP16: usize = 15_100_000_000; + + #[test] + fn small_model_gets_its_whole_context() { + // Qwen3-1.7B on a 24 GB card with 1.7 GB used elsewhere: weights ~1.1 GB. + // The 26k-token chat turn that timed out must fit. + let len = session_kv_len( + 40_960, + QWEN3_1_7B_KV, + 22 * GB, + GB + QWEN3_1_7B_FP16, + 4_600_000_000, + ); + assert_eq!(len, 40_960); + assert!(len > 26_147); + } + + #[test] + fn a_model_that_fills_the_card_keeps_the_old_floor() { + // Qwen3-8B: 4.8 GB weights + ~15 GB FP16 cache leave no room for more KV, + // so the session keeps the 2048 it always had rather than starving the cache. + let len = session_kv_len( + 40_960, + QWEN3_8B_KV, + 22 * GB, + 4_800_000_000 + QWEN3_8B_FP16, + 4_600_000_000, + ); + assert_eq!(len, 2048); + } + + #[test] + fn partial_room_is_used_and_never_exceeds_it() { + let room = 3 * GB; + let len = session_kv_len(1 << 20, QWEN3_1_7B_KV, 10 * GB, 5 * GB, 2 * GB); + assert!(len * QWEN3_1_7B_KV <= room); + assert!((len + 1) * QWEN3_1_7B_KV > room); + assert!(len > 2048); + } + + #[test] + fn floor_never_exceeds_a_short_context() { + assert_eq!(session_kv_len(512, QWEN3_1_7B_KV, 0, 0, 0), 512); + assert_eq!(session_kv_len(0, 0, 0, 0, 0), 1); + } + + /// K1: Qwen3-1.7B on a 24 GB card serves far past 4096 (the room after weights, + /// FP16 cache and reserve, split across the session KV and two batched slots: + /// ~18.8K positions), where serve used a fixed 4096 -- the > 4096 serve cells + /// overflowed the KV. With room for all of it (gx10, 128 GB) it is the whole 40960. + #[test] + fn serve_gives_a_small_model_far_more_than_4096() { + let len = serving_kv_len( + 40_960, + QWEN3_1_7B_KV, + 22 * GB, + GB + QWEN3_1_7B_FP16, + 4_600_000_000, + ); + assert!(len > 16_384, "{len}"); + assert!(len < 40_960, "{len}"); + assert_eq!( + serving_kv_len(40_960, QWEN3_1_7B_KV, 100 * GB, GB, 4_600_000_000), + 40_960 + ); + } + + /// K1: the budget pays for the batched slots too, so the session KV plus + /// `SERVE_BATCH_SLOTS` slots of the same length fit in the room. + #[test] + fn serve_budget_holds_the_batched_slots() { + let room = 3 * GB; + let len = serving_kv_len(1 << 20, QWEN3_1_7B_KV, 10 * GB, 5 * GB, 2 * GB); + let per_pos = QWEN3_1_7B_KV * (1 + SERVE_BATCH_SLOTS); + assert!(len * per_pos <= room); + assert!((len + 1) * per_pos > room); + assert!(len < session_kv_len(1 << 20, QWEN3_1_7B_KV, 10 * GB, 5 * GB, 2 * GB)); + } + + /// K1: a model that fills the card (Qwen3-8B + FP16 cache on 24 GB) keeps serve's + /// old 4096, never the session floor of 2048 -- no model serves less than before. + #[test] + fn serve_never_drops_below_its_old_length() { + let len = serving_kv_len( + 40_960, + QWEN3_8B_KV, + 22 * GB, + 4_800_000_000 + QWEN3_8B_FP16, + 4_600_000_000, + ); + assert_eq!(len, SERVE_FLOOR); + assert_eq!(serving_kv_len(2048, QWEN3_8B_KV, 0, 0, 0), 2048); + } +} diff --git a/crates/aprender-serve/src/gguf/inference/forward/dense_session.rs b/crates/aprender-serve/src/gguf/inference/forward/dense_session.rs index dabb3cf476..3d8790152d 100644 --- a/crates/aprender-serve/src/gguf/inference/forward/dense_session.rs +++ b/crates/aprender-serve/src/gguf/inference/forward/dense_session.rs @@ -551,6 +551,52 @@ fn grow_or_build( true } +/// Positions a serve turn can reach: the model's context, capped by the device +/// KV cache it was built with (`device_kv`, 0 when there is none) (D5, +/// ruling-3715-2010). +/// +/// Serve has no CPU copy to fall back to, so the device cache is the real +/// window. Reporting the model's context instead let a 27k-token turn through to +/// `reserve`, which failed as a GPU fault and reached the client as a 500. With +/// the cap as the session's context, a prompt past it is refused as +/// [`RealizarError::ContextLimitExceeded`] (a 400) and a large `max_tokens` is +/// clamped to what is left, the same rule as every other context window. +/// CUDA-free on purpose, so the rule is tested on every build. +#[must_use] +pub fn cap_context(context_length: usize, device_kv: usize) -> usize { + let context = context_length.max(1); + match device_kv { + 0 => context, + device => context.min(device), + } +} + #[cfg(test)] #[path = "dense_session_tests.rs"] mod tests; + +#[cfg(test)] +mod cap_context_tests { + use super::cap_context; + + #[test] + fn d5_the_device_kv_caps_a_longer_model_context() { + assert_eq!(cap_context(32768, 4096), 4096); + } + + #[test] + fn d5_a_shorter_model_context_is_kept() { + assert_eq!(cap_context(2048, 4096), 2048); + } + + #[test] + fn d5_no_device_cache_means_the_model_context() { + assert_eq!(cap_context(32768, 0), 32768); + } + + #[test] + fn d5_a_zero_model_context_still_holds_one_position() { + assert_eq!(cap_context(0, 0), 1); + assert_eq!(cap_context(0, 4096), 1); + } +} diff --git a/crates/aprender-serve/src/gguf/inference/forward/dense_session_borrowed.rs b/crates/aprender-serve/src/gguf/inference/forward/dense_session_borrowed.rs index e4091059ed..28389c8971 100644 --- a/crates/aprender-serve/src/gguf/inference/forward/dense_session_borrowed.rs +++ b/crates/aprender-serve/src/gguf/inference/forward/dense_session_borrowed.rs @@ -37,7 +37,9 @@ impl<'a> BorrowedCudaForward<'a> { pub fn new(model: &'a mut OwnedQuantizedModelCuda) -> Self { let config = &model.model().config; let arch = crate::tensor_names::normalize_architecture(&config.architecture); - let context_length = config.context_length.max(1); + // D5: the device KV cache is the real window; serve has no CPU copy to move to. + let context_length = + super::dense_session::cap_context(config.context_length, model.executor().max_kv_len()); let line = format!( "Backend: GPU ({}, {} MB VRAM)", model.device_name(), diff --git a/docs/audits/surface_audit.csv b/docs/audits/surface_audit.csv index 2e9e3091d6..a4b934c5ae 100644 --- a/docs/audits/surface_audit.csv +++ b/docs/audits/surface_audit.csv @@ -75,7 +75,7 @@ apr,POST /v1/chat/completions,4,x86_64-linux;nvidia-cuda,llama.cpp,yes,3,http-ap apr,"POST /v1/chat/completions (stream:true, SSE)",4,x86_64-linux;nvidia-cuda,llama.cpp,no,3,http-apr-serve,"crates/aprender-serve/src/api/openai_handlers.rs:755-787 (stream branches crates/aprender-serve/src/api/cuda_chat_backend.rs:129,275,405,520,986; ignore_eos fail-closed 12,362,471)",high apr,POST /v1/chat/completions/stream,4,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:139,high apr,POST /v1/completions,2,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:131,high -apr,"POST /v1/completions (stream:true, SSE)",2,x86_64-linux,llama.cpp,no,3,http-apr-serve,crates/aprender-serve/src/api/gpu_completions_handler.rs:398 (SSE :340-378),high +apr,"POST /v1/completions (stream:true, SSE)",2,x86_64-linux,llama.cpp,no,3,http-apr-serve,crates/aprender-serve/src/api/gpu_completions_handler.rs:405 (SSE :340-378),high apr,POST /v1/embeddings,4,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:142,high apr,POST /v1/explain,4,x86_64-linux,llama.cpp,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:145,high apr,POST /v1/gpu/warmup,7,x86_64-linux,vllm,yes,3,http-apr-serve,crates/aprender-serve/src/api/router.rs:148,medium diff --git a/docs/roadmaps/entries/PMAT-4616.yaml b/docs/roadmaps/entries/PMAT-4616.yaml new file mode 100644 index 0000000000..6363e80b88 --- /dev/null +++ b/docs/roadmaps/entries/PMAT-4616.yaml @@ -0,0 +1,19 @@ +- id: PMAT-4616 + github_issue: 4616 + item_type: task + title: 'D5: serve over device KV (4096) returns 500 — refuse with 400 context exceeds N; --context-length sizes KV for gate-named cells' + status: inprogress + priority: critical + assigned_to: aprender-2e-serve + created: 2026-09-28T20:47:39Z + updated: 2026-09-28T20:47:39Z + spec: null + acceptance_criteria: + - 'lambda RTX 4090, 5432-token chat turn: main returns 500, the branch returns 400 context exceeds 4096, the branch with --context-length 32768 returns 200' + - 'api/tests/serve_context_d5.rs: the refusal fires at prompt_tokens >= serving_context, and serving_context = min(model context, device KV)' + - 'CUDA chat, completions and stream all refuse before generating; ContextLimitExceeded maps to 400, never 500' + phases: [] + subtasks: [] + estimated_effort: null + labels: [] + notes: 'OPERATOR ruling D5 (ruling-3715). Child of #3715. The ladder part (--context-length on gate-named rungs) is 0d #4603.' diff --git a/docs/roadmaps/roadmap.yaml b/docs/roadmaps/roadmap.yaml index e7abff39da..20ae97a3ad 100644 --- a/docs/roadmaps/roadmap.yaml +++ b/docs/roadmaps/roadmap.yaml @@ -22211,3 +22211,22 @@ roadmap: labels: - kind:code notes: 'Operator ruling 2026-09-28 13:53Z (via the cop): own small PR, first in the queue, on main before the #4512 push.' +- id: PMAT-4616 + github_issue: 4616 + item_type: task + title: 'D5: serve over device KV (4096) returns 500 — refuse with 400 context exceeds N; --context-length sizes KV for gate-named cells' + status: inprogress + priority: critical + assigned_to: aprender-2e-serve + created: 2026-09-28T20:47:39Z + updated: 2026-09-28T20:47:39Z + spec: null + acceptance_criteria: + - 'lambda RTX 4090, 5432-token chat turn: main returns 500, the branch returns 400 context exceeds 4096, the branch with --context-length 32768 returns 200' + - 'api/tests/serve_context_d5.rs: the refusal fires at prompt_tokens >= serving_context, and serving_context = min(model context, device KV)' + - 'CUDA chat, completions and stream all refuse before generating; ContextLimitExceeded maps to 400, never 500' + phases: [] + subtasks: [] + estimated_effort: null + labels: [] + notes: 'OPERATOR ruling D5 (ruling-3715). Child of #3715. The ladder part (--context-length on gate-named rungs) is 0d #4603.'