diff --git a/oar-ocr-vl/build.rs b/oar-ocr-vl/build.rs index 0821cd0..fd74115 100644 --- a/oar-ocr-vl/build.rs +++ b/oar-ocr-vl/build.rs @@ -1,6 +1,6 @@ use std::process::Command; -const MIN_DFLASH_COMPUTE_CAP: u32 = 80; +const MIN_CUDA_COMPUTE_CAP: u32 = 80; fn parse_compute_cap(value: &str) -> Option<(String, u32)> { let value = value.trim().to_ascii_lowercase(); @@ -48,25 +48,25 @@ fn cuda_compute_arch() -> String { ) }); assert!( - base >= MIN_DFLASH_COMPUTE_CAP, - "HunyuanOCR DFlash CUDA kernels require compute capability 8.0 or newer; got CUDA_COMPUTE_CAP={value:?}" + base >= MIN_CUDA_COMPUTE_CAP, + "oar-ocr-vl CUDA kernels require compute capability 8.0 or newer; got CUDA_COMPUTE_CAP={value:?}" ); return format!("compute_{arch}"); } match detect_local_compute_cap() { - Some(base) if base >= MIN_DFLASH_COMPUTE_CAP => format!("compute_{base}"), + Some(base) if base >= MIN_CUDA_COMPUTE_CAP => format!("compute_{base}"), Some(base) => { println!( - "cargo:warning=detected GPU compute capability {base} is below the HunyuanOCR DFlash minimum; compiling forward-compatible compute_{MIN_DFLASH_COMPUTE_CAP} PTX" + "cargo:warning=detected GPU compute capability {base} is below the oar-ocr-vl CUDA kernel minimum; compiling forward-compatible compute_{MIN_CUDA_COMPUTE_CAP} PTX" ); - format!("compute_{MIN_DFLASH_COMPUTE_CAP}") + format!("compute_{MIN_CUDA_COMPUTE_CAP}") } None => { println!( - "cargo:warning=could not detect a CUDA GPU; compiling compute_{MIN_DFLASH_COMPUTE_CAP} PTX (set CUDA_COMPUTE_CAP to override for cross/headless builds)" + "cargo:warning=could not detect a CUDA GPU; compiling compute_{MIN_CUDA_COMPUTE_CAP} PTX (set CUDA_COMPUTE_CAP to override for cross/headless builds)" ); - format!("compute_{MIN_DFLASH_COMPUTE_CAP}") + format!("compute_{MIN_CUDA_COMPUTE_CAP}") } } } @@ -93,18 +93,18 @@ fn main() { .args(["--ptx", "--std=c++17", "-O3"]) .arg(format!("--gpu-architecture={cuda_arch}")) .arg("-o") - .arg(out_dir.join("hunyuan_dynamic_kv.ptx")) + .arg(out_dir.join("oar_vl_kernels.ptx")) .arg("src/hunyuanocr/dynamic_kv.cu") .output() .unwrap_or_else(|error| { panic!( - "failed to invoke {:?} for HunyuanOCR dynamic KV kernel; install the CUDA toolkit or set NVCC to the compiler path: {error}", + "failed to invoke {:?} for oar-ocr-vl CUDA kernels; install the CUDA toolkit or set NVCC to the compiler path: {error}", nvcc ) }); if !output.status.success() { panic!( - "{:?} failed for HunyuanOCR dynamic KV kernel ({cuda_arch}):\n{}", + "{:?} failed for oar-ocr-vl CUDA kernels ({cuda_arch}):\n{}", nvcc, String::from_utf8_lossy(&output.stderr) ); diff --git a/oar-ocr-vl/examples/glmocr.rs b/oar-ocr-vl/examples/glmocr.rs index a507bf7..5d20f90 100644 --- a/oar-ocr-vl/examples/glmocr.rs +++ b/oar-ocr-vl/examples/glmocr.rs @@ -27,6 +27,7 @@ use tracing::{error, info}; use oar_ocr_core::utils::load_image; use oar_ocr_vl::GlmOcr; use oar_ocr_vl::utils::parse_device; +use utils::token_fingerprint; #[derive(Parser)] #[command(name = "glmocr")] @@ -102,16 +103,18 @@ fn main() -> Result<(), Box> { let infer_start = Instant::now(); match model - .generate(&[rgb_img], &[args.prompt.as_str()], args.max_tokens) + .generate_tokens(&[rgb_img], &[args.prompt.as_str()], args.max_tokens) .into_iter() .next() { - Some(Ok(result)) => { + Some(Ok(tokens)) => { info!( - " Inference time: {:.2}ms", - infer_start.elapsed().as_secs_f64() * 1000.0 + " Inference time: {:.2}ms, tokens: {}, fingerprint: {:016x}", + infer_start.elapsed().as_secs_f64() * 1000.0, + tokens.len(), + token_fingerprint(&tokens) ); - println!("{}", result); + println!("{}", model.decode_tokens(&tokens)?); } Some(Err(e)) => error!(" Inference failed: {}", e), None => error!(" No result returned from model"), diff --git a/oar-ocr-vl/examples/hunyuanocr.rs b/oar-ocr-vl/examples/hunyuanocr.rs index 07bdbd3..cd5c379 100644 --- a/oar-ocr-vl/examples/hunyuanocr.rs +++ b/oar-ocr-vl/examples/hunyuanocr.rs @@ -29,6 +29,7 @@ use tracing::{error, info}; use oar_ocr_core::utils::load_image; use oar_ocr_vl::HunyuanOcr; use oar_ocr_vl::utils::parse_device; +use utils::token_fingerprint; #[derive(Parser)] #[command(name = "hunyuanocr")] @@ -70,12 +71,6 @@ struct Args { benchmark: bool, } -fn token_fingerprint(tokens: &[u32]) -> u64 { - tokens.iter().fold(0xcbf29ce484222325_u64, |hash, token| { - (hash ^ u64::from(*token)).wrapping_mul(0x100000001b3) - }) -} - fn main() -> Result<(), Box> { utils::init_tracing(); let args = Args::parse(); diff --git a/oar-ocr-vl/examples/mineru.rs b/oar-ocr-vl/examples/mineru.rs index c644a54..bcc6a84 100644 --- a/oar-ocr-vl/examples/mineru.rs +++ b/oar-ocr-vl/examples/mineru.rs @@ -34,6 +34,7 @@ use oar_ocr_vl::utils::{convert_otsl_to_html, truncate_repetitive_content}; use utils::mineru_layout::{ ContentBlock, LAYOUT_IMAGE_SIZE, LAYOUT_PROMPT, parse_layout_output, prepare_for_extract, }; +use utils::token_fingerprint; #[derive(Parser)] #[command(name = "mineru")] @@ -156,11 +157,17 @@ fn two_step_extract( LAYOUT_IMAGE_SIZE, imageops::FilterType::CatmullRom, ); - let layout = model - .generate(&[layout_image], &[LAYOUT_PROMPT], max_tokens) + let layout_tokens = model + .generate_tokens(&[layout_image], &[LAYOUT_PROMPT], max_tokens) .into_iter() .next() .ok_or("Layout detection returned no result")??; + info!( + " Layout tokens: {}, fingerprint: {:016x}", + layout_tokens.len(), + token_fingerprint(&layout_tokens) + ); + let layout = model.decode_tokens(&layout_tokens)?; if dump_layout { info!("Layout raw output:\n{}", layout); @@ -181,11 +188,18 @@ fn two_step_extract( for (i, (block_image, prompt)) in block_images.into_iter().zip(prompts.iter()).enumerate() { let idx = indices[i]; let output = model - .generate(&[block_image], &[prompt], max_tokens) + .generate_tokens(&[block_image], &[prompt], max_tokens) .into_iter() .next(); match output { - Some(Ok(content)) => { + Some(Ok(tokens)) => { + info!( + " Block {} tokens: {}, fingerprint: {:016x}", + idx, + tokens.len(), + token_fingerprint(&tokens) + ); + let content = model.decode_tokens(&tokens)?; let cleaned = truncate_repetitive_content(&content, 10, 10, 10); let content = if blocks[idx].block_type == "table" { convert_otsl_to_html(&cleaned) diff --git a/oar-ocr-vl/examples/paddleocr_vl.rs b/oar-ocr-vl/examples/paddleocr_vl.rs index bf65997..a8ec331 100644 --- a/oar-ocr-vl/examples/paddleocr_vl.rs +++ b/oar-ocr-vl/examples/paddleocr_vl.rs @@ -60,6 +60,7 @@ use tracing::{error, info}; use oar_ocr_core::utils::load_image; use oar_ocr_vl::utils::parse_device; use oar_ocr_vl::{PaddleOcrVl, PaddleOcrVlTask}; +use utils::token_fingerprint; /// Command-line arguments for the PaddleOCR-VL example #[derive(Parser)] @@ -174,17 +175,19 @@ fn main() -> Result<(), Box> { // Run inference let infer_start = Instant::now(); match model - .generate(&[rgb_img], &[task], args.max_tokens) + .generate_tokens(&[rgb_img], &[task], args.max_tokens) .pop() .unwrap() { - Ok(result) => { + Ok(tokens) => { let infer_duration = infer_start.elapsed(); info!( - " Inference time: {:.2}ms", - infer_duration.as_secs_f64() * 1000.0 + " Inference time: {:.2}ms, tokens: {}, fingerprint: {:016x}", + infer_duration.as_secs_f64() * 1000.0, + tokens.len(), + token_fingerprint(&tokens) ); - println!("{}", result); + println!("{}", model.decode_tokens(&tokens, task)?.1); } Err(e) => { error!(" Inference failed: {}", e); diff --git a/oar-ocr-vl/examples/utils/mod.rs b/oar-ocr-vl/examples/utils/mod.rs index 0c7c78f..4670eda 100644 --- a/oar-ocr-vl/examples/utils/mod.rs +++ b/oar-ocr-vl/examples/utils/mod.rs @@ -20,6 +20,14 @@ pub fn init_tracing() { .init(); } +/// Stable FNV-1a fingerprint used to compare generated token streams. +#[allow(dead_code)] +pub fn token_fingerprint(tokens: &[u32]) -> u64 { + tokens.iter().fold(0xcbf29ce484222325_u64, |hash, token| { + (hash ^ u64::from(*token)).wrapping_mul(0x100000001b3) + }) +} + #[allow(dead_code)] pub fn print_preview(tag: &str, text: &str) { let preview: String = text.chars().take(400).collect(); diff --git a/oar-ocr-vl/src/attention.rs b/oar-ocr-vl/src/attention.rs index 6ef1044..225af1b 100644 --- a/oar-ocr-vl/src/attention.rs +++ b/oar-ocr-vl/src/attention.rs @@ -21,6 +21,18 @@ use crate::utils::candle_to_ocr_processing; use candle_core::{D, DType, Device, IndexOp, Result, Tensor}; use oar_ocr_core::core::errors::OCRError; +use std::sync::OnceLock; + +fn grouped_query_attention_disabled() -> bool { + static DISABLED: OnceLock = OnceLock::new(); + *DISABLED.get_or_init(|| std::env::var_os("OAR_VL_DISABLE_GQA").is_some()) +} + +#[cfg(feature = "cuda")] +fn flash_attention_disabled() -> bool { + static DISABLED: OnceLock = OnceLock::new(); + *DISABLED.get_or_init(|| std::env::var_os("OAR_VL_DISABLE_FLASH_ATTN").is_some()) +} /// Helper function to handle Metal device computation. /// @@ -116,6 +128,11 @@ pub fn scaled_dot_product_attention_gqa( if num_kv_groups == 1 { return scaled_dot_product_attention(q, k, v, mask, scale, is_causal); } + if grouped_query_attention_disabled() { + let k = repeat_kv(k, num_kv_groups)?; + let v = repeat_kv(v, num_kv_groups)?; + return scaled_dot_product_attention(q, &k, &v, mask, scale, is_causal); + } let (batch, num_heads, query_len, head_dim) = q.dims4()?; let (k_batch, num_kv_heads, kv_len, k_head_dim) = k.dims4()?; @@ -181,7 +198,8 @@ pub fn flash_attention( causal: bool, ) -> Result> { #[cfg(feature = "cuda")] - if q.device().is_cuda() + if !flash_attention_disabled() + && q.device().is_cuda() && flash_attention_dtype_supported(q.dtype()) && k.dtype() == q.dtype() && v.dtype() == q.dtype() diff --git a/oar-ocr-vl/src/cuda_kernels.rs b/oar-ocr-vl/src/cuda_kernels.rs new file mode 100644 index 0000000..1b4b409 --- /dev/null +++ b/oar-ocr-vl/src/cuda_kernels.rs @@ -0,0 +1,452 @@ +use candle_core::backend::BackendStorage; +use candle_core::{CpuStorage, CustomOp1, CustomOp2, InplaceOp2, Layout, Result, Shape}; + +const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/oar_vl_kernels.ptx")); +pub(crate) const CUDA_KERNEL_MODULE: &str = "oar_vl_kernels"; +const PARTITIONS_PER_ROW: usize = 8; + +/// Stable row-wise argmax for F32 logits. Equal maxima keep the lowest +/// vocabulary id, matching the scalar CPU generation loops. +pub(crate) struct ArgmaxFirstF32; + +/// Stable row-wise argmax for BF16 logits. Values are compared after the same +/// exact BF16-to-F32 conversion used by the CPU fallback. +pub(crate) struct ArgmaxFirstBf16; + +/// Set a sparse list of vocabulary positions to negative infinity before +/// greedy selection. The same list is applied to every logits row. +pub(crate) struct MaskTokenIds; + +/// GPU row-wise multinomial/greedy selection with the selected token's +/// softmax probability. The result is F32 `[rows, 2]`: token id, confidence. +pub(crate) struct SampleWithConfidence { + pub(crate) temperature: f32, + pub(crate) greedy: bool, +} + +impl CustomOp1 for ArgmaxFirstF32 { + fn name(&self) -> &'static str { + "stable-argmax-first-f32" + } + + fn cpu_fwd(&self, _storage: &CpuStorage, _layout: &Layout) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("stable 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!("stable argmax requires contiguous F32 logits") + } + let (rows, vocab_size) = layout.shape().dims2()?; + if vocab_size == 0 { + candle_core::bail!("stable 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", CUDA_KERNEL_MODULE, 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 { + "stable-argmax-first-bf16" + } + + fn cpu_fwd(&self, _storage: &CpuStorage, _layout: &Layout) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("stable 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!("stable BF16 argmax requires contiguous logits") + } + let (rows, vocab_size) = layout.shape().dims2()?; + if vocab_size == 0 { + candle_core::bail!("stable 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", CUDA_KERNEL_MODULE, 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", + CUDA_KERNEL_MODULE, + 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 MaskTokenIds { + fn name(&self) -> &'static str { + "mask-token-ids" + } + + fn cpu_fwd( + &self, + _logits: &mut CpuStorage, + _logits_layout: &Layout, + _token_ids: &CpuStorage, + _token_ids_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("sparse token masking 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() + || token_ids.dtype() != candle_core::DType::U32 + { + candle_core::bail!( + "sparse token masking requires contiguous logits and contiguous U32 ids" + ) + } + let (rows, vocab_size) = logits_layout.shape().dims2()?; + let token_count = token_ids_layout.shape().dims1()?; + if token_count == 0 { + return Ok(()); + } + + let device = logits.device().clone(); + let token_ids = token_ids.as_cuda_slice::()?; + let token_ids = token_ids.slice(token_ids_layout.start_offset()..); + let count = rows * token_count; + match logits.dtype() { + candle_core::DType::BF16 => { + let logits = logits.as_cuda_slice_mut::()?; + let mut logits = logits.slice_mut(logits_layout.start_offset()..); + let function = device.get_or_load_custom_func( + "mask_token_ids_bf16", + CUDA_KERNEL_MODULE, + PTX, + )?; + let mut builder = function.builder(); + builder.arg(&mut logits); + builder.arg(&token_ids); + candle_core::builder_arg!( + builder, + rows as u32, + vocab_size as u32, + token_count as u32 + ); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + } + candle_core::DType::F32 => { + let logits = logits.as_cuda_slice_mut::()?; + let mut logits = logits.slice_mut(logits_layout.start_offset()..); + let function = device.get_or_load_custom_func( + "mask_token_ids_f32", + CUDA_KERNEL_MODULE, + PTX, + )?; + let mut builder = function.builder(); + builder.arg(&mut logits); + builder.arg(&token_ids); + candle_core::builder_arg!( + builder, + rows as u32, + vocab_size as u32, + token_count as u32 + ); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + } + dtype => candle_core::bail!( + "sparse token masking requires BF16 or F32 logits, got {dtype:?}" + ), + } + Ok(()) + } +} + +impl CustomOp2 for SampleWithConfidence { + fn name(&self) -> &'static str { + "sample-with-confidence" + } + + fn cpu_fwd( + &self, + _logits: &CpuStorage, + _logits_layout: &Layout, + _uniforms: &CpuStorage, + _uniforms_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("sample-with-confidence is CUDA-only") + } + + fn cuda_fwd( + &self, + logits: &candle_core::CudaStorage, + logits_layout: &Layout, + uniforms: &candle_core::CudaStorage, + uniforms_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() + || !uniforms_layout.is_contiguous() + || uniforms.dtype() != candle_core::DType::F32 + { + candle_core::bail!("sample-with-confidence requires contiguous logits and F32 uniforms") + } + let (rows, vocab_size) = logits_layout.shape().dims2()?; + if rows == 0 || vocab_size == 0 || uniforms_layout.shape().dims1()? != rows { + candle_core::bail!( + "sample-with-confidence shape mismatch: logits={:?}, uniforms={:?}", + logits_layout.shape(), + uniforms_layout.shape() + ) + } + if vocab_size >= (1 << 24) { + candle_core::bail!( + "sample-with-confidence requires vocabulary size < 2^24, got {vocab_size}" + ) + } + if !self.greedy && (!self.temperature.is_finite() || self.temperature <= 0.0) { + candle_core::bail!( + "sample-with-confidence requires a positive finite sampling temperature" + ) + } + + let device = logits.device().clone(); + let uniforms = uniforms.as_cuda_slice::()?; + let uniforms = uniforms.slice(uniforms_layout.start_offset()..); + let output = unsafe { device.alloc::(rows * 2) }?; + let function_name = match logits.dtype() { + candle_core::DType::BF16 => "sample_with_confidence_bf16", + candle_core::DType::F32 => "sample_with_confidence_f32", + dtype => candle_core::bail!( + "sample-with-confidence requires BF16 or F32 logits, got {dtype:?}" + ), + }; + let function = device.get_or_load_custom_func(function_name, CUDA_KERNEL_MODULE, PTX)?; + let inv_temperature = if self.greedy { + 1.0 + } else { + self.temperature.recip() + }; + let greedy = u32::from(self.greedy); + match logits.dtype() { + candle_core::DType::BF16 => { + let logits = logits.as_cuda_slice::()?; + let logits = logits.slice(logits_layout.start_offset()..); + let mut builder = function.builder(); + builder.arg(&logits); + builder.arg(&uniforms); + builder.arg(&output); + candle_core::builder_arg!( + builder, + rows as u32, + vocab_size as u32, + inv_temperature, + greedy + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + } + candle_core::DType::F32 => { + let logits = logits.as_cuda_slice::()?; + let logits = logits.slice(logits_layout.start_offset()..); + let mut builder = function.builder(); + builder.arg(&logits); + builder.arg(&uniforms); + builder.arg(&output); + candle_core::builder_arg!( + builder, + rows as u32, + vocab_size as u32, + inv_temperature, + greedy + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + } + _ => unreachable!(), + } + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + Shape::from_dims(&[rows, 2]), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use candle_core::{DType, Device, Tensor}; + + #[cfg(feature = "cuda")] + #[test] + fn stable_bf16_argmax_masks_ids_and_keeps_first_tie() -> Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let logits = Tensor::from_vec(vec![0.0f32, 10.0, 3.0, 10.0, 9.0, 1.0], (1, 6), &device)? + .to_dtype(DType::BF16)?; + let banned = Tensor::new(&[1u32, u32::MAX], &device)?; + logits.inplace_op2(&banned, &MaskTokenIds)?; + let token = logits + .apply_op1_no_bwd(&ArgmaxFirstBf16)? + .to_vec1::()?; + assert_eq!(token, vec![3]); + Ok(()) + } + + #[cfg(feature = "cuda")] + #[test] + fn stable_f32_argmax_masks_ids_for_every_row() -> Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let logits = Tensor::from_vec(vec![1.0f32, 8.0, 7.0, 0.0, 2.0, 9.0], (2, 3), &device)?; + let banned = Tensor::new(&[1u32], &device)?; + logits.inplace_op2(&banned, &MaskTokenIds)?; + let tokens = logits.apply_op1_no_bwd(&ArgmaxFirstF32)?.to_vec1::()?; + assert_eq!(tokens, vec![2, 2]); + Ok(()) + } + + #[cfg(feature = "cuda")] + #[test] + fn gpu_sampling_returns_stable_greedy_and_selected_probability() -> Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let logits = Tensor::from_vec(vec![2.0f32, 1.0, 0.0, 2.0, 2.0, 0.0], (2, 3), &device)? + .to_dtype(DType::BF16)?; + let uniforms = Tensor::zeros(2, DType::F32, &device)?; + let sampled = logits + .apply_op2_no_bwd( + &uniforms, + &SampleWithConfidence { + temperature: 1.0, + greedy: true, + }, + )? + .to_vec2::()?; + assert_eq!(sampled[0][0] as u32, 0); + // Equal maxima keep the first vocabulary id. + assert_eq!(sampled[1][0] as u32, 0); + assert!((sampled[0][1] - 0.665240).abs() < 1e-4); + assert!((sampled[1][1] - 0.468311).abs() < 1e-4); + Ok(()) + } + + #[cfg(feature = "cuda")] + #[test] + fn gpu_sampling_uses_cdf_uniforms() -> Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let logits = Tensor::from_vec(vec![0.0f32, 0.0, 0.0, 0.0], (2, 2), &device)?; + let uniforms = Tensor::new(&[0.1f32, 0.9], &device)?; + let sampled = logits + .apply_op2_no_bwd( + &uniforms, + &SampleWithConfidence { + temperature: 1.0, + greedy: false, + }, + )? + .to_vec2::()?; + assert_eq!(sampled[0][0] as u32, 0); + assert_eq!(sampled[1][0] as u32, 1); + assert!((sampled[0][1] - 0.5).abs() < 1e-6); + assert!((sampled[1][1] - 0.5).abs() < 1e-6); + Ok(()) + } +} diff --git a/oar-ocr-vl/src/decoder_graph.rs b/oar-ocr-vl/src/decoder_graph.rs new file mode 100644 index 0000000..33fab75 --- /dev/null +++ b/oar-ocr-vl/src/decoder_graph.rs @@ -0,0 +1,232 @@ +//! Shared single-token decoder CUDA-graph plumbing. +//! +//! Model-specific decoder code owns capture and replay because each model has +//! a different layer stack. This module centralizes the storage lifetime and +//! cache-bucket rules that must remain identical across those implementations. + +/// Select a bounded power-of-two KV-cache bucket for single-token decoding. +/// +/// Returning `None` means the request should stay on the eager path. A request +/// whose declared maximum exceeds `limit` may still use the largest bucket; +/// replay drops the graph and falls back to eager before the cache grows past +/// that bucket. +#[cfg(any(feature = "cuda", test))] +pub(crate) fn decoder_cache_capacity( + prompt_len: usize, + max_new_tokens: usize, + limit: usize, +) -> Option { + if max_new_tokens == 0 || prompt_len >= limit || limit == 0 { + return None; + } + let required = prompt_len.saturating_add(max_new_tokens).min(limit); + Some(required.max(1).next_power_of_two().min(limit)) +} + +/// Match eager decoder attention: a single query has no future token to mask, +/// while verification blocks must remain causal within the block. +#[cfg(any(feature = "cuda", test))] +pub(crate) const fn decoder_attention_is_causal(query_len: usize) -> bool { + query_len > 1 +} + +#[cfg(feature = "cuda")] +use crate::utils::candle_to_ocr_inference; +#[cfg(feature = "cuda")] +use candle_core::{CpuStorage, DType, Device, InplaceOp1, Layout, Tensor}; +#[cfg(feature = "cuda")] +use oar_ocr_core::core::OCRError; +#[cfg(feature = "cuda")] +use std::cell::RefCell; + +#[cfg(feature = "cuda")] +pub(crate) fn cuda_graph_error( + model_name: &str, + context: impl Into, + source: impl std::error::Error + Send + Sync + 'static, +) -> OCRError { + OCRError::Inference { + model_name: model_name.to_string(), + context: context.into(), + source: Box::new(source), + } +} + +#[cfg(feature = "cuda")] +pub(crate) fn sync_graph_tensor( + model_name: &str, + tensor: &Tensor, + operation: &'static str, +) -> Result<(), OCRError> { + tensor + .flatten_all() + .and_then(|x| x.narrow(0, 0, 1)) + .and_then(|x| x.to_dtype(DType::F32)) + .and_then(|x| x.to_vec1::()) + .map(|_| ()) + .map_err(|e| candle_to_ocr_inference(model_name, operation, e)) +} + +/// Persistent device/host pair used to update `[0, kv_len]` before replay. +/// +/// Keeping both allocations alive avoids constructing a temporary CUDA tensor +/// on every generated token. The page-locked host buffer also lets the tiny +/// copy remain ordered on the decoder stream without a whole-stream sync. +#[cfg(feature = "cuda")] +pub(crate) struct CudaGraphKvLengths { + tensor: Tensor, + host: RefCell>, +} + +#[cfg(feature = "cuda")] +struct CopyPinnedKvLengths<'a> { + source: &'a candle_core::cuda_backend::cudarc::driver::PinnedHostSlice, +} + +#[cfg(feature = "cuda")] +impl InplaceOp1 for CopyPinnedKvLengths<'_> { + fn name(&self) -> &'static str { + "copy-pinned-cuda-graph-kv-lengths" + } + + fn cpu_fwd(&self, _storage: &mut CpuStorage, _layout: &Layout) -> candle_core::Result<()> { + candle_core::bail!("CUDA-graph KV lengths are CUDA-only") + } + + fn cuda_fwd( + &self, + storage: &mut candle_core::CudaStorage, + layout: &Layout, + ) -> candle_core::Result<()> { + let Some((start, end)) = layout.contiguous_offsets() else { + candle_core::bail!("CUDA-graph KV lengths must be contiguous") + }; + if end.saturating_sub(start) != 2 { + candle_core::bail!("CUDA-graph KV lengths must contain two u32 values") + } + let device = storage.device.clone(); + let destination = storage.as_cuda_slice_mut::()?; + let mut destination = destination.slice_mut(start..end); + device.memcpy_htod(self.source, &mut destination) + } +} + +#[cfg(feature = "cuda")] +impl CudaGraphKvLengths { + pub(crate) fn new(initial_kv_len: usize, device: &Device) -> candle_core::Result { + use candle_core::cuda_backend::WrapErr; + + let Device::Cuda(cuda) = device else { + candle_core::bail!("CUDA-graph KV lengths require a CUDA device") + }; + let initial_kv_len = u32::try_from(initial_kv_len) + .map_err(|_| candle_core::Error::Msg("KV length exceeds u32".to_string()))?; + let tensor = Tensor::new(&[0u32, initial_kv_len], device)?; + let stream = cuda.cuda_stream(); + // SAFETY: both u32 elements are initialized immediately below before + // the page-locked allocation can be read or copied. + let mut host = unsafe { stream.context().alloc_pinned::(2) }.w()?; + let host_ptr = host.as_mut_ptr().w()?; + // SAFETY: `host` owns two properly aligned u32 slots. + unsafe { + host_ptr.write(0); + host_ptr.add(1).write(initial_kv_len); + } + Ok(Self { + tensor, + host: RefCell::new(host), + }) + } + + pub(crate) fn tensor(&self) -> &Tensor { + &self.tensor + } + + pub(crate) fn update(&self, kv_len: usize) -> candle_core::Result<()> { + use candle_core::cuda_backend::WrapErr; + + let kv_len = u32::try_from(kv_len) + .map_err(|_| candle_core::Error::Msg("KV length exceeds u32".to_string()))?; + let mut host = self.host.borrow_mut(); + let host_ptr = host.as_mut_ptr().w()?; + // SAFETY: `host` owns two properly aligned u32 slots, and waiting in + // `as_mut_ptr` makes the previous asynchronous copy safe to overwrite. + unsafe { + host_ptr.write(0); + host_ptr.add(1).write(kv_len); + } + self.tensor + .inplace_op1(&CopyPinnedKvLengths { source: &host }) + } +} + +#[cfg(feature = "cuda")] +impl std::fmt::Debug for CudaGraphKvLengths { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaGraphKvLengths").finish_non_exhaustive() + } +} + +/// Captured storage for a batch-1, query-length-1 decoder graph. +/// +/// The graph owns raw pointers into every tensor below, the model's fixed KV +/// storage, its weights, and the external LM head. Keep `graph` first so it is +/// destroyed before the tensors when this value is dropped. +#[cfg(feature = "cuda")] +pub(crate) struct SingleTokenDecoderCudaGraph { + pub(crate) graph: candle_core::cuda_backend::cudarc::driver::CudaGraph, + pub(crate) hidden_input: Tensor, + pub(crate) position_input: Tensor, + pub(crate) _query_lengths: Tensor, + pub(crate) kv_lengths: CudaGraphKvLengths, + pub(crate) logits_output: Tensor, + pub(crate) cache_len: usize, +} + +#[cfg(feature = "cuda")] +impl std::fmt::Debug for SingleTokenDecoderCudaGraph { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("SingleTokenDecoderCudaGraph") + .field("hidden", &self.hidden_input.shape()) + .field("cache_len", &self.cache_len) + .finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::{decoder_attention_is_causal, decoder_cache_capacity}; + + #[test] + fn cache_capacity_uses_bounded_power_of_two_buckets() { + const LIMIT: usize = 16_384; + assert_eq!(decoder_cache_capacity(1500, 256, LIMIT), Some(2048)); + assert_eq!(decoder_cache_capacity(2000, 4096, LIMIT), Some(8192)); + assert_eq!(decoder_cache_capacity(10_000, 20_000, LIMIT), Some(LIMIT)); + assert_eq!(decoder_cache_capacity(100, 0, LIMIT), None); + assert_eq!(decoder_cache_capacity(LIMIT, 1, LIMIT), None); + assert_eq!(decoder_cache_capacity(1, 1, 0), None); + } + + #[test] + fn single_token_decode_is_not_causal_but_verification_blocks_are() { + assert!(!decoder_attention_is_causal(1)); + assert!(decoder_attention_is_causal(2)); + assert!(decoder_attention_is_causal(16)); + } + + #[cfg(feature = "cuda")] + #[test] + fn persistent_kv_lengths_update_device_storage() -> candle_core::Result<()> { + use super::CudaGraphKvLengths; + use candle_core::Device; + + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let lengths = CudaGraphKvLengths::new(1, &device)?; + lengths.update(12_345)?; + assert_eq!(lengths.tensor().to_vec1::()?, [0, 12_345]); + Ok(()) + } +} diff --git a/oar-ocr-vl/src/glmocr/mod.rs b/oar-ocr-vl/src/glmocr/mod.rs index f471f10..8de0893 100644 --- a/oar-ocr-vl/src/glmocr/mod.rs +++ b/oar-ocr-vl/src/glmocr/mod.rs @@ -2,6 +2,8 @@ mod config; mod model; +#[cfg(feature = "cuda")] +mod mtp; mod processing; mod text; mod vision; diff --git a/oar-ocr-vl/src/glmocr/model.rs b/oar-ocr-vl/src/glmocr/model.rs index 1217b48..c4a47b7 100644 --- a/oar-ocr-vl/src/glmocr/model.rs +++ b/oar-ocr-vl/src/glmocr/model.rs @@ -1,4 +1,6 @@ use super::config::{EosTokenId, GlmOcrConfig, GlmOcrImageProcessorConfig}; +#[cfg(feature = "cuda")] +use super::mtp::GlmOcrMtpModel; use super::processing::{GlmOcrImageInputs, preprocess_image}; use super::text::GlmOcrTextModel; use super::vision::GlmOcrVisionModel; @@ -10,6 +12,14 @@ use oar_ocr_core::core::OCRError; use std::path::Path; use tokenizers::Tokenizer; +#[cfg(feature = "cuda")] +// Four recurrent proposals are fastest on the reference 4090 across text, +// formula, table, and long document outputs. The upstream-recommended three +// leaves useful acceptance on the table for Candle's batch-1 path. +const GLM_MTP_DRAFT_TOKENS: usize = 4; +#[cfg(feature = "cuda")] +const GLM_MTP_QUERY_LEN: usize = GLM_MTP_DRAFT_TOKENS + 1; + pub struct GlmOcr { device: Device, dtype: DType, @@ -17,6 +27,8 @@ pub struct GlmOcr { image_cfg: GlmOcrImageProcessorConfig, tokenizer: Tokenizer, text: GlmOcrTextModel, + #[cfg(feature = "cuda")] + mtp: Option, vision: GlmOcrVisionModel, lm_head: Linear, eos_token_ids: Vec, @@ -56,7 +68,22 @@ impl GlmOcr { }; let image_token_id = cfg.image_token_id; - let text = GlmOcrTextModel::load(&cfg.text_config, vb.pp("model").pp("language_model"))?; + let vb_language_model = vb.pp("model").pp("language_model"); + let text = GlmOcrTextModel::load(&cfg.text_config, vb_language_model.clone())?; + #[cfg(feature = "cuda")] + let mtp = if device.is_cuda() + && cfg.text_config.num_nextn_predict_layers == 1 + && std::env::var_os("OAR_GLMOCR_DISABLE_MTP").is_none() + && std::env::var_os("OAR_VL_DISABLE_SPECULATIVE").is_none() + { + Some(GlmOcrMtpModel::load( + &cfg.text_config, + cfg.text_config.num_hidden_layers, + vb_language_model, + )?) + } else { + None + }; let vision = GlmOcrVisionModel::load(&cfg.vision_config, vb.pp("model").pp("visual"))?; let lm_head = if cfg.text_config.tie_word_embeddings || cfg.tie_word_embeddings { @@ -82,6 +109,8 @@ impl GlmOcr { image_cfg, tokenizer, text, + #[cfg(feature = "cuda")] + mtp, vision, lm_head, eos_token_ids, @@ -217,71 +246,315 @@ impl GlmOcr { let rope_delta = max_pos + 1 - seq_len as i64; self.text.clear_kv_cache(); + #[cfg(feature = "cuda")] + let use_mtp = self.mtp.is_some() + && self.dtype == DType::BF16 + && max_new_tokens >= 8 + && std::env::var_os("OAR_VL_DISABLE_SPECULATIVE").is_none(); + #[cfg(not(feature = "cuda"))] + let use_mtp = false; + + #[cfg(feature = "cuda")] + if use_mtp { + let mtp = self.mtp.as_ref().expect("MTP availability checked"); + mtp.clear_kv_cache(); + if let Some(cache_len) = self.text.prepare_verification_cuda_graph( + seq_len, + max_new_tokens, + GLM_MTP_QUERY_LEN, + &self.lm_head, + )? { + mtp.prepare_cuda_graph(cache_len)?; + } else { + // An eager prefill may grow/reallocate MTP storage. Any + // graph captured for an earlier page would then retain + // stale device pointers. + mtp.disable_cuda_graph(); + } + } + if !use_mtp { + self.text + .prepare_ar_cuda_graph(seq_len, max_new_tokens, &self.lm_head)?; + } + let hidden = self.text.forward(&inputs_embeds, &position_ids, None)?; - let last = hidden.i((0, seq_len - 1, ..)).map_err(|e| { + #[cfg(feature = "cuda")] + if use_mtp { + let generated = self.generate_mtp_tokens( + &input_ids, + &position_ids, + &hidden, + rope_delta, + max_new_tokens, + self.mtp.as_ref().expect("MTP availability checked"), + )?; + results.push(generated); + continue; + } + + let generated = + self.generate_ar_tokens(&hidden, seq_len, rope_delta, max_new_tokens)?; + + results.push(generated); + } + + Ok(results) + } + + fn generate_ar_tokens( + &self, + prompt_hidden: &Tensor, + seq_len: usize, + rope_delta: i64, + max_new_tokens: usize, + ) -> Result, OCRError> { + let last = prompt_hidden.i((0, seq_len - 1, ..)).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: get last hidden", + e, + ) + })?; + let mut logits = self.logits_from_hidden(&last)?; + let mut generated = Vec::with_capacity(max_new_tokens); + + for (step, pos) in (seq_len as i64..).take(max_new_tokens).enumerate() { + let tok = logits + .argmax(D::Minus1) + .and_then(|t| t.to_scalar::()) + .map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: argmax", + e, + ) + })?; + if self.eos_token_ids.contains(&tok) { + break; + } + generated.push(tok); + if step + 1 == max_new_tokens { + break; + } + + let token = token_tensor(tok, &self.device)?; + let embed = self.text.embed(&token)?; + let pos_ids = text_position_ids(pos + rope_delta, 1, &self.device)?; + logits = self + .text + .forward_decode_logits(&embed, &pos_ids, &self.lm_head)?; + } + Ok(generated) + } + + #[cfg(feature = "cuda")] + fn generate_mtp_tokens( + &self, + prompt_ids: &[u32], + prompt_position_ids: &Tensor, + prompt_hidden: &Tensor, + rope_delta: i64, + max_new_tokens: usize, + mtp: &GlmOcrMtpModel, + ) -> Result, OCRError> { + let prompt_len = prompt_ids.len(); + let last = prompt_hidden.i((0, prompt_len - 1, ..)).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: get MTP target hidden", + e, + ) + })?; + let mut current_token = self + .logits_from_hidden(&last)? + .argmax(D::Minus1) + .and_then(|token| token.to_scalar::()) + .map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: get last hidden", + "GLM-OCR: initial MTP target argmax", e, ) })?; - let mut logits = self.logits_from_hidden(&last)?; - - let mut generated: Vec = Vec::new(); - - for (pos, _) in (seq_len as i64..).zip(0..max_new_tokens) { - let tok = logits - .argmax(D::Minus1) - .and_then(|t| t.to_scalar::()) - .map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: argmax", - e, - ) - })?; - if self.eos_token_ids.contains(&tok) { + if max_new_tokens == 0 || self.eos_token_ids.contains(¤t_token) { + return Ok(Vec::new()); + } + + // Match vLLM's EAGLE/MTP alignment: shift prompt token ids left, + // insert the target's certain next token at the tail, and pair that + // sequence with the target hidden states at the original positions. + let mut shifted_ids = Vec::with_capacity(prompt_len); + shifted_ids.extend_from_slice(&prompt_ids[1..]); + shifted_ids.push(current_token); + let shifted_ids = Tensor::from_vec(shifted_ids, (1, prompt_len), &self.device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "initial MTP shifted ids", e))?; + let (mtp_hidden, first_draft) = + mtp.sync_target_span(&shifted_ids, prompt_hidden, prompt_position_ids, true)?; + + let mut target_position = prompt_len as i64 + rope_delta; + let mut drafts = + self.complete_mtp_drafts(mtp, &mtp_hidden, &first_draft, target_position)?; + let mut target_cache_len = prompt_len; + let mut generated = Vec::with_capacity(max_new_tokens); + let mut rounds = 0usize; + let mut accepted_drafts = 0usize; + let mut accepted_by_position = [0usize; GLM_MTP_DRAFT_TOKENS]; + + loop { + if self.eos_token_ids.contains(¤t_token) { + break; + } + generated.push(current_token); + if generated.len() == max_new_tokens { + break; + } + + let mut query_ids = Vec::with_capacity(GLM_MTP_QUERY_LEN); + query_ids.push(current_token); + query_ids.extend_from_slice(&drafts); + let query = + Tensor::from_vec(query_ids.clone(), (1, GLM_MTP_QUERY_LEN), &self.device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP verification ids", e))?; + let query_embeds = self.text.embed(&query)?; + let query_positions = + text_position_ids(target_position, GLM_MTP_QUERY_LEN, &self.device)?; + let (target_hidden, target_tokens) = self.text.forward_verification_tokens( + &query_embeds, + &query_positions, + &self.lm_head, + )?; + let target_tokens = target_tokens.to_vec1::().map_err(|e| { + candle_to_ocr_inference("GLM-OCR", "copy MTP verification tokens", e) + })?; + if target_tokens.len() != GLM_MTP_QUERY_LEN { + return Err(OCRError::InvalidInput { + message: format!( + "GLM-OCR: verification returned {} tokens, expected {}", + target_tokens.len(), + GLM_MTP_QUERY_LEN + ), + }); + } + + rounds += 1; + let mut accepted = 0usize; + let mut stop = false; + while accepted < GLM_MTP_DRAFT_TOKENS && drafts[accepted] == target_tokens[accepted] { + let token = drafts[accepted]; + accepted += 1; + accepted_drafts += 1; + accepted_by_position[accepted - 1] += 1; + if self.eos_token_ids.contains(&token) { + stop = true; break; } - generated.push(tok); - - let token = Tensor::new(&[tok], &self.device) - .and_then(|t| t.reshape((1, 1))) - .map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: create next token", - e, - ) - })?; - let embed = self.text.embed(&token)?; - - let pos_val = pos + rope_delta; - let pos_ids = Tensor::new(vec![pos_val, pos_val, pos_val], &self.device) - .and_then(|t| t.reshape((3, 1, 1))) - .map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: create position ids", - e, - ) - })?; - - let hs = self.text.forward(&embed, &pos_ids, None)?; - let last = hs.i((0, 0, ..)).map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: get decode hidden", - e, - ) - })?; - logits = self.logits_from_hidden(&last)?; + generated.push(token); + if generated.len() == max_new_tokens { + stop = true; + break; + } + } + if stop { + break; } - results.push(generated); + let next_token = target_tokens[accepted]; + let keep = accepted + 1; + + // Target verification writes the whole block. Retain only the + // certain current token and its accepted draft prefix. The MTP + // recurrent tail is speculative too, so roll it back to the same + // pre-query base before synchronizing the accepted target span. + self.text.trim_kv_cache(target_cache_len + keep)?; + mtp.trim_kv_cache(target_cache_len)?; + target_cache_len += keep; + target_position += keep as i64; + current_token = next_token; + + if self.eos_token_ids.contains(¤t_token) { + break; + } + + let mut shifted_sync_ids = query_ids[1..keep].to_vec(); + shifted_sync_ids.push(current_token); + let shifted_sync_ids = Tensor::from_vec(shifted_sync_ids, (1, keep), &self.device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP sync ids", e))?; + let sync_hidden = target_hidden + .narrow(1, 0, keep) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP sync target hidden", e))?; + let sync_positions = query_positions + .narrow(2, 0, keep) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP sync positions", e))?; + let (mtp_hidden, first_draft) = + mtp.sync_target_span(&shifted_sync_ids, &sync_hidden, &sync_positions, false)?; + drafts = self.complete_mtp_drafts(mtp, &mtp_hidden, &first_draft, target_position)?; } - Ok(results) + if rounds > 0 { + tracing::debug!( + rounds, + accepted_drafts, + ?accepted_by_position, + mean_acceptance_length = 1.0 + accepted_drafts as f64 / rounds as f64, + "GLM-OCR MTP acceptance" + ); + } + Ok(generated) + } + + #[cfg(feature = "cuda")] + fn complete_mtp_drafts( + &self, + mtp: &GlmOcrMtpModel, + first_hidden: &Tensor, + first_tokens: &Tensor, + target_position: i64, + ) -> Result, OCRError> { + let seq_len = first_hidden + .dim(1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP result length", e))?; + let mut hidden = first_hidden + .narrow(1, seq_len - 1, 1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP final hidden state", e))?; + let token_count = first_tokens + .dim(0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP token count", e))?; + let mut token = first_tokens + .narrow(0, token_count - 1, 1) + .and_then(|token| token.reshape((1, 1))) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP first proposal", e))?; + let mut draft_tensors = Vec::with_capacity(GLM_MTP_DRAFT_TOKENS); + + while draft_tensors.len() < GLM_MTP_DRAFT_TOKENS { + // CUDA-graph replay overwrites its captured output storage. Keep + // each proposal in independent storage before launching the next + // recurrent step, otherwise earlier draft handles would silently + // observe the newest token. + draft_tensors.push( + token + .copy() + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "save MTP proposal", e))?, + ); + if draft_tensors.len() == GLM_MTP_DRAFT_TOKENS { + break; + } + let position = text_position_ids( + target_position + draft_tensors.len() as i64 - 1, + 1, + &self.device, + )?; + let (next_hidden, next_token) = mtp.predict_single(&token, &hidden, &position)?; + hidden = next_hidden; + token = next_token + .reshape((1, 1)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP proposal shape", e))?; + } + + let refs: Vec<&Tensor> = draft_tensors.iter().collect(); + Tensor::cat(&refs, 1) + .and_then(|drafts| drafts.flatten_all()) + .and_then(|drafts| drafts.to_vec1::()) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy MTP proposals", e)) } pub fn decode_tokens(&self, tokens: &[u32]) -> Result { @@ -466,6 +739,33 @@ impl GlmOcr { } } +fn token_tensor(token: u32, device: &Device) -> Result { + Tensor::new(&[token], device) + .and_then(|token| token.reshape((1, 1))) + .map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: create token tensor", + e, + ) + }) +} + +fn text_position_ids(start: i64, len: usize, device: &Device) -> Result { + let positions: Vec = (0..len).map(|offset| start + offset as i64).collect(); + let mut data = Vec::with_capacity(3 * len); + data.extend_from_slice(&positions); + data.extend_from_slice(&positions); + data.extend_from_slice(&positions); + Tensor::from_vec(data, (3, 1, len), device).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: create text position ids", + e, + ) + }) +} + fn build_prompt(instruction: &str) -> String { format!( "[gMASK]<|user|>\n<|begin_of_image|><|image|><|end_of_image|>{instruction}<|assistant|>\n" diff --git a/oar-ocr-vl/src/glmocr/mtp.rs b/oar-ocr-vl/src/glmocr/mtp.rs new file mode 100644 index 0000000..8e26e6e --- /dev/null +++ b/oar-ocr-vl/src/glmocr/mtp.rs @@ -0,0 +1,404 @@ +//! GLM-OCR multi-token predictor used for lossless greedy speculation. +//! +//! The checkpoint stores one trained MTP layer after the 16 target decoder +//! layers. It is recurrently reused to propose several tokens; the caller +//! chooses the verification width independently of the checkpoint layer. + +use super::config::GlmOcrTextConfig; +use super::text::{GlmOcrTextDecoderLayer, GlmOcrTextRotaryEmbedding}; +use crate::decoder_graph::{CudaGraphKvLengths, cuda_graph_error, sync_graph_tensor}; +use crate::utils::candle_to_ocr_inference; +use candle_core::{D, DType, Device, Tensor}; +use candle_nn::{ + Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, linear_no_bias, rms_norm, +}; +use oar_ocr_core::core::OCRError; +use std::cell::RefCell; + +struct GlmMtpCudaGraph { + // The graph owns device pointers into all tensors below. Drop it first. + graph: candle_core::cuda_backend::cudarc::driver::CudaGraph, + token_input: Tensor, + previous_hidden_input: Tensor, + position_input: Tensor, + _query_lengths: Tensor, + kv_lengths: CudaGraphKvLengths, + hidden_output: Tensor, + token_output: Tensor, + cache_len: usize, +} + +impl std::fmt::Debug for GlmMtpCudaGraph { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GlmMtpCudaGraph") + .field("cache_len", &self.cache_len) + .finish_non_exhaustive() + } +} + +#[derive(Debug)] +pub(super) struct GlmOcrMtpModel { + graph: RefCell>, + embed_tokens: Embedding, + enorm: RmsNorm, + hnorm: RmsNorm, + eh_proj: Linear, + layer: GlmOcrTextDecoderLayer, + shared_norm: RmsNorm, + shared_head: Linear, + rotary_emb: GlmOcrTextRotaryEmbedding, +} + +impl GlmOcrMtpModel { + pub(super) fn load( + cfg: &GlmOcrTextConfig, + layer_index: usize, + vb_language_model: VarBuilder, + ) -> Result { + let vb = vb_language_model.pp("layers").pp(layer_index); + let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens")) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "load MTP embeddings", e))?; + let enorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("enorm")) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "load MTP embedding norm", e))?; + let hnorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("hnorm")) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "load MTP hidden norm", e))?; + let eh_proj = linear_no_bias(cfg.hidden_size * 2, cfg.hidden_size, vb.pp("eh_proj")) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "load MTP fusion projection", e))?; + let layer = GlmOcrTextDecoderLayer::load(cfg, vb.clone())?; + let shared_norm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("shared_head").pp("norm"), + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "load MTP shared norm", e))?; + let shared_head = linear_no_bias( + cfg.hidden_size, + cfg.vocab_size, + vb.pp("shared_head").pp("head"), + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "load MTP shared head", e))?; + let rotary_emb = GlmOcrTextRotaryEmbedding::new(cfg, vb.device())?; + + Ok(Self { + graph: RefCell::new(None), + embed_tokens, + enorm, + hnorm, + eh_proj, + layer, + shared_norm, + shared_head, + rotary_emb, + }) + } + + fn fuse_inputs( + &self, + input_ids: &Tensor, + previous_hidden_states: &Tensor, + mask_first_position: bool, + ) -> Result { + let mut inputs_embeds = self + .embed_tokens + .forward(input_ids) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP embedding", e))?; + if mask_first_position { + let seq_len = inputs_embeds + .dim(1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP sequence length", e))?; + let hidden_size = inputs_embeds + .dim(2) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP hidden size", e))?; + let zero = Tensor::zeros( + (1, 1, hidden_size), + inputs_embeds.dtype(), + inputs_embeds.device(), + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP zero first embedding", e))?; + inputs_embeds = if seq_len == 1 { + zero + } else { + let tail = inputs_embeds + .narrow(1, 1, seq_len - 1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP embedding tail", e))?; + Tensor::cat(&[&zero, &tail], 1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP masked embeddings", e))? + }; + } + + let inputs_embeds = self + .enorm + .forward(&inputs_embeds) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP embedding norm", e))?; + let previous_hidden_states = self + .hnorm + .forward(previous_hidden_states) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP hidden norm", e))?; + let fused = Tensor::cat(&[&inputs_embeds, &previous_hidden_states], D::Minus1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP concatenate inputs", e))?; + self.eh_proj + .forward(&fused) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP fusion projection", e)) + } + + fn tokens_from_hidden(&self, hidden_states: &Tensor) -> Result { + let hidden_states = self + .shared_norm + .forward(hidden_states) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP shared norm", e))?; + self.shared_head + .forward(&hidden_states) + .and_then(|logits| logits.squeeze(0)) + .and_then(|logits| logits.argmax(D::Minus1)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP shared head argmax", e)) + } + + fn forward_tokens( + &self, + input_ids: &Tensor, + previous_hidden_states: &Tensor, + position_ids: &Tensor, + mask_first_position: bool, + ) -> Result<(Tensor, Tensor), OCRError> { + let hidden_states = + self.fuse_inputs(input_ids, previous_hidden_states, mask_first_position)?; + let (cos, sin) = self.rotary_emb.forward(&hidden_states, position_ids)?; + let hidden_states = self.layer.forward(&hidden_states, &cos, &sin, None)?; + let tokens = self.tokens_from_hidden(&hidden_states)?; + Ok((hidden_states, tokens)) + } + + fn forward_dynamic( + &self, + input_ids: &Tensor, + previous_hidden_states: &Tensor, + position_ids: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result<(Tensor, Tensor), OCRError> { + let hidden_states = self.fuse_inputs(input_ids, previous_hidden_states, false)?; + let (cos, sin) = self.rotary_emb.forward(&hidden_states, position_ids)?; + let hidden_states = + self.layer + .forward_dynamic(&hidden_states, &cos, &sin, query_lengths, kv_lengths)?; + let tokens = self.tokens_from_hidden(&hidden_states)?; + Ok((hidden_states, tokens)) + } + + /// Synchronize MTP with a target-model span and return its first proposal. + pub(super) fn sync_target_span( + &self, + shifted_input_ids: &Tensor, + target_hidden_states: &Tensor, + position_ids: &Tensor, + mask_first_position: bool, + ) -> Result<(Tensor, Tensor), OCRError> { + self.forward_tokens( + shifted_input_ids, + target_hidden_states, + position_ids, + mask_first_position, + ) + } + + /// Recurrently propose one more token from the preceding MTP hidden state. + pub(super) fn predict_single( + &self, + input_id: &Tensor, + previous_hidden_state: &Tensor, + position_ids: &Tensor, + ) -> Result<(Tensor, Tensor), OCRError> { + let kv_len = self.kv_cache_len().saturating_add(1); + if let Some(output) = + self.replay_cuda_graph(input_id, previous_hidden_state, position_ids, kv_len)? + { + return Ok(output); + } + self.forward_tokens(input_id, previous_hidden_state, position_ids, false) + } + + pub(super) fn prepare_cuda_graph(&self, cache_len: usize) -> Result<(), OCRError> { + if std::env::var_os("OAR_VL_DISABLE_CUDA_GRAPH").is_some() + || std::env::var_os("OAR_GLMOCR_DISABLE_CUDA_GRAPH").is_some() + { + self.invalidate_cuda_graph(); + return Ok(()); + } + if !self.embed_tokens.embeddings().device().is_cuda() + || !matches!( + self.embed_tokens.embeddings().dtype(), + DType::BF16 | DType::F16 + ) + { + self.invalidate_cuda_graph(); + return Ok(()); + } + let reusable = self + .graph + .borrow() + .as_ref() + .is_some_and(|graph| graph.cache_len == cache_len); + if reusable { + return Ok(()); + } + self.invalidate_cuda_graph(); + self.capture_cuda_graph(cache_len) + } + + fn capture_cuda_graph(&self, cache_len: usize) -> Result<(), OCRError> { + use candle_core::cuda_backend::cudarc::driver::sys::{ + CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, + }; + + let Device::Cuda(cuda) = self.embed_tokens.embeddings().device() else { + return Ok(()); + }; + let query_len = 1; + self.layer.prepare_dynamic_cache(query_len, cache_len)?; + let hidden_size = self + .embed_tokens + .embeddings() + .dim(1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP graph hidden size", e))?; + let device = self.embed_tokens.embeddings().device(); + let token_input = Tensor::zeros((1, 1), DType::U32, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP graph token", e))?; + let previous_hidden_input = Tensor::zeros( + (1, 1, hidden_size), + self.embed_tokens.embeddings().dtype(), + device, + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP graph hidden input", e))?; + let position_input = Tensor::zeros((3, 1, 1), DType::I64, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP graph positions", e))?; + let query_lengths = Tensor::new(&[0u32, 1u32], device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP query lengths", e))?; + let kv_lengths = CudaGraphKvLengths::new(query_len, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "MTP KV lengths", e))?; + let stream = cuda.cuda_stream(); + let _htod_cache = cuda.enable_cuda_graph_htod_cache(); + + let (_, warm_token) = self.forward_dynamic( + &token_input, + &previous_hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + sync_graph_tensor("GLM-OCR", &warm_token, "warm MTP CUDA graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| cuda_graph_error("GLM-OCR", "begin MTP CUDA graph capture", e))?; + let captured_output = self.forward_dynamic( + &token_input, + &previous_hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + ); + let (hidden_output, token_output) = match captured_output { + Ok(output) => output, + Err(error) => { + let _ = stream.end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + return Err(error); + } + }; + let graph = stream + .end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .map_err(|e| cuda_graph_error("GLM-OCR", "end MTP CUDA graph capture", e))? + .ok_or_else(|| OCRError::ConfigError { + message: "GLM-OCR MTP capture returned no graph".to_string(), + })?; + graph + .launch() + .map_err(|e| cuda_graph_error("GLM-OCR", "warm MTP CUDA graph", e))?; + sync_graph_tensor("GLM-OCR", &token_output, "sync MTP CUDA graph")?; + self.clear_kv_cache(); + *self.graph.borrow_mut() = Some(GlmMtpCudaGraph { + graph, + token_input, + previous_hidden_input, + position_input, + _query_lengths: query_lengths, + kv_lengths, + hidden_output, + token_output, + cache_len, + }); + Ok(()) + } + + fn replay_cuda_graph( + &self, + input_id: &Tensor, + previous_hidden_state: &Tensor, + position_ids: &Tensor, + kv_len: usize, + ) -> Result, OCRError> { + let captured_ref = self.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 input_id.shape() != captured.token_input.shape() + || previous_hidden_state.shape() != captured.previous_hidden_input.shape() + || position_ids.shape() != captured.position_input.shape() + { + return Ok(None); + } + captured + .token_input + .slice_set(input_id, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy MTP token", e))?; + captured + .previous_hidden_input + .slice_set(previous_hidden_state, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy MTP hidden", e))?; + captured + .position_input + .slice_set(position_ids, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy MTP positions", e))?; + captured + .kv_lengths + .update(kv_len) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "update MTP KV lengths", e))?; + captured + .graph + .launch() + .map_err(|e| cuda_graph_error("GLM-OCR", "launch MTP CUDA graph", e))?; + self.layer.set_kv_cache_len(kv_len)?; + Ok(Some(( + captured.hidden_output.clone(), + captured.token_output.clone(), + ))) + } + + fn invalidate_cuda_graph(&self) { + self.graph.borrow_mut().take(); + } + + pub(super) fn disable_cuda_graph(&self) { + self.invalidate_cuda_graph(); + } + + pub(super) fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { + self.layer.trim_kv_cache(len) + } + + pub(super) fn clear_kv_cache(&self) { + self.layer.clear_kv_cache(); + } + + pub(super) fn kv_cache_len(&self) -> usize { + self.layer.kv_cache_len() + } +} diff --git a/oar-ocr-vl/src/glmocr/text.rs b/oar-ocr-vl/src/glmocr/text.rs index 9805a4e..09c6f9d 100644 --- a/oar-ocr-vl/src/glmocr/text.rs +++ b/oar-ocr-vl/src/glmocr/text.rs @@ -1,5 +1,14 @@ use super::config::GlmOcrTextConfig; -use crate::attention::{repeat_kv, scaled_dot_product_attention}; +use crate::attention::{flash_attention, scaled_dot_product_attention_gqa}; +#[cfg(any(feature = "cuda", test))] +use crate::decoder_graph::decoder_cache_capacity; +#[cfg(feature = "cuda")] +use crate::decoder_graph::{ + CudaGraphKvLengths, SingleTokenDecoderCudaGraph, cuda_graph_error, decoder_attention_is_causal, + sync_graph_tensor, +}; +#[cfg(feature = "cuda")] +use crate::hunyuanocr::dynamic_kv::DynamicKvAppend; use crate::kv_trim::TrimmableKvCache; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing}; use candle_core::{D, DType, Device, IndexOp, Tensor}; @@ -9,6 +18,14 @@ use candle_nn::{ use oar_ocr_core::core::OCRError; use std::cell::RefCell; +#[cfg(any(feature = "cuda", test))] +const GLM_DECODE_CACHE_LEN: usize = 16_384; + +#[cfg(any(feature = "cuda", test))] +fn ar_cuda_graph_capacity(prompt_len: usize, max_new_tokens: usize) -> Option { + decoder_cache_capacity(prompt_len, max_new_tokens, GLM_DECODE_CACHE_LEN) +} + fn rotate_half_interleaved(x: &Tensor) -> Result { let (b, h, s, d) = x.dims4().map_err(|e| { candle_to_ocr_processing( @@ -67,7 +84,7 @@ fn rotate_half_interleaved(x: &Tensor) -> Result { }) } -fn apply_rotary_pos_emb( +pub(super) fn apply_rotary_pos_emb( q: &Tensor, k: &Tensor, cos: &Tensor, @@ -294,13 +311,13 @@ fn apply_mrope(freqs: &Tensor, mrope_section: &[usize]) -> Result, } impl GlmOcrTextRotaryEmbedding { - fn new(cfg: &GlmOcrTextConfig, device: &Device) -> Result { + pub(super) fn new(cfg: &GlmOcrTextConfig, device: &Device) -> Result { if cfg.rope_parameters.rope_type != "default" { return Err(OCRError::ConfigError { message: format!( @@ -349,7 +366,11 @@ impl GlmOcrTextRotaryEmbedding { }) } - fn forward(&self, x: &Tensor, position_ids: &Tensor) -> Result<(Tensor, Tensor), OCRError> { + pub(super) fn forward( + &self, + x: &Tensor, + position_ids: &Tensor, + ) -> Result<(Tensor, Tensor), OCRError> { let dtype = x.dtype(); let dims = position_ids.dims(); if dims.len() != 3 || dims[0] != 3 { @@ -571,13 +592,12 @@ impl GlmOcrTextAttention { }) } - fn forward( + fn project_qkv( &self, hidden_states: &Tensor, cos: &Tensor, sin: &Tensor, - attention_mask: Option<&Tensor>, - ) -> Result { + ) -> Result<(Tensor, Tensor, Tensor), OCRError> { let (b, seq_len, _) = hidden_states.dims3().map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, @@ -691,31 +711,61 @@ impl GlmOcrTextAttention { ) })?; - let (k, v) = self.kv_cache.borrow_mut().append(&k, &v).map_err(|e| { + Ok((q, k, v)) + } + + fn forward( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b, seq_len, _) = hidden_states.dims3().map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: kv_cache append", + "GLM-OCR: attn hidden dims3", e, ) })?; - let k = repeat_kv(&k, self.num_kv_groups).map_err(|e| { + let (q, k, v) = self.project_qkv(hidden_states, cos, sin)?; + + let (k, v) = self.kv_cache.borrow_mut().append(&k, &v).map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: repeat_kv k", + "GLM-OCR: kv_cache append", e, ) })?; - let v = repeat_kv(&v, self.num_kv_groups).map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "GLM-OCR: repeat_kv v", - e, + let flash = if b == 1 { + flash_attention(&q, &k, &v, self.scaling, seq_len > 1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text flash attention", e))? + } else { + None + }; + let attn = match flash { + Some(attn) => attn, + None => scaled_dot_product_attention_gqa( + &q, + &k, + &v, + attention_mask, + self.scaling, + true, + self.num_kv_groups, ) - })?; + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text grouped-query attention", e))?, + }; - let attn = scaled_dot_product_attention(&q, &k, &v, attention_mask, self.scaling, true) - .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text attention", e))?; + self.project_attention_output(&attn, b, seq_len) + } + fn project_attention_output( + &self, + attn: &Tensor, + batch: usize, + seq_len: usize, + ) -> Result { let attn = attn .transpose(1, 2) .map_err(|e| { @@ -725,7 +775,7 @@ impl GlmOcrTextAttention { e, ) })? - .reshape((b, seq_len, self.num_heads * self.head_dim)) + .reshape((batch, seq_len, self.num_heads * self.head_dim)) .map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, @@ -739,13 +789,105 @@ impl GlmOcrTextAttention { .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text o_proj forward", e)) } + #[cfg(feature = "cuda")] + fn prepare_dynamic_cache(&self, query_len: usize, cache_len: usize) -> Result<(), OCRError> { + let template = Tensor::zeros( + (1, self.num_kv_heads, query_len, self.head_dim), + self.k_proj.weight().dtype(), + self.k_proj.weight().device(), + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic KV template", e))?; + self.kv_cache + .borrow_mut() + .initialize_storage_with_capacity(&template, cache_len) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "initialize dynamic KV", e)) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let (batch, query_len, _) = hidden_states + .dims3() + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic attention hidden shape", e))?; + if batch != 1 { + return Err(OCRError::ConfigError { + message: "GLM-OCR CUDA-graph attention requires batch size 1".to_string(), + }); + } + let (q, k, v) = self.project_qkv(hidden_states, cos, sin)?; + 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: "GLM-OCR dynamic KV storage is not initialized".to_string(), + })?; + drop(cache); + let append = DynamicKvAppend { + query_len, + cache_len, + }; + cache_k + .inplace_op3(&k, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic key cache append", e))?; + cache_v + .inplace_op3(&v, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic value cache append", e))?; + + let q = q + .squeeze(0) + .and_then(|q| q.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic Q layout", e))?; + let cache_k = cache_k + .squeeze(0) + .and_then(|k| k.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic K layout", e))?; + let cache_v = cache_v + .squeeze(0) + .and_then(|v| v.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic V layout", e))?; + let attn = candle_flash_attn::flash_attn_varlen( + &q, + &cache_k, + &cache_v, + query_lengths, + kv_lengths, + query_len, + cache_len, + self.scaling as f32, + decoder_attention_is_causal(query_len), + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic flash attention", e))? + .transpose(0, 1) + .and_then(|attn| attn.unsqueeze(0)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic attention layout", e))?; + self.project_attention_output(&attn, batch, query_len) + } + fn clear_kv_cache(&self) { self.kv_cache.borrow_mut().reset(); } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + self.kv_cache.borrow().current_seq_len() + } + + #[cfg(feature = "cuda")] + fn set_kv_cache_len(&self, len: usize) -> Result<(), OCRError> { + self.kv_cache + .borrow_mut() + .set_current_len(len) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "set dynamic KV length", e)) + } } #[derive(Debug, Clone)] -struct GlmOcrTextDecoderLayer { +pub(super) struct GlmOcrTextDecoderLayer { self_attn: GlmOcrTextAttention, mlp: GlmOcrTextMLP, input_layernorm: RmsNorm, @@ -755,7 +897,7 @@ struct GlmOcrTextDecoderLayer { } impl GlmOcrTextDecoderLayer { - fn load(cfg: &GlmOcrTextConfig, vb: VarBuilder) -> Result { + pub(super) fn load(cfg: &GlmOcrTextConfig, vb: VarBuilder) -> Result { let self_attn = GlmOcrTextAttention::load(cfg, vb.clone())?; let mlp = GlmOcrTextMLP::load(cfg, vb.clone())?; let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm")) @@ -788,7 +930,7 @@ impl GlmOcrTextDecoderLayer { }) } - fn forward( + pub(super) fn forward( &self, hidden_states: &Tensor, cos: &Tensor, @@ -835,13 +977,120 @@ impl GlmOcrTextDecoderLayer { }) } - fn clear_kv_cache(&self) { + #[cfg(feature = "cuda")] + pub(super) fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let residual = hidden_states.clone(); + let hidden = self + .input_layernorm + .forward(hidden_states) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text input_layernorm forward", e))?; + let hidden = + self.self_attn + .forward_dynamic(&hidden, cos, sin, query_lengths, kv_lengths)?; + let hidden = self + .post_self_attn_layernorm + .forward(&hidden) + .map_err(|e| { + candle_to_ocr_inference("GLM-OCR", "text post_self_attn_layernorm forward", e) + })?; + let hidden = (&residual + &hidden).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: text residual add", + e, + ) + })?; + + let residual = hidden.clone(); + let hidden = self + .post_attention_layernorm + .forward(&hidden) + .map_err(|e| { + candle_to_ocr_inference("GLM-OCR", "text post_attention_layernorm forward", e) + })?; + let hidden = self.mlp.forward(&hidden)?; + let hidden = self.post_mlp_layernorm.forward(&hidden).map_err(|e| { + candle_to_ocr_inference("GLM-OCR", "text post_mlp_layernorm forward", e) + })?; + (&residual + &hidden).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "GLM-OCR: text residual add mlp", + e, + ) + }) + } + + pub(super) fn clear_kv_cache(&self) { self.self_attn.clear_kv_cache(); } + + #[cfg(feature = "cuda")] + pub(super) fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { + self.self_attn + .kv_cache + .borrow_mut() + .trim_to(len) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "trim dynamic KV cache", e)) + } + + #[cfg(feature = "cuda")] + pub(super) fn prepare_dynamic_cache( + &self, + query_len: usize, + cache_len: usize, + ) -> Result<(), OCRError> { + self.self_attn.prepare_dynamic_cache(query_len, cache_len) + } + + #[cfg(feature = "cuda")] + pub(super) fn kv_cache_len(&self) -> usize { + self.self_attn.kv_cache_len() + } + + #[cfg(feature = "cuda")] + pub(super) fn set_kv_cache_len(&self, len: usize) -> Result<(), OCRError> { + self.self_attn.set_kv_cache_len(len) + } } -#[derive(Debug, Clone)] +#[cfg(feature = "cuda")] +struct GlmVerificationCudaGraph { + // The graph owns device pointers into all tensors below. Drop it first. + graph: candle_core::cuda_backend::cudarc::driver::CudaGraph, + hidden_input: Tensor, + position_input: Tensor, + _query_lengths: Tensor, + kv_lengths: CudaGraphKvLengths, + hidden_output: Tensor, + token_output: Tensor, + cache_len: usize, + query_len: usize, +} + +#[cfg(feature = "cuda")] +impl std::fmt::Debug for GlmVerificationCudaGraph { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GlmVerificationCudaGraph") + .field("query_len", &self.query_len) + .field("cache_len", &self.cache_len) + .finish_non_exhaustive() + } +} + +#[derive(Debug)] pub struct GlmOcrTextModel { + #[cfg(feature = "cuda")] + decode_graph: RefCell>, + #[cfg(feature = "cuda")] + verification_graph: RefCell>, embed_tokens: Embedding, layers: Vec, norm: RmsNorm, @@ -863,6 +1112,10 @@ impl GlmOcrTextModel { .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text norm", e))?; Ok(Self { + #[cfg(feature = "cuda")] + decode_graph: RefCell::new(None), + #[cfg(feature = "cuda")] + verification_graph: RefCell::new(None), embed_tokens, layers, norm, @@ -887,18 +1140,537 @@ impl GlmOcrTextModel { attention_mask: Option<&Tensor>, ) -> Result { let (cos, sin) = self.rotary_emb.forward(inputs_embeds, position_ids)?; + self.forward_prepared(inputs_embeds, &cos, &sin, attention_mask) + } + + fn forward_prepared( + &self, + inputs_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { let mut hidden_states = inputs_embeds.clone(); for layer in &self.layers { - hidden_states = layer.forward(&hidden_states, &cos, &sin, attention_mask)?; + hidden_states = layer.forward(&hidden_states, cos, sin, attention_mask)?; } self.norm .forward(&hidden_states) .map_err(|e| candle_to_ocr_inference("GLM-OCR", "text norm forward", e)) } + fn project_logits(&self, hidden_states: &Tensor, lm_head: &Linear) -> Result { + lm_head + .forward(hidden_states) + .and_then(|logits| logits.i((0, 0, ..))) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "decode LM head", e)) + } + + #[cfg(feature = "cuda")] + fn project_all_logits( + &self, + hidden_states: &Tensor, + lm_head: &Linear, + ) -> Result { + lm_head + .forward(hidden_states) + .and_then(|logits| logits.squeeze(0)) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification LM head", e)) + } + + pub(crate) fn forward_decode_logits( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + lm_head: &Linear, + ) -> Result { + #[cfg(feature = "cuda")] + { + let kv_len = self.kv_cache_len().saturating_add(1); + if let Some(logits) = self.replay_cuda_graph(inputs_embeds, position_ids, kv_len)? { + return Ok(logits); + } + } + let hidden = self.forward(inputs_embeds, position_ids, None)?; + self.project_logits(&hidden, lm_head) + } + + /// Verify a fixed block of speculative tokens in one causal target pass. + /// + /// Returned hidden states are the target model's post-final-norm states; + /// GLM-OCR's MTP layer consumes exactly these states on its next sync pass. + #[cfg(feature = "cuda")] + pub(crate) fn forward_verification_tokens( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + lm_head: &Linear, + ) -> Result<(Tensor, Tensor), OCRError> { + let query_len = inputs_embeds + .dim(1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification query length", e))?; + let kv_len = self.kv_cache_len().saturating_add(query_len); + if let Some(output) = + self.replay_verification_cuda_graph(inputs_embeds, position_ids, kv_len)? + { + return Ok(output); + } + + let hidden = self.forward(inputs_embeds, position_ids, None)?; + let tokens = self + .project_all_logits(&hidden, lm_head)? + .argmax(D::Minus1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification argmax", e))?; + Ok((hidden, tokens)) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let (cos, sin) = self.rotary_emb.forward(inputs_embeds, position_ids)?; + let mut hidden_states = inputs_embeds.clone(); + for layer in &self.layers { + hidden_states = + layer.forward_dynamic(&hidden_states, &cos, &sin, query_lengths, kv_lengths)?; + } + self.norm + .forward(&hidden_states) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "dynamic text norm", e)) + } + + pub(crate) fn prepare_ar_cuda_graph( + &self, + prompt_len: usize, + max_new_tokens: usize, + lm_head: &Linear, + ) -> Result<(), OCRError> { + if std::env::var_os("OAR_VL_DISABLE_CUDA_GRAPH").is_some() + || std::env::var_os("OAR_GLMOCR_DISABLE_CUDA_GRAPH").is_some() + { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); + return Ok(()); + } + #[cfg(feature = "cuda")] + if self.embed_tokens.embeddings().device().is_cuda() + && matches!( + self.embed_tokens.embeddings().dtype(), + DType::BF16 | DType::F16 + ) + { + 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(GLM_DECODE_CACHE_LEN); + let reusable = self + .decode_graph + .borrow() + .as_ref() + .is_some_and(|graph| graph.cache_len >= required); + if reusable { + return Ok(()); + } + self.invalidate_cuda_graph(); + self.capture_cuda_graph(cache_len, lm_head)?; + } + let _ = prompt_len; + let _ = max_new_tokens; + let _ = lm_head; + Ok(()) + } + + /// Capture the target-model verification block used by greedy MTP. + #[cfg(feature = "cuda")] + pub(crate) fn prepare_verification_cuda_graph( + &self, + prompt_len: usize, + max_new_tokens: usize, + query_len: usize, + lm_head: &Linear, + ) -> Result, OCRError> { + if query_len == 0 + || std::env::var_os("OAR_VL_DISABLE_CUDA_GRAPH").is_some() + || std::env::var_os("OAR_GLMOCR_DISABLE_CUDA_GRAPH").is_some() + { + self.invalidate_cuda_graph(); + return Ok(None); + } + if !self.embed_tokens.embeddings().device().is_cuda() + || !matches!( + self.embed_tokens.embeddings().dtype(), + DType::BF16 | DType::F16 + ) + { + self.invalidate_cuda_graph(); + return Ok(None); + } + + // Verification can temporarily place a complete query block in KV + // beyond the user-visible output limit, so reserve one extra block. + let Some(cache_len) = + ar_cuda_graph_capacity(prompt_len, max_new_tokens.saturating_add(query_len)) + else { + self.invalidate_cuda_graph(); + return Ok(None); + }; + let required = prompt_len + .saturating_add(max_new_tokens) + .saturating_add(query_len) + .min(GLM_DECODE_CACHE_LEN); + let retained_cache_len = self + .verification_graph + .borrow() + .as_ref() + .filter(|graph| graph.query_len == query_len && graph.cache_len >= required) + .map(|graph| graph.cache_len); + if let Some(retained_cache_len) = retained_cache_len { + // The MTP graph shares this capacity contract. Returning the newly + // computed (possibly smaller) bucket would force it to recapture + // even though the retained target graph is still reusable. + return Ok(Some(retained_cache_len)); + } + + self.invalidate_cuda_graph(); + self.capture_verification_cuda_graph(cache_len, query_len, lm_head)?; + Ok(Some(cache_len)) + } + + #[cfg(feature = "cuda")] + fn capture_cuda_graph(&self, cache_len: usize, lm_head: &Linear) -> Result<(), OCRError> { + use candle_core::cuda_backend::cudarc::driver::sys::{ + CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, + }; + + if self.decode_graph.borrow().is_some() { + return Ok(()); + } + let Device::Cuda(cuda) = self.embed_tokens.embeddings().device() else { + return Ok(()); + }; + let query_len = 1; + for layer in &self.layers { + layer + .self_attn + .prepare_dynamic_cache(query_len, cache_len)?; + } + let hidden_size = self + .embed_tokens + .embeddings() + .dim(1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "graph hidden size", e))?; + let device = self.embed_tokens.embeddings().device(); + let hidden_input = Tensor::zeros( + (1, query_len, hidden_size), + self.embed_tokens.embeddings().dtype(), + device, + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "graph hidden input", e))?; + let position_input = Tensor::zeros((3, 1, query_len), DType::I64, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "graph position input", e))?; + let query_lengths = Tensor::new(&[0u32, query_len as u32], device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "graph query lengths", e))?; + let kv_lengths = CudaGraphKvLengths::new(query_len, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "graph KV lengths", e))?; + let stream = cuda.cuda_stream(); + let _htod_cache = cuda.enable_cuda_graph_htod_cache(); + + let warm = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + let warm_logits = self.project_logits(&warm, lm_head)?; + sync_graph_tensor("GLM-OCR", &warm_logits, "warm decoder CUDA graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| cuda_graph_error("GLM-OCR", "begin decoder CUDA graph capture", e))?; + let captured_output = (|| { + let hidden = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + self.project_logits(&hidden, lm_head) + })(); + let logits_output = match captured_output { + Ok(output) => output, + Err(error) => { + let _ = stream.end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + return Err(error); + } + }; + let graph = stream + .end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .map_err(|e| cuda_graph_error("GLM-OCR", "end decoder CUDA graph capture", e))? + .ok_or_else(|| OCRError::ConfigError { + message: "GLM-OCR decoder capture returned no graph".to_string(), + })?; + graph + .launch() + .map_err(|e| cuda_graph_error("GLM-OCR", "warm decoder CUDA graph", e))?; + sync_graph_tensor("GLM-OCR", &logits_output, "sync decoder CUDA graph")?; + self.clear_kv_cache(); + *self.decode_graph.borrow_mut() = Some(SingleTokenDecoderCudaGraph { + graph, + hidden_input, + position_input, + _query_lengths: query_lengths, + kv_lengths, + logits_output, + cache_len, + }); + Ok(()) + } + + #[cfg(feature = "cuda")] + fn capture_verification_cuda_graph( + &self, + cache_len: usize, + query_len: usize, + lm_head: &Linear, + ) -> Result<(), OCRError> { + use candle_core::cuda_backend::cudarc::driver::sys::{ + CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, + }; + + let Device::Cuda(cuda) = self.embed_tokens.embeddings().device() else { + return Ok(()); + }; + for layer in &self.layers { + layer.prepare_dynamic_cache(query_len, cache_len)?; + } + let hidden_size = self + .embed_tokens + .embeddings() + .dim(1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "graph hidden size", e))?; + let device = self.embed_tokens.embeddings().device(); + let hidden_input = Tensor::zeros( + (1, query_len, hidden_size), + self.embed_tokens.embeddings().dtype(), + device, + ) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification graph input", e))?; + let position_input = Tensor::zeros((3, 1, query_len), DType::I64, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification graph positions", e))?; + let query_lengths = Tensor::new(&[0u32, query_len as u32], device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification query lengths", e))?; + let kv_lengths = CudaGraphKvLengths::new(query_len, device) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "verification KV lengths", e))?; + let stream = cuda.cuda_stream(); + let _htod_cache = cuda.enable_cuda_graph_htod_cache(); + + let warm = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + let warm_tokens = self + .project_all_logits(&warm, lm_head)? + .argmax(D::Minus1) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "warm verification argmax", e))?; + sync_graph_tensor("GLM-OCR", &warm_tokens, "warm verification CUDA graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| cuda_graph_error("GLM-OCR", "begin verification graph capture", e))?; + let captured_output = (|| { + let hidden = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + let tokens = self + .project_all_logits(&hidden, lm_head)? + .argmax(D::Minus1) + .map_err(|e| { + candle_to_ocr_inference("GLM-OCR", "captured verification argmax", e) + })?; + Ok::<_, OCRError>((hidden, tokens)) + })(); + let (hidden_output, token_output) = match captured_output { + Ok(output) => output, + Err(error) => { + let _ = stream.end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + return Err(error); + } + }; + let graph = stream + .end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .map_err(|e| cuda_graph_error("GLM-OCR", "end verification graph capture", e))? + .ok_or_else(|| OCRError::ConfigError { + message: "GLM-OCR verification capture returned no graph".to_string(), + })?; + graph + .launch() + .map_err(|e| cuda_graph_error("GLM-OCR", "warm verification CUDA graph", e))?; + sync_graph_tensor("GLM-OCR", &token_output, "sync verification CUDA graph")?; + self.clear_kv_cache(); + *self.verification_graph.borrow_mut() = Some(GlmVerificationCudaGraph { + graph, + hidden_input, + position_input, + _query_lengths: query_lengths, + kv_lengths, + hidden_output, + token_output, + cache_len, + query_len, + }); + Ok(()) + } + + #[cfg(feature = "cuda")] + fn replay_cuda_graph( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + kv_len: usize, + ) -> Result, OCRError> { + 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() + || position_ids.shape() != captured.position_input.shape() + { + return Ok(None); + } + captured + .hidden_input + .slice_set(inputs_embeds, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy graph hidden", e))?; + captured + .position_input + .slice_set(position_ids, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy graph positions", e))?; + captured + .kv_lengths + .update(kv_len) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "update graph KV lengths", e))?; + captured + .graph + .launch() + .map_err(|e| cuda_graph_error("GLM-OCR", "launch decoder CUDA graph", e))?; + for layer in &self.layers { + layer.set_kv_cache_len(kv_len)?; + } + Ok(Some(captured.logits_output.clone())) + } + + #[cfg(feature = "cuda")] + fn replay_verification_cuda_graph( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + kv_len: usize, + ) -> Result, OCRError> { + let captured_ref = self.verification_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() + || position_ids.shape() != captured.position_input.shape() + { + return Ok(None); + } + captured + .hidden_input + .slice_set(inputs_embeds, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy verification hidden", e))?; + captured + .position_input + .slice_set(position_ids, 0, 0) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "copy verification positions", e))?; + captured + .kv_lengths + .update(kv_len) + .map_err(|e| candle_to_ocr_inference("GLM-OCR", "update verification KV lengths", e))?; + captured + .graph + .launch() + .map_err(|e| cuda_graph_error("GLM-OCR", "launch verification CUDA graph", e))?; + for layer in &self.layers { + layer.set_kv_cache_len(kv_len)?; + } + Ok(Some(( + captured.hidden_output.clone(), + captured.token_output.clone(), + ))) + } + + #[cfg(feature = "cuda")] + fn invalidate_cuda_graph(&self) { + self.decode_graph.borrow_mut().take(); + self.verification_graph.borrow_mut().take(); + } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + let len = self.layers.first().map_or(0, |layer| layer.kv_cache_len()); + debug_assert!(self.layers.iter().all(|layer| layer.kv_cache_len() == len)); + len + } + pub fn clear_kv_cache(&self) { for layer in &self.layers { layer.clear_kv_cache(); } } + + #[cfg(feature = "cuda")] + pub(crate) fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { + for layer in &self.layers { + layer.trim_kv_cache(len)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::{GLM_DECODE_CACHE_LEN, ar_cuda_graph_capacity}; + + #[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(GLM_DECODE_CACHE_LEN) + ); + assert_eq!(ar_cuda_graph_capacity(100, 0), None); + assert_eq!(ar_cuda_graph_capacity(GLM_DECODE_CACHE_LEN, 1), None); + } } diff --git a/oar-ocr-vl/src/hunyuanocr/dflash.rs b/oar-ocr-vl/src/hunyuanocr/dflash.rs index 8ec48e5..83d8775 100644 --- a/oar-ocr-vl/src/hunyuanocr/dflash.rs +++ b/oar-ocr-vl/src/hunyuanocr/dflash.rs @@ -7,10 +7,14 @@ #[cfg(feature = "cuda")] use super::dynamic_kv::{ - ArgmaxFirstBf16, DynamicPagedKvAppend, FusedAddRmsNormBf16, FusedRmsNormRopeBf16, - FusedRopeBf16, FusedSiluMulBf16, + DynamicPagedKvAppend, FusedAddRmsNormBf16, FusedRmsNormRopeBf16, FusedRopeBf16, + FusedSiluMulBf16, }; use crate::attention::{RotaryEmbedding, flash_attention, scaled_dot_product_attention_gqa}; +#[cfg(feature = "cuda")] +use crate::cuda_kernels::ArgmaxFirstBf16; +#[cfg(feature = "cuda")] +use crate::decoder_graph::{CudaGraphKvLengths, cuda_graph_error, sync_graph_tensor}; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; use candle_core::{D, DType, Device, Tensor}; use candle_nn::{Linear, Module, RmsNorm, VarBuilder, linear_no_bias, rms_norm}; @@ -721,29 +725,6 @@ impl DFlashMlp { } } -#[cfg(feature = "cuda")] -fn dflash_cuda_graph_error( - context: impl Into, - source: impl std::error::Error + Send + Sync + 'static, -) -> OCRError { - OCRError::Inference { - model_name: "HunyuanOCR DFlash".to_string(), - context: context.into(), - source: Box::new(source), - } -} - -#[cfg(feature = "cuda")] -fn sync_dflash_graph_tensor(tensor: &Tensor, operation: &'static str) -> Result<(), OCRError> { - tensor - .flatten_all() - .and_then(|x| x.narrow(0, 0, 1)) - .and_then(|x| x.to_dtype(DType::F32)) - .and_then(|x| x.to_vec1::()) - .map(|_| ()) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", operation, e)) -} - #[derive(Debug)] struct DFlashLayer { input_layernorm: RmsNorm, @@ -876,7 +857,7 @@ struct DFlashCudaGraph { cos_input: Tensor, sin_input: Tensor, _query_lengths: Tensor, - kv_lengths: Tensor, + kv_lengths: CudaGraphKvLengths, hidden_output: Tensor, proposals_output: Tensor, } @@ -1248,7 +1229,7 @@ impl DFlashModel { .map_err(|e| tensor_err("HunyuanOCR DFlash: full graph sin input", e))?; let query_lengths = Tensor::new(&[0u32, query_len as u32], &self.device) .map_err(|e| tensor_err("HunyuanOCR DFlash: full graph query lengths", e))?; - let kv_lengths = Tensor::new(&[0u32, query_len as u32], &self.device) + let kv_lengths = CudaGraphKvLengths::new(query_len, &self.device) .map_err(|e| tensor_err("HunyuanOCR DFlash: full graph KV lengths", e))?; let stream = cuda.cuda_stream(); let _htod_cache = cuda.enable_cuda_graph_htod_cache(); @@ -1258,21 +1239,27 @@ impl DFlashModel { &cos_input, &sin_input, &query_lengths, - &kv_lengths, + kv_lengths.tensor(), )?; let warm_proposals = self.proposals_from_hidden(&warm)?; - sync_dflash_graph_tensor(&warm_proposals, "warm full draft graph")?; + sync_graph_tensor( + "HunyuanOCR DFlash", + &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))?; + .map_err(|e| { + cuda_graph_error("HunyuanOCR DFlash", "begin full draft graph capture", e) + })?; let captured_output = (|| { let hidden = self.forward_queries_dynamic( &query_input, &cos_input, &sin_input, &query_lengths, - &kv_lengths, + kv_lengths.tensor(), )?; let proposals = self.proposals_from_hidden(&hidden)?; Ok::<_, OCRError>((hidden, proposals)) @@ -1290,14 +1277,18 @@ impl DFlashModel { .end_capture( CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ) - .map_err(|e| dflash_cuda_graph_error("end full draft graph capture", e))? + .map_err(|e| cuda_graph_error("HunyuanOCR DFlash", "end full draft graph capture", e))? .ok_or_else(|| OCRError::ConfigError { message: "HunyuanOCR DFlash full graph capture returned no graph".to_string(), })?; graph .launch() - .map_err(|e| dflash_cuda_graph_error("warm full draft graph", e))?; - sync_dflash_graph_tensor(&proposals_output, "sync full draft graph")?; + .map_err(|e| cuda_graph_error("HunyuanOCR DFlash", "warm full draft graph", e))?; + sync_graph_tensor( + "HunyuanOCR DFlash", + &proposals_output, + "sync full draft graph", + )?; self.clear_context(); *self.decode_graph.borrow_mut() = Some(DFlashCudaGraph { graph, @@ -1346,16 +1337,14 @@ impl DFlashModel { .sin_input .slice_set(sin, 0, 0) .map_err(|e| tensor_err("HunyuanOCR DFlash: copy full graph sin", e))?; - let lengths = Tensor::new(&[0u32, total_kv_len as u32], &self.device) - .map_err(|e| tensor_err("HunyuanOCR DFlash: create full graph KV lengths", e))?; captured .kv_lengths - .slice_set(&lengths, 0, 0) - .map_err(|e| tensor_err("HunyuanOCR DFlash: copy full graph KV lengths", e))?; + .update(total_kv_len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: update full graph KV lengths", e))?; captured .graph .launch() - .map_err(|e| dflash_cuda_graph_error("launch full draft graph", e))?; + .map_err(|e| cuda_graph_error("HunyuanOCR DFlash", "launch full draft graph", e))?; Ok(Some(captured.proposals_output.clone())) } diff --git a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu index d87a9ff..ce5e5f1 100644 --- a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu +++ b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu @@ -160,6 +160,45 @@ extern "C" __global__ void mark_repetition_history_u8( } } +// Apply a common sparse banned-token list to every logits row. This is used +// by exact greedy processors such as no-repeat-ngram without copying the full +// vocabulary to the host. +extern "C" __global__ void mask_token_ids_bf16( + __nv_bfloat16* logits, + const uint32_t* token_ids, + uint32_t rows, + uint32_t vocab_size, + uint32_t token_count) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = rows * token_count; + if (index >= count) { + return; + } + const uint32_t row = index / token_count; + const uint32_t token = token_ids[index - row * token_count]; + if (token < vocab_size) { + logits[row * vocab_size + token] = __float2bfloat16(-CUDART_INF_F); + } +} + +extern "C" __global__ void mask_token_ids_f32( + float* logits, + const uint32_t* token_ids, + uint32_t rows, + uint32_t vocab_size, + uint32_t token_count) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = rows * token_count; + if (index >= count) { + return; + } + const uint32_t row = index / token_count; + const uint32_t token = token_ids[index - row * token_count]; + if (token < vocab_size) { + logits[row * vocab_size + token] = -CUDART_INF_F; + } +} + extern "C" __global__ void argmax_first_bf16_stage1( const __nv_bfloat16* logits, float* partial_values, @@ -684,3 +723,177 @@ extern "C" __global__ void add_rmsnorm_bf16( output[element_count + row_offset + col] = __float2bfloat16_rn( scale * value * __bfloat162float(weight[col])); } + +template +static __device__ __forceinline__ float sampling_logit(const T* logits, + uint32_t index); + +template <> +__device__ __forceinline__ float sampling_logit(const float* logits, + uint32_t index) { + return logits[index]; +} + +template <> +__device__ __forceinline__ float sampling_logit<__nv_bfloat16>( + const __nv_bfloat16* logits, + uint32_t index) { + return __bfloat162float(logits[index]); +} + +// Select one token per logits row and return both the token id and its softmax +// probability. The output is F32 [rows, 2]; vocabulary ids below 2^24 are +// exactly representable, and MinerU's vocabulary is much smaller than that. +// +// Sampling uses a caller-provided U[0, 1) variate. This keeps RNG ownership on +// the host (and therefore preserves seeded/reproducible generation) while the +// O(rows*vocab) softmax/CDF work stays on the GPU. The scan is tiled so all 256 +// lanes evaluate exponentials in parallel instead of making one lane walk the +// full vocabulary. +template +static __device__ void sample_with_confidence_impl( + const T* logits, + const float* uniforms, + float* output, + uint32_t rows, + uint32_t vocab_size, + float inv_temperature, + uint32_t greedy, + float* scratch_values, + uint32_t* scratch_indices) { + const uint32_t row = blockIdx.x; + const uint32_t lane = threadIdx.x; + if (row >= rows) { + return; + } + + const uint32_t row_offset = row * vocab_size; + 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 = sampling_logit(logits, row_offset + col) * + inv_temperature; + if (value > best_value || (value == best_value && col < best_index)) { + best_value = value; + best_index = col; + } + } + + scratch_values[lane] = best_value; + scratch_indices[lane] = best_index; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + const float other_value = scratch_values[lane + width]; + const uint32_t other_index = scratch_indices[lane + width]; + if (other_value > scratch_values[lane] || + (other_value == scratch_values[lane] && + other_index < scratch_indices[lane])) { + scratch_values[lane] = other_value; + scratch_indices[lane] = other_index; + } + } + } + __syncthreads(); + const float max_value = scratch_values[0]; + const uint32_t argmax = + scratch_indices[0] == UINT32_MAX ? 0 : scratch_indices[0]; + + float local_sum = 0.0f; + for (uint32_t col = lane; col < vocab_size; col += blockDim.x) { + const float value = sampling_logit(logits, row_offset + col) * + inv_temperature; + local_sum += expf(value - max_value); + } + scratch_values[lane] = local_sum; + for (uint32_t width = blockDim.x / 2; width > 0; width >>= 1) { + __syncthreads(); + if (lane < width) { + scratch_values[lane] += scratch_values[lane + width]; + } + } + __syncthreads(); + const float sum = scratch_values[0]; + + uint32_t selected = argmax; + if (!greedy && isfinite(sum) && sum > 0.0f) { + const float threshold = uniforms[row] * sum; + float cumulative = 0.0f; + for (uint32_t base = 0; base < vocab_size; base += blockDim.x) { + const uint32_t col = base + lane; + const float weight = col < vocab_size + ? expf(sampling_logit(logits, row_offset + col) * inv_temperature - + max_value) + : 0.0f; + scratch_values[lane] = weight; + __syncthreads(); + + // Inclusive scan of this 256-token tile. + for (uint32_t offset = 1; offset < blockDim.x; offset <<= 1) { + const float add = lane >= offset ? scratch_values[lane - offset] : 0.0f; + __syncthreads(); + if (lane >= offset) { + scratch_values[lane] += add; + } + __syncthreads(); + } + + if (lane == 0) { + scratch_indices[0] = UINT32_MAX; + } + __syncthreads(); + if (col < vocab_size && cumulative + scratch_values[lane] > threshold) { + atomicMin(scratch_indices, col); + } + __syncthreads(); + const uint32_t crossing = scratch_indices[0]; + const float tile_sum = scratch_values[blockDim.x - 1]; + if (crossing != UINT32_MAX) { + selected = crossing; + break; + } + cumulative += tile_sum; + __syncthreads(); + } + } + + if (lane == 0) { + const float selected_value = + sampling_logit(logits, row_offset + selected) * inv_temperature; + const float confidence = isfinite(sum) && sum > 0.0f + ? expf(selected_value - max_value) / sum + : 1.0f; + output[row * 2] = static_cast(selected); + output[row * 2 + 1] = confidence; + } +} + +extern "C" __global__ void sample_with_confidence_f32( + const float* logits, + const float* uniforms, + float* output, + uint32_t rows, + uint32_t vocab_size, + float inv_temperature, + uint32_t greedy) { + __shared__ float scratch_values[256]; + __shared__ uint32_t scratch_indices[256]; + sample_with_confidence_impl(logits, uniforms, output, rows, vocab_size, + inv_temperature, greedy, scratch_values, + scratch_indices); +} + +extern "C" __global__ void sample_with_confidence_bf16( + const __nv_bfloat16* logits, + const float* uniforms, + float* output, + uint32_t rows, + uint32_t vocab_size, + float inv_temperature, + uint32_t greedy) { + __shared__ float scratch_values[256]; + __shared__ uint32_t scratch_indices[256]; + sample_with_confidence_impl(logits, uniforms, output, rows, vocab_size, + inv_temperature, greedy, scratch_values, + scratch_indices); +} diff --git a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs index 78e39b7..2bbeacc 100644 --- a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs +++ b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs @@ -1,12 +1,13 @@ +use crate::cuda_kernels::CUDA_KERNEL_MODULE; use candle_core::backend::BackendStorage; use candle_core::{ - CpuStorage, CustomOp1, CustomOp2, CustomOp3, InplaceOp2, InplaceOp3, Layout, Result, Shape, + CpuStorage, CustomOp2, CustomOp3, InplaceOp2, InplaceOp3, Layout, Result, Shape, }; -const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/hunyuan_dynamic_kv.ptx")); +const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/oar_vl_kernels.ptx")); const PARTITIONS_PER_ROW: usize = 8; -pub(super) struct DynamicKvAppend { +pub(crate) struct DynamicKvAppend { pub query_len: usize, pub cache_len: usize, } @@ -56,10 +57,6 @@ pub(super) struct RepetitionPenaltyF32 { pub penalty: f32, } -pub(super) struct ArgmaxFirstF32; - -pub(super) struct ArgmaxFirstBf16; - pub(super) struct MarkRepetitionHistoryU8; pub(super) struct DFlashRepetitionArgmaxBf16 { @@ -70,134 +67,6 @@ 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" @@ -244,7 +113,7 @@ impl InplaceOp2 for MarkRepetitionHistoryU8 { 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", + CUDA_KERNEL_MODULE, PTX, )?; let mut builder = function.builder(); @@ -312,7 +181,7 @@ impl CustomOp2 for RepetitionArgmaxBf16 { let output = unsafe { device.alloc::(rows) }?; let stage1 = device.get_or_load_custom_func( "repetition_argmax_bf16_stage1", - "hunyuan_dynamic_kv", + CUDA_KERNEL_MODULE, PTX, )?; let mut builder = stage1.builder(); @@ -337,7 +206,7 @@ impl CustomOp2 for RepetitionArgmaxBf16 { .w()?; let stage2 = device.get_or_load_custom_func( "dflash_repetition_argmax_stage2", - "hunyuan_dynamic_kv", + CUDA_KERNEL_MODULE, PTX, )?; let mut builder = stage2.builder(); @@ -431,7 +300,7 @@ impl CustomOp3 for DFlashRepetitionArgmaxBf16 { let output = unsafe { device.alloc::(rows) }?; let stage1 = device.get_or_load_custom_func( "dflash_repetition_argmax_bf16_stage1", - "hunyuan_dynamic_kv", + CUDA_KERNEL_MODULE, PTX, )?; let mut builder = stage1.builder(); @@ -458,7 +327,7 @@ impl CustomOp3 for DFlashRepetitionArgmaxBf16 { .w()?; let stage2 = device.get_or_load_custom_func( "dflash_repetition_argmax_stage2", - "hunyuan_dynamic_kv", + CUDA_KERNEL_MODULE, PTX, )?; let mut builder = stage2.builder(); @@ -537,7 +406,7 @@ impl InplaceOp2 for RepetitionPenaltyF32 { 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)?; + device.get_or_load_custom_func("repetition_penalty_f32", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&mut logits); builder.arg(&token_ids); @@ -630,8 +499,7 @@ impl CustomOp2 for FusedSiluMulBf16 { let up = up.slice(up_layout.start_offset()..); let count = gate_layout.shape().elem_count(); let output = unsafe { device.alloc::(count) }?; - let function = - device.get_or_load_custom_func("silu_mul_bf16", "hunyuan_dynamic_kv", PTX)?; + let function = device.get_or_load_custom_func("silu_mul_bf16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&gate); builder.arg(&up); @@ -715,7 +583,7 @@ impl CustomOp3 for FusedXdRopeRmsNormF16 { let output_count = if self.include_v { 2 * count } else { count }; let output = unsafe { device.alloc::(output_count) }?; let function = - device.get_or_load_custom_func("xdrope_rmsnorm_f16", "hunyuan_dynamic_kv", PTX)?; + device.get_or_load_custom_func("xdrope_rmsnorm_f16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&qkv); builder.arg(&cos_sin); @@ -815,7 +683,7 @@ impl CustomOp3 for FusedRmsNormRopeBf16 { let output_count = if self.include_v { 2 * count } else { count }; let output = unsafe { device.alloc::(output_count) }?; let function = - device.get_or_load_custom_func("rmsnorm_rope_bf16", "hunyuan_dynamic_kv", PTX)?; + device.get_or_load_custom_func("rmsnorm_rope_bf16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&qkv); builder.arg(&cos_sin); @@ -909,7 +777,7 @@ impl CustomOp3 for FusedAddRmsNormBf16 { let weight = weight.slice(weight_layout.start_offset()..); let output = unsafe { device.alloc::(2 * element_count) }?; let function = - device.get_or_load_custom_func("add_rmsnorm_bf16", "hunyuan_dynamic_kv", PTX)?; + device.get_or_load_custom_func("add_rmsnorm_bf16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&input); builder.arg(&delta); @@ -988,7 +856,7 @@ impl InplaceOp3 for FusedRopeBf16 { let input = input.slice(input_layout.start_offset()..); let cos_sin = cos_sin.slice(cos_sin_layout.start_offset()..); let count = heads * query_len * head_dim; - let function = device.get_or_load_custom_func("rope_bf16", "hunyuan_dynamic_kv", PTX)?; + let function = device.get_or_load_custom_func("rope_bf16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&mut output); builder.arg(&input); @@ -1059,7 +927,7 @@ impl InplaceOp3 for FusedXdRope { let qkv = qkv.slice(qkv_layout.start_offset()..); let cos_sin = cos_sin.slice(cos_sin_layout.start_offset()..); let count = heads * query_len * head_dim; - let function = device.get_or_load_custom_func("xdrope_bf16", "hunyuan_dynamic_kv", PTX)?; + let function = device.get_or_load_custom_func("xdrope_bf16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&mut output); builder.arg(&qkv); @@ -1133,7 +1001,7 @@ impl InplaceOp3 for DynamicKvAppend { let mut cache = cache.slice_mut(cache_layout.start_offset()..); let source = source.slice(source_layout.start_offset()..); let function = - device.get_or_load_custom_func($function, "hunyuan_dynamic_kv", PTX)?; + device.get_or_load_custom_func($function, CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&mut cache); builder.arg(&source); @@ -1214,7 +1082,7 @@ impl InplaceOp3 for DynamicPagedKvAppend { let lengths = lengths.slice(lengths_layout.start_offset()..); let count = num_heads * query_len * head_dim; let function = - device.get_or_load_custom_func("append_paged_kv_bf16", "hunyuan_dynamic_kv", PTX)?; + device.get_or_load_custom_func("append_paged_kv_bf16", CUDA_KERNEL_MODULE, PTX)?; let mut builder = function.builder(); builder.arg(&mut cache); builder.arg(&source); @@ -1234,6 +1102,7 @@ impl InplaceOp3 for DynamicPagedKvAppend { #[cfg(test)] mod tests { use super::*; + use crate::cuda_kernels::ArgmaxFirstBf16; use candle_core::{DType, Device, Tensor}; #[cfg(feature = "cuda")] diff --git a/oar-ocr-vl/src/hunyuanocr/llm.rs b/oar-ocr-vl/src/hunyuanocr/llm.rs index 2f6c56e..063a56f 100644 --- a/oar-ocr-vl/src/hunyuanocr/llm.rs +++ b/oar-ocr-vl/src/hunyuanocr/llm.rs @@ -6,6 +6,10 @@ use super::dynamic_kv::{ use crate::attention::{ RotaryEmbedding, flash_attention, repeat_kv, scaled_dot_product_attention, select_rope_sections, }; +#[cfg(feature = "cuda")] +use crate::decoder_graph::{ + CudaGraphKvLengths, cuda_graph_error, decoder_attention_is_causal, sync_graph_tensor, +}; use crate::kv_trim::TrimmableKvCache; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; #[cfg(feature = "cuda")] @@ -116,18 +120,6 @@ fn apply_xdrope_rotary_pos_emb( Ok((q_rot, k_rot)) } -#[cfg(feature = "cuda")] -fn cuda_graph_error( - context: impl Into, - source: impl std::error::Error + Send + Sync + 'static, -) -> OCRError { - OCRError::Inference { - model_name: "HunyuanOCR".to_string(), - context: context.into(), - source: Box::new(source), - } -} - #[derive(Debug)] struct HunyuanMlp { gate_up_proj: candle_nn::Linear, @@ -641,7 +633,7 @@ impl HunyuanAttention { query_len, cache_len, self.scaling as f32, - true, + decoder_attention_is_causal(query_len), ) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic flash attention", e))? .transpose(0, 1) @@ -683,7 +675,7 @@ struct TargetDecoderCudaGraph { cos_input: Tensor, sin_input: Tensor, _query_lengths: Tensor, - kv_lengths: Tensor, + kv_lengths: CudaGraphKvLengths, hidden_output: Tensor, logits_output: Tensor, aux_output: Option, @@ -702,17 +694,6 @@ impl std::fmt::Debug for TargetDecoderCudaGraph { } } -#[cfg(feature = "cuda")] -fn sync_graph_tensor(tensor: &Tensor, operation: &'static str) -> Result<(), OCRError> { - tensor - .flatten_all() - .and_then(|x| x.narrow(0, 0, 1)) - .and_then(|x| x.to_dtype(candle_core::DType::F32)) - .and_then(|x| x.to_vec1::()) - .map(|_| ()) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", operation, e)) -} - #[derive(Debug)] struct HunyuanDecoderLayer { self_attn: HunyuanAttention, @@ -1062,17 +1043,13 @@ impl HunyuanLlm { .sin_input .slice_set(sin, 0, 0) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "copy full graph sin", e))?; - let lengths = Tensor::new(&[0u32, kv_len as u32], inputs_embeds.device()).map_err(|e| { - candle_to_ocr_inference("HunyuanOCR", "create full graph KV lengths", e) + captured.kv_lengths.update(kv_len).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "update full graph KV lengths", e) })?; - captured - .kv_lengths - .slice_set(&lengths, 0, 0) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "copy full graph KV lengths", e))?; captured .graph .launch() - .map_err(|e| cuda_graph_error("launch full target decoder graph", e))?; + .map_err(|e| cuda_graph_error("HunyuanOCR", "launch full target decoder graph", e))?; for layer in &self.layers { layer.set_kv_cache_len(kv_len)?; } @@ -1385,7 +1362,7 @@ impl HunyuanLlm { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph query lengths", e))?; // 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()) + let kv_lengths = CudaGraphKvLengths::new(query_len, self.decode_cos.device()) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph KV lengths", e))?; let stream = cuda.cuda_stream(); let _htod_cache = cuda.enable_cuda_graph_htod_cache(); @@ -1395,22 +1372,24 @@ impl HunyuanLlm { &cos_input, &sin_input, &query_lengths, - &kv_lengths, + kv_lengths.tensor(), aux_layer_ids, )?; let warm_logits = self.project_logits(&warm.hidden_states)?; - sync_graph_tensor(&warm_logits, "warm full target decoder graph")?; + sync_graph_tensor("HunyuanOCR", &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))?; + .map_err(|e| { + cuda_graph_error("HunyuanOCR", "begin full target decoder graph capture", e) + })?; let captured_output = (|| { let output = self.forward_with_aux_dynamic( &hidden_input, &cos_input, &sin_input, &query_lengths, - &kv_lengths, + kv_lengths.tensor(), aux_layer_ids, )?; let logits = self.project_logits(&output.hidden_states)?; @@ -1429,15 +1408,21 @@ impl HunyuanLlm { .end_capture( CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ) - .map_err(|e| cuda_graph_error("end full target decoder graph capture", e))? + .map_err(|e| { + cuda_graph_error("HunyuanOCR", "end full target decoder graph capture", e) + })? .ok_or_else(|| OCRError::ConfigError { message: "HunyuanOCR full target decoder capture returned no graph".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(&logits_output, "sync full target decoder graph")?; + .map_err(|e| cuda_graph_error("HunyuanOCR", "warm full target decoder graph", e))?; + sync_graph_tensor( + "HunyuanOCR", + &logits_output, + "sync full target decoder graph", + )?; self.clear_kv_cache(); *self.decode_graph.borrow_mut() = Some(TargetDecoderCudaGraph { graph, diff --git a/oar-ocr-vl/src/hunyuanocr/mod.rs b/oar-ocr-vl/src/hunyuanocr/mod.rs index 4c96a81..f6221cb 100644 --- a/oar-ocr-vl/src/hunyuanocr/mod.rs +++ b/oar-ocr-vl/src/hunyuanocr/mod.rs @@ -7,7 +7,7 @@ mod config; mod dflash; #[cfg(feature = "cuda")] -mod dynamic_kv; +pub(crate) mod dynamic_kv; mod llm; mod model; mod processing; diff --git a/oar-ocr-vl/src/hunyuanocr/model.rs b/oar-ocr-vl/src/hunyuanocr/model.rs index 00540ce..4f44289 100644 --- a/oar-ocr-vl/src/hunyuanocr/model.rs +++ b/oar-ocr-vl/src/hunyuanocr/model.rs @@ -4,8 +4,7 @@ use super::config::{HunyuanOcrConfig, HunyuanOcrImageProcessorConfig, HunyuanOcr use super::dflash::DFlashModel; #[cfg(feature = "cuda")] use super::dynamic_kv::{ - ArgmaxFirstF32, DFlashRepetitionArgmaxBf16, MarkRepetitionHistoryU8, RepetitionArgmaxBf16, - RepetitionPenaltyF32, + DFlashRepetitionArgmaxBf16, MarkRepetitionHistoryU8, RepetitionArgmaxBf16, RepetitionPenaltyF32, }; use super::llm::HunyuanLlm; use super::processing::{HunyuanOcrImageInputs, preprocess_image}; @@ -13,6 +12,8 @@ use super::vision::HunyuanVisionModel; use crate::attention::{ combine_masks, create_causal_mask, create_generation_mask, create_left_padding_mask, }; +#[cfg(feature = "cuda")] +use crate::cuda_kernels::ArgmaxFirstF32; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing}; use candle_core::{D, DType, Device, IndexOp, Tensor}; use image::RgbImage; diff --git a/oar-ocr-vl/src/lib.rs b/oar-ocr-vl/src/lib.rs index 47b675b..ef7cf7c 100644 --- a/oar-ocr-vl/src/lib.rs +++ b/oar-ocr-vl/src/lib.rs @@ -38,6 +38,15 @@ pub mod utils; // Shared attention implementation pub mod attention; +// Small CUDA primitives shared by multiple VLM backends. Model-specific +// kernels remain in their respective modules, while stable sampling and other +// reusable decode operations live here. +#[cfg(feature = "cuda")] +pub(crate) mod cuda_kernels; + +// Shared lifetime and capacity plumbing for batch-1 decoder CUDA graphs. +pub(crate) mod decoder_graph; + // `TrimmableKvCache` backs the KV cache used by every model's attention // forward path. pub(crate) mod kv_trim; diff --git a/oar-ocr-vl/src/mineru/model.rs b/oar-ocr-vl/src/mineru/model.rs index 0186856..27ab66c 100644 --- a/oar-ocr-vl/src/mineru/model.rs +++ b/oar-ocr-vl/src/mineru/model.rs @@ -5,6 +5,8 @@ use super::vision::MinerUVisionModel; use crate::attention::{ combine_masks, create_causal_mask, create_generation_mask, create_left_padding_mask, }; +#[cfg(feature = "cuda")] +use crate::cuda_kernels::{ArgmaxFirstBf16, ArgmaxFirstF32, MaskTokenIds}; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing}; use candle_core::{DType, Device, IndexOp, Tensor}; use candle_nn::{Linear, Module, VarBuilder, linear_no_bias}; @@ -94,6 +96,8 @@ pub struct MinerU { temperature: f32, top_p: f32, top_k: usize, + #[cfg(feature = "cuda")] + gpu_greedy_sampling: bool, } #[derive(Debug, Deserialize)] @@ -260,6 +264,8 @@ impl MinerU { temperature, top_p, top_k, + #[cfg(feature = "cuda")] + gpu_greedy_sampling: std::env::var_os("OAR_MINERU_DISABLE_GPU_SAMPLING").is_none(), }) } @@ -512,32 +518,48 @@ impl MinerU { let position_ids = Tensor::cat(&pos_refs, 1) .map_err(|e| candle_to_ocr_inference("MinerU2.5", "stack pos", e))?; - let causal = create_causal_mask(max_seq_len, max_seq_len, self.dtype, &self.device) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "create causal", e))?; - let padding = create_left_padding_mask(&seq_lens, max_seq_len, self.dtype, &self.device) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "create padding", e))?; - let mask = combine_masks(&causal, &padding) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "combine masks", e))?; + let mask = if batch_size == 1 { + None + } else { + let causal = create_causal_mask(max_seq_len, max_seq_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "create causal", e))?; + let padding = + create_left_padding_mask(&seq_lens, max_seq_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "create padding", e))?; + Some( + combine_masks(&causal, &padding) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "combine masks", e))?, + ) + }; self.text.clear_kv_cache(); + if batch_size == 1 { + self.text + .prepare_ar_cuda_graph(max_seq_len, max_new_tokens, &self.lm_head)?; + } else { + // Batch-shaped prefill replaces the batch-1 KV backing storage. + // Drop the captured graph before those raw pointers become stale. + self.text.invalidate_ar_cuda_graph(); + } let hidden = self .text - .forward(&inputs_embeds, &position_ids, Some(&mask))?; - + .forward(&inputs_embeds, &position_ids, mask.as_ref())?; + + let last_hidden = hidden + .i((.., max_seq_len - 1, ..)) + .and_then(|hidden| hidden.contiguous()) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "get last hidden", e))?; + let batched_logits = self + .lm_head + .forward(&last_hidden) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "lm_head", e))?; let mut logits_list: Vec = Vec::with_capacity(batch_size); for i in 0..batch_size { - let last = hidden - .i((i, max_seq_len - 1, ..)) - .and_then(|t| t.unsqueeze(0)) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "get last hidden", e))?; - let logits = self - .lm_head - .forward(&last) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "lm_head", e))?; - let logits = logits - .squeeze(0) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "lm_head squeeze", e))?; - logits_list.push(logits); + logits_list.push( + batched_logits + .i(i) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "select logits", e))?, + ); } let mut generated: Vec> = vec![Vec::new(); batch_size]; @@ -555,11 +577,10 @@ impl MinerU { // Current KV cache length (grows during generation) let mut kv_len = max_seq_len; - for _ in 0..max_new_tokens { + for step in 0..max_new_tokens { if finished.iter().all(|&f| f) { break; } - let sampling_params = self.sampling_params(); let mut next_tokens: Vec = Vec::with_capacity(batch_size); for (i, logits) in logits_list.iter().enumerate() { @@ -580,6 +601,9 @@ impl MinerU { if finished.iter().all(|&f| f) { break; } + if step + 1 == max_new_tokens { + break; + } let tokens = Tensor::new(next_tokens, &self.device) .and_then(|t| t.reshape((batch_size, 1))) @@ -595,25 +619,35 @@ impl MinerU { // - Query position can attend to all non-padding positions in KV cache // - Padding positions (first pad_len tokens) should be masked kv_len += 1; - let gen_mask = create_generation_mask(&pad_lens, kv_len, self.dtype, &self.device) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "create gen mask", e))?; - - let hs = self.text.forward(&embeds, &pos, Some(&gen_mask))?; + let gen_mask = if batch_size == 1 { + None + } else { + Some( + create_generation_mask(&pad_lens, kv_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "create gen mask", e))?, + ) + }; logits_list.clear(); - for i in 0..batch_size { - let h = hs - .i((i, 0, ..)) - .and_then(|t| t.unsqueeze(0)) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "get hs", e))?; - let logits = self + if batch_size == 1 { + logits_list.push(self.text.forward_decode_logits( + &embeds, + &pos, + gen_mask.as_ref(), + &self.lm_head, + )?); + } else { + let hs = self.text.forward(&embeds, &pos, gen_mask.as_ref())?; + let batched_logits = self .lm_head - .forward(&h) + .forward(&hs) + .and_then(|t| t.squeeze(1)) .map_err(|e| candle_to_ocr_inference("MinerU2.5", "lm_head step", e))?; - let logits = logits - .squeeze(0) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "lm_head squeeze", e))?; - logits_list.push(logits); + for i in 0..batch_size { + logits_list.push(batched_logits.i(i).map_err(|e| { + candle_to_ocr_inference("MinerU2.5", "select step logits", e) + })?); + } } for (i, p) in positions.iter_mut().enumerate() { @@ -651,6 +685,8 @@ impl MinerU { temperature: self.temperature, top_p: self.top_p, top_k: self.top_k, + #[cfg(feature = "cuda")] + gpu_greedy_sampling: self.gpu_greedy_sampling, } } @@ -692,6 +728,8 @@ struct SamplingParams { temperature: f32, top_p: f32, top_k: usize, + #[cfg(feature = "cuda")] + gpu_greedy_sampling: bool, } fn select_next_token( @@ -699,6 +737,21 @@ fn select_next_token( history: &[u32], params: &SamplingParams, ) -> Result { + // The official MinerU generation config uses top_k=1 and + // repetition_penalty=1, so decoding is greedy even though do_sample=true. + // Keep the full vocabulary on-device for that common path. The custom + // reduction preserves the CPU loop's lowest-index tie break, and the + // sparse mask retains no-repeat-ngram semantics exactly. + #[cfg(feature = "cuda")] + if params.gpu_greedy_sampling + && logits.device().is_cuda() + && (!params.do_sample || params.top_k == 1) + && params.repetition_penalty <= 1.0 + && matches!(logits.dtype(), DType::BF16 | DType::F32) + { + return select_greedy_token_cuda(logits, history, params.no_repeat_ngram_size); + } + let logits = logits .to_dtype(DType::F32) .map_err(|e| candle_to_ocr_inference("MinerU2.5", "logits cast", e))? @@ -722,6 +775,41 @@ fn select_next_token( } } +#[cfg(feature = "cuda")] +fn select_greedy_token_cuda( + logits: &Tensor, + history: &[u32], + no_repeat_ngram_size: usize, +) -> Result { + let vocab_size = logits.elem_count(); + let logits = logits + .reshape((1, vocab_size)) + .and_then(|logits| logits.contiguous()) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "reshape GPU logits", e))?; + let banned = no_repeat_ngram_banned_tokens(history, no_repeat_ngram_size); + if !banned.is_empty() { + let banned = Tensor::new(banned, logits.device()) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "upload banned token ids", e))?; + logits.inplace_op2(&banned, &MaskTokenIds).map_err(|e| { + candle_to_ocr_inference("MinerU2.5", "apply GPU no-repeat-ngram mask", e) + })?; + } + let tokens = match logits.dtype() { + DType::BF16 => logits.apply_op1_no_bwd(&ArgmaxFirstBf16), + DType::F32 => logits.apply_op1_no_bwd(&ArgmaxFirstF32), + dtype => { + return Err(OCRError::ConfigError { + message: format!("MinerU2.5: unsupported GPU greedy logits dtype {dtype:?}"), + }); + } + } + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "stable GPU argmax", e))?; + tokens + .i(0) + .and_then(|token| token.to_scalar::()) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "copy selected token", e)) +} + fn apply_sampling_processors(logits: &mut [f32], history: &[u32], params: &SamplingParams) { apply_repetition_penalty(logits, history, params.repetition_penalty); apply_no_repeat_ngram(logits, history, params.no_repeat_ngram_size); @@ -861,8 +949,17 @@ fn sample_from_probs(probs: &[f32]) -> Option { } fn apply_no_repeat_ngram(logits: &mut [f32], history: &[u32], ngram_size: usize) { + for token in no_repeat_ngram_banned_tokens(history, ngram_size) { + let idx = token as usize; + if idx < logits.len() { + logits[idx] = f32::NEG_INFINITY; + } + } +} + +fn no_repeat_ngram_banned_tokens(history: &[u32], ngram_size: usize) -> Vec { if ngram_size <= 1 || history.len() < ngram_size { - return; + return Vec::new(); } let prefix_len = ngram_size - 1; let prefix_start = history.len() - prefix_len; @@ -873,12 +970,9 @@ fn apply_no_repeat_ngram(logits: &mut [f32], history: &[u32], ngram_size: usize) banned.insert(history[i + prefix_len]); } } - for token in banned { - let idx = token as usize; - if idx < logits.len() { - logits[idx] = f32::NEG_INFINITY; - } - } + let mut banned: Vec = banned.into_iter().collect(); + banned.sort_unstable(); + banned } fn expand_image_tokens( @@ -1027,6 +1121,43 @@ fn get_rope_index( mod tests { use super::*; + #[test] + fn no_repeat_ngram_collects_all_matching_continuations() { + assert_eq!( + no_repeat_ngram_banned_tokens(&[1, 2, 3, 1, 2, 4, 1, 2], 3), + vec![3, 4] + ); + assert!(no_repeat_ngram_banned_tokens(&[1, 2], 3).is_empty()); + } + + #[cfg(feature = "cuda")] + #[test] + fn mineru_cuda_greedy_matches_cpu_ties_and_no_repeat_ngram() -> candle_core::Result<()> { + let Ok(device) = Device::new_cuda(0) else { + return Ok(()); + }; + let mut values = vec![0.0f32; 1030]; + values[1] = 10.0; + values[1024] = 10.0; + let gpu_logits = Tensor::from_vec(values, 1030, &device)?.to_dtype(DType::BF16)?; + let cpu_logits = gpu_logits.to_device(&Device::Cpu)?; + let params = SamplingParams { + repetition_penalty: 1.0, + no_repeat_ngram_size: 3, + do_sample: true, + temperature: 0.01, + top_p: 0.001, + top_k: 1, + gpu_greedy_sampling: true, + }; + let history = [4, 5, 1, 8, 4, 5]; + let cpu = select_next_token(&cpu_logits, &history, ¶ms).unwrap(); + let gpu = select_next_token(&gpu_logits, &history, ¶ms).unwrap(); + assert_eq!(cpu, 1024); + assert_eq!(gpu, cpu); + Ok(()) + } + #[test] fn mineru_task_prompt_text_recognition_matches_official() { assert_eq!(MinerUTaskPrompt::Text.prompt(), "\nText Recognition:"); diff --git a/oar-ocr-vl/src/mineru/text.rs b/oar-ocr-vl/src/mineru/text.rs index a60c902..063d4e7 100644 --- a/oar-ocr-vl/src/mineru/text.rs +++ b/oar-ocr-vl/src/mineru/text.rs @@ -1,10 +1,21 @@ use super::config::MinerUConfig; use crate::attention::{ - RotaryEmbedding, repeat_kv, scaled_dot_product_attention, select_rope_sections, + RotaryEmbedding, flash_attention, scaled_dot_product_attention_gqa, select_rope_sections, }; +#[cfg(feature = "cuda")] +use crate::decoder_graph::decoder_cache_capacity; +#[cfg(feature = "cuda")] +use crate::decoder_graph::{ + CudaGraphKvLengths, SingleTokenDecoderCudaGraph, cuda_graph_error, decoder_attention_is_causal, + sync_graph_tensor, +}; +#[cfg(feature = "cuda")] +use crate::hunyuanocr::dynamic_kv::DynamicKvAppend; use crate::kv_trim::TrimmableKvCache; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; -use candle_core::Tensor; +#[cfg(feature = "cuda")] +use candle_core::{DType, Device}; +use candle_core::{IndexOp, Tensor}; use candle_nn::{ Embedding, Linear, Module, VarBuilder, embedding, linear, linear_no_bias, rms_norm, }; @@ -12,6 +23,9 @@ use oar_ocr_core::core::OCRError; use std::cell::RefCell; use std::sync::Arc; +#[cfg(feature = "cuda")] +const MINERU_DECODE_CACHE_LEN: usize = 16_384; + fn apply_multimodal_rotary_pos_emb( q: &Tensor, k: &Tensor, @@ -192,13 +206,12 @@ impl MinerUAttention { }) } - fn forward( + fn project_qkv( &self, hidden_states: &Tensor, cos: &Tensor, sin: &Tensor, - attention_mask: Option<&Tensor>, - ) -> Result { + ) -> Result<(Tensor, Tensor, Tensor), OCRError> { let (b, seq_len, _) = hidden_states .dims3() .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn hidden_states dims3", e))?; @@ -238,24 +251,59 @@ impl MinerUAttention { .contiguous() .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn v contiguous", e))?; + Ok((q, k, v)) + } + + fn forward( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + attention_mask: Option<&Tensor>, + ) -> Result { + let (b, seq_len, _) = hidden_states + .dims3() + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn hidden_states dims3", e))?; + let (q, k, v) = self.project_qkv(hidden_states, cos, sin)?; + let (k, v) = self .kv_cache .borrow_mut() .append(&k, &v) .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn kv_cache append", e))?; - let k = repeat_kv(&k, self.num_kv_groups) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "repeat kv k", e))?; - let v = repeat_kv(&v, self.num_kv_groups) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "repeat kv v", e))?; - let is_causal = attention_mask.is_none(); - let attn_output = - scaled_dot_product_attention(&q, &k, &v, attention_mask, self.scaling, is_causal) - .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn scaled_dot_product", e))?; + let flash = if b == 1 { + flash_attention(&q, &k, &v, self.scaling, seq_len > 1) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "flash attention", e))? + } else { + None + }; + let attn_output = match flash { + Some(attn) => attn, + None => scaled_dot_product_attention_gqa( + &q, + &k, + &v, + attention_mask, + self.scaling, + is_causal, + self.num_kv_groups, + ) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "grouped-query attention", e))?, + }; + self.project_attention_output(&attn_output, b, seq_len) + } + + fn project_attention_output( + &self, + attn_output: &Tensor, + batch: usize, + seq_len: usize, + ) -> Result { let attn_output = attn_output .transpose(1, 2) .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn output transpose", e))? - .reshape((b, seq_len, self.num_heads * self.head_dim)) + .reshape((batch, seq_len, self.num_heads * self.head_dim)) .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn output reshape", e))?; self.o_proj @@ -263,9 +311,101 @@ impl MinerUAttention { .map_err(|e| candle_to_ocr_inference("MinerU2.5", "attn o_proj", e)) } + #[cfg(feature = "cuda")] + fn prepare_dynamic_cache(&self, query_len: usize, cache_len: usize) -> Result<(), OCRError> { + let template = Tensor::zeros( + (1, self.num_kv_heads, query_len, self.head_dim), + self.k_proj.weight().dtype(), + self.k_proj.weight().device(), + ) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic KV template", e))?; + self.kv_cache + .borrow_mut() + .initialize_storage_with_capacity(&template, cache_len) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "initialize dynamic KV", e)) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let (batch, query_len, _) = hidden_states.dims3().map_err(|e| { + candle_to_ocr_inference("MinerU2.5", "dynamic attention hidden shape", e) + })?; + if batch != 1 { + return Err(OCRError::ConfigError { + message: "MinerU2.5 CUDA-graph attention requires batch size 1".to_string(), + }); + } + let (q, k, v) = self.project_qkv(hidden_states, cos, sin)?; + 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: "MinerU2.5 dynamic KV storage is not initialized".to_string(), + })?; + drop(cache); + let append = DynamicKvAppend { + query_len, + cache_len, + }; + cache_k + .inplace_op3(&k, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic key cache append", e))?; + cache_v + .inplace_op3(&v, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic value cache append", e))?; + + let q = q + .squeeze(0) + .and_then(|q| q.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic Q layout", e))?; + let cache_k = cache_k + .squeeze(0) + .and_then(|k| k.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic K layout", e))?; + let cache_v = cache_v + .squeeze(0) + .and_then(|v| v.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic V layout", e))?; + let attn = candle_flash_attn::flash_attn_varlen( + &q, + &cache_k, + &cache_v, + query_lengths, + kv_lengths, + query_len, + cache_len, + self.scaling as f32, + decoder_attention_is_causal(query_len), + ) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic flash attention", e))? + .transpose(0, 1) + .and_then(|attn| attn.unsqueeze(0)) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic attention layout", e))?; + self.project_attention_output(&attn, batch, query_len) + } + fn clear_kv_cache(&self) { self.kv_cache.borrow_mut().reset(); } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + self.kv_cache.borrow().current_seq_len() + } + + #[cfg(feature = "cuda")] + fn set_kv_cache_len(&self, len: usize) -> Result<(), OCRError> { + self.kv_cache + .borrow_mut() + .set_current_len(len) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "set dynamic KV length", e)) + } } pub struct MinerUDecoderLayer { @@ -333,12 +473,64 @@ impl MinerUDecoderLayer { }) } + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let residual = hidden_states.clone(); + let hidden_states = self + .input_layernorm + .forward(hidden_states) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "input_layernorm", e))?; + let hidden_states = + self.self_attn + .forward_dynamic(&hidden_states, cos, sin, query_lengths, kv_lengths)?; + let hidden_states = (&residual + &hidden_states).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "MinerU2.5: attn residual add failed", + e, + ) + })?; + + let residual = hidden_states.clone(); + let hidden_states = self + .post_attention_layernorm + .forward(&hidden_states) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "post_attention_layernorm", e))?; + let hidden_states = self.mlp.forward(&hidden_states)?; + (&residual + &hidden_states).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "MinerU2.5: mlp residual add failed", + e, + ) + }) + } + fn clear_kv_cache(&self) { self.self_attn.clear_kv_cache(); } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + self.self_attn.kv_cache_len() + } + + #[cfg(feature = "cuda")] + fn set_kv_cache_len(&self, len: usize) -> Result<(), OCRError> { + self.self_attn.set_kv_cache_len(len) + } } pub struct MinerUTextModel { + #[cfg(feature = "cuda")] + decode_graph: RefCell>, embed_tokens: Embedding, layers: Vec, norm: candle_nn::RmsNorm, @@ -367,6 +559,8 @@ impl MinerUTextModel { )?); Ok(Self { + #[cfg(feature = "cuda")] + decode_graph: RefCell::new(None), embed_tokens, layers, norm, @@ -399,6 +593,252 @@ impl MinerUTextModel { .map_err(|e| candle_to_ocr_inference("MinerU2.5", "norm forward", e)) } + fn project_logits(&self, hidden_states: &Tensor, lm_head: &Linear) -> Result { + lm_head + .forward(hidden_states) + .and_then(|logits| logits.i((0, 0, ..))) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "decode LM head", e)) + } + + pub(crate) fn forward_decode_logits( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + attention_mask: Option<&Tensor>, + lm_head: &Linear, + ) -> Result { + #[cfg(feature = "cuda")] + { + let kv_len = self.kv_cache_len().saturating_add(1); + if let Some(logits) = self.replay_cuda_graph(inputs_embeds, position_ids, kv_len)? { + return Ok(logits); + } + } + let hidden = self.forward(inputs_embeds, position_ids, attention_mask)?; + self.project_logits(&hidden, lm_head) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let (cos, sin) = self + .rotary_emb + .forward_multi_axis(position_ids, inputs_embeds.dtype())?; + let mut hidden_states = inputs_embeds.clone(); + for layer in &self.layers { + hidden_states = + layer.forward_dynamic(&hidden_states, &cos, &sin, query_lengths, kv_lengths)?; + } + self.norm + .forward(&hidden_states) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "dynamic norm", e)) + } + + pub(crate) fn prepare_ar_cuda_graph( + &self, + prompt_len: usize, + max_new_tokens: usize, + lm_head: &Linear, + ) -> Result<(), OCRError> { + if std::env::var_os("OAR_VL_DISABLE_CUDA_GRAPH").is_some() + || std::env::var_os("OAR_MINERU_DISABLE_CUDA_GRAPH").is_some() + { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); + return Ok(()); + } + #[cfg(feature = "cuda")] + if self.embed_tokens.embeddings().device().is_cuda() + && matches!( + self.embed_tokens.embeddings().dtype(), + DType::BF16 | DType::F16 + ) + { + let Some(cache_len) = + decoder_cache_capacity(prompt_len, max_new_tokens, MINERU_DECODE_CACHE_LEN) + else { + self.invalidate_cuda_graph(); + return Ok(()); + }; + let required = prompt_len + .saturating_add(max_new_tokens) + .min(MINERU_DECODE_CACHE_LEN); + let reusable = self + .decode_graph + .borrow() + .as_ref() + .is_some_and(|graph| graph.cache_len >= required); + if reusable { + return Ok(()); + } + self.invalidate_cuda_graph(); + self.capture_cuda_graph(cache_len, lm_head)?; + } + let _ = prompt_len; + let _ = max_new_tokens; + let _ = lm_head; + Ok(()) + } + + #[cfg(feature = "cuda")] + fn capture_cuda_graph(&self, cache_len: usize, lm_head: &Linear) -> Result<(), OCRError> { + use candle_core::cuda_backend::cudarc::driver::sys::{ + CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, + }; + + if self.decode_graph.borrow().is_some() { + return Ok(()); + } + let Device::Cuda(cuda) = self.embed_tokens.embeddings().device() else { + return Ok(()); + }; + let query_len = 1; + for layer in &self.layers { + layer + .self_attn + .prepare_dynamic_cache(query_len, cache_len)?; + } + let hidden_size = self + .embed_tokens + .embeddings() + .dim(1) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "graph hidden size", e))?; + let device = self.embed_tokens.embeddings().device(); + let hidden_input = Tensor::zeros( + (1, query_len, hidden_size), + self.embed_tokens.embeddings().dtype(), + device, + ) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "graph hidden input", e))?; + let position_input = Tensor::zeros((3, 1, query_len), DType::I64, device) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "graph position input", e))?; + let query_lengths = Tensor::new(&[0u32, query_len as u32], device) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "graph query lengths", e))?; + let kv_lengths = CudaGraphKvLengths::new(query_len, device) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "graph KV lengths", e))?; + let stream = cuda.cuda_stream(); + let _htod_cache = cuda.enable_cuda_graph_htod_cache(); + + let warm = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + let warm_logits = self.project_logits(&warm, lm_head)?; + sync_graph_tensor("MinerU2.5", &warm_logits, "warm decoder CUDA graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| cuda_graph_error("MinerU2.5", "begin decoder CUDA graph capture", e))?; + let captured_output = (|| { + let hidden = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + self.project_logits(&hidden, lm_head) + })(); + let logits_output = match captured_output { + Ok(output) => output, + Err(error) => { + let _ = stream.end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + return Err(error); + } + }; + let graph = stream + .end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .map_err(|e| cuda_graph_error("MinerU2.5", "end decoder CUDA graph capture", e))? + .ok_or_else(|| OCRError::ConfigError { + message: "MinerU2.5 decoder capture returned no graph".to_string(), + })?; + graph + .launch() + .map_err(|e| cuda_graph_error("MinerU2.5", "warm decoder CUDA graph", e))?; + sync_graph_tensor("MinerU2.5", &logits_output, "sync decoder CUDA graph")?; + self.clear_kv_cache(); + *self.decode_graph.borrow_mut() = Some(SingleTokenDecoderCudaGraph { + graph, + hidden_input, + position_input, + _query_lengths: query_lengths, + kv_lengths, + logits_output, + cache_len, + }); + Ok(()) + } + + #[cfg(feature = "cuda")] + fn replay_cuda_graph( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + kv_len: usize, + ) -> Result, OCRError> { + 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() + || position_ids.shape() != captured.position_input.shape() + { + return Ok(None); + } + captured + .hidden_input + .slice_set(inputs_embeds, 0, 0) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "copy graph hidden", e))?; + captured + .position_input + .slice_set(position_ids, 0, 0) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "copy graph positions", e))?; + captured + .kv_lengths + .update(kv_len) + .map_err(|e| candle_to_ocr_inference("MinerU2.5", "update graph KV lengths", e))?; + captured + .graph + .launch() + .map_err(|e| cuda_graph_error("MinerU2.5", "launch decoder CUDA graph", e))?; + for layer in &self.layers { + layer.set_kv_cache_len(kv_len)?; + } + Ok(Some(captured.logits_output.clone())) + } + + #[cfg(feature = "cuda")] + fn invalidate_cuda_graph(&self) { + self.decode_graph.borrow_mut().take(); + } + + pub(crate) fn invalidate_ar_cuda_graph(&self) { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); + } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + let len = self.layers.first().map_or(0, |layer| layer.kv_cache_len()); + debug_assert!(self.layers.iter().all(|layer| layer.kv_cache_len() == len)); + len + } + pub fn token_embedding_weight(&self) -> Tensor { self.embed_tokens.embeddings().clone() } diff --git a/oar-ocr-vl/src/mineru_diffusion/model.rs b/oar-ocr-vl/src/mineru_diffusion/model.rs index 15bae2b..9e9d102 100644 --- a/oar-ocr-vl/src/mineru_diffusion/model.rs +++ b/oar-ocr-vl/src/mineru_diffusion/model.rs @@ -13,6 +13,8 @@ use super::config::MinerUDiffusionConfig; use super::projector::VisionAbstractor; use super::text::{SdarKvCache, SdarModel}; +#[cfg(feature = "cuda")] +use crate::cuda_kernels::SampleWithConfidence; use crate::mineru::MinerUImageProcessorConfig; use crate::mineru::processing::preprocess_images; use crate::mineru::vision::MinerUVisionModel; @@ -490,6 +492,16 @@ fn sample_tokens_and_conf( top_p: f32, rng: &mut StdRng, ) -> Result<(Vec, Vec), OCRError> { + #[cfg(feature = "cuda")] + if matches!(logits.device(), Device::Cuda(_)) + && matches!(logits.dtype(), DType::BF16 | DType::F32) + && top_k == 0 + && top_p >= 1.0 + && std::env::var_os("OAR_MINERU_DIFFUSION_DISABLE_GPU_SAMPLING").is_none() + { + return sample_tokens_and_conf_cuda(logits, temperature, rng); + } + // logits: (1, block, vocab) -> rows of (vocab,) let logits = logits .squeeze(0) @@ -537,6 +549,52 @@ fn sample_tokens_and_conf( Ok((tokens, confs)) } +/// Fast path for the reference/default sampling configuration. The model's +/// logits remain on the GPU; only two F32 values per block position (token id +/// and selected probability) cross to the host. The CPU fallback above stays +/// available for top-k/top-p filters and as a numerical/debugging oracle via +/// `OAR_MINERU_DIFFUSION_DISABLE_GPU_SAMPLING=1`. +#[cfg(feature = "cuda")] +fn sample_tokens_and_conf_cuda( + logits: &Tensor, + temperature: f32, + rng: &mut StdRng, +) -> Result<(Vec, Vec), OCRError> { + let logits = logits + .squeeze(0) + .and_then(|t| t.contiguous()) + .map_err(|e| proc_err("CUDA sampling logits", e))?; + let rows = logits + .dim(0) + .map_err(|e| proc_err("CUDA sampling rows", e))?; + let greedy = temperature <= 0.0; + let uniforms = if greedy { + vec![0.0f32; rows] + } else { + (0..rows).map(|_| rng.random::()).collect() + }; + let uniforms = Tensor::from_vec(uniforms, rows, logits.device()) + .map_err(|e| proc_err("CUDA sampling uniforms", e))?; + let sampled = logits + .apply_op2_no_bwd( + &uniforms, + &SampleWithConfidence { + temperature, + greedy, + }, + ) + .and_then(|t| t.to_vec2::()) + .map_err(|e| proc_err("CUDA sample-with-confidence", e))?; + + let mut tokens = Vec::with_capacity(rows); + let mut confs = Vec::with_capacity(rows); + for row in sampled { + tokens.push(row[0] as u32); + confs.push(row[1]); + } + Ok((tokens, confs)) +} + /// Numerically stable softmax over a logits row. fn softmax(logits: &[f32]) -> Vec { let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max); diff --git a/oar-ocr-vl/src/mineru_diffusion/text.rs b/oar-ocr-vl/src/mineru_diffusion/text.rs index bc12104..f534efd 100644 --- a/oar-ocr-vl/src/mineru_diffusion/text.rs +++ b/oar-ocr-vl/src/mineru_diffusion/text.rs @@ -15,7 +15,7 @@ //! committed once (`store_kv = true`). use super::config::SdarConfig; -use crate::attention::{RotaryEmbedding, repeat_kv, scaled_dot_product_attention}; +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_nn::{ @@ -179,13 +179,21 @@ impl SdarAttention { } let n_rep = self.num_heads / self.num_kv_heads; - let k_rep = repeat_kv(&full_k, n_rep) - .map_err(|e| candle_to_ocr_inference("MinerU-Diffusion", "repeat_kv k", e))?; - let v_rep = repeat_kv(&full_v, n_rep) - .map_err(|e| candle_to_ocr_inference("MinerU-Diffusion", "repeat_kv v", e))?; - - let attn = scaled_dot_product_attention(&q, &k_rep, &v_rep, mask, self.scale, false) - .map_err(|e| candle_to_ocr_inference("MinerU-Diffusion", "attention", e))?; + let flash = if mask.is_none() { + flash_attention(&q, &full_k, &full_v, self.scale, false) + .map_err(|e| candle_to_ocr_inference("MinerU-Diffusion", "flash attention", e))? + } else { + None + }; + let attn = match flash { + Some(attn) => attn, + None => scaled_dot_product_attention_gqa( + &q, &full_k, &full_v, mask, self.scale, false, n_rep, + ) + .map_err(|e| { + candle_to_ocr_inference("MinerU-Diffusion", "grouped-query attention", e) + })?, + }; let attn = attn .transpose(1, 2) .map_err(|e| proc_err("SDAR attn out transpose", e))? diff --git a/oar-ocr-vl/src/paddleocr_vl/ernie.rs b/oar-ocr-vl/src/paddleocr_vl/ernie.rs index d624635..4a70400 100644 --- a/oar-ocr-vl/src/paddleocr_vl/ernie.rs +++ b/oar-ocr-vl/src/paddleocr_vl/ernie.rs @@ -1,14 +1,28 @@ use super::config::PaddleOcrVlConfig; use crate::attention::{ - RotaryEmbedding, repeat_kv, scaled_dot_product_attention, select_rope_sections, + RotaryEmbedding, flash_attention, repeat_kv, scaled_dot_product_attention, select_rope_sections, }; +#[cfg(feature = "cuda")] +use crate::decoder_graph::decoder_cache_capacity; +#[cfg(feature = "cuda")] +use crate::decoder_graph::{ + CudaGraphKvLengths, SingleTokenDecoderCudaGraph, cuda_graph_error, decoder_attention_is_causal, + sync_graph_tensor, +}; +#[cfg(feature = "cuda")] +use crate::hunyuanocr::dynamic_kv::DynamicKvAppend; use crate::kv_trim::TrimmableKvCache; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; -use candle_core::Tensor; -use candle_nn::Module; +#[cfg(feature = "cuda")] +use candle_core::{DType, Device}; +use candle_core::{IndexOp, Tensor}; +use candle_nn::{Linear, Module}; use oar_ocr_core::core::OCRError; use std::cell::RefCell; +#[cfg(feature = "cuda")] +const PADDLE_DECODE_CACHE_LEN: usize = 16_384; + fn apply_multimodal_rotary_pos_emb( q: &Tensor, k: &Tensor, @@ -201,13 +215,12 @@ impl Ernie4_5Attention { }) } - fn forward( + fn project_qkv( &self, hidden_states: &Tensor, cos: &Tensor, sin: &Tensor, - causal_mask: Option<&Tensor>, - ) -> Result { + ) -> Result<(Tensor, Tensor, Tensor), OCRError> { let (b, seq_len, _) = hidden_states .dims3() .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "attn hidden_states dims3", e))?; @@ -251,39 +264,74 @@ impl Ernie4_5Attention { .contiguous() .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "attn v contiguous", e))?; + Ok((q, k, v)) + } + + fn forward( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + causal_mask: Option<&Tensor>, + ) -> Result { + let (b, seq_len, _) = hidden_states + .dims3() + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "attn hidden_states dims3", e))?; + let (q, k, v) = self.project_qkv(hidden_states, cos, sin)?; + let (k_all, v_all) = self .kv_cache .borrow_mut() .append(&k, &v) .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "append kv_cache", e))?; - let key_states = repeat_kv(&k_all, self.num_kv_groups) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "repeat_kv key", e))?; - let value_states = repeat_kv(&v_all, self.num_kv_groups) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "repeat_kv value", e))?; - - let key_states = key_states.contiguous().map_err(|e| { - candle_to_ocr_inference("PaddleOCR-VL", "attn key_states contiguous", e) - })?; - let value_states = value_states.contiguous().map_err(|e| { - candle_to_ocr_inference("PaddleOCR-VL", "attn value_states contiguous", e) - })?; - - // Use unified attention implementation - let attn_output = scaled_dot_product_attention( - &q, - &key_states, - &value_states, - causal_mask, - self.scaling, - false, // not causal - mask is explicit - ) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "scaled_dot_product_attention", e))?; + let flash = if b == 1 { + flash_attention(&q, &k_all, &v_all, self.scaling, seq_len > 1) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "flash attention", e))? + } else { + None + }; + let is_causal = causal_mask.is_none(); + let attn_output = match flash { + Some(attn) => attn, + None => { + let key_states = repeat_kv(&k_all, self.num_kv_groups) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "repeat_kv key", e))?; + let value_states = repeat_kv(&v_all, self.num_kv_groups) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "repeat_kv value", e))?; + let key_states = key_states.contiguous().map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "attn key_states contiguous", e) + })?; + let value_states = value_states.contiguous().map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "attn value_states contiguous", e) + })?; + scaled_dot_product_attention( + &q, + &key_states, + &value_states, + causal_mask, + self.scaling, + is_causal, + ) + .map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "scaled_dot_product_attention", e) + })? + } + }; + + self.project_attention_output(&attn_output, b, seq_len) + } + fn project_attention_output( + &self, + attn_output: &Tensor, + batch: usize, + seq_len: usize, + ) -> Result { let attn_output = attn_output .transpose(1, 2) .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "attn out transpose", e))? - .reshape((b, seq_len, self.num_heads * self.head_dim)) + .reshape((batch, seq_len, self.num_heads * self.head_dim)) .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "attn out reshape", e))?; self.o_proj @@ -291,9 +339,101 @@ impl Ernie4_5Attention { .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "attn o_proj", e)) } + #[cfg(feature = "cuda")] + fn prepare_dynamic_cache(&self, query_len: usize, cache_len: usize) -> Result<(), OCRError> { + let template = Tensor::zeros( + (1, self.num_kv_heads, query_len, self.head_dim), + self.k_proj.weight().dtype(), + self.k_proj.weight().device(), + ) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic KV template", e))?; + self.kv_cache + .borrow_mut() + .initialize_storage_with_capacity(&template, cache_len) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "initialize dynamic KV", e)) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let (batch, query_len, _) = hidden_states.dims3().map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "dynamic attention hidden shape", e) + })?; + if batch != 1 { + return Err(OCRError::ConfigError { + message: "PaddleOCR-VL CUDA-graph attention requires batch size 1".to_string(), + }); + } + let (q, k, v) = self.project_qkv(hidden_states, cos, sin)?; + 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: "PaddleOCR-VL dynamic KV storage is not initialized".to_string(), + })?; + drop(cache); + let append = DynamicKvAppend { + query_len, + cache_len, + }; + cache_k + .inplace_op3(&k, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic key cache append", e))?; + cache_v.inplace_op3(&v, kv_lengths, &append).map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "dynamic value cache append", e) + })?; + + let q = q + .squeeze(0) + .and_then(|q| q.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic Q layout", e))?; + let cache_k = cache_k + .squeeze(0) + .and_then(|k| k.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic K layout", e))?; + let cache_v = cache_v + .squeeze(0) + .and_then(|v| v.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic V layout", e))?; + let attn = candle_flash_attn::flash_attn_varlen( + &q, + &cache_k, + &cache_v, + query_lengths, + kv_lengths, + query_len, + cache_len, + self.scaling as f32, + decoder_attention_is_causal(query_len), + ) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic flash attention", e))? + .transpose(0, 1) + .and_then(|attn| attn.unsqueeze(0)) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic attention layout", e))?; + self.project_attention_output(&attn, batch, query_len) + } + fn clear_kv_cache(&self) { self.kv_cache.borrow_mut().reset(); } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + self.kv_cache.borrow().current_seq_len() + } + + #[cfg(feature = "cuda")] + fn set_kv_cache_len(&self, len: usize) -> Result<(), OCRError> { + self.kv_cache + .borrow_mut() + .set_current_len(len) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "set dynamic KV length", e)) + } } #[derive(Debug)] @@ -355,13 +495,57 @@ impl Ernie4_5DecoderLayer { .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "llm mlp residual add", e)) } + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let residual = hidden_states.clone(); + let hidden_states = self + .input_layernorm + .forward(hidden_states) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "llm input_layernorm", e))?; + let attn_out = + self.self_attn + .forward_dynamic(&hidden_states, cos, sin, query_lengths, kv_lengths)?; + let hidden_states = (&residual + &attn_out) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "llm attn residual add", e))?; + + let residual = hidden_states.clone(); + let hidden_states = self + .post_attention_layernorm + .forward(&hidden_states) + .map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "llm post_attention_layernorm", e) + })?; + let mlp_out = self.mlp.forward(&hidden_states)?; + (&residual + &mlp_out) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "llm mlp residual add", e)) + } + fn clear_kv_cache(&self) { self.self_attn.clear_kv_cache(); } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + self.self_attn.kv_cache_len() + } + + #[cfg(feature = "cuda")] + fn set_kv_cache_len(&self, len: usize) -> Result<(), OCRError> { + self.self_attn.set_kv_cache_len(len) + } } #[derive(Debug)] pub struct Ernie4_5Model { + #[cfg(feature = "cuda")] + decode_graph: RefCell>, embed_tokens: candle_nn::Embedding, layers: Vec, norm: candle_nn::RmsNorm, @@ -384,6 +568,8 @@ impl Ernie4_5Model { .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "load final norm", e))?; let rotary = RotaryEmbedding::new_multi_axis(cfg.head_dim, cfg.rope_theta, 3, vb.device())?; Ok(Self { + #[cfg(feature = "cuda")] + decode_graph: RefCell::new(None), embed_tokens, layers, norm, @@ -407,15 +593,269 @@ impl Ernie4_5Model { .rotary .forward_multi_axis(position_ids, inputs_embeds.dtype())?; + self.forward_prepared(inputs_embeds, &cos, &sin, causal_mask) + } + + fn forward_prepared( + &self, + inputs_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + causal_mask: Option<&Tensor>, + ) -> Result { let mut hidden_states = inputs_embeds.clone(); for layer in &self.layers { - hidden_states = layer.forward(&hidden_states, &cos, &sin, causal_mask)?; + hidden_states = layer.forward(&hidden_states, cos, sin, causal_mask)?; } - let hidden_states = self - .norm + self.norm .forward(&hidden_states) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "llm final norm", e))?; - Ok(hidden_states) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "llm final norm", e)) + } + + fn project_logits(&self, hidden_states: &Tensor, lm_head: &Linear) -> Result { + lm_head + .forward(hidden_states) + .and_then(|logits| logits.i((0, 0, ..))) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "decode LM head", e)) + } + + pub(crate) fn forward_decode_logits( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + causal_mask: Option<&Tensor>, + lm_head: &Linear, + ) -> Result { + #[cfg(feature = "cuda")] + { + let kv_len = self.kv_cache_len().saturating_add(1); + if let Some(logits) = self.replay_cuda_graph(inputs_embeds, position_ids, kv_len)? { + return Ok(logits); + } + } + let hidden = self.forward(inputs_embeds, position_ids, causal_mask)?; + self.project_logits(&hidden, lm_head) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let (cos, sin) = self + .rotary + .forward_multi_axis(position_ids, inputs_embeds.dtype())?; + let mut hidden_states = inputs_embeds.clone(); + for layer in &self.layers { + hidden_states = + layer.forward_dynamic(&hidden_states, &cos, &sin, query_lengths, kv_lengths)?; + } + self.norm + .forward(&hidden_states) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "dynamic final norm", e)) + } + + pub(crate) fn prepare_ar_cuda_graph( + &self, + prompt_len: usize, + max_new_tokens: usize, + lm_head: &Linear, + ) -> Result<(), OCRError> { + if std::env::var_os("OAR_VL_DISABLE_CUDA_GRAPH").is_some() + || std::env::var_os("OAR_PADDLEOCR_VL_DISABLE_CUDA_GRAPH").is_some() + { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); + return Ok(()); + } + #[cfg(feature = "cuda")] + if self.embed_tokens.embeddings().device().is_cuda() + && matches!( + self.embed_tokens.embeddings().dtype(), + DType::BF16 | DType::F16 + ) + { + let Some(cache_len) = + decoder_cache_capacity(prompt_len, max_new_tokens, PADDLE_DECODE_CACHE_LEN) + else { + self.invalidate_cuda_graph(); + return Ok(()); + }; + let required = prompt_len + .saturating_add(max_new_tokens) + .min(PADDLE_DECODE_CACHE_LEN); + let reusable = self + .decode_graph + .borrow() + .as_ref() + .is_some_and(|graph| graph.cache_len >= required); + if reusable { + return Ok(()); + } + self.invalidate_cuda_graph(); + self.capture_cuda_graph(cache_len, lm_head)?; + } + let _ = prompt_len; + let _ = max_new_tokens; + let _ = lm_head; + Ok(()) + } + + #[cfg(feature = "cuda")] + fn capture_cuda_graph(&self, cache_len: usize, lm_head: &Linear) -> Result<(), OCRError> { + use candle_core::cuda_backend::cudarc::driver::sys::{ + CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, + }; + + if self.decode_graph.borrow().is_some() { + return Ok(()); + } + let Device::Cuda(cuda) = self.embed_tokens.embeddings().device() else { + return Ok(()); + }; + let query_len = 1; + for layer in &self.layers { + layer + .self_attn + .prepare_dynamic_cache(query_len, cache_len)?; + } + let hidden_size = self + .embed_tokens + .embeddings() + .dim(1) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "graph hidden size", e))?; + let device = self.embed_tokens.embeddings().device(); + let hidden_input = Tensor::zeros( + (1, query_len, hidden_size), + self.embed_tokens.embeddings().dtype(), + device, + ) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "graph hidden input", e))?; + let position_input = Tensor::zeros((3, 1, query_len), DType::I64, device) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "graph position input", e))?; + let query_lengths = Tensor::new(&[0u32, query_len as u32], device) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "graph query lengths", e))?; + let kv_lengths = CudaGraphKvLengths::new(query_len, device) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "graph KV lengths", e))?; + let stream = cuda.cuda_stream(); + let _htod_cache = cuda.enable_cuda_graph_htod_cache(); + + let warm = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + let warm_logits = self.project_logits(&warm, lm_head)?; + sync_graph_tensor("PaddleOCR-VL", &warm_logits, "warm decoder CUDA graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| cuda_graph_error("PaddleOCR-VL", "begin decoder CUDA graph capture", e))?; + let captured_output = (|| { + let hidden = self.forward_dynamic( + &hidden_input, + &position_input, + &query_lengths, + kv_lengths.tensor(), + )?; + self.project_logits(&hidden, lm_head) + })(); + let logits_output = match captured_output { + Ok(output) => output, + Err(error) => { + let _ = stream.end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ); + return Err(error); + } + }; + let graph = stream + .end_capture( + CUgraphInstantiate_flags_enum::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, + ) + .map_err(|e| cuda_graph_error("PaddleOCR-VL", "end decoder CUDA graph capture", e))? + .ok_or_else(|| OCRError::ConfigError { + message: "PaddleOCR-VL decoder capture returned no graph".to_string(), + })?; + graph + .launch() + .map_err(|e| cuda_graph_error("PaddleOCR-VL", "warm decoder CUDA graph", e))?; + sync_graph_tensor("PaddleOCR-VL", &logits_output, "sync decoder CUDA graph")?; + self.clear_kv_cache(); + *self.decode_graph.borrow_mut() = Some(SingleTokenDecoderCudaGraph { + graph, + hidden_input, + position_input, + _query_lengths: query_lengths, + kv_lengths, + logits_output, + cache_len, + }); + Ok(()) + } + + #[cfg(feature = "cuda")] + fn replay_cuda_graph( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + kv_len: usize, + ) -> Result, OCRError> { + 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() + || position_ids.shape() != captured.position_input.shape() + { + return Ok(None); + } + captured + .hidden_input + .slice_set(inputs_embeds, 0, 0) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "copy graph hidden", e))?; + captured + .position_input + .slice_set(position_ids, 0, 0) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "copy graph positions", e))?; + captured + .kv_lengths + .update(kv_len) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "update graph KV lengths", e))?; + captured + .graph + .launch() + .map_err(|e| cuda_graph_error("PaddleOCR-VL", "launch decoder CUDA graph", e))?; + for layer in &self.layers { + layer.set_kv_cache_len(kv_len)?; + } + Ok(Some(captured.logits_output.clone())) + } + + #[cfg(feature = "cuda")] + fn invalidate_cuda_graph(&self) { + self.decode_graph.borrow_mut().take(); + } + + pub(crate) fn invalidate_ar_cuda_graph(&self) { + #[cfg(feature = "cuda")] + self.invalidate_cuda_graph(); + } + + #[cfg(feature = "cuda")] + fn kv_cache_len(&self) -> usize { + let len = self.layers.first().map_or(0, |layer| layer.kv_cache_len()); + debug_assert!(self.layers.iter().all(|layer| layer.kv_cache_len() == len)); + len } pub fn clear_kv_cache(&self) { diff --git a/oar-ocr-vl/src/paddleocr_vl/model.rs b/oar-ocr-vl/src/paddleocr_vl/model.rs index 5701a0c..4591bb9 100644 --- a/oar-ocr-vl/src/paddleocr_vl/model.rs +++ b/oar-ocr-vl/src/paddleocr_vl/model.rs @@ -449,32 +449,51 @@ impl PaddleOcrVl { .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "stack pos", e))?; // 6. Create attention mask - let causal = create_causal_mask(max_seq_len, max_seq_len, self.dtype, &self.device) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "create causal", e))?; - let padding = create_left_padding_mask(&seq_lens, max_seq_len, self.dtype, &self.device) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "create padding", e))?; - let mask = combine_masks(&causal, &padding) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "combine masks", e))?; + let mask = if batch_size == 1 { + None + } else { + let causal = create_causal_mask(max_seq_len, max_seq_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "create causal", e))?; + let padding = + create_left_padding_mask(&seq_lens, max_seq_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "create padding", e))?; + Some( + combine_masks(&causal, &padding) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "combine masks", e))?, + ) + }; // 7. Prefill self.llm.clear_kv_cache(); + if batch_size == 1 { + self.llm + .prepare_ar_cuda_graph(max_seq_len, max_new_tokens, &self.lm_head)?; + } else { + // A batch-shaped prefill reallocates the per-layer KV storage. Drop + // a batch-1 graph before that happens because it owns raw pointers + // into the previous fixed storage. + self.llm.invalidate_ar_cuda_graph(); + } let hidden = self .llm - .forward(&inputs_embeds, &position_ids, Some(&mask))?; + .forward(&inputs_embeds, &position_ids, mask.as_ref())?; // 8. Get initial logits + let last_hidden = hidden + .i((.., max_seq_len - 1, ..)) + .and_then(|hidden| hidden.contiguous()) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "get last hidden", e))?; + let batched_logits = self + .lm_head + .forward(&last_hidden) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "lm_head", e))?; let mut logits_list: Vec = Vec::with_capacity(batch_size); for i in 0..batch_size { - let last = hidden - .i((i, max_seq_len - 1, ..)) - .and_then(|t| t.unsqueeze(0)) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "get last hidden", e))?; - let logits = self - .lm_head - .forward(&last) - .and_then(|t| t.squeeze(0)) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "lm_head", e))?; - logits_list.push(logits); + logits_list.push( + batched_logits + .i(i) + .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "select logits", e))?, + ); } // 9. Autoregressive decode @@ -492,11 +511,10 @@ impl PaddleOcrVl { let pad_lens: Vec = seq_lens.iter().map(|&len| max_seq_len - len).collect(); let mut kv_len = max_seq_len; - for _ in 0..max_new_tokens { + for step in 0..max_new_tokens { if finished.iter().all(|&f| f) { break; } - let mut next_tokens: Vec = Vec::with_capacity(batch_size); for (i, logits) in logits_list.iter().enumerate() { if finished[i] { @@ -519,6 +537,9 @@ impl PaddleOcrVl { if finished.iter().all(|&f| f) { break; } + if step + 1 == max_new_tokens { + break; + } // Batch forward let tokens = Tensor::new(next_tokens, &self.device) @@ -533,23 +554,36 @@ impl PaddleOcrVl { // 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("PaddleOCR-VL", "create gen mask", e))?; - - let hs = self.llm.forward(&embeds, &pos, Some(&gen_mask))?; + let gen_mask = if batch_size == 1 { + None + } else { + Some( + create_generation_mask(&pad_lens, kv_len, self.dtype, &self.device).map_err( + |e| candle_to_ocr_inference("PaddleOCR-VL", "create gen mask", e), + )?, + ) + }; logits_list.clear(); - for i in 0..batch_size { - let h = hs - .i((i, 0, ..)) - .and_then(|t| t.unsqueeze(0)) - .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "get hs", e))?; - let logits = self + if batch_size == 1 { + logits_list.push(self.llm.forward_decode_logits( + &embeds, + &pos, + gen_mask.as_ref(), + &self.lm_head, + )?); + } else { + let hs = self.llm.forward(&embeds, &pos, gen_mask.as_ref())?; + let batched_logits = self .lm_head - .forward(&h) - .and_then(|t| t.squeeze(0)) + .forward(&hs) + .and_then(|t| t.squeeze(1)) .map_err(|e| candle_to_ocr_inference("PaddleOCR-VL", "lm_head step", e))?; - logits_list.push(logits); + for i in 0..batch_size { + logits_list.push(batched_logits.i(i).map_err(|e| { + candle_to_ocr_inference("PaddleOCR-VL", "select step logits", e) + })?); + } } for (i, p) in positions.iter_mut().enumerate() {