diff --git a/openinfer-deepseek-v2-lite/src/scheduler.rs b/openinfer-deepseek-v2-lite/src/scheduler.rs index f307a0bd9..69a8b08a2 100644 --- a/openinfer-deepseek-v2-lite/src/scheduler.rs +++ b/openinfer-deepseek-v2-lite/src/scheduler.rs @@ -529,8 +529,8 @@ impl From for PendingRequest { max_tokens: req.max_tokens, lora_adapter: req.lora_adapter, token_tx: req.token_tx, - logprobs: req.logprobs, - echo: req.echo, + logprobs: req.logprobs.unwrap_or(0), + echo: req.prompt_logprobs.is_some(), } } } @@ -839,6 +839,14 @@ fn admission_decision(req: &PendingRequest, supported_context: usize) -> Admissi "DeepSeek-V2-Lite EP=2 mixed serving gate does not return logprobs yet".to_string(), ); } + // Honor-or-reject: the echo stub would emit all-None prompt logprobs, + // which silently strips the requested data (#720). + if req.echo { + return AdmissionDecision::Reject( + "DeepSeek-V2-Lite EP=2 mixed serving gate does not return prompt logprobs yet" + .to_string(), + ); + } if req.lora_adapter.is_some() { return AdmissionDecision::Reject( "DeepSeek-V2-Lite EP=2 mixed serving gate does not support LoRA adapters".to_string(), diff --git a/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs b/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs index e38156204..133ac3ef4 100644 --- a/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs +++ b/openinfer-deepseek-v2-lite/tests/e2e_ep2.rs @@ -631,8 +631,8 @@ fn run_mixed_serving_generation(model_path: &Path, model_path_label: &str) -> Re max_tokens, lora_adapter: None, token_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, }; receivers.push((id, token_rx)); requests.push(req); @@ -716,8 +716,8 @@ fn run_mixed_serving_position_fallback( max_tokens, lora_adapter: None, token_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, }; receivers.push((id, token_rx)); requests.push(req); @@ -782,8 +782,8 @@ fn run_mixed_serving_rejection_isolation( max_tokens: 4, lora_adapter: None, token_tx: invalid_tx, - logprobs: 1, - echo: false, + logprobs: Some(1), + prompt_logprobs: None, }; let (valid_tx, mut valid_rx) = TokenSink::standalone(); @@ -796,8 +796,8 @@ fn run_mixed_serving_rejection_isolation( max_tokens: 6, lora_adapter: None, token_tx: valid_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, }; submit_concurrently(handle, vec![invalid_req, valid_req])?; diff --git a/openinfer-dynamo-backend/src/engine.rs b/openinfer-dynamo-backend/src/engine.rs index 10958c767..5865f21ec 100644 --- a/openinfer-dynamo-backend/src/engine.rs +++ b/openinfer-dynamo-backend/src/engine.rs @@ -291,10 +291,11 @@ impl LLMEngine for OpeninferBackend { lora_adapter: None, token_tx: sink, // M1 does not surface per-token logprobs (the Dynamo `log_probs` - // slot stays None), so pin 0 rather than make openinfer pay the - // full-vocab O(V) logprob pass for a value we would then drop. - logprobs: 0, - echo: false, + // slot stays None), so leave logprobs disabled rather than make + // openinfer pay the full-vocab O(V) logprob pass for a value we + // would then drop. + logprobs: None, + prompt_logprobs: None, }; if handle.submit(req).is_err() { diff --git a/openinfer-engine/src/engine.rs b/openinfer-engine/src/engine.rs index d3c5e0dff..5319998fd 100644 --- a/openinfer-engine/src/engine.rs +++ b/openinfer-engine/src/engine.rs @@ -78,8 +78,17 @@ pub struct GenerateRequest { /// one engine share a single tagged output channel behind this sink (see /// [`TokenSink`]); the frontend demuxes by tag. pub token_tx: TokenSink, - pub logprobs: usize, - pub echo: bool, + /// Completion logprobs, mirroring the pinned vLLM contract: `None` + /// disables them, `Some(0)` requests the sampled token's logprob with no + /// additional top entries, and `Some(k)` adds the top-`k` alternatives. + /// (`-1` full-vocabulary requests are rejected by the frontend before a + /// request ever reaches the engine.) + pub logprobs: Option, + /// Prompt logprobs, same `None`/`Some(0)`/`Some(k)` semantics as + /// `logprobs` but independent of it. `Some(_)` makes the scheduler emit + /// [`TokenEvent::PromptTokens`] with one entry per prompt position; the + /// leading position carries `None` because it has no predecessor logits. + pub prompt_logprobs: Option, } #[derive(Clone, Debug, Eq, PartialEq)] diff --git a/openinfer-glm52/src/scheduler/admission.rs b/openinfer-glm52/src/scheduler/admission.rs index 10feb98c7..a4bb8e88d 100644 --- a/openinfer-glm52/src/scheduler/admission.rs +++ b/openinfer-glm52/src/scheduler/admission.rs @@ -59,8 +59,8 @@ fn validate_request(req: &GenerateRequest, max_model_len: usize) -> Result<(), S )); } } - if req.logprobs > 0 || req.echo { - return Err("GLM5.2 bring-up does not support logprobs/echo".to_owned()); + if req.logprobs.is_some() || req.prompt_logprobs.is_some() { + return Err("GLM5.2 bring-up does not support completion/prompt logprobs".to_owned()); } if req.lora_adapter.is_some() { return Err("GLM5.2 does not support LoRA adapters".to_owned()); diff --git a/openinfer-glm52/src/scheduler/testkit.rs b/openinfer-glm52/src/scheduler/testkit.rs index 3a1f7b88f..bad4cd8af 100644 --- a/openinfer-glm52/src/scheduler/testkit.rs +++ b/openinfer-glm52/src/scheduler/testkit.rs @@ -48,8 +48,8 @@ pub(super) fn request( max_tokens, lora_adapter: None, token_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, } } diff --git a/openinfer-kernels/build.rs b/openinfer-kernels/build.rs index af96b05b0..95a2dbfa8 100644 --- a/openinfer-kernels/build.rs +++ b/openinfer-kernels/build.rs @@ -1563,6 +1563,7 @@ fn main() { || stem == "flashinfer_sampling" || stem == "flashinfer_top1" || stem == "glm52_topk" + || stem == "logprobs" { for dir in &flashinfer.cccl { nvcc_args.extend(["-I".to_string(), dir.to_string_lossy().to_string()]); diff --git a/openinfer-kernels/csrc/shared/logprobs.cu b/openinfer-kernels/csrc/shared/logprobs.cu new file mode 100644 index 000000000..9168646b8 --- /dev/null +++ b/openinfer-kernels/csrc/shared/logprobs.cu @@ -0,0 +1,211 @@ +// Batched GPU logprobs reduction (#719). +// +// Replaces the per-row CPU path (extract_vec + full-vocab D2H + stream sync + +// three O(V) host passes) with: +// 1. `logprobs_lse_bf16_cuda` — one block per scored row; two-pass online +// log-sum-exp (f64 partial sums) plus the picked token's logprob. +// 2. `logprobs_topk_bf16_cuda` — vendored FlashInfer FilteredTopK with +// deterministic smallest-index tie-break plus an index-sort / +// stable-value-sort chain so output order exactly matches the host +// reference `token_logprob_from_row` (value desc, index asc). +// 3. `logprobs_gather_rows_bf16_cuda` — row-index gather so sparse row +// subsets can feed the contiguous FilteredTopK input layout. +// +// D2H per batch is O(rows * (k + 1)) instead of O(rows * V), with a single +// stream sync instead of one per row. + +#include "common.cuh" +#include "ffi_guard.cuh" + +#include + +#include +#include + +#define LOGPROBS_BLOCK 256 + +// --------------------------------------------------------------------------- +// logsumexp + picked-token logprob +// --------------------------------------------------------------------------- + +__global__ void logprobs_lse_kernel(const __nv_bfloat16* __restrict__ logits, + const unsigned int* __restrict__ row_indices, + const unsigned int* __restrict__ picked, + int vocab_size, float* __restrict__ out_lse, + float* __restrict__ out_picked_lp) { + const int scored = blockIdx.x; + const long long row = row_indices == nullptr ? scored : row_indices[scored]; + const __nv_bfloat16* x = logits + row * (long long)vocab_size; + + // Pass 1: block max. + float local_max = -INFINITY; + for (int i = threadIdx.x; i < vocab_size; i += LOGPROBS_BLOCK) { + local_max = fmaxf(local_max, __bfloat162float(x[i])); + } + local_max = warp_reduce_max(local_max); + + __shared__ float warp_max[LOGPROBS_BLOCK / WARP_SIZE]; + __shared__ double warp_sum[LOGPROBS_BLOCK / WARP_SIZE]; + const int warp = threadIdx.x / WARP_SIZE; + const int lane = threadIdx.x % WARP_SIZE; + if (lane == 0) { + warp_max[warp] = local_max; + } + __syncthreads(); + + float row_max = warp_max[0]; + for (int w = 1; w < LOGPROBS_BLOCK / WARP_SIZE; ++w) { + row_max = fmaxf(row_max, warp_max[w]); + } + + // Pass 2: f64 partial sums of exp(x - max), matching the f64 accumulation of + // the host reference. Row is L2-hot from pass 1. + double local_sum = 0.0; + for (int i = threadIdx.x; i < vocab_size; i += LOGPROBS_BLOCK) { + local_sum += (double)expf(__bfloat162float(x[i]) - row_max); + } + // f64 warp reduction via 64-bit shuffle. + for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) { + local_sum += __shfl_down_sync(0xffffffff, local_sum, offset); + } + if (lane == 0) { + warp_sum[warp] = local_sum; + } + __syncthreads(); + + if (threadIdx.x == 0) { + double total = 0.0; + for (int w = 0; w < LOGPROBS_BLOCK / WARP_SIZE; ++w) { + total += warp_sum[w]; + } + const float lse = row_max + (float)log(total); + out_lse[scored] = lse; + const float picked_val = __bfloat162float(x[picked[scored]]); + out_picked_lp[scored] = picked_val - lse; + } +} + +extern "C" int logprobs_lse_bf16_cuda(const __nv_bfloat16* logits, + const unsigned int* row_indices, + const unsigned int* picked, int num_rows, + int vocab_size, float* out_lse, + float* out_picked_lp, cudaStream_t stream) { + OPENINFER_FFI_GUARD_BEGIN + if (logits == nullptr || picked == nullptr || out_lse == nullptr || + out_picked_lp == nullptr) { + return static_cast(CUDA_ERROR_INVALID_VALUE); + } + if (num_rows <= 0 || vocab_size <= 0) { + return static_cast(CUDA_ERROR_INVALID_VALUE); + } + logprobs_lse_kernel<<>>( + logits, row_indices, picked, vocab_size, out_lse, out_picked_lp); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "logprobs_lse_bf16_cuda: launch failed: %s\n", + cudaGetErrorString(err)); + return static_cast(CUDA_ERROR_LAUNCH_FAILED); + } + return static_cast(CUDA_SUCCESS); + OPENINFER_FFI_GUARD_END(-1) +} + +// --------------------------------------------------------------------------- +// Row-index gather (sparse subset -> contiguous [num_rows, vocab] bf16) +// --------------------------------------------------------------------------- + +__global__ void logprobs_gather_rows_kernel(const uint4* __restrict__ logits, + const unsigned int* __restrict__ row_indices, + uint4* __restrict__ out, int vec4_per_row) { + const long long row = row_indices[blockIdx.x]; + const uint4* src = logits + row * (long long)vec4_per_row; + uint4* dst = out + (long long)blockIdx.x * vec4_per_row; + for (int i = threadIdx.x; i < vec4_per_row; i += LOGPROBS_BLOCK) { + dst[i] = src[i]; + } +} + +extern "C" int logprobs_gather_rows_bf16_cuda(const __nv_bfloat16* logits, + const unsigned int* row_indices, + __nv_bfloat16* out, int num_rows, + int vocab_size, cudaStream_t stream) { + OPENINFER_FFI_GUARD_BEGIN + if (logits == nullptr || row_indices == nullptr || out == nullptr) { + return static_cast(CUDA_ERROR_INVALID_VALUE); + } + if (num_rows <= 0 || vocab_size <= 0 || vocab_size % 8 != 0) { + // % 8: rows are copied as 16B uint4 chunks (8 bf16). Vocab sizes in-tree + // are all multiples of 8; anything else keeps the caller on the CPU path. + return static_cast(CUDA_ERROR_INVALID_VALUE); + } + logprobs_gather_rows_kernel<<>>( + reinterpret_cast(logits), row_indices, + reinterpret_cast(out), vocab_size / 8); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "logprobs_gather_rows_bf16_cuda: launch failed: %s\n", + cudaGetErrorString(err)); + return static_cast(CUDA_ERROR_LAUNCH_FAILED); + } + return static_cast(CUDA_SUCCESS); + OPENINFER_FFI_GUARD_END(-1) +} + +// --------------------------------------------------------------------------- +// Deterministic top-k (FilteredTopK, smallest-index tie-break) +// --------------------------------------------------------------------------- + +extern "C" int logprobs_topk_bf16_cuda(const __nv_bfloat16* logits, int num_rows, + int vocab_size, int top_k, + __nv_bfloat16* out_values, int* out_indices, + cudaStream_t stream) { + OPENINFER_FFI_GUARD_BEGIN + if (logits == nullptr || out_values == nullptr || out_indices == nullptr) { + return static_cast(CUDA_ERROR_INVALID_VALUE); + } + if (num_rows <= 0 || vocab_size <= 0 || top_k <= 0) { + return static_cast(CUDA_ERROR_INVALID_VALUE); + } + // FilteredTopK needs ~128KB dynamic smem (Hopper+). Report unsupported so + // the caller can fall back to the host path instead of crashing. + if (!flashinfer::sampling::CanImplementFilteredTopK()) { + return static_cast(CUDA_ERROR_NOT_SUPPORTED); + } + + auto* input = const_cast<__nv_bfloat16*>(logits); + cudaError_t err = flashinfer::sampling::FilteredTopK<__nv_bfloat16, int>( + input, out_indices, out_values, nullptr, static_cast(num_rows), + static_cast(top_k), static_cast(vocab_size), + /*deterministic=*/true, flashinfer::sampling::TopKTieBreak::Small, stream, + /*dsa_graph_safe=*/true); + if (err != cudaSuccess) { + fprintf(stderr, "logprobs_topk_bf16_cuda: FilteredTopK failed: %s\n", + cudaGetErrorString(err)); + return static_cast(CUDA_ERROR_LAUNCH_FAILED); + } + // FilteredTopK emits >pivot winners in atomicAdd race order; only selection + // (with smallest-index tie-break) is deterministic. Restore a canonical + // (value desc, index asc) order to match the CPU reference exactly: + // stable value-descending sort preceded by an index-ascending sort. + err = flashinfer::sampling::LaunchSortTopKByIndex< + flashinfer::sampling::FilteredTopKMode::Plain, __nv_bfloat16, int>( + out_indices, out_values, nullptr, 0, nullptr, nullptr, + static_cast(num_rows), static_cast(top_k), + static_cast(vocab_size), stream); + if (err != cudaSuccess) { + fprintf(stderr, "logprobs_topk_bf16_cuda: LaunchSortTopKByIndex failed: %s\n", + cudaGetErrorString(err)); + return static_cast(CUDA_ERROR_LAUNCH_FAILED); + } + err = flashinfer::sampling::StableSortTopKByValue<__nv_bfloat16, int>( + out_indices, out_values, static_cast(num_rows), + static_cast(top_k), static_cast(vocab_size), stream); + if (err != cudaSuccess) { + fprintf(stderr, + "logprobs_topk_bf16_cuda: StableSortTopKByValue failed: %s\n", + cudaGetErrorString(err)); + return static_cast(CUDA_ERROR_LAUNCH_FAILED); + } + return static_cast(CUDA_SUCCESS); + OPENINFER_FFI_GUARD_END(-1) +} diff --git a/openinfer-kernels/src/ffi/shared.rs b/openinfer-kernels/src/ffi/shared.rs index eb1caafe5..d31ca17c3 100644 --- a/openinfer-kernels/src/ffi/shared.rs +++ b/openinfer-kernels/src/ffi/shared.rs @@ -743,3 +743,47 @@ unsafe extern "C" { unsafe extern "C" { pub fn openinfer_kernels_last_error() -> *const std::os::raw::c_char; } + +// #719: batched device-side logprobs reduction over bf16 logits rows. +// Replaces per-row full-vocab D2H + host O(V) passes with O(rows * (k + 1)) +// D2H. Semantics match `openinfer_sample::token_logprob_from_row`: +// fp32 log-sum-exp, top-k ordered (value desc, index asc) with +// smallest-index tie-break selection. +unsafe extern "C" { + /// Per-row log-sum-exp + picked-token logprob. `row_indices == nullptr` + /// scores rows `0..num_rows` contiguously. + pub fn logprobs_lse_bf16_cuda( + logits: *const Half, + row_indices: *const u32, + picked: *const u32, + num_rows: i32, + vocab_size: i32, + out_lse: *mut f32, + out_picked_lp: *mut f32, + stream: CUstream, + ) -> i32; + + /// Gather `num_rows` logits rows by index into a contiguous + /// [num_rows, vocab_size] block for the FilteredTopK layout. + pub fn logprobs_gather_rows_bf16_cuda( + logits: *const Half, + row_indices: *const u32, + out: *mut Half, + num_rows: i32, + vocab_size: i32, + stream: CUstream, + ) -> i32; + + /// Deterministic top-k over a contiguous [num_rows, vocab_size] bf16 + /// block, output ordered (value desc, index asc). Returns + /// CUDA_ERROR_NOT_SUPPORTED when the GPU cannot run FilteredTopK. + pub fn logprobs_topk_bf16_cuda( + logits: *const Half, + num_rows: i32, + vocab_size: i32, + top_k: i32, + out_values: *mut Half, + out_indices: *mut i32, + stream: CUstream, + ) -> i32; +} diff --git a/openinfer-kernels/src/ops.rs b/openinfer-kernels/src/ops.rs index 06cd99d10..5067990ab 100644 --- a/openinfer-kernels/src/ops.rs +++ b/openinfer-kernels/src/ops.rs @@ -12,6 +12,7 @@ mod glm52; #[cfg(feature = "kimi-k2")] mod kimi_k2; mod linear; +mod logprobs; mod lora; mod norm; mod sampling; @@ -61,6 +62,9 @@ pub use linear::{ gemm_token_range_into_checked, gemv, linear, numeric_policy, per_token_served, pin_served, reset_numeric_policy_counters, set_numeric_policy, }; +pub use logprobs::{ + logprobs_gather_rows_bf16_into, logprobs_lse_bf16_into, logprobs_topk_bf16_into, +}; pub use lora::{ LoraDecodeGroupedProjection, lora_decode_fused_delta_group3_into, lora_decode_fused_delta_into, pack_lora_b_rows_into, diff --git a/openinfer-kernels/src/ops/logprobs.rs b/openinfer-kernels/src/ops/logprobs.rs new file mode 100644 index 000000000..bd565ad39 --- /dev/null +++ b/openinfer-kernels/src/ops/logprobs.rs @@ -0,0 +1,183 @@ +//! Batched GPU logprobs reduction (#719): log-sum-exp, picked-token logprob, +//! and deterministic top-k over bf16 logits rows, replacing the per-row +//! full-vocab D2H and host O(V) passes. + +use anyhow::{Result, ensure}; +use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut}; + +use crate::ffi; +use crate::tensor::{DeviceContext, HiddenStates}; + +/// `CUDA_ERROR_NOT_SUPPORTED` — FilteredTopK needs Hopper-class smem. +const CUDA_ERROR_NOT_SUPPORTED: i32 = 801; + +/// log-sum-exp + picked-token logprob for `num_rows` scored rows. +/// +/// `row_indices == None` scores rows `0..num_rows` contiguously; otherwise +/// row `i` of the output scores logits row `row_indices[i]`. `picked[r]` is +/// the sampled token of scored row `r`. +pub fn logprobs_lse_bf16_into( + ctx: &DeviceContext, + logits: &HiddenStates, + row_indices: Option<&CudaSlice>, + picked: &CudaSlice, + num_rows: usize, + out_lse: &mut CudaSlice, + out_picked_lp: &mut CudaSlice, +) -> Result<()> { + ensure!(num_rows > 0, "logprobs lse requires num_rows > 0"); + ensure!( + num_rows <= logits.seq_len || row_indices.is_some(), + "logprobs lse: num_rows {num_rows} exceeds logits seq_len {} without row indices", + logits.seq_len + ); + if let Some(indices) = row_indices { + ensure!( + indices.len() >= num_rows, + "logprobs lse row_indices too small: have {}, need {num_rows}", + indices.len() + ); + } + ensure!( + picked.len() >= num_rows, + "logprobs lse picked too small: have {}, need {num_rows}", + picked.len() + ); + ensure!( + out_lse.len() >= num_rows && out_picked_lp.len() >= num_rows, + "logprobs lse outputs too small: have {}/{}, need {num_rows}", + out_lse.len(), + out_picked_lp.len() + ); + + let (logits_ptr, _lg) = logits.data.device_ptr(&ctx.stream); + let (indices_ptr, _ig) = match row_indices { + Some(indices) => { + let (ptr, guard) = indices.device_ptr(&ctx.stream); + (ptr as *const u32, Some(guard)) + } + None => (std::ptr::null(), None), + }; + let (picked_ptr, _pg) = picked.device_ptr(&ctx.stream); + let (lse_ptr, _og) = out_lse.device_ptr_mut(&ctx.stream); + let (picked_lp_ptr, _pg2) = out_picked_lp.device_ptr_mut(&ctx.stream); + + let result = unsafe { + ffi::logprobs_lse_bf16_cuda( + logits_ptr as *const ffi::Half, + indices_ptr, + picked_ptr as *const u32, + num_rows as i32, + logits.hidden_dim as i32, + lse_ptr as *mut f32, + picked_lp_ptr as *mut f32, + ctx.stream.cu_stream(), + ) + }; + ensure!( + result == 0, + "logprobs lse launch failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + Ok(()) +} + +/// Gather `num_rows` logits rows by index into a contiguous +/// [num_rows, hidden_dim] bf16 buffer for the FilteredTopK layout. +pub fn logprobs_gather_rows_bf16_into( + ctx: &DeviceContext, + logits: &HiddenStates, + row_indices: &CudaSlice, + num_rows: usize, + out: &mut CudaSlice, +) -> Result<()> { + ensure!(num_rows > 0, "logprobs gather requires num_rows > 0"); + ensure!( + row_indices.len() >= num_rows, + "logprobs gather row_indices too small: have {}, need {num_rows}", + row_indices.len() + ); + ensure!( + out.len() >= num_rows * logits.hidden_dim, + "logprobs gather output too small: have {}, need {}", + out.len(), + num_rows * logits.hidden_dim + ); + + let (logits_ptr, _lg) = logits.data.device_ptr(&ctx.stream); + let (indices_ptr, _ig) = row_indices.device_ptr(&ctx.stream); + let (out_ptr, _og) = out.device_ptr_mut(&ctx.stream); + + let result = unsafe { + ffi::logprobs_gather_rows_bf16_cuda( + logits_ptr as *const ffi::Half, + indices_ptr as *const u32, + out_ptr as *mut ffi::Half, + num_rows as i32, + logits.hidden_dim as i32, + ctx.stream.cu_stream(), + ) + }; + ensure!( + result == 0, + "logprobs gather launch failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + Ok(()) +} + +/// Deterministic top-k over a contiguous [num_rows, hidden_dim] bf16 logits +/// block (FilteredTopK, smallest-index tie-break). +/// +/// Returns `Ok(false)` when the GPU cannot run FilteredTopK (pre-Hopper +/// smem), in which case the caller should fall back to the host path. +pub fn logprobs_topk_bf16_into( + ctx: &DeviceContext, + logits: &CudaSlice, + num_rows: usize, + hidden_dim: usize, + top_k: usize, + out_values: &mut CudaSlice, + out_indices: &mut CudaSlice, +) -> Result { + ensure!(num_rows > 0, "logprobs top-k requires num_rows > 0"); + ensure!(top_k > 0, "logprobs top-k requires top_k > 0"); + ensure!( + logits.len() >= num_rows * hidden_dim, + "logprobs top-k input too small: have {}, need {}", + logits.len(), + num_rows * hidden_dim + ); + ensure!( + out_values.len() >= num_rows * top_k && out_indices.len() >= num_rows * top_k, + "logprobs top-k outputs too small: have {}/{}, need {}", + out_values.len(), + out_indices.len(), + num_rows * top_k + ); + + let (logits_ptr, _lg) = logits.device_ptr(&ctx.stream); + let (values_ptr, _vg) = out_values.device_ptr_mut(&ctx.stream); + let (indices_ptr, _ig) = out_indices.device_ptr_mut(&ctx.stream); + + let result = unsafe { + ffi::logprobs_topk_bf16_cuda( + logits_ptr as *const ffi::Half, + num_rows as i32, + hidden_dim as i32, + top_k as i32, + values_ptr as *mut ffi::Half, + indices_ptr as *mut i32, + ctx.stream.cu_stream(), + ) + }; + if result == CUDA_ERROR_NOT_SUPPORTED { + return Ok(false); + } + ensure!( + result == 0, + "logprobs top-k launch failed with error {result}{}", + crate::ops::ffi_exception_message(result) + ); + Ok(true) +} diff --git a/openinfer-kimi-k2/src/batch_decode_trace.rs b/openinfer-kimi-k2/src/batch_decode_trace.rs index d8f0ac57c..52f016eca 100644 --- a/openinfer-kimi-k2/src/batch_decode_trace.rs +++ b/openinfer-kimi-k2/src/batch_decode_trace.rs @@ -133,8 +133,8 @@ pub fn trace_runtime_decode_kernel_calls( max_tokens: 2, lora_adapter: None, token_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, })?; receivers.push(std::thread::spawn(move || -> Result<()> { while let Some((_, event)) = token_rx.blocking_recv() { diff --git a/openinfer-kimi-k2/src/runner/executor/tp8_dp1.rs b/openinfer-kimi-k2/src/runner/executor/tp8_dp1.rs index 15473198e..7aa09f16b 100644 --- a/openinfer-kimi-k2/src/runner/executor/tp8_dp1.rs +++ b/openinfer-kimi-k2/src/runner/executor/tp8_dp1.rs @@ -36,7 +36,7 @@ impl ForwardExecutor for Tp8Dp1ForwardExecutor { if self.workers.is_empty() { bail!("Kimi TP8 executor has no rank workers"); } - ensure_no_logprobs_tp8(row.logprobs > 0)?; + ensure_no_logprobs_tp8(row.logprobs.is_some())?; ensure_greedy_tp8(&row.sampling)?; // Every rank replicates the KV pool with identical geometry, so the // same page assignment is broadcast to all eight workers. @@ -104,7 +104,7 @@ impl ForwardExecutor for Tp8Dp1ForwardExecutor { slots.len() ); } - ensure_no_logprobs_tp8(rows.iter().any(|r| r.logprobs > 0))?; + ensure_no_logprobs_tp8(rows.iter().any(|r| r.logprobs.is_some()))?; for row in rows { ensure_greedy_tp8(&row.sampling)?; } diff --git a/openinfer-kimi-k2/src/runner/scheduler.rs b/openinfer-kimi-k2/src/runner/scheduler.rs index b2198e26f..bba270158 100644 --- a/openinfer-kimi-k2/src/runner/scheduler.rs +++ b/openinfer-kimi-k2/src/runner/scheduler.rs @@ -765,8 +765,8 @@ mod tests { max_tokens, lora_adapter: None, token_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, }; (req, token_rx) } @@ -852,7 +852,7 @@ mod tests { let mut scheduler = test_scheduler(&calls, test_pool()); let (mut echo_req, mut token_rx) = request_with_channel(vec![11, 22], 4); - echo_req.echo = true; + echo_req.prompt_logprobs = Some(1); scheduler.handle_request_batch(vec![echo_req]); @@ -867,8 +867,8 @@ mod tests { panic!("expected Rejected event"); }; assert!( - message.contains("echo"), - "rejection names the unsupported field: {message}" + message.contains("prompt logprobs"), + "rejection names the unsupported feature: {message}" ); } diff --git a/openinfer-kimi-k2/src/runner/scheduler/dp.rs b/openinfer-kimi-k2/src/runner/scheduler/dp.rs index 5b6265a2b..f409388a2 100644 --- a/openinfer-kimi-k2/src/runner/scheduler/dp.rs +++ b/openinfer-kimi-k2/src/runner/scheduler/dp.rs @@ -1102,8 +1102,8 @@ mod tests { max_tokens, lora_adapter: None, token_tx, - logprobs: 0, - echo: false, + logprobs: None, + prompt_logprobs: None, } } @@ -1213,7 +1213,7 @@ mod tests { assert_eq!(positions, vec![6]); assert_eq!(slots, vec![MAX_BATCH_PER_DP - 1]); assert_eq!(rows.len(), 1); - assert_eq!(rows[0].logprobs, 0); + assert_eq!(rows[0].logprobs, None); // 6 KV tokens + this step's append fit one 16-token page. assert_eq!(kv_pages.rows(), 1); assert_eq!(kv_pages.row(0).expect("row 0").len(), 1); @@ -1236,7 +1236,7 @@ mod tests { append_position: 5, slot: 3, options: KimiRowOptions { - logprobs: 4, + logprobs: Some(4), sampling: SamplingParams::default(), }, pages: vec![5], @@ -1265,9 +1265,9 @@ mod tests { assert_eq!(token_ids, vec![11, 99]); assert_eq!(positions, vec![5, 0]); assert_eq!(slots, vec![3, 7]); - assert_eq!(rows[0].logprobs, 4); + assert_eq!(rows[0].logprobs, Some(4)); assert!(rows[0].sampling.is_greedy()); - assert_eq!(rows[1].logprobs, 0); + assert_eq!(rows[1].logprobs, None); assert!(!rows[1].sampling.is_greedy()); assert!((rows[1].sampling.temperature - 0.8).abs() < f32::EPSILON); assert_eq!(kv_pages.rows(), 2); @@ -1357,7 +1357,7 @@ mod tests { max_tokens: 16, last_token: 7, options: KimiRowOptions { - logprobs: 0, + logprobs: None, sampling: SamplingParams { ignore_eos: true, ..SamplingParams::default() diff --git a/openinfer-kimi-k2/src/runner/scheduler/lifecycle.rs b/openinfer-kimi-k2/src/runner/scheduler/lifecycle.rs index 734ea3925..daa0f1bc2 100644 --- a/openinfer-kimi-k2/src/runner/scheduler/lifecycle.rs +++ b/openinfer-kimi-k2/src/runner/scheduler/lifecycle.rs @@ -114,14 +114,14 @@ fn unschedulable_verdict(req: &GenerateRequest) -> Option "Kimi-K2 forward requires at least one prompt token".to_string(), )); } - // Honor-or-reject (#236): prompt echo needs per-position prompt logprobs, - // which the prefill path does not compute (lm_head runs on the last - // position only). Reject instead of silently returning a response with - // the echo stripped. - if req.echo { + // Honor-or-reject (#236): prompt logprobs need per-position prompt + // logits, which the prefill path does not compute (lm_head runs on the + // last position only). Reject instead of silently returning a response + // with the prompt logprobs stripped. + if req.prompt_logprobs.is_some() { return Some(UnschedulableVerdict::Reject( - "echo is not supported on the Kimi-K2 serving path: prompt \ - logprobs are not computed; set echo=false" + "prompt logprobs are not supported on the Kimi-K2 serving path: \ + all-position prompt logits are not computed" .to_string(), )); } diff --git a/openinfer-kimi-k2/src/runner/worker.rs b/openinfer-kimi-k2/src/runner/worker.rs index 51fc2352e..9cdab98b0 100644 --- a/openinfer-kimi-k2/src/runner/worker.rs +++ b/openinfer-kimi-k2/src/runner/worker.rs @@ -152,7 +152,7 @@ pub(super) struct KimiRankWeightLoadReport { /// and how many logprobs to report for it. #[derive(Clone, Copy, Debug, Default)] pub(crate) struct KimiRowOptions { - pub(crate) logprobs: usize, + pub(crate) logprobs: Option, pub(crate) sampling: openinfer_core::sampler::SamplingParams, } @@ -170,7 +170,7 @@ pub(super) struct KimiOneTokenForwardReport { pub moe_layers_executed: usize, /// Exact log-softmax of the picked token plus the top-K, computed on the /// host from the full-vocab logits row. `Some` only when the request - /// asked for logprobs (`GenerateRequest::logprobs > 0`); the serving + /// asked for logprobs (`GenerateRequest::logprobs` is `Some`); the serving /// path never pays for it. pub logprob: Option, } diff --git a/openinfer-kimi-k2/src/runner/worker/state.rs b/openinfer-kimi-k2/src/runner/worker/state.rs index e832392e3..44b80d9c8 100644 --- a/openinfer-kimi-k2/src/runner/worker/state.rs +++ b/openinfer-kimi-k2/src/runner/worker/state.rs @@ -394,7 +394,7 @@ impl KimiRankThreadState { picks[sampling_row.row].0 = *token; } - let host_logits = if rows.iter().any(|r| r.logprobs > 0) { + let host_logits = if rows.iter().any(|r| r.logprobs.is_some()) { ensure!( cache.vocab_start == 0 && cache.vocab_rows == KIMI_K2_VOCAB, "Kimi logprobs require an unsharded vocab (TP1); a vocab shard's \ @@ -412,11 +412,13 @@ impl KimiRankThreadState { let mut reports = Vec::with_capacity(active_len); for (row, (local_next, local_top_logit_f32)) in picks.into_iter().enumerate() { let logprob = match &host_logits { - Some(host) if rows[row].logprobs > 0 => openinfer_sample::token_logprob_from_row( - &host[row * cache.vocab_rows..(row + 1) * cache.vocab_rows], - local_next, - rows[row].logprobs, - ), + Some(host) if rows[row].logprobs.is_some() => { + openinfer_sample::token_logprob_from_row( + &host[row * cache.vocab_rows..(row + 1) * cache.vocab_rows], + local_next, + rows[row].logprobs.unwrap_or(0), + ) + } _ => None, }; reports.push(KimiOneTokenForwardReport { @@ -649,7 +651,7 @@ impl KimiRankThreadState { .with_context(|| format!("Kimi rank {rank} prefill sampling"))?; local_next = sampled[0]; } - let logprob = if row.logprobs > 0 { + let logprob = if row.logprobs.is_some() { ensure!( cache.vocab_start == 0 && cache.vocab_rows == KIMI_K2_VOCAB, "Kimi logprobs require an unsharded vocab (TP1); a vocab \ @@ -659,7 +661,7 @@ impl KimiRankThreadState { .stream .clone_dtoh(&logits.data) .with_context(|| format!("Kimi rank {rank} D2H prefill logits"))?; - openinfer_sample::token_logprob_from_row(&host, local_next, row.logprobs) + openinfer_sample::token_logprob_from_row(&host, local_next, row.logprobs.unwrap_or(0)) } else { None }; diff --git a/openinfer-kimi-k2/tests/vllm_golden_gate.rs b/openinfer-kimi-k2/tests/vllm_golden_gate.rs index 4af9dad25..00022c213 100644 --- a/openinfer-kimi-k2/tests/vllm_golden_gate.rs +++ b/openinfer-kimi-k2/tests/vllm_golden_gate.rs @@ -342,8 +342,8 @@ fn submit( max_tokens, lora_adapter: None, token_tx: tx, - logprobs, - echo: false, + logprobs: (logprobs > 0).then_some(logprobs), + prompt_logprobs: None, }) .expect("submit to kimi engine"); PendingRequest { label, rx } diff --git a/openinfer-qwen3/src/bin/qwen3_decode_context.rs b/openinfer-qwen3/src/bin/qwen3_decode_context.rs index fa4990f70..5365d1322 100644 --- a/openinfer-qwen3/src/bin/qwen3_decode_context.rs +++ b/openinfer-qwen3/src/bin/qwen3_decode_context.rs @@ -185,8 +185,8 @@ fn prefill_one( prompt.to_vec(), 4096, params, - 0, - false, + None, + None, )]; let result = executor.execute_prefill(PrefillPlan { requests: &requests, @@ -203,7 +203,7 @@ fn decode_one_step( params: SamplingParams, rng: &mut StdRng, ) -> Result { - let requests = [DecodeStepItem::new(request_id, *token, params, 0)]; + let requests = [DecodeStepItem::new(request_id, *token, params, None)]; let start = Instant::now(); let result = executor.execute_decode(DecodePlan { requests: &requests, @@ -242,7 +242,7 @@ fn decode_batch_step( ) -> Result { let requests: Vec = batch .iter() - .map(|&(request_id, token)| DecodeStepItem::new(request_id, token, params, 0)) + .map(|&(request_id, token)| DecodeStepItem::new(request_id, token, params, None)) .collect(); let start = Instant::now(); let result = executor.execute_decode(DecodePlan { diff --git a/openinfer-qwen3/src/executor.rs b/openinfer-qwen3/src/executor.rs index c6cceddf9..40a8ed6d2 100644 --- a/openinfer-qwen3/src/executor.rs +++ b/openinfer-qwen3/src/executor.rs @@ -4,6 +4,8 @@ use std::thread; use anyhow::{Context, Result, ensure}; use crossbeam_channel as channel; +use cudarc::driver::CudaSlice; +use half::bf16; use crate::batch_decode::DecodeGraphUse; use crate::batch_decode_buffers::{BATCH_BUCKETS, BatchDecodeBuffers}; @@ -18,7 +20,7 @@ use openinfer_core::engine::{ use openinfer_core::kv_pool::KvLayout; use openinfer_core::ops; use openinfer_core::sampler::SamplingParams; -use openinfer_core::tensor::{DeviceContext, DeviceVec, HiddenStates}; +use openinfer_core::tensor::{DeviceContext, HiddenStates}; use openinfer_kv_cache::{ KvBlockGuard, KvBuffer, KvCacheEvent, KvCacheManager, KvView, LoadReservation, PrefixProbe, RegisteredBlock, @@ -61,8 +63,11 @@ pub struct PrefillStepItem { pub(crate) prompt_tokens: Vec, pub(crate) max_output_tokens: usize, pub(crate) params: SamplingParams, - pub(crate) logprobs: usize, - pub(crate) echo: bool, + /// Completion logprob top-k count (`None` = disabled, `Some(0)` = scored + /// token only). + pub(crate) logprobs: Option, + /// Prompt logprob top-k count; `Some(_)` needs all-position logits. + pub(crate) prompt_logprobs: Option, pub(crate) lora_adapter: Option, /// Leading prompt tokens whose KV came from the prefix cache. /// Set by the executor after matching; the forward pass only computes @@ -85,8 +90,8 @@ impl PrefillStepItem { prompt_tokens: Vec, max_output_tokens: usize, params: SamplingParams, - logprobs: usize, - echo: bool, + logprobs: Option, + prompt_logprobs: Option, ) -> Self { let chunk_tokens = prompt_tokens.len(); Self { @@ -95,7 +100,7 @@ impl PrefillStepItem { max_output_tokens, params, logprobs, - echo, + prompt_logprobs, lora_adapter: None, cached_tokens: 0, chunk_budget: usize::MAX, @@ -135,7 +140,7 @@ pub struct DecodeStepItem { pub(crate) request_id: RequestId, pub(crate) token_id: u32, pub(crate) params: SamplingParams, - pub(crate) logprobs: usize, + pub(crate) logprobs: Option, pub(crate) lora_adapter: Option, } @@ -144,7 +149,7 @@ impl DecodeStepItem { request_id: RequestId, token_id: u32, params: SamplingParams, - logprobs: usize, + logprobs: Option, ) -> Self { Self { request_id, @@ -162,6 +167,260 @@ impl DecodeStepItem { } } +/// Largest top-k the device logprobs path serves; FlashInfer FilteredTopK +/// caps at 2048 (shared-memory index buffer). Larger requests take the host +/// path per row. +const LOGPROBS_DEVICE_TOPK_MAX: usize = 2048; + +/// One scored row of a batched logprobs call: logits arena row, the picked +/// token, and the per-row top-k budget. +struct LogprobsJob { + row: u32, + picked: u32, + top_k: usize, +} + +/// Device buffers for the batched logprobs reduction (#719), grown on demand +/// and reused across steps so the decode hot path stays allocation-free. +struct LogprobsScratch { + row_indices: CudaSlice, + picked: CudaSlice, + lse: CudaSlice, + picked_lp: CudaSlice, + /// [max_rows x vocab] gathered block feeding FilteredTopK (k > 0 only). + gathered: CudaSlice, + topk_values: CudaSlice, + topk_indices: CudaSlice, + max_rows: usize, + max_k: usize, +} + +impl LogprobsScratch { + fn new(ctx: &DeviceContext) -> Result { + Ok(Self { + row_indices: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + picked: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + lse: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + picked_lp: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + gathered: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + topk_values: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + topk_indices: ctx + .stream + .alloc_zeros(1) + .map_err(|e| anyhow::anyhow!("LogprobsScratch alloc failed: {e}"))?, + max_rows: 0, + max_k: 0, + }) + } + + fn ensure_capacity( + &mut self, + ctx: &DeviceContext, + vocab: usize, + rows: usize, + k: usize, + ) -> Result<()> { + let rows = rows.max(1); + let grew_rows = rows > self.max_rows; + if grew_rows { + self.row_indices = ctx + .stream + .alloc_zeros(rows) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + self.picked = ctx + .stream + .alloc_zeros(rows) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + self.lse = ctx + .stream + .alloc_zeros(rows) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + self.picked_lp = ctx + .stream + .alloc_zeros(rows) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + self.gathered = ctx + .stream + .alloc_zeros(rows * vocab) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + self.max_rows = rows; + } + let k = k.max(1); + if grew_rows || k > self.max_k { + self.max_k = self.max_k.max(k); + self.topk_values = ctx + .stream + .alloc_zeros(self.max_rows * self.max_k) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + self.topk_indices = ctx + .stream + .alloc_zeros(self.max_rows * self.max_k) + .map_err(|e| anyhow::anyhow!("LogprobsScratch grow failed: {e}"))?; + } + Ok(()) + } +} + +/// Per-row host logprobs (full-vocab D2H + CPU passes) — the pre-#719 path, +/// kept as the fallback for device-unsupported cases and as the golden +/// reference in tests. +fn logprobs_host_row( + ctx: &DeviceContext, + logits: &HiddenStates, + job: &LogprobsJob, +) -> Result> { + let row = ops::extract_vec(ctx, logits, job.row as usize)?; + let row_f32 = row.to_host(ctx)?; + Ok(openinfer_sample::token_logprob_from_row( + &row_f32, job.picked, job.top_k, + )) +} + +/// Batched device-side logprobs over `logits` rows (#719): one LSE pass for +/// every job, one FilteredTopK over the gathered rows when any job wants +/// alternatives, then a single D2H of O(rows * (k + 1)) and one stream sync — +/// replacing the per-row full-vocab D2H + host passes. Per-row fallbacks to +/// the host path cover top_k above the device cap and GPUs without the +/// FilteredTopK shared-memory budget. +fn extract_logprobs_batch( + ctx: &DeviceContext, + scratch: &mut LogprobsScratch, + logits: &HiddenStates, + jobs: &[LogprobsJob], +) -> Result>> { + let n = jobs.len(); + let mut results: Vec> = (0..n).map(|_| None).collect(); + if n == 0 { + return Ok(results); + } + let vocab = logits.hidden_dim; + + // Host path for jobs beyond the device top-k cap. + let mut device_jobs: Vec = Vec::with_capacity(n); + for (i, job) in jobs.iter().enumerate() { + if job.top_k > LOGPROBS_DEVICE_TOPK_MAX { + results[i] = logprobs_host_row(ctx, logits, job)?; + } else { + device_jobs.push(i); + } + } + if device_jobs.is_empty() { + return Ok(results); + } + + let m = device_jobs.len(); + let k_dev = device_jobs + .iter() + .map(|&i| jobs[i].top_k) + .max() + .unwrap_or(0) + .min(vocab); + scratch.ensure_capacity(ctx, vocab, m, k_dev)?; + + let rows: Vec = device_jobs.iter().map(|&i| jobs[i].row).collect(); + let picked: Vec = device_jobs.iter().map(|&i| jobs[i].picked).collect(); + ctx.stream + .memcpy_htod(&rows, &mut scratch.row_indices) + .map_err(|e| anyhow::anyhow!("logprobs row indices H2D failed: {e}"))?; + ctx.stream + .memcpy_htod(&picked, &mut scratch.picked) + .map_err(|e| anyhow::anyhow!("logprobs picked H2D failed: {e}"))?; + + openinfer_kernels::ops::logprobs_lse_bf16_into( + ctx, + logits, + Some(&scratch.row_indices), + &scratch.picked, + m, + &mut scratch.lse, + &mut scratch.picked_lp, + )?; + + let mut topk_ok = false; + if k_dev > 0 { + openinfer_kernels::ops::logprobs_gather_rows_bf16_into( + ctx, + logits, + &scratch.row_indices, + m, + &mut scratch.gathered, + )?; + topk_ok = openinfer_kernels::ops::logprobs_topk_bf16_into( + ctx, + &scratch.gathered, + m, + vocab, + k_dev, + &mut scratch.topk_values, + &mut scratch.topk_indices, + )?; + } + + let lse_host = ctx + .stream + .clone_dtoh(&scratch.lse) + .map_err(|e| anyhow::anyhow!("logprobs lse D2H failed: {e}"))?; + let picked_lp_host = ctx + .stream + .clone_dtoh(&scratch.picked_lp) + .map_err(|e| anyhow::anyhow!("logprobs picked D2H failed: {e}"))?; + let topk_host = if topk_ok { + Some(( + ctx.stream + .clone_dtoh(&scratch.topk_values) + .map_err(|e| anyhow::anyhow!("logprobs top-k values D2H failed: {e}"))?, + ctx.stream + .clone_dtoh(&scratch.topk_indices) + .map_err(|e| anyhow::anyhow!("logprobs top-k indices D2H failed: {e}"))?, + )) + } else { + None + }; + ctx.sync()?; + + for (out_idx, &job_idx) in device_jobs.iter().enumerate() { + let job = &jobs[job_idx]; + if job.top_k > 0 && !topk_ok { + // GPU cannot run FilteredTopK; keep this row on the host path. + results[job_idx] = logprobs_host_row(ctx, logits, job)?; + continue; + } + let lse = lse_host[out_idx]; + let top_logprobs = match &topk_host { + Some((values, indices)) if job.top_k > 0 => { + let base = out_idx * k_dev; + (0..job.top_k.min(vocab)) + .map(|t| (indices[base + t] as u32, values[base + t].to_f32() - lse)) + .collect() + } + _ => Vec::new(), + }; + results[job_idx] = Some(TokenLogprob { + logprob: picked_lp_host[out_idx], + top_logprobs, + }); + } + Ok(results) +} + fn build_prefill_request_results( lane: &mut LocalQwen3Lane, requests: &[PrefillStepItem], @@ -170,31 +429,67 @@ fn build_prefill_request_results( all_position_logits: Option<&HiddenStates>, compute_prompt_logprobs: bool, ) -> Result> { + // Pass 1: collect device-batchable logprobs jobs (#719) — first-token + // rows from `logits`, prompt positions from `all_position_logits`. + let mut first_jobs = Vec::new(); + let mut prompt_jobs = Vec::new(); let mut token_offset = 0usize; + for (i, req) in requests.iter().enumerate() { + if req.is_final_chunk() && req.logprobs.is_some() { + first_jobs.push(LogprobsJob { + row: i as u32, + picked: tokens[i], + top_k: req.logprobs.unwrap_or(0), + }); + } + if compute_prompt_logprobs && all_position_logits.is_some() { + if let Some(prompt_top_k) = req.prompt_logprobs { + for j in 1..req.prompt_tokens.len() { + prompt_jobs.push(LogprobsJob { + row: (token_offset + j - 1) as u32, + picked: req.prompt_tokens[j], + top_k: prompt_top_k, + }); + } + } + } + token_offset += req.chunk_tokens; + } + let mut first_lp = extract_logprobs_batch( + lane.model.device_ctx(), + &mut lane.logprobs_scratch, + logits, + &first_jobs, + )? + .into_iter(); + let mut prompt_lp = match all_position_logits { + Some(all_logits) if !prompt_jobs.is_empty() => extract_logprobs_batch( + lane.model.device_ctx(), + &mut lane.logprobs_scratch, + all_logits, + &prompt_jobs, + )? + .into_iter(), + _ => Vec::new().into_iter(), + }; + + // Pass 2: assemble per-request outputs in the original wire shape. let mut outputs = Vec::with_capacity(requests.len()); for (i, req) in requests.iter().enumerate() { let completed = req.is_final_chunk(); let first_token = tokens[i]; - let first_token_logprob = if completed && req.logprobs > 0 { - let logits_i = ops::extract_vec(lane.model.device_ctx(), logits, i)?; - Some(lane.extract_logprobs(&logits_i, first_token, req.logprobs)?) + let first_token_logprob = if completed && req.logprobs.is_some() { + first_lp.next().flatten() } else { None }; - let prompt_logprobs = if req.echo { + let prompt_logprobs = if req.prompt_logprobs.is_some() { if compute_prompt_logprobs { let mut echo_logprobs = Vec::with_capacity(req.prompt_tokens.len()); echo_logprobs.push(None); - if let Some(all_logits) = all_position_logits { - for j in 1..req.prompt_tokens.len() { - let prev_pos = token_offset + j - 1; - let target_token = req.prompt_tokens[j]; - echo_logprobs.push(lane.extract_prompt_logprobs( - all_logits, - prev_pos, - target_token, - req.logprobs, - )); + if all_position_logits.is_some() { + for _ in 1..req.prompt_tokens.len() { + echo_logprobs.push(prompt_lp.next().flatten()); } } else { for _ in 1..req.prompt_tokens.len() { @@ -208,7 +503,6 @@ fn build_prefill_request_results( } else { None }; - token_offset += req.chunk_tokens; outputs.push(PrefillRequestResult { request_id: req.request_id, first_token, @@ -229,12 +523,28 @@ fn build_decode_request_results( row_offset: usize, tokens: &[u32], ) -> Result> { + let jobs: Vec = requests + .iter() + .enumerate() + .filter(|(_, req)| req.logprobs.is_some()) + .map(|(i, req)| LogprobsJob { + row: (row_offset + i) as u32, + picked: tokens[row_offset + i], + top_k: req.logprobs.unwrap_or(0), + }) + .collect(); + let mut lp = extract_logprobs_batch( + lane.model.device_ctx(), + &mut lane.logprobs_scratch, + logits, + &jobs, + )? + .into_iter(); let mut outputs = Vec::with_capacity(requests.len()); for (i, req) in requests.iter().enumerate() { let token = tokens[row_offset + i]; - let logprob = if req.logprobs > 0 { - let logits_i = ops::extract_vec(lane.model.device_ctx(), logits, row_offset + i)?; - Some(lane.extract_logprobs(&logits_i, token, req.logprobs)?) + let logprob = if req.logprobs.is_some() { + lp.next().flatten() } else { None }; @@ -264,12 +574,28 @@ fn build_batch_decode_request_results( &mut lane.sample_scratch, )?; + let jobs: Vec = requests + .iter() + .enumerate() + .filter(|(_, req)| req.logprobs.is_some()) + .map(|(i, req)| LogprobsJob { + row: i as u32, + picked: tokens[i], + top_k: req.logprobs.unwrap_or(0), + }) + .collect(); + let mut lp = extract_logprobs_batch( + lane.model.device_ctx(), + &mut lane.logprobs_scratch, + &lane.bufs.logits, + &jobs, + )? + .into_iter(); let mut outputs = Vec::with_capacity(requests.len()); for (i, req) in requests.iter().enumerate() { let token = tokens[i]; - let logprob = if req.logprobs > 0 { - let logits_i = ops::extract_vec(lane.model.device_ctx(), &lane.bufs.logits, i)?; - Some(lane.extract_logprobs(&logits_i, token, req.logprobs)?) + let logprob = if req.logprobs.is_some() { + lp.next().flatten() } else { None }; @@ -1745,9 +2071,10 @@ impl Qwen3Executor { req.max_output_tokens, req.lora_adapter.as_deref(), ); - // Echo needs logits for every prompt position; cached positions - // are never forwarded, so echo requests prefill from scratch. - if self.prefix_cache_enabled && !req.echo { + // Prompt-logprobs requests need logits for every prompt position; + // cached positions are never forwarded, so they prefill from + // scratch. + if self.prefix_cache_enabled && req.prompt_logprobs.is_none() { req.cached_tokens = rkv.match_and_add_prefix(self.kv_mgr.pool())?; } self.request_kvs.insert(req.request_id, rkv); @@ -1762,9 +2089,10 @@ impl Qwen3Executor { .expect("inserted above"); req.chunk_start = rkv.kv_position(); let remaining = req.prompt_tokens.len() - req.chunk_start; - // Echo must produce all-position logits in a single forward, so it is - // exempt from chunking (the scheduler never splits echo requests). - req.chunk_tokens = if req.echo { + // Prompt-logprobs requests must produce all-position logits in a + // single forward, so they are exempt from chunking (the scheduler + // never splits them). + req.chunk_tokens = if req.prompt_logprobs.is_some() { remaining } else { remaining.min(req.chunk_budget) @@ -3041,6 +3369,9 @@ struct LocalQwen3Lane { kv_buffer: KvBuffer, layout: KvLayout, sample_scratch: openinfer_sample::SampleScratch, + /// Batched logprobs scratch (#719), grown on demand; keeps the decode + /// hot path allocation-free after the first logprobs step. + logprobs_scratch: LogprobsScratch, /// Request-local decode steps handed to `select_batch`, reused across /// steps to keep the sampling hot path allocation-free. All zeros until /// the scheduler wires generated counts through (sampling-parity 1b). @@ -3114,12 +3445,14 @@ impl LocalQwen3Lane { model.config().vocab_size, max_bucket, )?; + let logprobs_scratch = LogprobsScratch::new(model.device_ctx())?; Ok(Self { model, kv_buffer, layout, bufs, sample_scratch, + logprobs_scratch, steps_buf: Vec::new(), max_prefill_tokens, inflight_prefill: None, @@ -3302,32 +3635,6 @@ impl LocalQwen3Lane { ) } - fn extract_logprobs( - &self, - logits: &DeviceVec, - sampled_token: u32, - top_k: usize, - ) -> Result { - let logits_f32 = logits.to_host(self.model.device_ctx())?; - openinfer_sample::token_logprob_from_row(&logits_f32, sampled_token, top_k) - .ok_or_else(|| anyhow::anyhow!("logprobs computation failed")) - } - - fn extract_prompt_logprobs( - &self, - all_logits: &HiddenStates, - prev_pos: usize, - target_token: u32, - top_k: usize, - ) -> Option { - openinfer_core::ops::extract_vec(self.model.device_ctx(), all_logits, prev_pos) - .ok() - .and_then(|logits_vec| { - let logits_f32 = logits_vec.to_host(self.model.device_ctx()).ok()?; - openinfer_sample::token_logprob_from_row(&logits_f32, target_token, top_k) - }) - } - fn execute_prefill( &mut self, prompts: &[&[u32]], diff --git a/openinfer-qwen3/src/executor/dflash_prefill.rs b/openinfer-qwen3/src/executor/dflash_prefill.rs index be8440d96..1242d70e9 100644 --- a/openinfer-qwen3/src/executor/dflash_prefill.rs +++ b/openinfer-qwen3/src/executor/dflash_prefill.rs @@ -1,17 +1,20 @@ //! DFlash prefill-capture eligibility predicates. //! //! A request can seed the DFlash draft only if its prefill produces clean target -//! hidden states: no LoRA, no prefix-cache hit, no echo, no logprobs. Sampling -//! params are irrelevant here — the prompt's hidden states are -//! sampling-independent, and sampled-verify (#512) speculates the full -//! sampling surface, so capture is as valid for a sampled request as for a -//! greedy one. +//! hidden states: no LoRA, no prefix-cache hit, no prompt logprobs, no +//! completion logprobs. Sampling params are irrelevant here — the prompt's +//! hidden states are sampling-independent, and sampled-verify (#512) speculates +//! the full sampling surface, so capture is as valid for a sampled request as +//! for a greedy one. use super::{PrefillStepItem, RequestId}; /// Whether a prefill request is eligible to capture DFlash target context. pub(super) fn dflash_prefill_supported(req: &PrefillStepItem) -> bool { - req.lora_adapter.is_none() && req.cached_tokens == 0 && req.logprobs == 0 && !req.echo + req.lora_adapter.is_none() + && req.cached_tokens == 0 + && req.logprobs.is_none() + && req.prompt_logprobs.is_none() } /// Eligible AND continuous: either the first chunk, or a later chunk whose diff --git a/openinfer-qwen3/src/scheduler.rs b/openinfer-qwen3/src/scheduler.rs index 045027bec..ad47c7f5a 100644 --- a/openinfer-qwen3/src/scheduler.rs +++ b/openinfer-qwen3/src/scheduler.rs @@ -49,8 +49,9 @@ pub(super) struct ActiveRequestState { pub(super) max_tokens: usize, pub(super) prompt_len: usize, pub(super) params: SamplingParams, - /// Number of top logprobs to return (0 = disabled). - pub(super) logprobs: usize, + /// Completion logprob top-k count (`None` = disabled, `Some(0)` = scored + /// token only). + pub(super) logprobs: Option, } #[derive(Clone)] @@ -61,8 +62,8 @@ pub(super) struct PendingRequest { pub(super) params: SamplingParams, pub(super) max_tokens: usize, pub(super) token_tx: TokenSink, - pub(super) logprobs: usize, - pub(super) echo: bool, + pub(super) logprobs: Option, + pub(super) prompt_logprobs: Option, pub(super) queued_at_unix_s: Option, /// Whether this request has already been offered to async KV prefetch. /// Offered at most once; a no-hit offer leaves the request in the normal @@ -90,7 +91,7 @@ impl PendingRequest { max_tokens: req.max_tokens, token_tx: req.token_tx, logprobs: req.logprobs, - echo: req.echo, + prompt_logprobs: req.prompt_logprobs, queued_at_unix_s: req.queued_at_unix_s, prefetch_offered: false, prefill_pos: 0, @@ -102,12 +103,18 @@ impl PendingRequest { fn remaining_prompt_tokens(&self) -> usize { self.prompt_tokens.len() - self.prefill_pos } + + /// Prompt-logprobs ("echo") requests need logits for every prompt + /// position in one forward: no prefix-cache reuse, no chunked prefill. + pub(super) fn wants_prompt_logprobs(&self) -> bool { + self.prompt_logprobs.is_some() + } } /// Pull the next prefill step set off the front of `prefilling`, capping the /// step's total forwarded tokens at `max_prefill_tokens`. Each taken request -/// gets its per-step chunk recorded in `step_chunk`. Echo requests need -/// logits for every prompt position in one forward, so they only run when +/// gets its per-step chunk recorded in `step_chunk`. Prompt-logprobs requests +/// need logits for every prompt position in one forward, so they only run when /// their whole remainder fits the profiled prefill bound. Under request-local /// chunking, a request takes `min(remaining, max_prefill_tokens)` whole or skips /// the step, so its chunk boundaries depend only on its own length and are @@ -122,7 +129,7 @@ fn take_prefill_chunks( let mut i = 0; while i < prefilling.len() && budget > 0 { let remaining = prefilling[i].remaining_prompt_tokens(); - let chunk = if prefilling[i].echo { + let chunk = if prefilling[i].wants_prompt_logprobs() { if remaining > budget { i += 1; continue; @@ -386,12 +393,12 @@ fn reclaim_ready_prefetch( /// request that doesn't start a load (pure GPU hit, miss, or block pressure) /// stays in `deferred`, flagged so it isn't re-probed next tick. /// -/// Echo requests are never offered: their prefill forwards the whole prompt to -/// recover prompt logprobs and so skips `match_and_add_prefix` (see -/// `execute_prefill`). Prefetched blocks would never be matched/reused — they -/// would only park restored KV that admission credits but prefill can't spend, -/// starving the request under tight budgets. Leaving `prefetch_offered` unset -/// for echo is harmless: the `!req.echo` guard keeps them from being probed. +/// Prompt-logprobs requests are never offered: their prefill forwards the +/// whole prompt to recover prompt logprobs and so skips `match_and_add_prefix` +/// (see `execute_prefill`). Prefetched blocks would never be matched/reused — +/// they would only park restored KV that admission credits but prefill can't +/// spend, starving the request under tight budgets. Leaving `prefetch_offered` +/// unset for them is harmless: the guard keeps them from being probed. fn offer_prefetch( executor: &mut E, deferred: &mut Vec, @@ -402,7 +409,7 @@ fn offer_prefetch( ) { let mut keep = Vec::with_capacity(deferred.len()); for mut req in deferred.drain(..) { - if !req.prefetch_offered && !req.echo { + if !req.prefetch_offered && !req.wants_prompt_logprobs() { req.prefetch_offered = true; if executor.begin_kv_prefetch( req.request_id, @@ -1044,7 +1051,7 @@ fn active_future_blocks(active: &[ActiveRequestState], block_size: usize) -> usi } fn echo_exceeds_prefill_bound(req: &PendingRequest, max_prefill_tokens: usize) -> bool { - req.echo && req.prompt_tokens.len() > max_prefill_tokens + req.wants_prompt_logprobs() && req.prompt_tokens.len() > max_prefill_tokens } /// Free blocks already promised to admitted requests (active decode growth + @@ -1192,7 +1199,7 @@ fn send_rejection(req: &PendingRequest, reason: RejectReason) { req.max_tokens ), RejectReason::EchoPrefillTokens { limit } => format!( - "echo request prompt exceeds the profiled prefill limit of {} tokens: prompt_tokens={}", + "prompt-logprobs request prompt exceeds the profiled prefill limit of {} tokens: prompt_tokens={}", limit, req.prompt_tokens.len() ), diff --git a/openinfer-qwen3/src/scheduler/plan.rs b/openinfer-qwen3/src/scheduler/plan.rs index 538a4cad5..fa399aac3 100644 --- a/openinfer-qwen3/src/scheduler/plan.rs +++ b/openinfer-qwen3/src/scheduler/plan.rs @@ -50,23 +50,24 @@ pub(super) fn build_next_plan( pending: Vec, speculative: bool, ) -> Option { - // echo+logprobs requests need all-position logits, which the unified forward - // does not compute (it passes all_position_logits=None). And under DFlash - // speculation, an eligible request must capture its target hidden context - // during prefill — the unified forward skips that capture, so a request - // prefilled via Unified would never become draft-ready and DFlash would - // silently no-op for it forever. Either way, route pending through a + // Prompt-logprobs requests need all-position logits, which the unified + // forward does not compute (it passes all_position_logits=None). And under + // DFlash speculation, an eligible request must capture its target hidden + // context during prefill — the unified forward skips that capture, so a + // request prefilled via Unified would never become draft-ready and DFlash + // would silently no-op for it forever. Either way, route pending through a // dedicated prefill step instead of degrading silently. - let needs_prompt_logprobs = pending.iter().any(|r| r.echo && r.logprobs > 0); + let needs_prompt_logprobs = pending.iter().any(PendingRequest::wants_prompt_logprobs); // Deliberately a loose superset of the real capture eligibility - // (`dflash_prefill_supported`, which also needs `cached_tokens == 0 && !echo`): - // over-routing an ineligible request to a dedicated prefill only costs one - // fusion, but under-routing a capture-eligible one into Unified would silently - // break its readiness. Never tighten this into the dangerous direction. + // (`dflash_prefill_supported`, which also needs `cached_tokens == 0 && no + // prompt logprobs`): over-routing an ineligible request to a dedicated + // prefill only costs one fusion, but under-routing a capture-eligible one + // into Unified would silently break its readiness. Never tighten this into + // the dangerous direction. let needs_dflash_capture = speculative && pending .iter() - .any(|r| r.lora_adapter.is_none() && r.logprobs == 0); + .any(|r| r.lora_adapter.is_none() && r.logprobs.is_none()); if !pending.is_empty() && have_active && !needs_prompt_logprobs && !needs_dflash_capture { Some(ExecutionPlan::Unified { pending }) } else if !pending.is_empty() { @@ -89,7 +90,7 @@ pub(super) fn execute_plan( let scheduled_at_unix_s = openinfer_core::engine::unix_now_s(); let indices: Vec = (0..pending.len()).collect(); let requests = build_prefill_items(&pending, &indices); - let any_echo = pending.iter().any(|req| req.echo); + let any_echo = pending.iter().any(PendingRequest::wants_prompt_logprobs); let mut result = executor.execute_prefill(PrefillPlan { requests: &requests, echo: any_echo, @@ -166,7 +167,7 @@ pub(super) fn should_speculative_decode( && active.iter().all(|req| { executor.speculative_request_ready(req.request_id) && req.lora_adapter.is_none() - && req.logprobs == 0 + && req.logprobs.is_none() }) } @@ -214,7 +215,7 @@ fn build_prefill_items(pending: &[PendingRequest], indices: &[usize]) -> Vec
 3 blocks; future growth = 2
         prompt_len: 16,
         params: SamplingParams::default(),
