diff --git a/oar-ocr-vl/src/hunyuanocr/dflash.rs b/oar-ocr-vl/src/hunyuanocr/dflash.rs index 581aa02..8ec48e5 100644 --- a/oar-ocr-vl/src/hunyuanocr/dflash.rs +++ b/oar-ocr-vl/src/hunyuanocr/dflash.rs @@ -7,12 +7,12 @@ #[cfg(feature = "cuda")] use super::dynamic_kv::{ - DynamicPagedKvAppend, FusedAddRmsNormBf16, FusedRmsNormRopeBf16, FusedRopeBf16, - FusedSiluMulBf16, + ArgmaxFirstBf16, DynamicPagedKvAppend, FusedAddRmsNormBf16, FusedRmsNormRopeBf16, + FusedRopeBf16, FusedSiluMulBf16, }; use crate::attention::{RotaryEmbedding, flash_attention, scaled_dot_product_attention_gqa}; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; -use candle_core::{DType, Device, Tensor}; +use candle_core::{D, DType, Device, Tensor}; use candle_nn::{Linear, Module, RmsNorm, VarBuilder, linear_no_bias, rms_norm}; use oar_ocr_core::core::OCRError; use serde::Deserialize; @@ -57,14 +57,12 @@ impl DFlashConfig { } fn validate(&self) -> Result<(), OCRError> { - if self.block_size == 0 + if self.block_size < 2 || self.num_hidden_layers == 0 || self.dflash_config.target_layer_ids.is_empty() { return Err(OCRError::ConfigError { - message: - "HunyuanOCR DFlash: block size, layer count, and target layers must be non-zero" - .to_string(), + message: "HunyuanOCR DFlash: block size must be at least 2, and layer count and target layers must be non-zero".to_string(), }); } if !self @@ -115,6 +113,14 @@ impl ContextKv { self.len = 0; } + #[cfg(feature = "cuda")] + fn reset_graph_storage(&mut self) { + self.storage = None; + self.block_table = None; + self.capacity = CONTEXT_KV_INITIAL_CAPACITY; + self.reset(); + } + fn ensure_capacity( &mut self, template_k: &Tensor, @@ -695,11 +701,7 @@ impl DFlashMlp { .narrow(2, self.intermediate_size, self.intermediate_size) .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP up", e))?; #[cfg(feature = "cuda")] - let hidden = if gate.device().is_cuda() - && gate.dtype() == DType::BF16 - && gate.is_contiguous() - && up.is_contiguous() - { + let hidden = if gate.device().is_cuda() && gate.dtype() == DType::BF16 { gate.apply_op2_no_bwd(&up, &FusedSiluMulBf16) .map_err(|e| tensor_err("HunyuanOCR DFlash: fused MLP SiLU*up", e))? } else { @@ -875,7 +877,8 @@ struct DFlashCudaGraph { sin_input: Tensor, _query_lengths: Tensor, kv_lengths: Tensor, - output: Tensor, + hidden_output: Tensor, + proposals_output: Tensor, } #[cfg(feature = "cuda")] @@ -883,7 +886,8 @@ impl std::fmt::Debug for DFlashCudaGraph { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("DFlashCudaGraph") .field("query_input", &self.query_input.shape()) - .field("output", &self.output.shape()) + .field("hidden_output", &self.hidden_output.shape()) + .field("proposals_output", &self.proposals_output.shape()) .finish_non_exhaustive() } } @@ -901,6 +905,7 @@ pub(crate) struct DFlashModel { rope_cos: Tensor, rope_sin: Tensor, context_kv_proj: Linear, + lm_head_weight_t: Tensor, caches: RefCell>, dtype: DType, device: Device, @@ -911,10 +916,25 @@ impl DFlashModel { model_dir: impl AsRef, dtype: DType, device: &Device, + lm_head_weight: &Tensor, ) -> Result { let model_dir = model_dir.as_ref(); let cfg = DFlashConfig::from_path(model_dir.join("config.json"))?; cfg.validate()?; + let (target_vocab_size, target_hidden_size) = lm_head_weight + .dims2() + .map_err(|e| tensor_err("HunyuanOCR DFlash: target LM head shape", e))?; + if target_vocab_size != cfg.vocab_size || target_hidden_size != cfg.hidden_size { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR DFlash: target LM head shape {target_vocab_size}x{target_hidden_size} does not match draft vocab/hidden {}x{}", + cfg.vocab_size, cfg.hidden_size + ), + }); + } + let lm_head_weight_t = lm_head_weight + .transpose(0, 1) + .map_err(|e| tensor_err("HunyuanOCR DFlash: transpose target LM head", e))?; let files = crate::utils::collect_safetensors(model_dir, "HunyuanOCR DFlash")?; // SAFETY: the checkpoint files remain mapped for the lifetime of the tensors. let vb = unsafe { @@ -966,21 +986,12 @@ impl DFlashModel { rope_cos, rope_sin, context_kv_proj, + lm_head_weight_t, caches: RefCell::new(caches), dtype, device: device.clone(), }; - // `DynamicPagedKvAppend` (the graph's per-layer append kernel) only - // accepts BF16; capturing under another dtype override would fail - // model load entirely even though the eager DFlash path below - // supports F16/F32. - #[cfg(feature = "cuda")] - if device.is_cuda() - && dtype == DType::BF16 - && std::env::var_os("OAR_HUNYUAN_DISABLE_CUDA_GRAPH").is_none() - { - model.capture_cuda_graph()?; - } + model.prepare_cuda_graph()?; Ok(model) } @@ -988,6 +999,38 @@ impl DFlashModel { &self.cfg } + pub(crate) fn prepare_cuda_graph(&self) -> Result<(), OCRError> { + #[cfg(feature = "cuda")] + { + if std::env::var_os("OAR_HUNYUAN_DISABLE_CUDA_GRAPH").is_some() { + self.invalidate_cuda_graph(); + return Ok(()); + } + // `DynamicPagedKvAppend` accepts BF16 only. Other dtype + // overrides retain the eager path instead of failing model load. + if self.device.is_cuda() + && self.dtype == DType::BF16 + && self.decode_graph.borrow().is_none() + { + let mut caches = self.caches.borrow_mut(); + if caches.iter().any(|cache| cache.len != 0) { + return Ok(()); + } + // A long-context eager fallback may have grown the physical + // cache beyond the graph's fixed 16K stride. At the next page + // boundary it is safe to shrink and capture fresh pointers. + for cache in caches.iter_mut() { + if cache.capacity != CONTEXT_KV_INITIAL_CAPACITY { + cache.reset_graph_storage(); + } + } + drop(caches); + self.capture_cuda_graph()?; + } + } + Ok(()) + } + #[cfg(feature = "cuda")] fn invalidate_cuda_graph(&self) { // A captured graph owns raw pointers into the fixed-size cache. Once @@ -1137,6 +1180,26 @@ impl DFlashModel { .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "dynamic final norm", e)) } + fn proposals_from_hidden(&self, hidden: &Tensor) -> Result { + let num_spec = self.cfg.block_size.saturating_sub(1); + let draft_rows = hidden + .narrow(1, 1, num_spec) + .and_then(|tensor| tensor.squeeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: draft mask rows", e))?; + let logits = draft_rows + .matmul(&self.lm_head_weight_t) + .map_err(|e| tensor_err("HunyuanOCR DFlash: graph draft logits", e))?; + #[cfg(feature = "cuda")] + if logits.device().is_cuda() && logits.dtype() == DType::BF16 { + return logits + .apply_op1_no_bwd(&ArgmaxFirstBf16) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused draft argmax", e)); + } + logits + .argmax(D::Minus1) + .map_err(|e| tensor_err("HunyuanOCR DFlash: draft argmax", e)) + } + #[cfg(feature = "cuda")] fn capture_cuda_graph(&self) -> Result<(), OCRError> { use candle_core::cuda_backend::cudarc::driver::sys::{ @@ -1197,18 +1260,24 @@ impl DFlashModel { &query_lengths, &kv_lengths, )?; - sync_dflash_graph_tensor(&warm, "warm full draft graph")?; + let warm_proposals = self.proposals_from_hidden(&warm)?; + sync_dflash_graph_tensor(&warm_proposals, "warm full draft graph")?; stream .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) .map_err(|e| dflash_cuda_graph_error("begin full draft graph capture", e))?; - let output = match self.forward_queries_dynamic( - &query_input, - &cos_input, - &sin_input, - &query_lengths, - &kv_lengths, - ) { + let captured_output = (|| { + let hidden = self.forward_queries_dynamic( + &query_input, + &cos_input, + &sin_input, + &query_lengths, + &kv_lengths, + )?; + let proposals = self.proposals_from_hidden(&hidden)?; + Ok::<_, OCRError>((hidden, proposals)) + })(); + let (hidden_output, proposals_output) = match captured_output { Ok(output) => output, Err(error) => { let _ = stream.end_capture( @@ -1228,7 +1297,7 @@ impl DFlashModel { graph .launch() .map_err(|e| dflash_cuda_graph_error("warm full draft graph", e))?; - sync_dflash_graph_tensor(&output, "sync full draft graph")?; + sync_dflash_graph_tensor(&proposals_output, "sync full draft graph")?; self.clear_context(); *self.decode_graph.borrow_mut() = Some(DFlashCudaGraph { graph, @@ -1237,7 +1306,8 @@ impl DFlashModel { sin_input, _query_lengths: query_lengths, kv_lengths, - output, + hidden_output, + proposals_output, }); Ok(()) } @@ -1286,12 +1356,13 @@ impl DFlashModel { .graph .launch() .map_err(|e| dflash_cuda_graph_error("launch full draft graph", e))?; - Ok(Some(captured.output.clone())) + Ok(Some(captured.proposals_output.clone())) } - /// Run the bonus+mask query block. Returns post-norm hidden states for all - /// query positions; the caller samples only rows `1..` (the mask rows). - pub(crate) fn forward_queries(&self, query_embeds: &Tensor) -> Result { + /// Run the bonus+mask query block and project all mask rows into draft + /// proposals. CUDA graph replay includes both the shared LM head and the + /// argmax so they do not incur per-round host dispatch. + pub(crate) fn forward_proposals(&self, query_embeds: &Tensor) -> Result { #[cfg(feature = "cuda")] let _cuda_htod_cache = match &self.device { Device::Cuda(device) => Some(device.enable_cuda_graph_htod_cache()), @@ -1324,9 +1395,11 @@ impl DFlashModel { for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { hidden = layer.forward(&hidden, &cos, &sin, cos_sin.as_ref(), cache)?; } - self.norm + let hidden = self + .norm .forward(&hidden) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "final norm", e)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "final norm", e))?; + self.proposals_from_hidden(&hidden) } pub(crate) fn clear_context(&self) { diff --git a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu index cdd006d..d87a9ff 100644 --- a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu +++ b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu @@ -1,5 +1,6 @@ #include #include +#include #include // Copy one fixed-size verification block into a preallocated KV cache. The @@ -90,17 +91,338 @@ extern "C" __global__ void silu_mul_bf16( const __nv_bfloat16* gate, const __nv_bfloat16* up, __nv_bfloat16* output, - uint32_t count) { + uint32_t count, + uint32_t ncols, + uint32_t gate_row_stride, + uint32_t up_row_stride) { const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; if (index >= count) { return; } + const uint32_t row = index / ncols; + const uint32_t col = index - row * ncols; + const uint32_t gate_index = row * gate_row_stride + col; + const uint32_t up_index = row * up_row_stride + col; const __nv_bfloat16 one = __float2bfloat16_rn(1.0f); - const __nv_bfloat16 negative = -gate[index]; + const __nv_bfloat16 gate_value = gate[gate_index]; + const __nv_bfloat16 negative = -gate_value; const __nv_bfloat16 exponential = hexp(negative); const __nv_bfloat16 denominator = one + exponential; - const __nv_bfloat16 activated = gate[index] / denominator; - output[index] = activated * up[index]; + const __nv_bfloat16 activated = gate_value / denominator; + output[index] = activated * up[up_index]; +} + +// Apply HuggingFace's repetition-penalty rule sparsely. token_ids contains one +// deduplicated row per logits row; UINT32_MAX pads rows to a common width. The +// logits are F32 so this has the same rounding points as the previous CPU +// implementation while avoiding a full-vocabulary device-to-host copy. +extern "C" __global__ void repetition_penalty_f32( + float* logits, + const uint32_t* token_ids, + uint32_t token_stride, + uint32_t token_rows, + uint32_t rows, + uint32_t vocab_size, + float penalty) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = rows * token_stride; + if (index >= count) { + return; + } + const uint32_t row = index / token_stride; + const uint32_t col = index - row * token_stride; + const uint32_t token_row = token_rows == 1 ? 0 : row; + const uint32_t token = token_ids[token_row * token_stride + col]; + if (token >= vocab_size) { + return; + } + const uint32_t offset = row * vocab_size + token; + const float value = logits[offset]; + logits[offset] = value > 0.0f ? __fdiv_rn(value, penalty) + : __fmul_rn(value, penalty); +} + +// Maintain a device-resident presence map for generated tokens. Repetition +// penalty is set-based, so duplicate ids may race while writing the same byte +// value without changing the result. +extern "C" __global__ void mark_repetition_history_u8( + uint8_t* history, + const uint32_t* token_ids, + uint32_t token_count, + uint32_t vocab_size) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= token_count) { + return; + } + const uint32_t token = token_ids[index]; + if (token < vocab_size) { + history[token] = 1; + } +} + +extern "C" __global__ void argmax_first_bf16_stage1( + const __nv_bfloat16* logits, + float* partial_values, + uint32_t* partial_indices, + uint32_t rows, + uint32_t vocab_size, + uint32_t partitions_per_row) { + const uint32_t row = blockIdx.x / partitions_per_row; + const uint32_t partition = blockIdx.x - row * partitions_per_row; + const uint32_t lane = threadIdx.x; + if (row >= rows) { + return; + } + + float best_value = -CUDART_INF_F; + uint32_t best_index = UINT32_MAX; + const uint32_t global_lane = partition * blockDim.x + lane; + const uint32_t global_stride = partitions_per_row * blockDim.x; + for (uint32_t col = global_lane; col < vocab_size; col += global_stride) { + const float value = __bfloat162float(logits[row * vocab_size + col]); + if (value > best_value || (value == best_value && col < best_index)) { + best_value = value; + best_index = col; + } + } + + __shared__ float values[256]; + __shared__ uint32_t indices[256]; + values[lane] = best_value; + indices[lane] = best_index; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + const float other_value = values[lane + width]; + const uint32_t other_index = indices[lane + width]; + if (other_value > values[lane] || + (other_value == values[lane] && other_index < indices[lane])) { + values[lane] = other_value; + indices[lane] = other_index; + } + } + } + if (lane == 0) { + partial_values[blockIdx.x] = values[0]; + partial_indices[blockIdx.x] = indices[0]; + } +} + +// Fuse BF16-to-F32 conversion, set-based repetition penalty, and stable +// argmax for autoregressive sampling. +extern "C" __global__ void repetition_argmax_bf16_stage1( + const __nv_bfloat16* logits, + const uint8_t* history, + float* partial_values, + uint32_t* partial_indices, + uint32_t rows, + uint32_t vocab_size, + uint32_t partitions_per_row, + float penalty) { + const uint32_t row = blockIdx.x / partitions_per_row; + const uint32_t partition = blockIdx.x - row * partitions_per_row; + const uint32_t lane = threadIdx.x; + if (row >= rows) { + return; + } + + float best_value = -CUDART_INF_F; + uint32_t best_index = UINT32_MAX; + const uint32_t global_lane = partition * blockDim.x + lane; + const uint32_t global_stride = partitions_per_row * blockDim.x; + for (uint32_t col = global_lane; col < vocab_size; col += global_stride) { + float value = __bfloat162float(logits[row * vocab_size + col]); + if (history[row * vocab_size + col] != 0) { + value = value > 0.0f ? __fdiv_rn(value, penalty) + : __fmul_rn(value, penalty); + } + if (value > best_value || (value == best_value && col < best_index)) { + best_value = value; + best_index = col; + } + } + + __shared__ float values[256]; + __shared__ uint32_t indices[256]; + values[lane] = best_value; + indices[lane] = best_index; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + const float other_value = values[lane + width]; + const uint32_t other_index = indices[lane + width]; + if (other_value > values[lane] || + (other_value == values[lane] && other_index < indices[lane])) { + values[lane] = other_value; + indices[lane] = other_index; + } + } + } + if (lane == 0) { + partial_values[blockIdx.x] = values[0]; + partial_indices[blockIdx.x] = indices[0]; + } +} + +// DFlash row r sees the common generated history plus proposals [0, r); +// equal maxima keep the lowest vocabulary id. +extern "C" __global__ void dflash_repetition_argmax_bf16_stage1( + const __nv_bfloat16* logits, + const uint8_t* history, + const uint32_t* proposals, + float* partial_values, + uint32_t* partial_indices, + uint32_t proposal_count, + uint32_t rows, + uint32_t vocab_size, + uint32_t partitions_per_row, + float penalty) { + const uint32_t row = blockIdx.x / partitions_per_row; + const uint32_t partition = blockIdx.x - row * partitions_per_row; + const uint32_t lane = threadIdx.x; + if (row >= rows) { + return; + } + + extern __shared__ uint32_t proposal_ids[]; + for (uint32_t i = lane; i < proposal_count; i += blockDim.x) { + proposal_ids[i] = proposals[i]; + } + __syncthreads(); + + const uint32_t prefix_len = row < proposal_count ? row : proposal_count; + float best_value = -CUDART_INF_F; + uint32_t best_index = UINT32_MAX; + const uint32_t global_lane = partition * blockDim.x + lane; + const uint32_t global_stride = partitions_per_row * blockDim.x; + for (uint32_t col = global_lane; col < vocab_size; col += global_stride) { + bool seen = history[col] != 0; + for (uint32_t i = 0; !seen && i < prefix_len; ++i) { + seen = proposal_ids[i] == col; + } + float value = __bfloat162float(logits[row * vocab_size + col]); + if (seen) { + value = value > 0.0f ? __fdiv_rn(value, penalty) + : __fmul_rn(value, penalty); + } + if (value > best_value || (value == best_value && col < best_index)) { + best_value = value; + best_index = col; + } + } + + __shared__ float values[256]; + __shared__ uint32_t indices[256]; + values[lane] = best_value; + indices[lane] = best_index; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + const float other_value = values[lane + width]; + const uint32_t other_index = indices[lane + width]; + if (other_value > values[lane] || + (other_value == values[lane] && other_index < indices[lane])) { + values[lane] = other_value; + indices[lane] = other_index; + } + } + } + if (lane == 0) { + partial_values[blockIdx.x] = values[0]; + partial_indices[blockIdx.x] = indices[0]; + } +} + +extern "C" __global__ void dflash_repetition_argmax_stage2( + const float* partial_values, + const uint32_t* partial_indices, + uint32_t* output, + uint32_t partitions_per_row, + uint32_t rows) { + const uint32_t row = blockIdx.x; + const uint32_t lane = threadIdx.x; + if (row >= rows) { + return; + } + float best_value = -CUDART_INF_F; + uint32_t best_index = UINT32_MAX; + for (uint32_t partition = lane; partition < partitions_per_row; + partition += blockDim.x) { + const float value = + partial_values[row * partitions_per_row + partition]; + const uint32_t index = + partial_indices[row * partitions_per_row + partition]; + if (value > best_value || + (value == best_value && index < best_index)) { + best_value = value; + best_index = index; + } + } + __shared__ float values[32]; + __shared__ uint32_t indices[32]; + values[lane] = best_value; + indices[lane] = best_index; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + const float other_value = values[lane + width]; + const uint32_t other_index = indices[lane + width]; + if (other_value > values[lane] || + (other_value == values[lane] && other_index < indices[lane])) { + values[lane] = other_value; + indices[lane] = other_index; + } + } + } + if (lane == 0) { + output[row] = indices[0] == UINT32_MAX ? 0 : indices[0]; + } +} + +// Candle's parallel argmax does not preserve the lowest vocabulary index when +// equal maxima land in different reduction lanes. BF16 logits produce real +// ties, while the previous CPU loop deliberately kept the first index. Reduce +// (value, index) pairs so GPU repetition penalty remains token-exact. +extern "C" __global__ void argmax_first_f32( + const float* logits, + uint32_t* output, + uint32_t rows, + uint32_t vocab_size) { + const uint32_t row = blockIdx.x; + const uint32_t lane = threadIdx.x; + if (row >= rows) { + return; + } + + float best_value = -CUDART_INF_F; + uint32_t best_index = UINT32_MAX; + for (uint32_t col = lane; col < vocab_size; col += blockDim.x) { + const float value = logits[row * vocab_size + col]; + if (value > best_value || (value == best_value && col < best_index)) { + best_value = value; + best_index = col; + } + } + + __shared__ float values[256]; + __shared__ uint32_t indices[256]; + values[lane] = best_value; + indices[lane] = best_index; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + const float other_value = values[lane + width]; + const uint32_t other_index = indices[lane + width]; + if (other_value > values[lane] || + (other_value == values[lane] && other_index < indices[lane])) { + values[lane] = other_value; + indices[lane] = other_index; + } + } + } + if (lane == 0) { + output[row] = indices[0] == UINT32_MAX ? 0 : indices[0]; + } } // Projected QKV is laid out [token, q+k+v]. Produce the contiguous diff --git a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs index 0e6e6b2..78e39b7 100644 --- a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs +++ b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs @@ -1,7 +1,10 @@ use candle_core::backend::BackendStorage; -use candle_core::{CpuStorage, CustomOp2, CustomOp3, InplaceOp3, Layout, Result, Shape}; +use candle_core::{ + CpuStorage, CustomOp1, CustomOp2, CustomOp3, InplaceOp2, InplaceOp3, Layout, Result, Shape, +}; const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/hunyuan_dynamic_kv.ptx")); +const PARTITIONS_PER_ROW: usize = 8; pub(super) struct DynamicKvAppend { pub query_len: usize, @@ -49,6 +52,508 @@ pub(super) struct FusedAddRmsNormBf16 { pub(super) struct FusedSiluMulBf16; +pub(super) struct RepetitionPenaltyF32 { + pub penalty: f32, +} + +pub(super) struct ArgmaxFirstF32; + +pub(super) struct ArgmaxFirstBf16; + +pub(super) struct MarkRepetitionHistoryU8; + +pub(super) struct DFlashRepetitionArgmaxBf16 { + pub penalty: f32, +} + +pub(super) struct RepetitionArgmaxBf16 { + pub penalty: f32, +} + +impl CustomOp1 for ArgmaxFirstF32 { + fn name(&self) -> &'static str { + "hunyuan-argmax-first-f32" + } + + fn cpu_fwd(&self, _storage: &CpuStorage, _layout: &Layout) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("deterministic argmax is CUDA-only") + } + + fn cuda_fwd( + &self, + storage: &candle_core::CudaStorage, + layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !layout.is_contiguous() || storage.dtype() != candle_core::DType::F32 { + candle_core::bail!("deterministic argmax requires contiguous F32 logits") + } + let (rows, vocab_size) = layout.shape().dims2()?; + if vocab_size == 0 { + candle_core::bail!("deterministic argmax does not support an empty vocabulary") + } + let device = storage.device().clone(); + let logits = storage.as_cuda_slice::()?; + let logits = logits.slice(layout.start_offset()..); + let output = unsafe { device.alloc::(rows) }?; + let function = + device.get_or_load_custom_func("argmax_first_f32", "hunyuan_dynamic_kv", PTX)?; + let mut builder = function.builder(); + builder.arg(&logits); + builder.arg(&output); + candle_core::builder_arg!(builder, rows as u32, vocab_size as u32); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + Shape::from_dims(&[rows]), + )) + } +} + +impl CustomOp1 for ArgmaxFirstBf16 { + fn name(&self) -> &'static str { + "hunyuan-argmax-first-bf16" + } + + fn cpu_fwd(&self, _storage: &CpuStorage, _layout: &Layout) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("deterministic BF16 argmax is CUDA-only") + } + + fn cuda_fwd( + &self, + storage: &candle_core::CudaStorage, + layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !layout.is_contiguous() || storage.dtype() != candle_core::DType::BF16 { + candle_core::bail!("deterministic BF16 argmax requires contiguous logits") + } + let (rows, vocab_size) = layout.shape().dims2()?; + if vocab_size == 0 { + candle_core::bail!("deterministic BF16 argmax does not support an empty vocabulary") + } + let device = storage.device().clone(); + let logits = storage.as_cuda_slice::()?; + let logits = logits.slice(layout.start_offset()..); + let partial_count = rows * PARTITIONS_PER_ROW; + let partial_values = unsafe { device.alloc::(partial_count) }?; + let partial_indices = unsafe { device.alloc::(partial_count) }?; + let output = unsafe { device.alloc::(rows) }?; + let stage1 = device.get_or_load_custom_func( + "argmax_first_bf16_stage1", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = stage1.builder(); + builder.arg(&logits); + builder.arg(&partial_values); + builder.arg(&partial_indices); + candle_core::builder_arg!( + builder, + rows as u32, + vocab_size as u32, + PARTITIONS_PER_ROW as u32 + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (partial_count as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + let stage2 = device.get_or_load_custom_func( + "dflash_repetition_argmax_stage2", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = stage2.builder(); + builder.arg(&partial_values); + builder.arg(&partial_indices); + builder.arg(&output); + candle_core::builder_arg!(builder, PARTITIONS_PER_ROW as u32, rows as u32); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + Shape::from_dims(&[rows]), + )) + } +} + +impl InplaceOp2 for MarkRepetitionHistoryU8 { + fn name(&self) -> &'static str { + "hunyuan-mark-repetition-history-u8" + } + + fn cpu_fwd( + &self, + _history: &mut CpuStorage, + _history_layout: &Layout, + _token_ids: &CpuStorage, + _token_ids_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("repetition-history marking is CUDA-only") + } + + fn cuda_fwd( + &self, + history: &mut candle_core::CudaStorage, + history_layout: &Layout, + token_ids: &candle_core::CudaStorage, + token_ids_layout: &Layout, + ) -> Result<()> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !history_layout.is_contiguous() + || !token_ids_layout.is_contiguous() + || history.dtype() != candle_core::DType::U8 + || token_ids.dtype() != candle_core::DType::U32 + { + candle_core::bail!( + "repetition-history marking requires contiguous U8 history and U32 ids" + ) + } + let vocab_size = history_layout.shape().elem_count(); + let token_count = token_ids_layout.shape().dims1()?; + if token_count == 0 { + return Ok(()); + } + let device = history.device().clone(); + let history = history.as_cuda_slice_mut::()?; + let token_ids = token_ids.as_cuda_slice::()?; + let mut history = history.slice_mut(history_layout.start_offset()..); + let token_ids = token_ids.slice(token_ids_layout.start_offset()..); + let function = device.get_or_load_custom_func( + "mark_repetition_history_u8", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = function.builder(); + builder.arg(&mut history); + builder.arg(&token_ids); + candle_core::builder_arg!(builder, token_count as u32, vocab_size as u32); + unsafe { builder.launch(LaunchConfig::for_num_elems(token_count as u32)) }.w()?; + Ok(()) + } +} + +impl CustomOp2 for RepetitionArgmaxBf16 { + fn name(&self) -> &'static str { + "hunyuan-repetition-argmax-bf16" + } + + fn cpu_fwd( + &self, + _logits: &CpuStorage, + _logits_layout: &Layout, + _history: &CpuStorage, + _history_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("fused repetition argmax is CUDA-only") + } + + fn cuda_fwd( + &self, + logits: &candle_core::CudaStorage, + logits_layout: &Layout, + history: &candle_core::CudaStorage, + history_layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !logits_layout.is_contiguous() || !history_layout.is_contiguous() { + candle_core::bail!("fused repetition argmax requires contiguous tensors") + } + let (rows, vocab_size) = logits_layout.shape().dims2()?; + if history_layout.shape().dims() != [rows, vocab_size] + || logits.dtype() != candle_core::DType::BF16 + || history.dtype() != candle_core::DType::U8 + || !self.penalty.is_finite() + || self.penalty < 1.0 + { + candle_core::bail!( + "fused repetition argmax mismatch logits={:?}/{:?} history={:?}/{:?} penalty={}", + logits_layout.shape(), + logits.dtype(), + history_layout.shape(), + history.dtype(), + self.penalty + ) + } + + let device = logits.device().clone(); + let logits = logits.as_cuda_slice::()?; + let history = history.as_cuda_slice::()?; + let logits = logits.slice(logits_layout.start_offset()..); + let history = history.slice(history_layout.start_offset()..); + let partial_count = rows * PARTITIONS_PER_ROW; + let partial_values = unsafe { device.alloc::(partial_count) }?; + let partial_indices = unsafe { device.alloc::(partial_count) }?; + let output = unsafe { device.alloc::(rows) }?; + let stage1 = device.get_or_load_custom_func( + "repetition_argmax_bf16_stage1", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = stage1.builder(); + builder.arg(&logits); + builder.arg(&history); + builder.arg(&partial_values); + builder.arg(&partial_indices); + candle_core::builder_arg!( + builder, + rows as u32, + vocab_size as u32, + PARTITIONS_PER_ROW as u32, + self.penalty + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (partial_count as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + let stage2 = device.get_or_load_custom_func( + "dflash_repetition_argmax_stage2", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = stage2.builder(); + builder.arg(&partial_values); + builder.arg(&partial_indices); + builder.arg(&output); + candle_core::builder_arg!(builder, PARTITIONS_PER_ROW as u32, rows as u32); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + Shape::from_dims(&[rows]), + )) + } +} + +impl CustomOp3 for DFlashRepetitionArgmaxBf16 { + fn name(&self) -> &'static str { + "hunyuan-dflash-repetition-argmax-bf16" + } + + fn cpu_fwd( + &self, + _logits: &CpuStorage, + _logits_layout: &Layout, + _history: &CpuStorage, + _history_layout: &Layout, + _proposals: &CpuStorage, + _proposals_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("fused DFlash repetition argmax is CUDA-only") + } + + fn cuda_fwd( + &self, + logits: &candle_core::CudaStorage, + logits_layout: &Layout, + history: &candle_core::CudaStorage, + history_layout: &Layout, + proposals: &candle_core::CudaStorage, + proposals_layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !logits_layout.is_contiguous() + || !history_layout.is_contiguous() + || !proposals_layout.is_contiguous() + { + candle_core::bail!("fused DFlash repetition argmax requires contiguous tensors") + } + let (rows, vocab_size) = logits_layout.shape().dims2()?; + let history_size = history_layout.shape().dims1()?; + let proposal_count = proposals_layout.shape().dims1()?; + if rows != proposal_count + 1 + || history_size != vocab_size + || logits.dtype() != candle_core::DType::BF16 + || history.dtype() != candle_core::DType::U8 + || proposals.dtype() != candle_core::DType::U32 + || !self.penalty.is_finite() + || self.penalty < 1.0 + { + candle_core::bail!( + "fused DFlash repetition argmax mismatch logits={:?}/{:?} history={:?}/{:?} proposals={:?}/{:?} penalty={}", + logits_layout.shape(), + logits.dtype(), + history_layout.shape(), + history.dtype(), + proposals_layout.shape(), + proposals.dtype(), + self.penalty + ) + } + + let device = logits.device().clone(); + let logits = logits.as_cuda_slice::()?; + let history = history.as_cuda_slice::()?; + let proposals = proposals.as_cuda_slice::()?; + let logits = logits.slice(logits_layout.start_offset()..); + let history = history.slice(history_layout.start_offset()..); + let proposals = proposals.slice(proposals_layout.start_offset()..); + let partial_count = rows * PARTITIONS_PER_ROW; + let partial_values = unsafe { device.alloc::(partial_count) }?; + let partial_indices = unsafe { device.alloc::(partial_count) }?; + let output = unsafe { device.alloc::(rows) }?; + let stage1 = device.get_or_load_custom_func( + "dflash_repetition_argmax_bf16_stage1", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = stage1.builder(); + builder.arg(&logits); + builder.arg(&history); + builder.arg(&proposals); + builder.arg(&partial_values); + builder.arg(&partial_indices); + candle_core::builder_arg!( + builder, + proposal_count as u32, + rows as u32, + vocab_size as u32, + PARTITIONS_PER_ROW as u32, + self.penalty + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (partial_count as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: (proposal_count * std::mem::size_of::()) as u32, + }) + } + .w()?; + let stage2 = device.get_or_load_custom_func( + "dflash_repetition_argmax_stage2", + "hunyuan_dynamic_kv", + PTX, + )?; + let mut builder = stage2.builder(); + builder.arg(&partial_values); + builder.arg(&partial_indices); + builder.arg(&output); + candle_core::builder_arg!(builder, PARTITIONS_PER_ROW as u32, rows as u32); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + Shape::from_dims(&[rows]), + )) + } +} + +impl InplaceOp2 for RepetitionPenaltyF32 { + fn name(&self) -> &'static str { + "hunyuan-repetition-penalty-f32" + } + + fn cpu_fwd( + &self, + _logits: &mut CpuStorage, + _logits_layout: &Layout, + _token_ids: &CpuStorage, + _token_ids_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("fused repetition penalty is CUDA-only") + } + + fn cuda_fwd( + &self, + logits: &mut candle_core::CudaStorage, + logits_layout: &Layout, + token_ids: &candle_core::CudaStorage, + token_ids_layout: &Layout, + ) -> Result<()> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !logits_layout.is_contiguous() || !token_ids_layout.is_contiguous() { + candle_core::bail!("repetition penalty requires contiguous tensors") + } + let (rows, vocab_size) = logits_layout.shape().dims2()?; + let (token_rows, token_stride) = token_ids_layout.shape().dims2()?; + if !matches!(token_rows, 1) && rows != token_rows + || logits.dtype() != candle_core::DType::F32 + || token_ids.dtype() != candle_core::DType::U32 + || !self.penalty.is_finite() + || self.penalty < 1.0 + { + candle_core::bail!( + "repetition penalty shape/dtype mismatch logits={:?}/{:?} ids={:?}/{:?} penalty={}", + logits_layout.shape(), + logits.dtype(), + token_ids_layout.shape(), + token_ids.dtype(), + self.penalty + ) + } + if token_stride == 0 { + return Ok(()); + } + + let device = logits.device().clone(); + let logits = logits.as_cuda_slice_mut::()?; + let token_ids = token_ids.as_cuda_slice::()?; + let mut logits = logits.slice_mut(logits_layout.start_offset()..); + let token_ids = token_ids.slice(token_ids_layout.start_offset()..); + let count = rows * token_stride; + let function = + device.get_or_load_custom_func("repetition_penalty_f32", "hunyuan_dynamic_kv", PTX)?; + let mut builder = function.builder(); + builder.arg(&mut logits); + builder.arg(&token_ids); + candle_core::builder_arg!( + builder, + token_stride as u32, + token_rows as u32, + rows as u32, + vocab_size as u32, + self.penalty + ); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + Ok(()) + } +} + impl CustomOp2 for FusedSiluMulBf16 { fn name(&self) -> &'static str { "hunyuan-fused-silu-mul-bf16" @@ -74,14 +579,50 @@ impl CustomOp2 for FusedSiluMulBf16 { use candle_core::cuda_backend::WrapErr; use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; - if !gate_layout.is_contiguous() - || !up_layout.is_contiguous() - || gate_layout.shape() != up_layout.shape() + if gate_layout.shape() != up_layout.shape() || gate.dtype() != candle_core::DType::BF16 || up.dtype() != candle_core::DType::BF16 { - candle_core::bail!("fused SiLU*up requires matching contiguous BF16 tensors") + candle_core::bail!("fused SiLU*up requires matching BF16 tensors") } + let dims = gate_layout.shape().dims(); + let Some(&ncols) = dims.last() else { + candle_core::bail!("fused SiLU*up does not support scalar tensors") + }; + if ncols == 0 { + candle_core::bail!("fused SiLU*up does not support empty rows") + } + // A last-dimension narrow of fused [gate, up] projections has unit + // column stride but a row stride of 2*ncols. Accept that layout + // directly so the fusion does not need two materializing copies. + let flat_row_stride = |layout: &Layout| -> Option { + let dims = layout.shape().dims(); + let strides = layout.stride(); + if strides.last().copied()? != 1 { + return None; + } + if dims.len() == 1 { + return Some(dims[0]); + } + let row_dim = dims.len() - 2; + let row_stride = strides[row_dim]; + let mut expected = row_stride.checked_mul(dims[row_dim])?; + for dim in (0..row_dim).rev() { + if dims[dim] > 1 && strides[dim] != expected { + return None; + } + expected = expected.checked_mul(dims[dim])?; + } + Some(row_stride) + }; + let gate_row_stride = flat_row_stride(gate_layout).ok_or_else(|| { + candle_core::Error::Msg(format!( + "unsupported fused SiLU gate layout {gate_layout:?}" + )) + })?; + let up_row_stride = flat_row_stride(up_layout).ok_or_else(|| { + candle_core::Error::Msg(format!("unsupported fused SiLU up layout {up_layout:?}")) + })?; let device = gate.device().clone(); let gate = gate.as_cuda_slice::()?; let up = up.as_cuda_slice::()?; @@ -95,7 +636,13 @@ impl CustomOp2 for FusedSiluMulBf16 { builder.arg(&gate); builder.arg(&up); builder.arg(&output); - candle_core::builder_arg!(builder, count as u32); + candle_core::builder_arg!( + builder, + count as u32, + ncols as u32, + gate_row_stride as u32, + up_row_stride as u32 + ); unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; Ok(( candle_core::CudaStorage::wrap_cuda_slice(output, device), @@ -683,3 +1230,56 @@ impl InplaceOp3 for DynamicPagedKvAppend { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::{DType, Device, Tensor}; + + #[cfg(feature = "cuda")] + #[test] + fn deterministic_bf16_argmax_keeps_first_index() -> Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let mut values = vec![0.0f32; 2 * 1030]; + values[1] = 10.0; + values[1024] = 10.0; + values[1030 + 2] = 7.0; + values[1030 + 1025] = 7.0; + let logits = Tensor::from_vec(values, (2, 1030), &device)?.to_dtype(DType::BF16)?; + let tokens = logits + .apply_op1_no_bwd(&ArgmaxFirstBf16)? + .to_vec1::()?; + assert_eq!(tokens, vec![1, 2]); + Ok(()) + } + + #[test] + fn fused_silu_mul_supports_strided_projection_halves() -> Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let gate_up = Tensor::from_vec( + vec![ + -2.0f32, -1.0, 0.5, 2.0, 3.0, -4.0, // row 0: gate, up + 1.5, -0.5, 4.0, -3.0, 2.0, 0.25, // row 1: gate, up + ], + (1, 2, 6), + &device, + )? + .to_dtype(DType::BF16)?; + let gate = gate_up.narrow(2, 0, 3)?; + let up = gate_up.narrow(2, 3, 3)?; + assert!(!gate.is_contiguous()); + assert!(!up.is_contiguous()); + + let actual = gate.apply_op2_no_bwd(&up, &FusedSiluMulBf16)?; + let expected = (candle_nn::ops::silu(&gate)? * up)?; + assert_eq!( + actual.to_vec3::()?, + expected.to_vec3::()? + ); + Ok(()) + } +} diff --git a/oar-ocr-vl/src/hunyuanocr/llm.rs b/oar-ocr-vl/src/hunyuanocr/llm.rs index 401a0f6..2f6c56e 100644 --- a/oar-ocr-vl/src/hunyuanocr/llm.rs +++ b/oar-ocr-vl/src/hunyuanocr/llm.rs @@ -10,7 +10,7 @@ use crate::kv_trim::TrimmableKvCache; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; #[cfg(feature = "cuda")] use candle_core::Device; -use candle_core::Tensor; +use candle_core::{D, Tensor}; use candle_nn::Module; use oar_ocr_core::core::OCRError; use std::cell::RefCell; @@ -22,6 +22,22 @@ fn target_cuda_graph_dtype_supported(dtype: candle_core::DType) -> bool { matches!(dtype, candle_core::DType::F16 | candle_core::DType::BF16) } +#[cfg(any(feature = "cuda", test))] +fn ar_cuda_graph_capacity(prompt_len: usize, max_new_tokens: usize) -> Option { + if max_new_tokens == 0 || prompt_len >= DECODE_ROPE_CACHE_LEN { + return None; + } + let required = prompt_len + .saturating_add(max_new_tokens) + .min(DECODE_ROPE_CACHE_LEN); + Some( + required + .max(1) + .next_power_of_two() + .min(DECODE_ROPE_CACHE_LEN), + ) +} + /// Apply XDRoPE to `(q, k)` using already-section-mixed F32 `cos`/`sin`. /// /// The section-mix (`select_rope_sections`) and the F32 cast of cos/sin are @@ -114,9 +130,9 @@ fn cuda_graph_error( #[derive(Debug)] struct HunyuanMlp { - gate_proj: candle_nn::Linear, - up_proj: candle_nn::Linear, + gate_up_proj: candle_nn::Linear, down_proj: candle_nn::Linear, + intermediate_size: usize, } impl HunyuanMlp { @@ -130,28 +146,29 @@ impl HunyuanMlp { let down_proj = candle_nn::linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj")) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "load down_proj", e))?; + let gate_up_weight = Tensor::cat(&[gate_proj.weight(), up_proj.weight()], 0) + .and_then(|weight| weight.contiguous()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fuse gate/up weights", e))?; Ok(Self { - gate_proj, - up_proj, + gate_up_proj: candle_nn::Linear::new(gate_up_weight, None), down_proj, + intermediate_size: cfg.intermediate_size, }) } fn forward(&self, xs: &Tensor) -> Result { - let gate = self - .gate_proj - .forward(xs) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate_proj", e))?; - let up = self - .up_proj + let gate_up = self + .gate_up_proj .forward(xs) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp up_proj", e))?; + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate_up_proj", e))?; + let gate = gate_up + .narrow(D::Minus1, 0, self.intermediate_size) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate", e))?; + let up = gate_up + .narrow(D::Minus1, self.intermediate_size, self.intermediate_size) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp up", e))?; #[cfg(feature = "cuda")] - let prod = if gate.device().is_cuda() - && gate.dtype() == candle_core::DType::BF16 - && gate.is_contiguous() - && up.is_contiguous() - { + let prod = if gate.device().is_cuda() && gate.dtype() == candle_core::DType::BF16 { gate.apply_op2_no_bwd(&up, &FusedSiluMulBf16) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused mlp silu*up", e))? } else { @@ -553,7 +570,7 @@ impl HunyuanAttention { } #[cfg(feature = "cuda")] - fn prepare_dynamic_cache(&self, query_len: usize) -> Result<(), OCRError> { + fn prepare_dynamic_cache(&self, query_len: usize, cache_len: usize) -> Result<(), OCRError> { let device = self.qkv_proj.weight().device(); let template = Tensor::zeros( (1, self.num_kv_heads, query_len, self.head_dim), @@ -563,7 +580,7 @@ impl HunyuanAttention { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic KV template", e))?; self.kv_cache .borrow_mut() - .initialize_storage(&template) + .initialize_storage_with_capacity(&template, cache_len) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "initialize dynamic KV", e)) } @@ -586,16 +603,15 @@ impl HunyuanAttention { .to_string(), }); } - let (cache_k, cache_v) = - self.kv_cache - .borrow() - .storage() - .ok_or_else(|| OCRError::ConfigError { - message: "HunyuanOCR dynamic KV storage is not initialized".to_string(), - })?; + let cache = self.kv_cache.borrow(); + let cache_len = cache.storage_capacity(); + let (cache_k, cache_v) = cache.storage().ok_or_else(|| OCRError::ConfigError { + message: "HunyuanOCR dynamic KV storage is not initialized".to_string(), + })?; + drop(cache); let append = DynamicKvAppend { query_len, - cache_len: DECODE_ROPE_CACHE_LEN, + cache_len, }; cache_k .inplace_op3(k, kv_lengths, &append) @@ -623,7 +639,7 @@ impl HunyuanAttention { query_lengths, kv_lengths, query_len, - DECODE_ROPE_CACHE_LEN, + cache_len, self.scaling as f32, true, ) @@ -669,8 +685,10 @@ struct TargetDecoderCudaGraph { _query_lengths: Tensor, kv_lengths: Tensor, hidden_output: Tensor, - aux_output: Tensor, + logits_output: Tensor, + aux_output: Option, aux_layer_ids: Vec, + cache_len: usize, } #[cfg(feature = "cuda")] @@ -679,6 +697,7 @@ impl std::fmt::Debug for TargetDecoderCudaGraph { f.debug_struct("TargetDecoderCudaGraph") .field("hidden", &self.hidden_input.shape()) .field("aux_layer_ids", &self.aux_layer_ids) + .field("cache_len", &self.cache_len) .finish_non_exhaustive() } } @@ -838,13 +857,17 @@ impl HunyuanDecoderLayer { /// DFlash target interface. pub(crate) struct HunyuanLlmOutput { pub hidden_states: Tensor, + /// LM-head output is populated when decode runs through a CUDA graph. + /// Prefill/eager paths leave it absent to avoid projecting every prompt + /// position into the full vocabulary. + pub logits: Option, pub aux_hidden_states: Option, } #[derive(Debug)] pub struct HunyuanLlm { #[cfg(feature = "cuda")] - dflash_decode_graph: RefCell>, + decode_graph: RefCell>, embed_tokens: candle_nn::Embedding, layers: Vec, norm: candle_nn::RmsNorm, @@ -897,7 +920,7 @@ impl HunyuanLlm { Ok(Self { #[cfg(feature = "cuda")] - dflash_decode_graph: RefCell::new(None), + decode_graph: RefCell::new(None), embed_tokens, layers, norm, @@ -974,7 +997,7 @@ impl HunyuanLlm { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode sequence length", e))?; if start + len > DECODE_ROPE_CACHE_LEN { #[cfg(feature = "cuda")] - self.invalidate_dflash_cuda_graph(); + self.invalidate_cuda_graph(); let positions: Vec = (start..start + len) .flat_map(|position| [position as i64; 4]) .collect(); @@ -995,7 +1018,7 @@ impl HunyuanLlm { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode rope sin slice", e))?; #[cfg(feature = "cuda")] if let Some(output) = - self.replay_dflash_cuda_graph(inputs_embeds, &cos, &sin, start + len, aux_layer_ids)? + self.replay_cuda_graph(inputs_embeds, &cos, &sin, start + len, aux_layer_ids)? { return Ok(output); } @@ -1003,7 +1026,7 @@ impl HunyuanLlm { } #[cfg(feature = "cuda")] - fn replay_dflash_cuda_graph( + fn replay_cuda_graph( &self, inputs_embeds: &Tensor, cos: &Tensor, @@ -1011,14 +1034,15 @@ impl HunyuanLlm { kv_len: usize, aux_layer_ids: &[usize], ) -> Result, OCRError> { - if kv_len > DECODE_ROPE_CACHE_LEN { - self.invalidate_dflash_cuda_graph(); - return Ok(None); - } - let captured_ref = self.dflash_decode_graph.borrow(); + let captured_ref = self.decode_graph.borrow(); let Some(captured) = captured_ref.as_ref() else { return Ok(None); }; + if kv_len > captured.cache_len { + drop(captured_ref); + self.invalidate_cuda_graph(); + return Ok(None); + } if inputs_embeds.shape() != captured.hidden_input.shape() || cos.shape() != captured.cos_input.shape() || sin.shape() != captured.sin_input.shape() @@ -1054,7 +1078,8 @@ impl HunyuanLlm { } Ok(Some(HunyuanLlmOutput { hidden_states: captured.hidden_output.clone(), - aux_hidden_states: Some(captured.aux_output.clone()), + logits: Some(captured.logits_output.clone()), + aux_hidden_states: captured.aux_output.clone(), })) } @@ -1071,8 +1096,14 @@ impl HunyuanLlm { let incoming_len = inputs_embeds.dim(1).map_err(|e| { candle_to_ocr_inference("HunyuanOCR", "prepared sequence length", e) })?; - if self.kv_cache_len().saturating_add(incoming_len) > DECODE_ROPE_CACHE_LEN { - self.invalidate_dflash_cuda_graph(); + let required = self.kv_cache_len().saturating_add(incoming_len); + let graph_capacity = self + .decode_graph + .borrow() + .as_ref() + .map(|graph| graph.cache_len); + if graph_capacity.is_some_and(|capacity| required > capacity) { + self.invalidate_cuda_graph(); } } // Decode invokes the same non-contiguous elementwise layouts thousands @@ -1131,6 +1162,7 @@ impl HunyuanLlm { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm final norm", e))?; Ok(HunyuanLlmOutput { hidden_states, + logits: None, aux_hidden_states, }) } @@ -1167,20 +1199,46 @@ impl HunyuanLlm { aux.push(hidden_states.clone()); } } - let refs: Vec<&Tensor> = aux.iter().collect(); - let aux_hidden_states = Tensor::cat(&refs, 2).map_err(|e| { - candle_to_ocr_inference("HunyuanOCR", "dynamic target auxiliary states", e) - })?; + let aux_hidden_states = if aux.is_empty() { + None + } else { + let refs: Vec<&Tensor> = aux.iter().collect(); + Some(Tensor::cat(&refs, 2).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "dynamic target auxiliary states", e) + })?) + }; let hidden_states = self .norm .forward(&hidden_states) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic target final norm", e))?; Ok(HunyuanLlmOutput { hidden_states, - aux_hidden_states: Some(aux_hidden_states), + logits: None, + aux_hidden_states, }) } + #[cfg(feature = "cuda")] + fn project_logits(&self, hidden_states: &Tensor) -> Result { + let (batch_size, sequence_length, hidden_size) = hidden_states + .dims3() + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "graph hidden shape", e))?; + let weight_t = self + .embed_tokens + .embeddings() + .transpose(0, 1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "transpose LM head", e))?; + let vocab_size = weight_t + .dim(1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "LM head vocab size", e))?; + hidden_states + .reshape((batch_size * sequence_length, hidden_size)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "flatten graph hidden", e))? + .matmul(&weight_t) + .and_then(|logits| logits.reshape((batch_size, sequence_length, vocab_size))) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "project graph logits", e)) + } + pub fn clear_kv_cache(&self) { for layer in &self.layers { layer.clear_kv_cache(); @@ -1193,6 +1251,8 @@ impl HunyuanLlm { aux_layer_ids: &[usize], ) -> Result<(), OCRError> { if std::env::var_os("OAR_HUNYUAN_DISABLE_CUDA_GRAPH").is_some() { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); return Ok(()); } // Dynamic target attention pins an F16 KV cache for graph replay. @@ -1202,23 +1262,85 @@ impl HunyuanLlm { if self.decode_cos.device().is_cuda() && target_cuda_graph_dtype_supported(self.embed_tokens.embeddings().dtype()) { + let reusable = self.decode_graph.borrow().as_ref().is_some_and(|graph| { + graph.hidden_input.dim(1).ok() == Some(query_len) + && graph.cache_len == DECODE_ROPE_CACHE_LEN + && graph.aux_layer_ids == aux_layer_ids + }); + if reusable { + return Ok(()); + } + self.invalidate_cuda_graph(); let cos = self.decode_cos.narrow(2, 0, query_len).map_err(|e| { candle_to_ocr_inference("HunyuanOCR", "graph decode cos template", e) })?; let sin = self.decode_sin.narrow(2, 0, query_len).map_err(|e| { candle_to_ocr_inference("HunyuanOCR", "graph decode sin template", e) })?; - self.capture_dflash_cuda_graph(query_len, &cos, &sin, aux_layer_ids)?; + self.capture_cuda_graph(query_len, DECODE_ROPE_CACHE_LEN, &cos, &sin, aux_layer_ids)?; } let _ = query_len; let _ = aux_layer_ids; Ok(()) } + /// Lazily capture the single-token autoregressive decoder. Capacity is + /// selected per workload and rounded up so short requests do not reserve + /// the full 16K-token graph cache. A later larger request transparently + /// recaptures with a larger bucket. + pub(crate) fn prepare_ar_cuda_graph( + &self, + prompt_len: usize, + max_new_tokens: usize, + ) -> Result<(), OCRError> { + if std::env::var_os("OAR_HUNYUAN_DISABLE_CUDA_GRAPH").is_some() + || std::env::var_os("OAR_HUNYUAN_DISABLE_AR_CUDA_GRAPH").is_some() + { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); + return Ok(()); + } + if max_new_tokens == 0 { + return Ok(()); + } + #[cfg(feature = "cuda")] + if self.decode_cos.device().is_cuda() + && target_cuda_graph_dtype_supported(self.embed_tokens.embeddings().dtype()) + { + let Some(cache_len) = ar_cuda_graph_capacity(prompt_len, max_new_tokens) else { + self.invalidate_cuda_graph(); + return Ok(()); + }; + let required = prompt_len + .saturating_add(max_new_tokens) + .min(DECODE_ROPE_CACHE_LEN); + let reusable = self.decode_graph.borrow().as_ref().is_some_and(|graph| { + graph.hidden_input.dim(1).ok() == Some(1) + && graph.cache_len >= required + && graph.aux_layer_ids.is_empty() + }); + if reusable { + return Ok(()); + } + self.invalidate_cuda_graph(); + let cos = self.decode_cos.narrow(2, 0, 1).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "AR graph decode cos template", e) + })?; + let sin = self.decode_sin.narrow(2, 0, 1).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "AR graph decode sin template", e) + })?; + self.capture_cuda_graph(1, cache_len, &cos, &sin, &[])?; + } + let _ = prompt_len; + let _ = max_new_tokens; + Ok(()) + } + #[cfg(feature = "cuda")] - fn capture_dflash_cuda_graph( + fn capture_cuda_graph( &self, query_len: usize, + cache_len: usize, cos_template: &Tensor, sin_template: &Tensor, aux_layer_ids: &[usize], @@ -1227,20 +1349,16 @@ impl HunyuanLlm { CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, }; - if self.dflash_decode_graph.borrow().is_some() { + if self.decode_graph.borrow().is_some() { return Ok(()); } let Device::Cuda(cuda) = self.decode_cos.device() else { return Ok(()); }; - if aux_layer_ids.is_empty() { - return Err(OCRError::ConfigError { - message: "HunyuanOCR DFlash CUDA graph requires auxiliary target layers" - .to_string(), - }); - } for layer in &self.layers { - layer.self_attn.prepare_dynamic_cache(query_len)?; + layer + .self_attn + .prepare_dynamic_cache(query_len, cache_len)?; } let hidden_size = self .embed_tokens @@ -1265,7 +1383,7 @@ impl HunyuanLlm { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph sin input", e))?; let query_lengths = Tensor::new(&[0u32, query_len as u32], self.decode_cos.device()) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph query lengths", e))?; - // These must not share storage: query length stays fixed at 16 while + // These must not share storage: query length stays fixed while // the cumulative KV length is overwritten before every graph replay. let kv_lengths = Tensor::new(&[0u32, query_len as u32], self.decode_cos.device()) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph KV lengths", e))?; @@ -1280,19 +1398,25 @@ impl HunyuanLlm { &kv_lengths, aux_layer_ids, )?; - sync_graph_tensor(&warm.hidden_states, "warm full target decoder graph")?; + let warm_logits = self.project_logits(&warm.hidden_states)?; + sync_graph_tensor(&warm_logits, "warm full target decoder graph")?; stream .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) .map_err(|e| cuda_graph_error("begin full target decoder graph capture", e))?; - let output = match self.forward_with_aux_dynamic( - &hidden_input, - &cos_input, - &sin_input, - &query_lengths, - &kv_lengths, - aux_layer_ids, - ) { + let captured_output = (|| { + let output = self.forward_with_aux_dynamic( + &hidden_input, + &cos_input, + &sin_input, + &query_lengths, + &kv_lengths, + aux_layer_ids, + )?; + let logits = self.project_logits(&output.hidden_states)?; + Ok::<_, OCRError>((output, logits)) + })(); + let (output, logits_output) = match captured_output { Ok(output) => output, Err(error) => { let _ = stream.end_capture( @@ -1309,18 +1433,13 @@ impl HunyuanLlm { .ok_or_else(|| OCRError::ConfigError { message: "HunyuanOCR full target decoder capture returned no graph".to_string(), })?; - let aux_output = output - .aux_hidden_states - .ok_or_else(|| OCRError::ConfigError { - message: "HunyuanOCR full target decoder graph produced no auxiliary states" - .to_string(), - })?; + let aux_output = output.aux_hidden_states; graph .launch() .map_err(|e| cuda_graph_error("warm full target decoder graph", e))?; - sync_graph_tensor(&output.hidden_states, "sync full target decoder graph")?; + sync_graph_tensor(&logits_output, "sync full target decoder graph")?; self.clear_kv_cache(); - *self.dflash_decode_graph.borrow_mut() = Some(TargetDecoderCudaGraph { + *self.decode_graph.borrow_mut() = Some(TargetDecoderCudaGraph { graph, hidden_input, cos_input, @@ -1328,18 +1447,20 @@ impl HunyuanLlm { _query_lengths: query_lengths, kv_lengths, hidden_output: output.hidden_states, + logits_output, aux_output, aux_layer_ids: aux_layer_ids.to_vec(), + cache_len, }); Ok(()) } #[cfg(feature = "cuda")] - fn invalidate_dflash_cuda_graph(&self) { + fn invalidate_cuda_graph(&self) { // Eager cache growth reallocates the storage whose raw pointers were // captured by the graph. Dropping it before growth also prevents a // stale replay after the logical cache is reset for another document. - self.dflash_decode_graph.borrow_mut().take(); + self.decode_graph.borrow_mut().take(); } /// Drop any captured target-decoder CUDA graph. @@ -1352,7 +1473,7 @@ impl HunyuanLlm { /// could replay the graph against that freed memory. pub(crate) fn invalidate_target_cuda_graph(&self) { #[cfg(feature = "cuda")] - self.invalidate_dflash_cuda_graph(); + self.invalidate_cuda_graph(); } pub(crate) fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { @@ -1371,7 +1492,7 @@ impl HunyuanLlm { #[cfg(test)] mod tests { - use super::target_cuda_graph_dtype_supported; + use super::{DECODE_ROPE_CACHE_LEN, ar_cuda_graph_capacity, target_cuda_graph_dtype_supported}; use candle_core::DType; #[test] @@ -1381,4 +1502,16 @@ mod tests { assert!(!target_cuda_graph_dtype_supported(DType::F32)); assert!(!target_cuda_graph_dtype_supported(DType::F64)); } + + #[test] + fn ar_graph_capacity_uses_bounded_power_of_two_buckets() { + assert_eq!(ar_cuda_graph_capacity(1500, 256), Some(2048)); + assert_eq!(ar_cuda_graph_capacity(2000, 4096), Some(8192)); + assert_eq!( + ar_cuda_graph_capacity(10_000, 20_000), + Some(DECODE_ROPE_CACHE_LEN) + ); + assert_eq!(ar_cuda_graph_capacity(100, 0), None); + assert_eq!(ar_cuda_graph_capacity(DECODE_ROPE_CACHE_LEN, 1), None); + } } diff --git a/oar-ocr-vl/src/hunyuanocr/model.rs b/oar-ocr-vl/src/hunyuanocr/model.rs index 4ab9e70..00540ce 100644 --- a/oar-ocr-vl/src/hunyuanocr/model.rs +++ b/oar-ocr-vl/src/hunyuanocr/model.rs @@ -2,6 +2,11 @@ use super::config::{HunyuanOcrConfig, HunyuanOcrImageProcessorConfig, HunyuanOcrVersion}; use super::dflash::DFlashModel; +#[cfg(feature = "cuda")] +use super::dynamic_kv::{ + ArgmaxFirstF32, DFlashRepetitionArgmaxBf16, MarkRepetitionHistoryU8, RepetitionArgmaxBf16, + RepetitionPenaltyF32, +}; use super::llm::HunyuanLlm; use super::processing::{HunyuanOcrImageInputs, preprocess_image}; use super::vision::HunyuanVisionModel; @@ -60,7 +65,7 @@ fn load_generation_eos_ids(model_dir: &Path) -> Option> { /// applying the penalty per *occurrence* would compound to `penalty^k` for a /// token repeated `k` times and quickly suppresses legitimate high-frequency /// tokens like `` in a structured HTML page. We dedup before applying. -fn argmax_with_repetition_penalty( +fn argmax_with_repetition_penalty_cpu( logits: &Tensor, seen: &[u32], penalty: f32, @@ -92,6 +97,273 @@ fn argmax_with_repetition_penalty( Ok(best_idx as u32) } +/// Incrementally maintained unique generated-token history. Repetition +/// penalty is presence-based, so keeping another full ordered copy and +/// sorting/deduplicating it for every decode step is unnecessary. +struct RepetitionHistory { + present: Vec, + unique: Vec, +} + +impl RepetitionHistory { + fn new(vocab_size: usize) -> Self { + Self { + present: vec![false; vocab_size], + unique: Vec::new(), + } + } + + fn insert(&mut self, token: u32) { + let index = token as usize; + if index < self.present.len() && !self.present[index] { + self.present[index] = true; + self.unique.push(token); + } + } + + fn contains(&self, token: u32) -> bool { + self.present.get(token as usize).copied().unwrap_or(false) + } + + fn ids(&self) -> &[u32] { + &self.unique + } + + fn is_empty(&self) -> bool { + self.unique.is_empty() + } +} + +/// Apply a common unique history to every row plus a small unique per-row +/// suffix. DFlash uses the generated history as the common set and only the +/// preceding proposal tokens as row suffixes, avoiding 16 copies of the full +/// history on every verification round. +fn batched_argmax_with_unique_repetition_parts( + logits: &Tensor, + common: &[u32], + row_extras: &[&[u32]], + penalty: f32, +) -> Result, OCRError> { + let vocab_size = logits + .dim(D::Minus1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "repetition penalty vocab", e))?; + let rows = logits.elem_count() / vocab_size; + if row_extras.len() != rows { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR: repetition penalty got {} suffix rows for {rows} logits rows", + row_extras.len() + ), + }); + } + let logits = logits + .reshape((rows, vocab_size)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "reshape penalty logits", e))?; + + #[cfg(feature = "cuda")] + if logits.device().is_cuda() { + let adjusted = logits + .to_dtype(DType::F32) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "cast repetition logits", e))?; + let common: Vec = common + .iter() + .copied() + .filter(|&token| (token as usize) < vocab_size) + .collect(); + if !common.is_empty() { + let common_len = common.len(); + let token_ids = + Tensor::from_vec(common, (1, common_len), logits.device()).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "upload common repetition ids", e) + })?; + adjusted + .inplace_op2(&token_ids, &RepetitionPenaltyF32 { penalty }) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "GPU common repetition penalty", e) + })?; + } + + let token_stride = row_extras + .iter() + .map(|extra| extra.len()) + .max() + .unwrap_or(0); + if token_stride > 0 { + let mut token_ids = vec![u32::MAX; rows * token_stride]; + for (row, extra) in row_extras.iter().enumerate() { + let start = row * token_stride; + let mut write = start; + for &token in *extra { + if (token as usize) < vocab_size { + token_ids[write] = token; + write += 1; + } + } + } + let token_ids = Tensor::from_vec(token_ids, (rows, token_stride), logits.device()) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "upload row repetition ids", e) + })?; + adjusted + .inplace_op2(&token_ids, &RepetitionPenaltyF32 { penalty }) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "GPU row repetition penalty", e) + })?; + } + return adjusted + .apply_op1_no_bwd(&ArgmaxFirstF32) + .and_then(|tokens| tokens.to_vec1::()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "GPU penalized argmax", e)); + } + + let mut tokens = Vec::with_capacity(rows); + for (row, extra) in row_extras.iter().enumerate() { + let row_logits = logits + .i((row, ..)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "penalty logits row", e))?; + let mut seen = Vec::with_capacity(common.len() + extra.len()); + seen.extend_from_slice(common); + seen.extend_from_slice(extra); + tokens.push(argmax_with_repetition_penalty_cpu( + &row_logits, + &seen, + penalty, + )?); + } + Ok(tokens) +} + +fn dflash_argmax_with_host_proposals( + logits: &Tensor, + repetition_history: &RepetitionHistory, + proposals: &[u32], + penalty: f32, +) -> Result, OCRError> { + let mut row_extras = Vec::with_capacity(proposals.len() + 1); + let mut proposal_prefix = Vec::with_capacity(proposals.len()); + for prefix_len in 0..=proposals.len() { + row_extras.push(proposal_prefix.clone()); + if let Some(&proposal) = proposals.get(prefix_len) + && !repetition_history.contains(proposal) + && !proposal_prefix.contains(&proposal) + { + proposal_prefix.push(proposal); + } + } + let extra_refs: Vec<&[u32]> = row_extras.iter().map(Vec::as_slice).collect(); + batched_argmax_with_unique_repetition_parts( + logits, + repetition_history.ids(), + &extra_refs, + penalty, + ) +} + +#[cfg(feature = "cuda")] +fn dflash_argmax_with_repetition_penalty( + logits: &Tensor, + history: &Tensor, + proposals: &Tensor, + penalty: f32, +) -> Result, OCRError> { + let vocab_size = logits + .dim(D::Minus1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "penalty vocab", e))?; + let rows = logits.elem_count() / vocab_size; + let proposal_count = proposals + .dims1() + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "proposal count", e))?; + if rows != proposal_count + 1 { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR DFlash: {rows} target rows do not match {proposal_count} proposals" + ), + }); + } + logits + .reshape((rows, vocab_size)) + .and_then(|logits| { + logits.apply_op3_no_bwd(history, proposals, &DFlashRepetitionArgmaxBf16 { penalty }) + }) + .and_then(|tokens| tokens.to_vec1::()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "penalized argmax", e)) +} + +#[cfg(feature = "cuda")] +fn mark_repetition_history(history: &Tensor, token_ids: &[u32]) -> Result<(), OCRError> { + if token_ids.is_empty() { + return Ok(()); + } + let token_ids = Tensor::new(token_ids, history.device()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "upload repetition history ids", e))?; + history + .inplace_op2(&token_ids, &MarkRepetitionHistoryU8) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mark repetition history", e)) +} + +#[cfg(feature = "cuda")] +fn argmax_with_device_repetition_history( + logits: &Tensor, + history: &Tensor, + penalty: f32, +) -> Result { + let vocab_size = logits + .dim(D::Minus1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "device penalty vocab", e))?; + logits + .reshape((1, vocab_size)) + .and_then(|logits| logits.apply_op2_no_bwd(history, &RepetitionArgmaxBf16 { penalty })) + .and_then(|tokens| tokens.i(0)) + .and_then(|tokens| tokens.to_scalar::()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "device penalized argmax", e)) +} + +/// Apply repetition penalty to one or more rows and return one greedy token per +/// row. CUDA keeps both the sparse penalty update and argmax on-device; only +/// the resulting token ids cross back to the host. Other backends retain the +/// reference CPU implementation above. +#[cfg(test)] +fn batched_argmax_with_repetition_penalty( + logits: &Tensor, + seen_rows: &[&[u32]], + penalty: f32, +) -> Result, OCRError> { + let vocab_size = logits + .dim(D::Minus1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "repetition penalty vocab", e))?; + let rows = logits.elem_count() / vocab_size; + if seen_rows.len() != rows { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR: repetition penalty got {} history rows for {rows} logits rows", + seen_rows.len() + ), + }); + } + let mut unique_rows = Vec::with_capacity(rows); + for seen in seen_rows { + let mut unique: Vec = seen + .iter() + .copied() + .filter(|&token| (token as usize) < vocab_size) + .collect(); + unique.sort_unstable(); + unique.dedup(); + unique_rows.push(unique); + } + let unique_refs: Vec<&[u32]> = unique_rows.iter().map(Vec::as_slice).collect(); + batched_argmax_with_unique_repetition_parts(logits, &[], &unique_refs, penalty) +} + +fn argmax_with_unique_repetition_penalty( + logits: &Tensor, + seen: &[u32], + penalty: f32, +) -> Result { + batched_argmax_with_unique_repetition_parts(logits, seen, &[&[]], penalty) + .map(|tokens| tokens[0]) +} + pub struct HunyuanOcr { device: Device, dtype: DType, @@ -193,7 +465,12 @@ impl HunyuanOcr { message: "HunyuanOCR: DFlash requires the 1.5 checkpoint".to_string(), }); } - let dflash = DFlashModel::from_dir(dflash_dir, model.dtype, &model.device)?; + let dflash = DFlashModel::from_dir( + dflash_dir, + model.dtype, + &model.device, + &model.llm.token_embedding_weight(), + )?; if dflash.config().hidden_size != model.cfg.hidden_size || dflash.config().vocab_size != model.cfg.vocab_size { @@ -535,6 +812,24 @@ impl HunyuanOcr { // 5. Prefill self.llm.clear_kv_cache(); let active_dflash = if batch_size == 1 { + if let Some(dflash) = self.dflash.as_ref() { + // A previous multi-image request or an over-capacity decode + // may have invalidated either graph. Recreate them before + // prefill; steady state is an inexpensive capacity check. + dflash.prepare_cuda_graph()?; + let target_layers: Vec = dflash + .config() + .dflash_config + .target_layer_ids + .iter() + .map(|id| id + 1) + .collect(); + self.llm + .prepare_dflash_cuda_graphs(dflash.config().block_size, &target_layers)?; + } else { + self.llm + .prepare_ar_cuda_graph(max_seq_len, max_new_tokens)?; + } self.dflash.as_ref() } else { // Multi-image requests disable DFlash and prefill a batch>1 @@ -542,9 +837,7 @@ impl HunyuanOcr { // captured target CUDA graph points at. Drop the graph before // that reallocation so a later single-image DFlash decode can't // replay it against freed memory. - if self.dflash.is_some() { - self.llm.invalidate_target_cuda_graph(); - } + self.llm.invalidate_target_cuda_graph(); None }; let hidden = if let Some(dflash) = active_dflash { @@ -600,6 +893,23 @@ impl HunyuanOcr { // 7. Autoregressive decode let mut generated: Vec> = vec![Vec::new(); batch_size]; + let mut repetition_histories: Vec = (0..batch_size) + .map(|_| RepetitionHistory::new(self.cfg.vocab_size)) + .collect(); + #[cfg(feature = "cuda")] + let repetition_history_device = if batch_size == 1 + && self.device.is_cuda() + && self.dtype == DType::BF16 + && self.repetition_penalty > 1.0 + { + Some( + Tensor::zeros((1, self.cfg.vocab_size), DType::U8, &self.device).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "allocate AR repetition history", e) + })?, + ) + } else { + None + }; let mut finished: Vec = vec![false; batch_size]; let mut positions: Vec = seq_lens.iter().map(|&len| len as i64).collect(); @@ -626,23 +936,50 @@ impl HunyuanOcr { // can spiral into runaway-repeat loops on large-context // inputs (observed on chart_01.jpg, seq≈11584, producing // 33K chars of synthetic Mermaid node IDs `BZ, BZW, …`). - let tok = if self.repetition_penalty > 1.0 && !generated[i].is_empty() { - argmax_with_repetition_penalty( - logits, - &generated[i], - self.repetition_penalty as f32, - )? - } else { - logits - .argmax(D::Minus1) - .and_then(|t| t.to_scalar::()) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "argmax", e))? - }; + let tok = + if self.repetition_penalty > 1.0 && !repetition_histories[i].is_empty() { + #[cfg(feature = "cuda")] + { + if let Some(history) = repetition_history_device.as_ref() { + argmax_with_device_repetition_history( + logits, + history, + self.repetition_penalty as f32, + )? + } else { + argmax_with_unique_repetition_penalty( + logits, + repetition_histories[i].ids(), + self.repetition_penalty as f32, + )? + } + } + #[cfg(not(feature = "cuda"))] + { + argmax_with_unique_repetition_penalty( + logits, + repetition_histories[i].ids(), + self.repetition_penalty as f32, + )? + } + } else { + logits + .argmax(D::Minus1) + .and_then(|t| t.to_scalar::()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "argmax", e))? + }; if self.stop_token_ids.contains(&tok) { finished[i] = true; } else { generated[i].push(tok); + repetition_histories[i].insert(tok); + #[cfg(feature = "cuda")] + if i == 0 + && let Some(history) = repetition_history_device.as_ref() + { + mark_repetition_history(history, &[tok])?; + } } next_tokens.push(tok); } @@ -658,26 +995,44 @@ impl HunyuanOcr { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "create tokens", e))?; let embeds = self.llm.embed(&tokens)?; - // Build 4-axis position IDs for decode step - let pos_data: Vec = positions.iter().flat_map(|&p| [p, p, p, p]).collect(); - let pos = Tensor::new(pos_data, &self.device) - .and_then(|t| t.reshape((4, batch_size, 1))) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "create pos", e))?; - - // Mask out left-padding positions in the KV cache for this step. kv_len += 1; - let gen_mask = create_generation_mask(&pad_lens, kv_len, self.dtype, &self.device) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "create gen mask", e))?; - - let hs = self.llm.forward(&embeds, &pos, Some(&gen_mask))?; + let (hs, graph_logits) = if batch_size == 1 { + // There is no left padding and a one-token query cannot see a + // future key. Reuse the precomputed scalar-position RoPE cache + // and avoid rebuilding four-axis positions plus an unused + // generation mask on every decode step. + let output = + self.llm + .forward_with_aux_decode(&embeds, positions[0] as usize, None, &[])?; + (output.hidden_states, output.logits) + } else { + // Build 4-axis position IDs for the true-batch decode path. + let pos_data: Vec = positions.iter().flat_map(|&p| [p, p, p, p]).collect(); + let pos = Tensor::new(pos_data, &self.device) + .and_then(|t| t.reshape((4, batch_size, 1))) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "create pos", e))?; + + // Mask out left-padding positions in the KV cache for this + // step when prompt lengths differ across the batch. + let gen_mask = create_generation_mask(&pad_lens, kv_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "create gen mask", e))?; + (self.llm.forward(&embeds, &pos, Some(&gen_mask))?, None) + }; logits_list.clear(); - for i in 0..batch_size { - let h = hs - .i((i, 0, ..)) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "get hs", e))?; - let logits = self.logits_from_hidden(&h)?; + if let Some(logits) = graph_logits { + let logits = logits + .i((0, 0, ..)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "get graph logits", e))?; logits_list.push(logits); + } else { + for i in 0..batch_size { + let h = hs + .i((i, 0, ..)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "get hs", e))?; + let logits = self.logits_from_hidden(&h)?; + logits_list.push(logits); + } } for (i, p) in positions.iter_mut().enumerate() { @@ -691,17 +1046,6 @@ impl HunyuanOcr { Ok(generated) } - fn greedy_token(&self, logits: &Tensor, seen: &[u32]) -> Result { - if self.repetition_penalty > 1.0 && !seen.is_empty() { - argmax_with_repetition_penalty(logits, seen, self.repetition_penalty as f32) - } else { - logits - .argmax(D::Minus1) - .and_then(|t| t.to_scalar::()) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "argmax", e)) - } - } - /// Greedy DFlash speculative decoding for one document. The target model /// remains authoritative: accepted draft prefixes plus the first target /// recovery/bonus token are exactly equivalent to autoregressive greedy @@ -719,7 +1063,10 @@ impl HunyuanOcr { return Ok(Vec::new()); } - let first = self.greedy_token(&initial_logits, &[])?; + let first = initial_logits + .argmax(D::Minus1) + .and_then(|token| token.to_scalar::()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "initial argmax", e))?; if self.stop_token_ids.contains(&first) { self.llm.clear_kv_cache(); dflash.clear_context(); @@ -742,6 +1089,22 @@ impl HunyuanOcr { .map(|id| id + 1) .collect(); let mut generated = vec![first]; + let mut repetition_history = RepetitionHistory::new(self.cfg.vocab_size); + repetition_history.insert(first); + #[cfg(feature = "cuda")] + let repetition_history_device = if self.device.is_cuda() + && self.dtype == DType::BF16 + && self.repetition_penalty > 1.0 + { + let history = + Tensor::zeros((self.cfg.vocab_size,), DType::U8, &self.device).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "allocate repetition history", e) + })?; + mark_repetition_history(&history, &[first])?; + Some(history) + } else { + None + }; let mut draft_rounds = 0usize; let mut accepted_draft_tokens = 0usize; @@ -759,18 +1122,10 @@ impl HunyuanOcr { .and_then(|t| t.reshape((1, num_spec + 1))) .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "query ids", e))?; let query_embeds = self.llm.embed(&query_ids)?; - let draft_hidden = dflash.forward_queries(&query_embeds)?; - let draft_rows = draft_hidden - .narrow(1, 1, num_spec) - .and_then(|t| t.squeeze(0)) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "draft mask rows", e))?; - let draft_logits = self.logits_from_hidden_batch(&draft_rows)?; - // All draft rows are independent. Keep their sampling as one GPU - // operation; target-side processors remain authoritative during - // verification and therefore preserve the target distribution. - let proposals = draft_logits - .argmax(D::Minus1) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "draft argmax", e))?; + // All draft rows are independent. The CUDA graph includes their + // shared LM head and argmax; target-side processors remain + // authoritative during verification. + let proposals = dflash.forward_proposals(&query_embeds)?; // The target verifies [bonus, proposal_1, ..., proposal_N] in one // causal pass. Row i predicts proposal i+1; the final row predicts @@ -818,49 +1173,92 @@ impl HunyuanOcr { "HunyuanOCR DFlash: target verification did not return auxiliary states" .to_string(), })?; - let target_rows = output.hidden_states.squeeze(0).map_err(|e| { - candle_to_ocr_inference("HunyuanOCR DFlash", "target verify rows", e) - })?; - let target_logits = self.logits_from_hidden_batch(&target_rows)?; - let target_tokens = if self.repetition_penalty <= 1.0 { - Some( - target_logits - .argmax(D::Minus1) - .and_then(|t| t.to_vec1::()) - .map_err(|e| { - candle_to_ocr_inference("HunyuanOCR DFlash", "batched target argmax", e) - })?, - ) + let target_logits = if let Some(logits) = output.logits.as_ref() { + logits.squeeze(0).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "graph target logits", e) + })? } else { + let target_rows = output.hidden_states.squeeze(0).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "target verify rows", e) + })?; + self.logits_from_hidden_batch(&target_rows)? + }; + // CPU backends need proposal values to build each row's history. + // CUDA keeps them device-resident through target sampling so the + // target-token transfer is the round's only blocking D2H copy. + let proposals_host = if self.device.is_cuda() + && (self.repetition_penalty <= 1.0 || self.dtype == DType::BF16) + { None + } else { + Some(proposals.to_vec1::().map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "draft proposals", e) + })?) + }; + let target_tokens = if self.repetition_penalty <= 1.0 { + target_logits + .argmax(D::Minus1) + .and_then(|t| t.to_vec1::()) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "batched target argmax", e) + })? + } else { + #[cfg(feature = "cuda")] + { + if self.device.is_cuda() && self.dtype == DType::BF16 { + dflash_argmax_with_repetition_penalty( + &target_logits, + repetition_history_device.as_ref().ok_or_else(|| { + OCRError::ConfigError { + message: "HunyuanOCR DFlash: missing CUDA repetition history" + .to_string(), + } + })?, + &proposals, + self.repetition_penalty as f32, + )? + } else { + dflash_argmax_with_host_proposals( + &target_logits, + &repetition_history, + proposals_host + .as_deref() + .expect("non-CUDA proposals were copied to the host"), + self.repetition_penalty as f32, + )? + } + } + #[cfg(not(feature = "cuda"))] + { + dflash_argmax_with_host_proposals( + &target_logits, + &repetition_history, + proposals_host + .as_deref() + .expect("non-CUDA proposals were copied to the host"), + self.repetition_penalty as f32, + )? + } + }; + // Target sampling above has already synchronized CUDA. This small + // proposal copy is now ready immediately and does not stall the + // draft/target pipeline a second time. + let proposals = match proposals_host { + Some(proposals) => proposals, + None => proposals.to_vec1::().map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "draft proposals", e) + })?, }; - let proposals = proposals - .to_vec1::() - .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "draft proposals", e))?; let mut accepted = 0usize; let mut recovery = None; - let mut seen = if self.repetition_penalty > 1.0 { - generated.clone() - } else { - Vec::new() - }; for (index, &proposal) in proposals.iter().enumerate() { - let target_token = match &target_tokens { - Some(tokens) => tokens[index], - None => { - let logits = target_logits.i((index, ..)).map_err(|e| { - candle_to_ocr_inference("HunyuanOCR DFlash", "target verify logits", e) - })?; - self.greedy_token(&logits, &seen)? - } - }; + let target_token = target_tokens[index]; if target_token != proposal { recovery = Some(target_token); break; } accepted += 1; - seen.push(proposal); } accepted_draft_tokens += accepted; @@ -871,15 +1269,7 @@ impl HunyuanOcr { // of the authoritative history before the recovery token. (1 + accepted, tokens) } else { - let target_bonus = match &target_tokens { - Some(tokens) => tokens[num_spec], - None => { - let bonus_logits = target_logits.i((num_spec, ..)).map_err(|e| { - candle_to_ocr_inference("HunyuanOCR DFlash", "target bonus logits", e) - })?; - self.greedy_token(&bonus_logits, &seen)? - } - }; + let target_bonus = target_tokens[num_spec]; let mut tokens = proposals; tokens.push(target_bonus); (num_spec + 1, tokens) @@ -901,6 +1291,8 @@ impl HunyuanOcr { dflash.append_context(&accepted_aux)?; let mut stop = false; + #[cfg(feature = "cuda")] + let mut new_history_tokens = Vec::with_capacity(next_tokens.len()); for token in next_tokens { if self.stop_token_ids.contains(&token) { stop = true; @@ -911,10 +1303,17 @@ impl HunyuanOcr { break; } generated.push(token); + repetition_history.insert(token); + #[cfg(feature = "cuda")] + new_history_tokens.push(token); } if stop { break; } + #[cfg(feature = "cuda")] + if let Some(history) = repetition_history_device.as_ref() { + mark_repetition_history(history, &new_history_tokens)?; + } debug_assert_eq!(self.llm.kv_cache_len(), dflash.context_len()); debug_assert!(self.llm.kv_cache_len() >= prompt_len); @@ -1144,6 +1543,96 @@ fn build_position_ids( mod tests { use super::*; + fn assert_repetition_penalty_tokens(device: &Device) { + let logits = Tensor::from_vec( + vec![ + 10.0f32, 9.0, 8.0, -1.0, // duplicate + out-of-range ids + -1.0, -2.0, 3.0, 2.5, // positive penalized winner + -1.0, -1.1, -1.2, -1.3, // negative penalized winner + ], + (3, 4), + device, + ) + .unwrap(); + let row0 = [0u32, 0, 99]; + let row1 = [2u32, 2]; + let row2 = [0u32]; + let histories: [&[u32]; 3] = [&row0, &row1, &row2]; + let tokens = batched_argmax_with_repetition_penalty(&logits, &histories, 2.0).unwrap(); + assert_eq!(tokens, vec![1, 3, 1]); + } + + #[test] + fn repetition_penalty_cpu_matches_huggingface_semantics() { + assert_repetition_penalty_tokens(&Device::Cpu); + } + + #[cfg(feature = "cuda")] + #[test] + fn repetition_penalty_cuda_matches_cpu() { + let Ok(device) = Device::new_cuda(0) else { + return; + }; + assert_repetition_penalty_tokens(&device); + + // Candle's generic CUDA argmax can choose index 1024 here because the + // equal maxima live in different reduction lanes. The reference CPU + // scan, and our deterministic kernel, must keep the first index. + let mut tied = vec![0.0f32; 1030]; + tied[1] = 10.0; + tied[1024] = 10.0; + let logits = Tensor::from_vec(tied, (1, 1030), &device).unwrap(); + let history = [2u32]; + let tokens = batched_argmax_with_repetition_penalty(&logits, &[&history], 1.08).unwrap(); + assert_eq!(tokens, vec![1]); + + let mut tied = vec![0.0f32; 1030]; + tied[1] = 10.0; + tied[1024] = 10.0; + let logits = Tensor::from_vec(tied, (1, 1030), &device) + .unwrap() + .to_dtype(DType::BF16) + .unwrap(); + let history = Tensor::zeros((1, 1030), DType::U8, &device).unwrap(); + mark_repetition_history(&history, &[2]).unwrap(); + let token = argmax_with_device_repetition_history(&logits, &history, 1.08).unwrap(); + assert_eq!(token, 1); + + let logits = Tensor::from_vec( + vec![10.0f32, 9.0, 0.0, 0.0, 8.0, 7.0, 0.0, 0.0], + (2, 4), + &device, + ) + .unwrap(); + let empty: &[u32] = &[]; + let tokens = + batched_argmax_with_unique_repetition_parts(&logits, &[0], &[empty, empty], 2.0) + .unwrap(); + assert_eq!(tokens, vec![1, 1]); + + // Row r sees only proposal ids [0, r), with ids already present in + // the common history and duplicate proposals penalized exactly once. + let logits = Tensor::from_vec( + vec![ + 10.0f32, 9.0, 5.2, 0.0, // common history only + 10.0, 9.0, 5.2, 0.0, // common + proposal 1 + 10.0, 20.0, 7.0, 0.0, // duplicate proposal 1 stays once + 10.0, 20.0, 7.0, 0.0, // proposal 0 is already common + ], + (4, 4), + &device, + ) + .unwrap() + .to_dtype(DType::BF16) + .unwrap(); + let proposals = Tensor::new(&[1u32, 1, 0], &device).unwrap(); + let history = Tensor::zeros((4,), DType::U8, &device).unwrap(); + mark_repetition_history(&history, &[0, 0, 99]).unwrap(); + let tokens = + dflash_argmax_with_repetition_penalty(&logits, &history, &proposals, 2.0).unwrap(); + assert_eq!(tokens, vec![1, 2, 1, 1]); + } + #[test] fn v15_prompt_omits_legacy_empty_system_token() { let prompt = build_prompt("read", HunyuanOcrVersion::V1_5); diff --git a/oar-ocr-vl/src/kv_trim.rs b/oar-ocr-vl/src/kv_trim.rs index 3826b1d..50d9946 100644 --- a/oar-ocr-vl/src/kv_trim.rs +++ b/oar-ocr-vl/src/kv_trim.rs @@ -199,14 +199,54 @@ impl TrimmableKvCache { /// This is used by CUDA-graph decode paths whose device kernel writes at a /// runtime offset while the host only tracks the resulting logical length. pub fn initialize_storage(&mut self, template: &Tensor) -> Result<()> { - if self.storage.is_none() { + self.initialize_storage_with_capacity(template, self.configured_capacity) + } + + /// Ensure fixed-capacity storage of exactly `capacity` tokens exists. + /// CUDA graphs capture the backing pointers and the physical head stride, + /// so a larger existing allocation cannot be reused with a smaller + /// captured `cache_len` value. + pub fn initialize_storage_with_capacity( + &mut self, + template: &Tensor, + capacity: usize, + ) -> Result<()> { + if capacity == 0 { + return Err(candle_core::Error::Msg( + "TrimmableKvCache capacity must be non-zero".into(), + )); + } + let reusable = self.storage.as_ref().is_some_and(|(storage_k, storage_v)| { + self.capacity == capacity + && storage_k.dtype() == template.dtype() + && storage_v.dtype() == template.dtype() + && storage_k.device().same_device(template.device()) + && storage_v.device().same_device(template.device()) + && storage_k.dims().len() == template.dims().len() + && storage_k + .dims() + .iter() + .zip(template.dims()) + .enumerate() + .all(|(dim, (stored, new))| { + if dim == self.cat_dim { + *stored == capacity + } else { + stored == new + } + }) + && storage_v.dims() == storage_k.dims() + }); + if !reusable { let mut shape = template.dims().to_vec(); - shape[self.cat_dim] = self.configured_capacity; - self.capacity = self.configured_capacity; + shape[self.cat_dim] = capacity; + self.capacity = capacity; self.storage = Some(( Tensor::zeros(shape.as_slice(), template.dtype(), template.device())?, Tensor::zeros(shape.as_slice(), template.dtype(), template.device())?, )); + self.kv = None; + self.cur_len = 0; } Ok(()) } @@ -246,6 +286,11 @@ impl TrimmableKvCache { self.configured_capacity } + /// Physical sequence capacity of the current backing storage. + pub fn storage_capacity(&self) -> usize { + self.capacity + } + pub fn reset(&mut self) { self.kv = None; self.cur_len = 0; @@ -331,6 +376,29 @@ mod tests { Ok(()) } + #[test] + fn fixed_storage_uses_exact_requested_capacity() -> Result<()> { + let mut c = TrimmableKvCache::new(2, 64); + let template = Tensor::zeros((1, 2, 1, 4), DType::F32, &dev())?; + c.initialize_storage_with_capacity(&template, 8)?; + assert_eq!(c.storage_capacity(), 8); + assert_eq!(c.storage().unwrap().0.dims(), &[1, 2, 8, 4]); + + c.set_current_len(4)?; + c.initialize_storage_with_capacity(&template, 16)?; + assert_eq!(c.storage_capacity(), 16); + assert_eq!(c.current_seq_len(), 0); + assert_eq!(c.storage().unwrap().0.dims(), &[1, 2, 16, 4]); + Ok(()) + } + + #[test] + fn fixed_storage_rejects_zero_capacity() { + let mut c = TrimmableKvCache::new(2, 64); + let template = Tensor::zeros((1, 2, 1, 4), DType::F32, &dev()).unwrap(); + assert!(c.initialize_storage_with_capacity(&template, 0).is_err()); + } + #[test] fn incompatible_storage_starts_a_new_logical_cache() -> Result<()> { let mut c = TrimmableKvCache::new(2, 64);