-        logprobs: 0,
+        logprobs: None,
     }];
 
     let mk = |id: u64, prompt_len, max_tokens| {
@@ -405,7 +405,7 @@ fn admission_respects_decode_batch_capacity() {
             max_tokens: 2,
             prompt_len: 16,
             params: SamplingParams::default(),
-            logprobs: 0,
+            logprobs: None,
         });
     }
     let pending = PendingRequest::from_scheduler_request(RequestId(64), request(16, 1).0);
@@ -523,7 +523,7 @@ fn echo_requests_run_only_when_their_prompt_fits_the_prefill_bound() {
     let mk_echo = |id: u64, prompt_len| {
         let (req, _rx) = request(prompt_len, 1);
         let mut pending = PendingRequest::from_scheduler_request(RequestId(id), req);
-        pending.echo = true;
+        pending.prompt_logprobs = Some(2);
         pending
     };
     let mk = |id: u64, prompt_len| {
@@ -565,7 +565,7 @@ fn oversized_echo_request_is_rejected_at_admission() {
     let active: [ActiveRequestState; 0] = [];
     let mk_echo = |id: u64, prompt_len| {
         let (mut req, _rx) = request(prompt_len, 1);
-        req.echo = true;
+        req.prompt_logprobs = Some(2);
         PendingRequest::from_scheduler_request(RequestId(id), req)
     };
     let mk = |id: u64, prompt_len| {
@@ -757,8 +757,8 @@ fn pending(request_id: u64, echo: bool) -> PendingRequest {
         params: SamplingParams::default(),
         max_tokens: 1,
         token_tx,
-        logprobs: 0,
-        echo,
+        logprobs: None,
+        prompt_logprobs: echo.then_some(2),
         queued_at_unix_s: None,
         prefetch_offered: false,
         prefill_pos: 0,
@@ -804,8 +804,8 @@ fn request(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         },
         token_rx,
     )
@@ -945,7 +945,7 @@ fn retiring_multiple_active_requests_tolerates_unsorted_indices() {
             max_tokens: 2,
             prompt_len: 16,
             params: SamplingParams::default(),
-            logprobs: 0,
+            logprobs: None,
         });
         executor
             .ensure_request_tokens(request_id, 16)
@@ -1097,7 +1097,7 @@ fn spec_active(
             ignore_eos,
             ..SamplingParams::default()
         },
-        logprobs: 0,
+        logprobs: None,
     }
 }
 
diff --git a/openinfer-qwen3/tests/batch_invariance_attention_path.rs b/openinfer-qwen3/tests/batch_invariance_attention_path.rs
index fc8806f9a..09c92671c 100644
--- a/openinfer-qwen3/tests/batch_invariance_attention_path.rs
+++ b/openinfer-qwen3/tests/batch_invariance_attention_path.rs
@@ -36,8 +36,8 @@ fn pitem(id: RequestId, prompt: Vec) -> PrefillStepItem {
         prompt,
         MAX_OUTPUT_TOKENS,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
@@ -87,7 +87,7 @@ fn a_decode_at_batch(
         id_a,
         a_first,
         SamplingParams::default(),
-        LOGPROBS,
+        Some(LOGPROBS),
     )];
     let mut cofill_ids = Vec::with_capacity(n_cofill);
     for i in 0..n_cofill {
@@ -97,7 +97,7 @@ fn a_decode_at_batch(
             id,
             f_first,
             SamplingParams::default(),
-            LOGPROBS,
+            Some(LOGPROBS),
         ));
         cofill_ids.push(id);
     }
diff --git a/openinfer-qwen3/tests/batch_invariance_decode_gemm_graph.rs b/openinfer-qwen3/tests/batch_invariance_decode_gemm_graph.rs
index f4c08e554..0bc0c6194 100644
--- a/openinfer-qwen3/tests/batch_invariance_decode_gemm_graph.rs
+++ b/openinfer-qwen3/tests/batch_invariance_decode_gemm_graph.rs
@@ -49,8 +49,8 @@ fn pitem(id: RequestId, prompt: Vec) -> PrefillStepItem {
         prompt,
         MAX_OUTPUT_TOKENS,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
@@ -79,7 +79,12 @@ fn a_first_and_decode(ex: &mut Qwen3Executor, n_requests: usize) -> (u32, Vec<(u
         .iter()
         .zip(&pr.requests)
         .map(|((id, _), req)| {
-            DecodeStepItem::new(*id, req.first_token, SamplingParams::default(), LOGPROBS)
+            DecodeStepItem::new(
+                *id,
+                req.first_token,
+                SamplingParams::default(),
+                Some(LOGPROBS),
+            )
         })
         .collect();
     let dr = ex
diff --git a/openinfer-qwen3/tests/batch_invariance_decode_splitkv_graph.rs b/openinfer-qwen3/tests/batch_invariance_decode_splitkv_graph.rs
index fdb84b7c1..2598b2852 100644
--- a/openinfer-qwen3/tests/batch_invariance_decode_splitkv_graph.rs
+++ b/openinfer-qwen3/tests/batch_invariance_decode_splitkv_graph.rs
@@ -37,8 +37,8 @@ fn pitem(id: RequestId, prompt: Vec) -> PrefillStepItem {
         prompt,
         MAX_OUTPUT_TOKENS,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
@@ -77,8 +77,8 @@ fn a_decode_cobatched_with(
     let b_first = prefill_chunked(ex, id_b, b_prompt);
     // Decode A+B together (batch 2, A row 0); B's KV length sets the batch max_seq_len.
     let ditems = vec![
-        DecodeStepItem::new(id_a, a_first, SamplingParams::default(), LOGPROBS),
-        DecodeStepItem::new(id_b, b_first, SamplingParams::default(), LOGPROBS),
+        DecodeStepItem::new(id_a, a_first, SamplingParams::default(), Some(LOGPROBS)),
+        DecodeStepItem::new(id_b, b_first, SamplingParams::default(), Some(LOGPROBS)),
     ];
     let dr = ex
         .execute_decode(DecodePlan {
diff --git a/openinfer-qwen3/tests/batch_invariance_endtoend.rs b/openinfer-qwen3/tests/batch_invariance_endtoend.rs
index 1cf58702e..a1a63a950 100644
--- a/openinfer-qwen3/tests/batch_invariance_endtoend.rs
+++ b/openinfer-qwen3/tests/batch_invariance_endtoend.rs
@@ -36,8 +36,8 @@ fn item(id: RequestId, prompt: Vec) -> PrefillStepItem {
         prompt,
         MAX_OUTPUT_TOKENS,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
diff --git a/openinfer-qwen3/tests/batch_invariance_envelope.rs b/openinfer-qwen3/tests/batch_invariance_envelope.rs
index 241b31f14..1baae1bb9 100644
--- a/openinfer-qwen3/tests/batch_invariance_envelope.rs
+++ b/openinfer-qwen3/tests/batch_invariance_envelope.rs
@@ -31,8 +31,8 @@ fn prefill_first(ex: &mut Qwen3Executor, id: RequestId, prompt: &[u32]) -> u32 {
             prompt.to_vec(),
             64,
             SamplingParams::default(),
-            0,
-            false,
+            None,
+            None,
         )],
         echo: false,
         sample_seed: 0,
@@ -74,10 +74,15 @@ fn pin_serves_production_envelope_without_fallback() {
                     synth(pf, 100 + i as u64),
                     64,
                     SamplingParams::default(),
-                    0,
-                    false,
+                    None,
+                    None,
+                )],
+                decode_requests: &[DecodeStepItem::new(
+                    id_d,
+                    t,
+                    SamplingParams::default(),
+                    Some(64),
                 )],
-                decode_requests: &[DecodeStepItem::new(id_d, t, SamplingParams::default(), 64)],
                 sample_seed: 0,
             })
             .unwrap_or_else(|e| panic!("Unified N={} bailed: {e}", pf + 1));
@@ -95,7 +100,7 @@ fn pin_serves_production_envelope_without_fallback() {
         .collect();
     let items: Vec = dec
         .iter()
-        .map(|&(id, tok)| DecodeStepItem::new(id, tok, SamplingParams::default(), 0))
+        .map(|&(id, tok)| DecodeStepItem::new(id, tok, SamplingParams::default(), None))
         .collect();
     reset_numeric_policy_counters();
     let _ = ex
@@ -119,7 +124,7 @@ fn pin_serves_production_envelope_without_fallback() {
     let id_big = RequestId::new(40000);
     let decode_items: Vec = decoders
         .iter()
-        .map(|&(id, tok)| DecodeStepItem::new(id, tok, SamplingParams::default(), 0))
+        .map(|&(id, tok)| DecodeStepItem::new(id, tok, SamplingParams::default(), None))
         .collect();
     reset_numeric_policy_counters();
     let _ = ex
@@ -129,8 +134,8 @@ fn pin_serves_production_envelope_without_fallback() {
                 synth(1024, 7),
                 64,
                 SamplingParams::default(),
-                0,
-                false,
+                None,
+                None,
             )],
             decode_requests: &decode_items,
             sample_seed: 0,
diff --git a/openinfer-qwen3/tests/batch_invariance_output.rs b/openinfer-qwen3/tests/batch_invariance_output.rs
index 13d430cf5..cd823c188 100644
--- a/openinfer-qwen3/tests/batch_invariance_output.rs
+++ b/openinfer-qwen3/tests/batch_invariance_output.rs
@@ -95,8 +95,8 @@ impl Harness {
                 max_tokens: output.0,
                 lora_adapter: None,
                 token_tx,
-                logprobs: output.1,
-                echo: false,
+                logprobs: (output.1 > 0).then_some(output.1),
+                prompt_logprobs: None,
             })
             .expect("submit failed");
         (tag, cutoff)
diff --git a/openinfer-qwen3/tests/batch_invariance_unified.rs b/openinfer-qwen3/tests/batch_invariance_unified.rs
index d0c129bbc..d074234de 100644
--- a/openinfer-qwen3/tests/batch_invariance_unified.rs
+++ b/openinfer-qwen3/tests/batch_invariance_unified.rs
@@ -29,8 +29,8 @@ fn prefill_first(ex: &mut Qwen3Executor, id: RequestId, prompt: &[u32]) -> u32 {
             prompt.to_vec(),
             64,
             SamplingParams::default(),
-            0,
-            false,
+            None,
+            None,
         )],
         echo: false,
         sample_seed: 0,
@@ -91,10 +91,15 @@ fn unified_decode_row(
                 chunk,
                 64,
                 SamplingParams::default(),
-                0,
-                false,
+                None,
+                None,
+            )],
+            decode_requests: &[DecodeStepItem::new(
+                id_a,
+                t0,
+                SamplingParams::default(),
+                Some(64),
             )],
-            decode_requests: &[DecodeStepItem::new(id_a, t0, SamplingParams::default(), 64)],
             sample_seed: 0,
         })
         .expect("unified");
@@ -111,7 +116,12 @@ fn pure_decode_row(ex: &mut Qwen3Executor, p: &[u32], id_dec: u64) -> Row {
     reset_numeric_policy_counters();
     let dr = ex
         .execute_decode(DecodePlan {
-            requests: &[DecodeStepItem::new(id, t0, SamplingParams::default(), 64)],
+            requests: &[DecodeStepItem::new(
+                id,
+                t0,
+                SamplingParams::default(),
+                Some(64),
+            )],
             sample_seed: 0,
         })
         .expect("decode");
@@ -252,7 +262,7 @@ fn prefill_each(
 fn decode_items(decoders: &[(RequestId, u32)]) -> Vec {
     decoders
         .iter()
-        .map(|&(id, tok)| DecodeStepItem::new(id, tok, SamplingParams::default(), 64))
+        .map(|&(id, tok)| DecodeStepItem::new(id, tok, SamplingParams::default(), Some(64)))
         .collect()
 }
 
@@ -293,8 +303,8 @@ fn unified_decode_batch(
                 mixed_prompt(PAST_SPLIT_CAP + 1),
                 64,
                 SamplingParams::default(),
-                0,
-                false,
+                None,
+                None,
             )],
             decode_requests: &decode_items,
             sample_seed: 0,
diff --git a/openinfer-qwen3/tests/cached_tokens_usage.rs b/openinfer-qwen3/tests/cached_tokens_usage.rs
index 2bfc9bfa0..293f99735 100644
--- a/openinfer-qwen3/tests/cached_tokens_usage.rs
+++ b/openinfer-qwen3/tests/cached_tokens_usage.rs
@@ -50,8 +50,8 @@ fn run_and_capture_cached(handle: &EngineHandle, prompt_tokens: Vec) -> usi
             max_tokens: 4,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/context_window.rs b/openinfer-qwen3/tests/context_window.rs
index 7db75c76d..4440ac17d 100644
--- a/openinfer-qwen3/tests/context_window.rs
+++ b/openinfer-qwen3/tests/context_window.rs
@@ -61,8 +61,8 @@ fn generate_text(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
@@ -111,8 +111,8 @@ fn oversized_prompt_is_rejected_with_context_length_error() {
             max_tokens: 8,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/context_window_in_window.rs b/openinfer-qwen3/tests/context_window_in_window.rs
index eac0f202b..37286b9f7 100644
--- a/openinfer-qwen3/tests/context_window_in_window.rs
+++ b/openinfer-qwen3/tests/context_window_in_window.rs
@@ -75,8 +75,8 @@ fn in_window_prompt_past_old_rope_table_is_served() {
             max_tokens: 1,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/dflash_sampled_equivalence.rs b/openinfer-qwen3/tests/dflash_sampled_equivalence.rs
index bccb064e4..e6a296986 100644
--- a/openinfer-qwen3/tests/dflash_sampled_equivalence.rs
+++ b/openinfer-qwen3/tests/dflash_sampled_equivalence.rs
@@ -190,8 +190,8 @@ fn sample_wave(
                     max_tokens: POSITIONS,
                     lora_adapter: None,
                     token_tx,
-                    logprobs: 0,
-                    echo: false,
+                    logprobs: None,
+                    prompt_logprobs: None,
                 })
                 .expect("submit failed");
             rx
diff --git a/openinfer-qwen3/tests/dflash_speculative_gate.rs b/openinfer-qwen3/tests/dflash_speculative_gate.rs
index b61a98b41..5d36b0b0e 100644
--- a/openinfer-qwen3/tests/dflash_speculative_gate.rs
+++ b/openinfer-qwen3/tests/dflash_speculative_gate.rs
@@ -150,8 +150,8 @@ fn generate(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs,
-            echo: false,
+            logprobs: (logprobs > 0).then_some(logprobs),
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
@@ -193,8 +193,8 @@ fn generate_concurrent(handle: &EngineHandle, requests: Vec<(Vec, usize)>)
                     max_tokens,
                     lora_adapter: None,
                     token_tx,
-                    logprobs: 0,
-                    echo: false,
+                    logprobs: None,
+                    prompt_logprobs: None,
                 })
                 .expect("submit failed");
             rx
@@ -246,8 +246,8 @@ fn prefill_next(handle: &EngineHandle, context: Vec, logprobs: usize) -> St
             max_tokens: 1,
             lora_adapter: None,
             token_tx,
-            logprobs,
-            echo: true,
+            logprobs: Some(logprobs),
+            prompt_logprobs: Some(logprobs),
         })
         .expect("submit failed");
 
@@ -706,8 +706,8 @@ fn dflash_request_in_draft_headroom_is_rejected_not_panicked() {
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/dflash_speculative_perf.rs b/openinfer-qwen3/tests/dflash_speculative_perf.rs
index 717ec5dea..b814977ce 100644
--- a/openinfer-qwen3/tests/dflash_speculative_perf.rs
+++ b/openinfer-qwen3/tests/dflash_speculative_perf.rs
@@ -92,8 +92,8 @@ fn timed_generate(handle: &EngineHandle, prompt_tokens: Vec) -> (usize, Dur
             max_tokens: GENERATED_TOKENS,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/hf_golden_gate.rs b/openinfer-qwen3/tests/hf_golden_gate.rs
index e31218c04..a1bb909a0 100644
--- a/openinfer-qwen3/tests/hf_golden_gate.rs
+++ b/openinfer-qwen3/tests/hf_golden_gate.rs
@@ -371,13 +371,13 @@ fn prefill_item(id: RequestId, prompt: Vec) -> PrefillStepItem {
         prompt,
         MAX_OUTPUT_TOKENS,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
 fn decode_item(id: RequestId, fed: u32) -> DecodeStepItem {
-    DecodeStepItem::new(id, fed, SamplingParams::default(), LOGPROBS)
+    DecodeStepItem::new(id, fed, SamplingParams::default(), Some(LOGPROBS))
 }
 
 /// Teacher-force the golden sequences `seqs` through `ex` and fold every
diff --git a/openinfer-qwen3/tests/kv_offload_cpu_hit.rs b/openinfer-qwen3/tests/kv_offload_cpu_hit.rs
index 8dd5e4d2c..b1891a54a 100644
--- a/openinfer-qwen3/tests/kv_offload_cpu_hit.rs
+++ b/openinfer-qwen3/tests/kv_offload_cpu_hit.rs
@@ -69,8 +69,8 @@ fn prefill_item(id: u64, prompt: &[u32]) -> PrefillStepItem {
         prompt.to_vec(),
         MAX_OUTPUT,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
diff --git a/openinfer-qwen3/tests/lora_golden_gate.rs b/openinfer-qwen3/tests/lora_golden_gate.rs
index 1f05f6176..da1f00c5c 100644
--- a/openinfer-qwen3/tests/lora_golden_gate.rs
+++ b/openinfer-qwen3/tests/lora_golden_gate.rs
@@ -313,14 +313,14 @@ fn prefill_item(id: RequestId, prompt: Vec, lora: bool) -> PrefillStepItem
         prompt,
         MAX_OUTPUT_TOKENS,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
     .with_lora_adapter(lora.then(|| ADAPTER_NAME.to_string()))
 }
 
 fn decode_item(id: RequestId, fed: u32, lora: bool) -> DecodeStepItem {
-    DecodeStepItem::new(id, fed, SamplingParams::default(), LOGPROBS)
+    DecodeStepItem::new(id, fed, SamplingParams::default(), Some(LOGPROBS))
         .with_lora_adapter(lora.then(|| ADAPTER_NAME.to_string()))
 }
 
diff --git a/openinfer-qwen3/tests/lora_smoke.rs b/openinfer-qwen3/tests/lora_smoke.rs
index 21fbd139a..4c4c34015 100644
--- a/openinfer-qwen3/tests/lora_smoke.rs
+++ b/openinfer-qwen3/tests/lora_smoke.rs
@@ -101,8 +101,8 @@ fn generate_tokens(
             max_tokens,
             lora_adapter,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/prefix_cache.rs b/openinfer-qwen3/tests/prefix_cache.rs
index 990deaf09..ed1f4a4ad 100644
--- a/openinfer-qwen3/tests/prefix_cache.rs
+++ b/openinfer-qwen3/tests/prefix_cache.rs
@@ -77,13 +77,18 @@ fn prefill_item(id: u64, prompt: &[u32]) -> PrefillStepItem {
         prompt.to_vec(),
         MAX_OUTPUT,
         SamplingParams::default(),
-        LOGPROBS,
-        false,
+        Some(LOGPROBS),
+        None,
     )
 }
 
 fn decode_item(id: u64, fed: u32) -> DecodeStepItem {
-    DecodeStepItem::new(RequestId::new(id), fed, SamplingParams::default(), LOGPROBS)
+    DecodeStepItem::new(
+        RequestId::new(id),
+        fed,
+        SamplingParams::default(),
+        Some(LOGPROBS),
+    )
 }
 
 fn top_logprobs(lp: Option<&TokenLogprob>) -> Vec<(u32, f32)> {
diff --git a/openinfer-qwen3/tests/sampling_behavior.rs b/openinfer-qwen3/tests/sampling_behavior.rs
index c931806e0..e626275b8 100644
--- a/openinfer-qwen3/tests/sampling_behavior.rs
+++ b/openinfer-qwen3/tests/sampling_behavior.rs
@@ -53,8 +53,8 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec, params: SamplingPara
             max_tokens: GENERATED_TOKENS,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen3/tests/scheduler_robustness.rs b/openinfer-qwen3/tests/scheduler_robustness.rs
index 56ff9f66e..c3ea50607 100644
--- a/openinfer-qwen3/tests/scheduler_robustness.rs
+++ b/openinfer-qwen3/tests/scheduler_robustness.rs
@@ -64,8 +64,8 @@ fn generate_text(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
@@ -132,8 +132,8 @@ fn scheduler_survives_consumer_drop() {
             max_tokens: 10,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
     std::thread::sleep(Duration::from_millis(500));
diff --git a/openinfer-qwen3/tests/tp_concurrent_decode.rs b/openinfer-qwen3/tests/tp_concurrent_decode.rs
index 96740142e..4de3c100d 100644
--- a/openinfer-qwen3/tests/tp_concurrent_decode.rs
+++ b/openinfer-qwen3/tests/tp_concurrent_decode.rs
@@ -126,8 +126,8 @@ fn tp2_graph_dump_when_available_and_concurrent_decode_complete() {
                     max_tokens: 24 + (i % 4) * 24,
                     lora_adapter: None,
                     token_tx,
-                    logprobs: 0,
-                    echo: false,
+                    logprobs: None,
+                    prompt_logprobs: None,
                 })
                 .expect("submit failed");
             rx
diff --git a/openinfer-qwen35-4b/src/executor.rs b/openinfer-qwen35-4b/src/executor.rs
index c0cb947be..67f4c42fd 100644
--- a/openinfer-qwen35-4b/src/executor.rs
+++ b/openinfer-qwen35-4b/src/executor.rs
@@ -35,11 +35,11 @@ impl RequestId {
 pub struct PrefillStepItem {
     pub(crate) request_id: RequestId,
     pub(crate) prompt_tokens: Vec,
-    pub(crate) logprobs: usize,
+    pub(crate) logprobs: Option,
 }
 
 impl PrefillStepItem {
-    pub fn new(request_id: RequestId, prompt_tokens: Vec, logprobs: usize) -> Self {
+    pub fn new(request_id: RequestId, prompt_tokens: Vec, logprobs: Option) -> Self {
         Self {
             request_id,
             prompt_tokens,
@@ -52,11 +52,11 @@ impl PrefillStepItem {
 pub struct DecodeStepItem {
     pub(crate) request_id: RequestId,
     pub(crate) token_id: u32,
-    pub(crate) logprobs: usize,
+    pub(crate) logprobs: Option,
 }
 
 impl DecodeStepItem {
-    pub fn new(request_id: RequestId, token_id: u32, logprobs: usize) -> Self {
+    pub fn new(request_id: RequestId, token_id: u32, logprobs: Option) -> Self {
         Self {
             request_id,
             token_id,
@@ -174,7 +174,8 @@ impl Qwen35Executor {
             self.model
                 .batch_prefill_logits(&prompts, &mut kv_states, &mut recurrent_refs)?;
 
-        let requested_logprobs: Vec = plan.requests.iter().map(|req| req.logprobs).collect();
+        let requested_logprobs: Vec> =
+            plan.requests.iter().map(|req| req.logprobs).collect();
         let cpu_logits =
             snapshot_requested_logprobs(self.model.device_ctx(), &logits, &requested_logprobs)?;
         let tokens =
@@ -184,7 +185,11 @@ impl Qwen35Executor {
         for (i, (req, kv)) in plan.requests.iter().zip(kv_states).enumerate() {
             let first_token = tokens[i];
             let first_token_logprob = cpu_logits[i].as_ref().and_then(|row| {
-                openinfer_sample::token_logprob_from_row(row, first_token, req.logprobs)
+                openinfer_sample::token_logprob_from_row(
+                    row,
+                    first_token,
+                    req.logprobs.unwrap_or(0),
+                )
             });
             let slot_idx = self.active.len();
             self.graph_state.copy_state_to_slot(
@@ -229,7 +234,8 @@ impl Qwen35Executor {
         self.model
             .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state)?;
 
-        let requested_logprobs: Vec = plan.requests.iter().map(|req| req.logprobs).collect();
+        let requested_logprobs: Vec> =
+            plan.requests.iter().map(|req| req.logprobs).collect();
         let cpu_logits = snapshot_requested_logprobs(
             self.model.device_ctx(),
             &self.graph_state.buffers.logits,
@@ -246,9 +252,9 @@ impl Qwen35Executor {
         let mut results = Vec::with_capacity(plan.requests.len());
         for (i, req) in plan.requests.iter().enumerate() {
             let token = tokens[i];
-            let logprob = cpu_logits[i]
-                .as_ref()
-                .and_then(|row| openinfer_sample::token_logprob_from_row(row, token, req.logprobs));
+            let logprob = cpu_logits[i].as_ref().and_then(|row| {
+                openinfer_sample::token_logprob_from_row(row, token, req.logprobs.unwrap_or(0))
+            });
             results.push(DecodeRequestResult {
                 request_id: req.request_id,
                 token,
diff --git a/openinfer-qwen35-4b/src/logprobs.rs b/openinfer-qwen35-4b/src/logprobs.rs
index fbbfceb4e..c747c4af7 100644
--- a/openinfer-qwen35-4b/src/logprobs.rs
+++ b/openinfer-qwen35-4b/src/logprobs.rs
@@ -4,7 +4,7 @@ use openinfer_core::tensor::{DeviceContext, HiddenStates};
 pub(crate) fn snapshot_requested_logprobs(
     ctx: &DeviceContext,
     logits: &HiddenStates,
-    requested_top_k: &[usize],
+    requested_top_k: &[Option],
 ) -> Result>>> {
     anyhow::ensure!(
         requested_top_k.len() <= logits.seq_len,
@@ -12,15 +12,15 @@ pub(crate) fn snapshot_requested_logprobs(
         requested_top_k.len(),
         logits.seq_len
     );
-    if !requested_top_k.iter().any(|&top_k| top_k > 0) {
+    if !requested_top_k.iter().any(Option::is_some) {
         return Ok(vec![None; requested_top_k.len()]);
     }
 
     requested_top_k
         .iter()
         .enumerate()
-        .map(|(i, &top_k)| {
-            if top_k == 0 {
+        .map(|(i, top_k)| {
+            if top_k.is_none() {
                 Ok(None)
             } else {
                 let row = crate::ops::extract_vec(ctx, logits, i)?;
diff --git a/openinfer-qwen35-4b/src/scheduler.rs b/openinfer-qwen35-4b/src/scheduler.rs
index be624fb9f..fdf9bf9f3 100644
--- a/openinfer-qwen35-4b/src/scheduler.rs
+++ b/openinfer-qwen35-4b/src/scheduler.rs
@@ -50,8 +50,9 @@ struct ActiveRequest35 {
     max_tokens: usize,
     prompt_len: usize,
     params: SamplingParams,
-    /// Number of top logprobs to return (0 = disabled).
-    logprobs: usize,
+    /// Completion logprob top-k count (`None` = disabled, `Some(0)` = scored
+    /// token only).
+    logprobs: Option,
 }
 
 /// A request whose prompt is being prefilled across multiple scheduler steps.
@@ -320,7 +321,7 @@ impl SingleGpuBackend {
             pending.len(),
             "Qwen3.5 prefill logits rows must preserve pending request order"
         );
-        let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect();
+        let requested_logprobs: Vec> = pending.iter().map(|r| r.logprobs).collect();
         let cpu_logits =
             snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?;
         let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect();
@@ -340,7 +341,7 @@ impl SingleGpuBackend {
                     openinfer_sample::token_logprob_from_row(
                         &logits_f32,
                         tokens[i],
-                        pending[i].logprobs,
+                        pending[i].logprobs.unwrap_or(0),
                     )
                 })
             })
@@ -353,7 +354,7 @@ impl SingleGpuBackend {
         active: &[ActiveRequest35],
         rng: &mut StdRng,
     ) -> Result<(Vec, Vec>)> {
-        let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect();
+        let requested_logprobs: Vec> = active.iter().map(|r| r.logprobs).collect();
         let cpu_logits = snapshot_requested_logprobs(
             self.model.device_ctx(),
             &self.graph_state.buffers.logits,
@@ -375,7 +376,7 @@ impl SingleGpuBackend {
                     openinfer_sample::token_logprob_from_row(
                         &logits_f32,
                         tokens[i],
-                        active[i].logprobs,
+                        active[i].logprobs.unwrap_or(0),
                     )
                 })
             })
@@ -790,6 +791,23 @@ fn scheduler_loop(
 
         // 3. Admit new prompts. In-flight prefills reserve their promotion slot
         //    and future KV growth, so shrink the slot/page budgets accordingly
+        //
+        // Honor-or-reject: prompt logprobs need all-position logits, which the
+        // Qwen3.5 prefill path does not produce yet. Reject loudly at intake
+        // rather than returning empty prompt logprobs (same policy as Kimi-K2,
+        // #236).
+        pending.retain(|req| {
+            if req.prompt_logprobs.is_none() {
+                return true;
+            }
+            let _ = req.token_tx.send(TokenEvent::Rejected {
+                message: "prompt logprobs are not supported on the Qwen3.5 serving path yet"
+                    .to_string(),
+                prompt_tokens: req.prompt_tokens.len(),
+                completion_tokens: 0,
+            });
+            false
+        });
         let active_budget: Vec = active
             .iter()
             .map(|req| ActiveKvBudget {
@@ -1391,14 +1409,6 @@ fn promote_or_requeue(
         let first_token = tokens[i];
         let logprob = logprobs[i].clone();
 
-        if req.echo {
-            let echo_logprobs = vec![None; req.prompt_tokens.len()];
-            let _ = req.token_tx.send(TokenEvent::PromptTokens {
-                ids: req.prompt_tokens.clone(),
-                logprobs: echo_logprobs,
-            });
-        }
-
         if !req.params.ignore_eos && backend.is_stop_token(first_token) {
             debug!(
                 "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}",
diff --git a/openinfer-qwen35-4b/src/scheduler/tests.rs b/openinfer-qwen35-4b/src/scheduler/tests.rs
index 679eb0259..548009767 100644
--- a/openinfer-qwen35-4b/src/scheduler/tests.rs
+++ b/openinfer-qwen35-4b/src/scheduler/tests.rs
@@ -15,8 +15,8 @@ fn send_rejection_reports_kv_lifetime_request_tokens() {
         max_tokens: 65,
         lora_adapter: None,
         token_tx,
-        logprobs: 0,
-        echo: false,
+        logprobs: None,
+        prompt_logprobs: None,
     };
 
     send_rejection(&req, RejectReason::KvBudget);
@@ -99,8 +99,8 @@ fn tp2_scheduler_chunked_prefill_then_decode_smoke() {
             max_tokens: 3,
             lora_adapter: None,
             token_tx,
-            logprobs: 1,
-            echo: false,
+            logprobs: Some(1),
+            prompt_logprobs: None,
         })
         .expect("submit TP scheduler request");
 
@@ -148,8 +148,8 @@ fn send_rejection_reports_context_window_limit() {
         max_tokens: 17,
         lora_adapter: None,
         token_tx,
-        logprobs: 0,
-        echo: false,
+        logprobs: None,
+        prompt_logprobs: None,
     };
 
     send_rejection(&req, RejectReason::ContextLength { limit: 32 });
diff --git a/openinfer-qwen35-4b/src/tp_executor.rs b/openinfer-qwen35-4b/src/tp_executor.rs
index e7dd7577b..22d3ee867 100644
--- a/openinfer-qwen35-4b/src/tp_executor.rs
+++ b/openinfer-qwen35-4b/src/tp_executor.rs
@@ -110,7 +110,7 @@ pub struct Qwen35TpExecutor {
 pub struct TpPrefillChunkItem {
     request_id: RequestId,
     prompt_tokens: Vec,
-    logprobs: usize,
+    logprobs: Option,
     sampling_params: SamplingParams,
     finish_prefill: bool,
 }
@@ -119,7 +119,7 @@ impl TpPrefillChunkItem {
     pub fn new(
         request_id: RequestId,
         prompt_tokens: Vec,
-        logprobs: usize,
+        logprobs: Option,
         finish_prefill: bool,
     ) -> Self {
         Self {
@@ -134,7 +134,7 @@ impl TpPrefillChunkItem {
     pub fn new_with_sampling(
         request_id: RequestId,
         prompt_tokens: Vec,
-        logprobs: usize,
+        logprobs: Option,
         sampling_params: SamplingParams,
         finish_prefill: bool,
     ) -> Self {
@@ -152,7 +152,7 @@ impl TpPrefillChunkItem {
 pub struct TpDecodeStepItem {
     request_id: RequestId,
     token_id: u32,
-    logprobs: usize,
+    logprobs: Option,
     sampling_params: SamplingParams,
 }
 
@@ -160,7 +160,7 @@ impl TpDecodeStepItem {
     pub fn new(
         request_id: RequestId,
         token_id: u32,
-        logprobs: usize,
+        logprobs: Option,
         sampling_params: SamplingParams,
     ) -> Self {
         Self {
@@ -1042,7 +1042,7 @@ impl TpWorkerState {
         )?;
         let first_token = tokens[0];
         let first_token_logprob = cpu_logits[0].as_ref().and_then(|row| {
-            openinfer_sample::token_logprob_from_row(row, first_token, chunk.logprobs)
+            openinfer_sample::token_logprob_from_row(row, first_token, chunk.logprobs.unwrap_or(0))
         });
         Ok(PrefillRequestResult {
             request_id: chunk.request_id,
@@ -1113,7 +1113,11 @@ impl TpWorkerState {
                 )?;
                 let token = tokens[0];
                 let logprob = cpu_logits[0].as_ref().and_then(|row| {
-                    openinfer_sample::token_logprob_from_row(row, token, request.logprobs)
+                    openinfer_sample::token_logprob_from_row(
+                        row,
+                        token,
+                        request.logprobs.unwrap_or(0),
+                    )
                 });
                 primary_results.push(DecodeRequestResult {
                     request_id: request.request_id,
diff --git a/openinfer-qwen35-4b/tests/chunked_prefill.rs b/openinfer-qwen35-4b/tests/chunked_prefill.rs
index 4437ab4a5..c4de7ffa6 100644
--- a/openinfer-qwen35-4b/tests/chunked_prefill.rs
+++ b/openinfer-qwen35-4b/tests/chunked_prefill.rs
@@ -66,8 +66,8 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec) -> (Vec, Finish
             max_tokens: GENERATED_TOKENS,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-qwen35-4b/tests/e2e_scheduler.rs b/openinfer-qwen35-4b/tests/e2e_scheduler.rs
index 655d684e7..bd76c266b 100644
--- a/openinfer-qwen35-4b/tests/e2e_scheduler.rs
+++ b/openinfer-qwen35-4b/tests/e2e_scheduler.rs
@@ -130,8 +130,8 @@ fn generate_tokens_with_logprobs(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs,
-            echo: false,
+            logprobs: (logprobs > 0).then_some(logprobs),
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
@@ -220,8 +220,8 @@ fn expect_context_window_rejection(handle: &EngineHandle, max_context_tokens: us
             max_tokens: 1,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit over-context request");
 
@@ -431,8 +431,8 @@ fn run_full_scheduler_e2e(
                     max_tokens: case.max_new_tokens,
                     lora_adapter: None,
                     token_tx,
-                    logprobs: 0,
-                    echo: false,
+                    logprobs: None,
+                    prompt_logprobs: None,
                 })
                 .expect("submit failed");
             receivers.push((case.name.to_string(), 0, token_rx));
@@ -471,8 +471,8 @@ fn run_full_scheduler_e2e(
                     max_tokens: 8,
                     lora_adapter: None,
                     token_tx,
-                    logprobs,
-                    echo: false,
+                    logprobs: (logprobs > 0).then_some(logprobs),
+                    prompt_logprobs: None,
                 })
                 .expect("submit failed");
             receivers.push((name, logprobs, token_rx));
@@ -512,8 +512,8 @@ fn run_full_scheduler_e2e(
                 max_tokens: 10,
                 lora_adapter: None,
                 token_tx,
-                logprobs: 0,
-                echo: false,
+                logprobs: None,
+                prompt_logprobs: None,
             })
             .expect("submit failed");
         std::thread::sleep(std::time::Duration::from_millis(500));
diff --git a/openinfer-qwen35-4b/tests/hf_golden_gate.rs b/openinfer-qwen35-4b/tests/hf_golden_gate.rs
index 14bf81113..919abb8eb 100644
--- a/openinfer-qwen35-4b/tests/hf_golden_gate.rs
+++ b/openinfer-qwen35-4b/tests/hf_golden_gate.rs
@@ -425,11 +425,11 @@ fn report_fixture_shape(golden: &Golden) {
 }
 
 fn prefill_item(id: RequestId, prompt: Vec) -> PrefillStepItem {
-    PrefillStepItem::new(id, prompt, LOGPROBS)
+    PrefillStepItem::new(id, prompt, Some(LOGPROBS))
 }
 
 fn decode_item(id: RequestId, fed: u32) -> DecodeStepItem {
-    DecodeStepItem::new(id, fed, LOGPROBS)
+    DecodeStepItem::new(id, fed, Some(LOGPROBS))
 }
 
 fn run(g: &Golden, ex: &mut Qwen35Executor, seqs: &[usize], batched: bool) -> (Stats, Vec) {
diff --git a/openinfer-qwen35-4b/tests/sampling_behavior.rs b/openinfer-qwen35-4b/tests/sampling_behavior.rs
index 22f526218..b3e767ea7 100644
--- a/openinfer-qwen35-4b/tests/sampling_behavior.rs
+++ b/openinfer-qwen35-4b/tests/sampling_behavior.rs
@@ -53,8 +53,8 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec, params: SamplingPara
             max_tokens: GENERATED_TOKENS,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .expect("submit failed");
 
diff --git a/openinfer-sample/tests/logprobs_golden.rs b/openinfer-sample/tests/logprobs_golden.rs
new file mode 100644
index 000000000..5c5222e52
--- /dev/null
+++ b/openinfer-sample/tests/logprobs_golden.rs
@@ -0,0 +1,248 @@
+//! Golden tests for the batched device logprobs reduction (#719).
+//!
+//! The device path (`logprobs_lse_bf16_into`, `logprobs_gather_rows_bf16_into`,
+//! and `logprobs_topk_bf16_into`) must reproduce
+//! `openinfer_sample::token_logprob_from_row` semantics over the same bf16
+//! logits: identical top-k token ids in (value desc, id asc) order, and
+//! logprobs within a small tolerance — the host accumulates `exp` in f64
+//! sequentially while the device block-reduces f64 partials, so the LSE can
+//! differ by a few fp32 ULPs. Requires a GPU.
+
+use half::bf16;
+use openinfer_kernels::ops::{
+    logprobs_gather_rows_bf16_into, logprobs_lse_bf16_into, logprobs_topk_bf16_into,
+};
+use openinfer_kernels::tensor::{DeviceContext, HiddenStates};
+use openinfer_sample::token_logprob_from_row;
+
+/// Logprob tolerance: covers fp32-ULP LSE divergence between the sequential
+/// host sum and the device tree reduction (values are O(10); fp32 eps ~ 1e-7).
+const LP_TOLERANCE: f32 = 1e-4;
+
+/// Deterministic xorshift64 PRNG so the test needs no extra deps.
+struct Rng(u64);
+
+impl Rng {
+    fn next_f32(&mut self, lo: f32, hi: f32) -> f32 {
+        self.0 ^= self.0 << 13;
+        self.0 ^= self.0 >> 7;
+        self.0 ^= self.0 << 17;
+        let unit = (self.0 >> 40) as f32 / (1u64 << 24) as f32;
+        lo + unit * (hi - lo)
+    }
+}
+
+fn make_arena(ctx: &DeviceContext, rows: &[Vec]) -> HiddenStates {
+    let vocab = rows[0].len();
+    assert!(rows.iter().all(|r| r.len() == vocab), "ragged rows");
+    let mut hs = HiddenStates::zeros(ctx, vocab, rows.len()).unwrap();
+    let flat: Vec = rows
+        .iter()
+        .flat_map(|r| r.iter().map(|&x| bf16::from_f32(x)))
+        .collect();
+    ctx.stream.memcpy_htod(&flat, &mut hs.data).unwrap();
+    ctx.sync().unwrap();
+    hs
+}
+
+/// One scored job: (arena row, picked token, top_k).
+struct Job {
+    row: u32,
+    picked: u32,
+    top_k: usize,
+}
+
+/// Run the device batch exactly the way the executor does and assemble
+/// host-side `(picked logprob, top logprobs)` results.
+fn run_device_batch(
+    ctx: &DeviceContext,
+    hs: &HiddenStates,
+    jobs: &[Job],
+) -> Vec<(f32, Vec<(u32, f32)>)> {
+    let m = jobs.len();
+    let vocab = hs.hidden_dim;
+    let k_dev = jobs.iter().map(|j| j.top_k).max().unwrap_or(0).min(vocab);
+
+    let rows: Vec = jobs.iter().map(|j| j.row).collect();
+    let picked: Vec = jobs.iter().map(|j| j.picked).collect();
+    let mut row_indices = ctx.stream.alloc_zeros::(m).unwrap();
+    let mut picked_dev = ctx.stream.alloc_zeros::(m).unwrap();
+    ctx.stream.memcpy_htod(&rows, &mut row_indices).unwrap();
+    ctx.stream.memcpy_htod(&picked, &mut picked_dev).unwrap();
+
+    let mut lse = ctx.stream.alloc_zeros::(m).unwrap();
+    let mut picked_lp = ctx.stream.alloc_zeros::(m).unwrap();
+    logprobs_lse_bf16_into(
+        ctx,
+        hs,
+        Some(&row_indices),
+        &picked_dev,
+        m,
+        &mut lse,
+        &mut picked_lp,
+    )
+    .unwrap();
+
+    let (values, indices) = if k_dev > 0 {
+        let mut gathered = ctx.stream.alloc_zeros::(m * vocab).unwrap();
+        logprobs_gather_rows_bf16_into(ctx, hs, &row_indices, m, &mut gathered).unwrap();
+        let mut values = ctx.stream.alloc_zeros::(m * k_dev).unwrap();
+        let mut indices = ctx.stream.alloc_zeros::(m * k_dev).unwrap();
+        let ok = logprobs_topk_bf16_into(
+            ctx,
+            &gathered,
+            m,
+            vocab,
+            k_dev,
+            &mut values,
+            &mut indices,
+        )
+        .unwrap();
+        assert!(ok, "FilteredTopK unsupported on this GPU");
+        (values, indices)
+    } else {
+        (
+            ctx.stream.alloc_zeros::(1).unwrap(),
+            ctx.stream.alloc_zeros::(1).unwrap(),
+        )
+    };
+
+    let lse_h = ctx.stream.clone_dtoh(&lse).unwrap();
+    let picked_lp_h = ctx.stream.clone_dtoh(&picked_lp).unwrap();
+    let values_h = ctx.stream.clone_dtoh(&values).unwrap();
+    let indices_h = ctx.stream.clone_dtoh(&indices).unwrap();
+    ctx.sync().unwrap();
+
+    jobs.iter()
+        .enumerate()
+        .map(|(out_idx, job)| {
+            let lse = lse_h[out_idx];
+            let top = if job.top_k > 0 {
+                let base = out_idx * k_dev;
+                (0..job.top_k)
+                    .map(|t| (indices_h[base + t] as u32, values_h[base + t].to_f32() - lse))
+                    .collect()
+            } else {
+                Vec::new()
+            };
+            (picked_lp_h[out_idx], top)
+        })
+        .collect()
+}
+
+/// Host reference for one arena row, reading the bf16 values back so both
+/// sides see bit-identical inputs.
+fn host_reference(
+    ctx: &DeviceContext,
+    hs: &HiddenStates,
+    jobs: &[Job],
+) -> Vec<(f32, Vec<(u32, f32)>)> {
+    let all = ctx.stream.clone_dtoh(&hs.data).unwrap();
+    ctx.sync().unwrap();
+    let vocab = hs.hidden_dim;
+    jobs.iter()
+        .map(|job| {
+            let row: Vec = all[job.row as usize * vocab..(job.row as usize + 1) * vocab]
+                .iter()
+                .map(|x| x.to_f32())
+                .collect();
+            let r = token_logprob_from_row(&row, job.picked, job.top_k).unwrap();
+            (r.logprob, r.top_logprobs)
+        })
+        .collect()
+}
+
+fn assert_matches(device: &[(f32, Vec<(u32, f32)>)], reference: &[(f32, Vec<(u32, f32)>)]) {
+    assert_eq!(device.len(), reference.len());
+    for (i, (d, r)) in device.iter().zip(reference).enumerate() {
+        assert!(
+            (d.0 - r.0).abs() <= LP_TOLERANCE,
+            "job {i}: picked logprob {} vs reference {}",
+            d.0,
+            r.0
+        );
+        let d_ids: Vec = d.1.iter().map(|t| t.0).collect();
+        let r_ids: Vec = r.1.iter().map(|t| t.0).collect();
+        assert_eq!(d_ids, r_ids, "job {i}: top-k token ids differ");
+        for (t, (dt, rt)) in d.1.iter().zip(&r.1).enumerate() {
+            assert!(
+                (dt.1 - rt.1).abs() <= LP_TOLERANCE,
+                "job {i} top {t}: logprob {} vs reference {}",
+                dt.1,
+                rt.1
+            );
+        }
+    }
+}
+
+#[test]
+fn golden_random_rows_match_host_reference() {
+    let ctx = DeviceContext::new().unwrap();
+    let vocab = 151_936;
+    let mut rng = Rng(0x5EED_1234_5678_9ABC);
+    let rows: Vec> = (0..8)
+        .map(|_| (0..vocab).map(|_| rng.next_f32(-20.0, 20.0)).collect())
+        .collect();
+    let arena = make_arena(&ctx, &rows);
+
+    let jobs = vec![
+        Job { row: 0, picked: 1234, top_k: 5 },
+        Job { row: 3, picked: 0, top_k: 1 },
+        Job { row: 7, picked: vocab as u32 - 1, top_k: 20 },
+        Job { row: 5, picked: 999, top_k: 0 },
+    ];
+    let device = run_device_batch(&ctx, &arena, &jobs);
+    let reference = host_reference(&ctx, &arena, &jobs);
+    assert_matches(&device, &reference);
+}
+
+#[test]
+fn golden_tie_heavy_rows_keep_index_order() {
+    let ctx = DeviceContext::new().unwrap();
+    let vocab = 4096;
+    // Only a handful of distinct values: massive ties, including at the top-k
+    // boundary — the device must keep (value desc, token id asc) order and the
+    // smallest-id tie-break selection, exactly like the host insertion pass.
+    let mut rng = Rng(0xDEAD_BEEF_0BAD_F00D);
+    let rows: Vec> = (0..4)
+        .map(|_| {
+            (0..vocab)
+                .map(|_| [3.0f32, 1.5, 0.0, -2.5][(rng.next_f32(0.0, 4.0) as usize) % 4])
+                .collect()
+        })
+        .collect();
+    let arena = make_arena(&ctx, &rows);
+
+    let jobs = vec![
+        Job { row: 0, picked: 4000, top_k: 3 },
+        Job { row: 1, picked: 7, top_k: 100 },
+        Job { row: 2, picked: 1, top_k: 1 },
+        Job { row: 3, picked: 2048, top_k: 2048 },
+    ];
+    let device = run_device_batch(&ctx, &arena, &jobs);
+    let reference = host_reference(&ctx, &arena, &jobs);
+    assert_matches(&device, &reference);
+}
+
+#[test]
+fn golden_chosen_only_rows_match() {
+    let ctx = DeviceContext::new().unwrap();
+    let vocab = 32_000;
+    let mut rng = Rng(0x0BAD_5EED_CAFE_F00D);
+    let rows: Vec> = (0..3)
+        .map(|_| (0..vocab).map(|_| rng.next_f32(-10.0, 10.0)).collect())
+        .collect();
+    let arena = make_arena(&ctx, &rows);
+
+    // top_k = 0: chosen-token logprob only, empty alternatives.
+    let jobs: Vec = (0..3)
+        .map(|i| Job {
+            row: i,
+            picked: (i * 977 + 13) % vocab as u32,
+            top_k: 0,
+        })
+        .collect();
+    let device = run_device_batch(&ctx, &arena, &jobs);
+    let reference = host_reference(&ctx, &arena, &jobs);
+    assert_matches(&device, &reference);
+}
diff --git a/openinfer-server/src/bin/bench_serving/decode.rs b/openinfer-server/src/bin/bench_serving/decode.rs
index 8ecc2afb3..15aa94f70 100644
--- a/openinfer-server/src/bin/bench_serving/decode.rs
+++ b/openinfer-server/src/bin/bench_serving/decode.rs
@@ -286,8 +286,8 @@ fn measure_decode_stream(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .map_err(|e| anyhow!("scheduler submit failed: {e}"))?;
 
diff --git a/openinfer-server/src/bin/bench_serving/exec.rs b/openinfer-server/src/bin/bench_serving/exec.rs
index 13b8421ca..e1642f6fd 100644
--- a/openinfer-server/src/bin/bench_serving/exec.rs
+++ b/openinfer-server/src/bin/bench_serving/exec.rs
@@ -131,8 +131,8 @@ pub(crate) fn run_scheduler_stream(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .map_err(|e| anyhow::anyhow!("scheduler submit failed: {e}"))?;
 
diff --git a/openinfer-server/src/bin/glm52_step_bench.rs b/openinfer-server/src/bin/glm52_step_bench.rs
index e3e74943b..9520fc5a9 100644
--- a/openinfer-server/src/bin/glm52_step_bench.rs
+++ b/openinfer-server/src/bin/glm52_step_bench.rs
@@ -259,8 +259,8 @@ fn run_stream(
             max_tokens,
             lora_adapter: None,
             token_tx,
-            logprobs: 0,
-            echo: false,
+            logprobs: None,
+            prompt_logprobs: None,
         })
         .map_err(|e| anyhow::anyhow!("scheduler submit failed: {e}"))?;
 
diff --git a/openinfer-sim/src/lib.rs b/openinfer-sim/src/lib.rs
index cc501a7e8..bfcacf785 100644
--- a/openinfer-sim/src/lib.rs
+++ b/openinfer-sim/src/lib.rs
@@ -102,16 +102,32 @@ async fn run_simulated_request(req: GenerateRequest, config: SimulatedEngineConf
         return;
     }
 
-    if req.echo
-        && req
+    if let Some(top_k) = req.prompt_logprobs {
+        let logprobs = req
+            .prompt_tokens
+            .iter()
+            .enumerate()
+            .map(|(index, &id)| {
+                // The leading prompt token has no predecessor logits.
+                if index == 0 {
+                    return None;
+                }
+                Some(TokenLogprob {
+                    logprob: -0.5,
+                    top_logprobs: dummy_top_logprobs(id, top_k),
+                })
+            })
+            .collect();
+        if req
             .token_tx
             .send(TokenEvent::PromptTokens {
                 ids: req.prompt_tokens.clone(),
-                logprobs: vec![None; req.prompt_tokens.len()],
+                logprobs,
             })
             .is_err()
-    {
-        return;
+        {
+            return;
+        }
     }
 
     let script = &config.scripted_completion;
@@ -130,15 +146,15 @@ async fn run_simulated_request(req: GenerateRequest, config: SimulatedEngineConf
             tokio::time::sleep(config.tpot()).await;
         }
 
-        let logprob = (req.logprobs > 0).then_some(TokenLogprob {
-            logprob: 0.0,
-            top_logprobs: Vec::new(),
-        });
         let id = if script.is_empty() {
             fake_token_id(&req.prompt_tokens, index, config.fallback_token_id)
         } else {
             script[index]
         };
+        let logprob = req.logprobs.map(|top_k| TokenLogprob {
+            logprob: -0.25,
+            top_logprobs: dummy_top_logprobs(id, top_k),
+        });
         if req
             .token_tx
             .send(TokenEvent::Token { id, logprob })
@@ -168,6 +184,13 @@ fn fake_token_id(prompt_tokens: &[u32], index: usize, fallback_token_id: u32) ->
     prompt_tokens[index % prompt_tokens.len()]
 }
 
+/// `top_k` fabricated alternatives that never collide with the scored token.
+fn dummy_top_logprobs(scored_id: u32, top_k: usize) -> Vec<(u32, f32)> {
+    (0..top_k)
+        .map(|j| (scored_id.wrapping_add(j as u32 + 1), -1.0 - j as f32))
+        .collect()
+}
+
 fn duration_from_ms(ms: f64) -> Duration {
     Duration::from_secs_f64(ms / 1000.0)
 }
@@ -212,8 +235,8 @@ mod tests {
                 max_tokens: 8,
                 lora_adapter: None,
                 token_tx,
-                logprobs: 0,
-                echo: false,
+                logprobs: None,
+                prompt_logprobs: None,
             },
             config,
         )
@@ -257,8 +280,8 @@ mod tests {
                 max_tokens: 2,
                 lora_adapter: None,
                 token_tx,
-                logprobs: 0,
-                echo: false,
+                logprobs: None,
+                prompt_logprobs: None,
             },
             config,
         )
@@ -309,8 +332,8 @@ mod tests {
                 max_tokens: 3,
                 lora_adapter: None,
                 token_tx,
-                logprobs: 1,
-                echo: false,
+                logprobs: Some(1),
+                prompt_logprobs: None,
             },
             config,
         )
diff --git a/openinfer-sim/tests/frontend_e2e.rs b/openinfer-sim/tests/frontend_e2e.rs
index 567ad5646..d0707f7e3 100644
--- a/openinfer-sim/tests/frontend_e2e.rs
+++ b/openinfer-sim/tests/frontend_e2e.rs
@@ -746,6 +746,227 @@ async fn simulated_frontend_metadata_contract_is_executable() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn echo_with_logprobs_returns_prompt_and_completion_logprobs() -> Result<()> {
+    let server = SimServer::spawn().await?;
+    let client = test_client()?;
+    // Non-streaming echo=true + logprobs asks the engine for prompt logprobs
+    // via the vLLM lowering (prompt_logprobs := logprobs). Previously the
+    // bridge dropped the request field and response assembly 500'd.
+    let mut body = completion_body(&server.model_name, false);
+    body["echo"] = json!(true);
+    body["logprobs"] = json!(2);
+    body["return_tokens_as_token_ids"] = json!(true);
+
+    let completion = post_completion_body(&client, &server.base_url, &body).await?;
+    let choice = &completion["choices"][0];
+    let logprobs = &choice["logprobs"];
+    let tokens = logprobs["tokens"]
+        .as_array()
+        .ok_or_else(|| anyhow!("echo logprobs missing tokens: {completion}"))?;
+    // 2 prompt tokens + 3 completion tokens in one concatenated payload.
+    assert_eq!(
+        tokens.len(),
+        5,
+        "prompt + completion positions: {completion}"
+    );
+    assert!(
+        logprobs["token_logprobs"][0].is_null(),
+        "leading prompt position has no predecessor logprob: {completion}"
+    );
+    assert!(logprobs["top_logprobs"][0].is_null());
+    for index in 1..5 {
+        assert!(
+            logprobs["token_logprobs"][index].is_number(),
+            "position {index} must carry the scored token's logprob: {completion}"
+        );
+        assert_eq!(
+            logprobs["top_logprobs"][index]
+                .as_object()
+                .map(serde_json::Map::len),
+            Some(3),
+            "position {index} must carry scored + top-2 entries: {completion}"
+        );
+    }
+    // echo lowers to prompt_logprobs, so the choice-level field is set too.
+    let prompt_logprobs = choice["prompt_logprobs"]
+        .as_array()
+        .ok_or_else(|| anyhow!("echo response missing prompt_logprobs: {completion}"))?;
+    assert_eq!(prompt_logprobs.len(), 2);
+    assert!(prompt_logprobs[0].is_null());
+    assert!(prompt_logprobs[1].is_object());
+
+    server.shutdown().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn explicit_prompt_logprobs_returns_scored_prompt_positions() -> Result<()> {
+    let server = SimServer::spawn().await?;
+    let client = test_client()?;
+    // The previously-500 case: an explicit multi-token prompt_logprobs request
+    // must come back with one map per prompt position instead of an error.
+    let mut body = completion_body(&server.model_name, false);
+    body["prompt"] = json!([1, 2, 1]);
+    body["prompt_logprobs"] = json!(2);
+    body["return_tokens_as_token_ids"] = json!(true);
+
+    let completion = post_completion_body(&client, &server.base_url, &body).await?;
+    let choice = &completion["choices"][0];
+    let prompt_logprobs = choice["prompt_logprobs"]
+        .as_array()
+        .ok_or_else(|| anyhow!("missing prompt_logprobs: {completion}"))?;
+    assert_eq!(
+        prompt_logprobs.len(),
+        3,
+        "one entry per prompt token: {completion}"
+    );
+    assert!(
+        prompt_logprobs[0].is_null(),
+        "leading position: {completion}"
+    );
+    for (index, position) in prompt_logprobs.iter().enumerate().skip(1) {
+        let map = position
+            .as_object()
+            .ok_or_else(|| anyhow!("position {index} must be a map: {completion}"))?;
+        assert_eq!(
+            map.len(),
+            3,
+            "scored + top-2 at position {index}: {completion}"
+        );
+        let scored = format!("token_id:{}", [1, 2, 1][index]);
+        assert!(
+            map.contains_key(&scored),
+            "position {index} must contain its scored token {scored}: {completion}"
+        );
+    }
+    // Completion logprobs were not requested.
+    assert!(
+        choice.get("logprobs").is_none_or(Value::is_null),
+        "no completion logprobs without a logprobs request: {completion}"
+    );
+
+    server.shutdown().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn explicit_zero_logprobs_returns_scored_token_only() -> Result<()> {
+    let server = SimServer::spawn().await?;
+    let client = test_client()?;
+    // logprobs=0 is a real request (scored token, no alternatives) — not the
+    // disabled value. Response assembly used to expect a logprob the engine
+    // never emitted.
+    let mut body = completion_body(&server.model_name, false);
+    body["logprobs"] = json!(0);
+    body["return_tokens_as_token_ids"] = json!(true);
+
+    let completion = post_completion_body(&client, &server.base_url, &body).await?;
+    let logprobs = &completion["choices"][0]["logprobs"];
+    let token_logprobs = logprobs["token_logprobs"]
+        .as_array()
+        .ok_or_else(|| anyhow!("logprobs=0 missing token_logprobs: {completion}"))?;
+    assert_eq!(
+        token_logprobs.len(),
+        3,
+        "one per completion token: {completion}"
+    );
+    let top_logprobs = logprobs["top_logprobs"].as_array().unwrap();
+    for (index, position) in top_logprobs.iter().enumerate() {
+        assert!(
+            token_logprobs[index].is_number(),
+            "position {index}: {completion}"
+        );
+        assert_eq!(
+            position.as_object().map(serde_json::Map::len),
+            Some(1),
+            "logprobs=0 yields exactly the scored entry at position {index}: {completion}"
+        );
+    }
+
+    server.shutdown().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn explicit_zero_prompt_logprobs_returns_scored_token_only() -> Result<()> {
+    let server = SimServer::spawn().await?;
+    let client = test_client()?;
+    let mut body = completion_body(&server.model_name, false);
+    body["prompt_logprobs"] = json!(0);
+    body["return_tokens_as_token_ids"] = json!(true);
+
+    let completion = post_completion_body(&client, &server.base_url, &body).await?;
+    let prompt_logprobs = completion["choices"][0]["prompt_logprobs"]
+        .as_array()
+        .ok_or_else(|| anyhow!("prompt_logprobs=0 missing payload: {completion}"))?;
+    assert_eq!(prompt_logprobs.len(), 2);
+    assert!(prompt_logprobs[0].is_null());
+    assert_eq!(
+        prompt_logprobs[1].as_object().map(serde_json::Map::len),
+        Some(1),
+        "prompt_logprobs=0 yields exactly the scored entry: {completion}"
+    );
+
+    server.shutdown().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn single_token_prompt_logprobs_answers_from_the_request() -> Result<()> {
+    let server = SimServer::spawn().await?;
+    let client = test_client()?;
+    let mut body = completion_body(&server.model_name, false);
+    body["prompt"] = json!([1]);
+    body["prompt_logprobs"] = json!(1);
+    body["return_tokens_as_token_ids"] = json!(true);
+
+    let completion = post_completion_body(&client, &server.base_url, &body).await?;
+    let prompt_logprobs = completion["choices"][0]["prompt_logprobs"]
+        .as_array()
+        .ok_or_else(|| anyhow!("single-token prompt_logprobs missing: {completion}"))?;
+    assert_eq!(prompt_logprobs.len(), 1);
+    assert!(prompt_logprobs[0].is_null());
+
+    server.shutdown().await
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn full_vocabulary_logprobs_fails_before_gpu_work() -> Result<()> {
+    let server = SimServer::spawn().await?;
+    let client = test_client()?;
+    // The -1 full-vocabulary sentinel used to be silently degraded to
+    // "disabled"; now the bridge rejects the request at admission instead.
+    let mut body = completion_body(&server.model_name, false);
+    body["prompt_logprobs"] = json!(-1);
+
+    let response = client
+        .post(format!("{}/v1/completions", server.base_url))
+        .header(reqwest::header::CONTENT_TYPE, "application/json")
+        .body(body.to_string())
+        .send()
+        .await?;
+    let status = response.status();
+    let response_body = response.text().await?;
+    if status.is_success() {
+        bail!("prompt_logprobs=-1 unexpectedly succeeded: {response_body}");
+    }
+
+    // The endpoint stays healthy for well-formed requests afterwards.
+    assert_non_streaming_completion_has_output(&client, &server.base_url, &server.model_name)
+        .await?;
+    server.shutdown().await
+}
+
+async fn post_completion_body(client: &Client, base_url: &str, body: &Value) -> Result {
+    client
+        .post(format!("{base_url}/v1/completions"))
+        .header(reqwest::header::CONTENT_TYPE, "application/json")
+        .body(body.to_string())
+        .send()
+        .await?
+        .error_for_status()?
+        .json()
+        .await
+        .context("failed to parse completion response")
+}
+
 async fn assert_models_endpoint(client: &Client, base_url: &str, model_name: &str) -> Result<()> {
     let models: Value = client
         .get(format!("{base_url}/v1/models"))
@@ -992,5 +1213,5 @@ const TINY_TOKENIZER_CONFIG_JSON: &str = r#"{
 const TINY_CONFIG_JSON: &str = r#"{
   "model_type": "openinfer_sim",
   "max_position_embeddings": 128,
-  "vocab_size": 3
+  "vocab_size": 16
 }"#;
diff --git a/openinfer-vllm-frontend/src/bridge.rs b/openinfer-vllm-frontend/src/bridge.rs
index 4417da758..6d1a8864d 100644
--- a/openinfer-vllm-frontend/src/bridge.rs
+++ b/openinfer-vllm-frontend/src/bridge.rs
@@ -32,7 +32,7 @@ use openinfer_engine::engine::{
 
 use crate::wire::{
     convert_finish_reason, convert_sampling, lora_adapter_from_sampling_params, requested_logprobs,
-    to_wire_position_logprobs,
+    requested_prompt_logprobs, to_wire_position_logprobs,
 };
 
 pub(crate) struct LocalEngineBridge {
@@ -319,7 +319,7 @@ impl LocalEngineBridge {
                 output_tx,
                 request_id,
                 EngineCoreFinishReason::Error,
-                None,
+                Some(StopReason::Text(unsupported)),
                 None,
                 None,
             )?;
@@ -357,7 +357,7 @@ impl LocalEngineBridge {
                 lora_adapter,
                 token_tx,
                 logprobs: requested_logprobs(&sampling_params),
-                echo: false,
+                prompt_logprobs: requested_prompt_logprobs(&sampling_params),
             })
             .context("failed to submit request to scheduler")?;
 
@@ -474,6 +474,7 @@ fn reduce_request(
     let mut token_ids: Vec = Vec::new();
     let mut positions: Vec = Vec::new();
     let mut has_logprobs = false;
+    let mut prompt_logprobs: Option = None;
     let mut finish_reason: Option = None;
     let mut stop_reason: Option = None;
     let mut terminated = false;
@@ -518,8 +519,31 @@ fn reduce_request(
                     });
                 }
             }
-            TokenEvent::PromptTokens { .. } => {
-                // Prompt logprobs are intentionally deferred for this bridge.
+            TokenEvent::PromptTokens { ids, logprobs } => {
+                // The wire payload carries one position per *scored* prompt
+                // token (every prompt token except the leading one, which has
+                // no predecessor logits); the frontend prepends the leading
+                // `None` itself. A single-token prompt scores nothing, so we
+                // publish no payload and let the frontend's one-token branch
+                // answer instead.
+                if ids.len() > 1 {
+                    let positions: Vec = ids
+                        .iter()
+                        .zip(logprobs.iter())
+                        .skip(1)
+                        .filter_map(|(&id, logprob)| {
+                            let position = to_wire_position_logprobs(id, logprob.clone());
+                            if position.is_none() {
+                                warn!(
+                                    "request {request_id}: missing prompt logprob for a \
+                                     scored position; dropping it"
+                                );
+                            }
+                            position
+                        })
+                        .collect();
+                    prompt_logprobs = Some(Logprobs { positions });
+                }
             }
             TokenEvent::Finished {
                 finish_reason: fr, ..
@@ -544,7 +568,7 @@ fn reduce_request(
         }
     }
 
-    if token_ids.is_empty() && !terminated {
+    if token_ids.is_empty() && prompt_logprobs.is_none() && !terminated {
         return (None, false);
     }
 
@@ -553,6 +577,7 @@ fn reduce_request(
         request_id.to_string(),
         token_ids,
         logprobs,
+        prompt_logprobs.map(MaybeWireLogprobs::Direct),
         finish_reason,
         stop_reason,
         state.first_token_events.take(),
@@ -633,6 +658,7 @@ fn send_terminal_output(
                 request_id.clone(),
                 Vec::new(),
                 None,
+                None,
                 Some(finish_reason),
                 stop_reason,
                 events,
@@ -688,6 +714,7 @@ fn engine_output(
     request_id: String,
     new_token_ids: Vec,
     new_logprobs: Option,
+    new_prompt_logprobs_tensors: Option,
     finish_reason: Option,
     stop_reason: Option,
     events: Option>,
@@ -697,7 +724,7 @@ fn engine_output(
         request_id,
         new_token_ids,
         new_logprobs,
-        new_prompt_logprobs_tensors: None,
+        new_prompt_logprobs_tensors,
         pooling_output: None,
         finish_reason,
         stop_reason,
diff --git a/openinfer-vllm-frontend/src/bridge/tests.rs b/openinfer-vllm-frontend/src/bridge/tests.rs
index b112d9617..f22ac25b5 100644
--- a/openinfer-vllm-frontend/src/bridge/tests.rs
+++ b/openinfer-vllm-frontend/src/bridge/tests.rs
@@ -452,6 +452,124 @@ fn rejected_request_is_reported_as_error() {
     );
 }
 
+/// `PromptTokens` becomes the prompt-logprobs wire payload: one position per
+/// *scored* prompt token (the leading prompt token is excluded — the frontend
+/// prepends its `None` itself), each position led by the actual prompt token
+/// at that offset. The payload rides the same coalesced output as the burst's
+/// tokens and lands exactly once.
+#[test]
+fn prompt_tokens_become_prompt_logprobs_payload() {
+    let mut d = Demux::new();
+    d.add("req-prompt");
+    d.emit(
+        "req-prompt",
+        TokenEvent::PromptTokens {
+            ids: vec![9, 8, 7],
+            logprobs: vec![
+                None,
+                Some(TokenLogprob {
+                    logprob: -0.3,
+                    top_logprobs: vec![(8, -0.3), (1, -2.0)],
+                }),
+                Some(TokenLogprob {
+                    logprob: -1.1,
+                    top_logprobs: vec![(7, -1.1), (2, -1.4)],
+                }),
+            ],
+        },
+    );
+    d.emit(
+        "req-prompt",
+        TokenEvent::Token {
+            id: 42,
+            logprob: None,
+        },
+    );
+    assert!(d.drain());
+
+    let batch = d.next_output().expect("coalesced output");
+    assert_eq!(batch.outputs.len(), 1);
+    let output = &batch.outputs[0];
+    assert_eq!(output.new_token_ids, vec![42]);
+    let direct = match output
+        .new_prompt_logprobs_tensors
+        .as_ref()
+        .expect("prompt logprobs payload")
+    {
+        MaybeWireLogprobs::Direct(direct) => direct,
+        MaybeWireLogprobs::Wire(_) => panic!("expected direct prompt logprobs"),
+    };
+    assert_eq!(direct.positions.len(), 2, "leading prompt token excluded");
+    let first = &direct.positions[0].entries;
+    assert_eq!(
+        first[0].token_id, 8,
+        "scored prompt token leads its position"
+    );
+    assert_eq!(
+        first.len(),
+        3,
+        "chosen + full top-k (rectangular wire width, sampled token keeps its slot)"
+    );
+    let second = &direct.positions[1].entries;
+    assert_eq!(second[0].token_id, 7);
+    assert!(d.next_output().is_none());
+}
+
+/// A lone `PromptTokens` burst (no sampled token yet) still flushes the
+/// payload immediately instead of parking it — the frontend reads prompt
+/// logprobs off the request's first output.
+#[test]
+fn lone_prompt_tokens_flush_their_payload() {
+    let mut d = Demux::new();
+    d.add("req-lone");
+    d.emit(
+        "req-lone",
+        TokenEvent::PromptTokens {
+            ids: vec![5, 6],
+            logprobs: vec![
+                None,
+                Some(TokenLogprob {
+                    logprob: -0.7,
+                    top_logprobs: Vec::new(),
+                }),
+            ],
+        },
+    );
+    assert!(d.drain());
+
+    let batch = d.next_output().expect("payload-only output");
+    let output = &batch.outputs[0];
+    assert!(output.new_token_ids.is_empty());
+    assert!(output.new_prompt_logprobs_tensors.is_some());
+    assert!(output.finish_reason.is_none());
+}
+
+/// A single-token prompt scores no position: no payload is published (the
+/// frontend answers the one-token case from the request itself).
+#[test]
+fn single_token_prompt_publishes_no_payload() {
+    let mut d = Demux::new();
+    d.add("req-single");
+    d.emit(
+        "req-single",
+        TokenEvent::PromptTokens {
+            ids: vec![9],
+            logprobs: vec![None],
+        },
+    );
+    d.emit(
+        "req-single",
+        TokenEvent::Token {
+            id: 3,
+            logprob: None,
+        },
+    );
+    assert!(d.drain());
+
+    let batch = d.next_output().expect("token output");
+    assert!(batch.outputs[0].new_prompt_logprobs_tensors.is_none());
+}
+
 /// The scheduler-stats task turns each load-watch snapshot into a stats-only
 /// batch (no request outputs, no finished set) with the queue gauges and the
 /// fractional KV usage the frontend records into Prometheus, sends the current
diff --git a/openinfer-vllm-frontend/src/wire.rs b/openinfer-vllm-frontend/src/wire.rs
index 81016da4e..5aff4b20a 100644
--- a/openinfer-vllm-frontend/src/wire.rs
+++ b/openinfer-vllm-frontend/src/wire.rs
@@ -25,10 +25,11 @@ pub(crate) fn to_wire_position_logprobs(
         logprob: lp.logprob,
         rank: 1,
     });
+    // The msgpack ndarray is rectangular (vLLM's LogprobsTensors is
+    // [positions, max_num_logprobs + 1]), so every position must carry the
+    // full top-k even when the sampled token already appears in it — the
+    // duplicate collapses back out when clients build per-position dicts.
     for (index, (alt_id, alt_logprob)) in lp.top_logprobs.into_iter().enumerate() {
-        if alt_id == token_id {
-            continue;
-        }
         entries.push(WireTokenLogprob {
             token_id: alt_id,
             logprob: alt_logprob,
@@ -107,14 +108,42 @@ pub(crate) fn unsupported_sampling(params: &EngineCoreSamplingParams) -> Option<
             params.repetition_penalty
         ));
     }
+    for (field, value) in [
+        ("logprobs", params.logprobs),
+        ("prompt_logprobs", params.prompt_logprobs),
+    ] {
+        match value {
+            // `-1` means the full vocabulary in the vLLM contract; fail loud
+            // instead of silently degrading it to "disabled".
+            Some(-1) => {
+                return Some(format!(
+                    "{field}=-1 (full-vocabulary logprobs) is not supported yet; \
+                     request a finite top-k count instead"
+                ));
+            }
+            Some(value) if value < -1 => {
+                return Some(format!("{field}={value} is invalid (must be >= -1)"));
+            }
+            _ => {}
+        }
+    }
     None
 }
 
-pub(crate) fn requested_logprobs(params: &EngineCoreSamplingParams) -> usize {
-    params
-        .logprobs
-        .and_then(|value| usize::try_from(value).ok())
-        .unwrap_or(0)
+/// Map the pinned contract's `Option` logprob counts onto the engine's
+/// `Option`: `None` stays disabled, `Some(0)` stays "scored token
+/// only", and `Some(k)` stays top-`k`. Negative values are rejected upstream
+/// by [`unsupported_sampling`], so they are unreachable here.
+fn logprob_count(value: Option) -> Option {
+    value.map(|value| usize::try_from(value).expect("negative logprobs rejected upstream"))
+}
+
+pub(crate) fn requested_logprobs(params: &EngineCoreSamplingParams) -> Option {
+    logprob_count(params.logprobs)
+}
+
+pub(crate) fn requested_prompt_logprobs(params: &EngineCoreSamplingParams) -> Option {
+    logprob_count(params.prompt_logprobs)
 }
 
 pub(crate) fn lora_adapter_from_sampling_params(
@@ -221,6 +250,49 @@ mod tests {
         params.presence_penalty = 0.0;
         params.repetition_penalty = 1.2;
         assert!(unsupported_sampling(¶ms).is_some());
+        params.repetition_penalty = 1.0;
+    }
+
+    #[test]
+    fn unsupported_sampling_rejects_full_vocabulary_and_invalid_logprobs() {
+        let mut params = EngineCoreSamplingParams::for_test();
+        assert_eq!(unsupported_sampling(¶ms), None);
+
+        params.logprobs = Some(-1);
+        let message = unsupported_sampling(¶ms).expect("-1 must be rejected");
+        assert!(message.contains("full-vocabulary"), "{message}");
+
+        params.logprobs = Some(-2);
+        let message = unsupported_sampling(¶ms).expect("<-1 must be rejected");
+        assert!(message.contains("invalid"), "{message}");
+
+        params.logprobs = Some(0);
+        assert_eq!(unsupported_sampling(¶ms), None);
+
+        params.prompt_logprobs = Some(-1);
+        let message = unsupported_sampling(¶ms).expect("prompt -1 must be rejected");
+        assert!(message.contains("prompt_logprobs"), "{message}");
+        params.prompt_logprobs = Some(3);
+        assert_eq!(unsupported_sampling(¶ms), None);
+    }
+
+    #[test]
+    fn requested_logprobs_preserves_the_option_contract() {
+        let mut params = EngineCoreSamplingParams::for_test();
+        assert_eq!(requested_logprobs(¶ms), None);
+        assert_eq!(requested_prompt_logprobs(¶ms), None);
+
+        // Some(0) requests the scored token's logprob with no top entries —
+        // distinct from the disabled value, not an alias for it.
+        params.logprobs = Some(0);
+        params.prompt_logprobs = Some(0);
+        assert_eq!(requested_logprobs(¶ms), Some(0));
+        assert_eq!(requested_prompt_logprobs(¶ms), Some(0));
+
+        params.logprobs = Some(5);
+        params.prompt_logprobs = Some(2);
+        assert_eq!(requested_logprobs(¶ms), Some(5));
+        assert_eq!(requested_prompt_logprobs(¶ms), Some(2));
     }
 
     #[test]
@@ -258,14 +330,47 @@ mod tests {
             MaybeWireLogprobs::Wire(_) => panic!("expected Direct logprobs"),
         };
         assert_eq!(direct.positions.len(), 1);
+        // Rectangular wire shape: the sampled token keeps its top-k slot too,
+        // so every position stays max_num_logprobs + 1 wide.
         let entries = &direct.positions[0].entries;
-        assert_eq!(entries.len(), 2);
+        assert_eq!(entries.len(), 3);
         assert_eq!(entries[0].token_id, 7);
         assert_logprob_eq(entries[0].logprob, -0.5);
         assert_eq!(entries[0].rank, 1);
-        assert_eq!(entries[1].token_id, 42);
-        assert_logprob_eq(entries[1].logprob, -1.5);
-        assert_eq!(entries[1].rank, 2);
+        assert_eq!(entries[1].token_id, 7);
+        assert_logprob_eq(entries[1].logprob, -0.5);
+        assert_eq!(entries[1].rank, 1);
+        assert_eq!(entries[2].token_id, 42);
+        assert_logprob_eq(entries[2].logprob, -1.5);
+        assert_eq!(entries[2].rank, 2);
+    }
+
+    #[test]
+    fn to_wire_logprobs_positions_have_uniform_width() {
+        // Regression test for the engine-core msgpack encode crash: a prompt
+        // logprobs batch whose sampled tokens fall inside the top-k for some
+        // positions and outside it for others must stay rectangular.
+        let sampled_in_topk = to_wire_position_logprobs(
+            7,
+            Some(TokenLogprob {
+                logprob: -0.5,
+                top_logprobs: vec![(7, -0.5), (42, -1.5)],
+            }),
+        )
+        .expect("position");
+        let sampled_outside_topk = to_wire_position_logprobs(
+            1,
+            Some(TokenLogprob {
+                logprob: -14.0,
+                top_logprobs: vec![(7, -0.5), (42, -1.5)],
+            }),
+        )
+        .expect("position");
+        assert_eq!(
+            sampled_in_topk.entries.len(),
+            sampled_outside_topk.entries.len()
+        );
+        assert_eq!(sampled_in_topk.entries.len(), 3);
     }
 
     #[test]