diff --git a/oar-ocr-vl/Cargo.toml b/oar-ocr-vl/Cargo.toml index be1e21a..1226759 100644 --- a/oar-ocr-vl/Cargo.toml +++ b/oar-ocr-vl/Cargo.toml @@ -24,6 +24,7 @@ download-binaries = ["oar-ocr-core/download-binaries"] # When enabled, turns on Candle's CUDA backend for GPU acceleration. cuda = [ "candle-core/cuda", + "dep:candle-flash-attn", "candle-nn/cuda", "candle-transformers/cuda", "oar-ocr-core/cuda", @@ -39,9 +40,11 @@ metal = [ [dependencies] candle-core = "0.11.0" +candle-flash-attn = { version = "0.11.0", optional = true } candle-nn = "0.11.0" candle-transformers = "0.11.0" html-escape = "0.2" +half = "2" image.workspace = true oar-ocr-core.workspace = true once_cell = "1.19" diff --git a/oar-ocr-vl/README.md b/oar-ocr-vl/README.md index 749ebdb..93e610f 100644 --- a/oar-ocr-vl/README.md +++ b/oar-ocr-vl/README.md @@ -11,7 +11,7 @@ This crate provides native Rust inference for document VLMs using [Candle](https | [PaddleOCR-VL](https://huggingface.co/PaddlePaddle/PaddleOCR-VL) | 0.9B | SOTA document parsing VLM supporting 109 languages, text, tables, formulas, and 11 chart types | | [PaddleOCR-VL-1.5](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.5) | 0.9B | Next-gen PaddleOCR-VL with 94.5% on OmniDocBench v1.5, adds text spotting and seal recognition | | [PaddleOCR-VL-1.6](https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.6) | 1.0B | Region-aware refinement on top of PaddleOCR-VL-1.5; 96.33% on OmniDocBench v1.6 (SOTA), drop-in compatible with the 1.5 loader | -| [HunyuanOCR 1.5](https://huggingface.co/tencent/HunyuanOCR) | Lightweight | End-to-end OCR VLM for multilingual document parsing, text spotting, and information extraction (archived 1.0 weights also supported) | +| [HunyuanOCR 1.5](https://huggingface.co/tencent/HunyuanOCR) | 1.0B | End-to-end OCR VLM for multilingual document parsing, text spotting, and information extraction (archived 1.0 weights also supported) | | [GLM-OCR](https://huggingface.co/zai-org/GLM-OCR) | 0.9B | #1 on OmniDocBench v1.5 (94.62), optimized for real-world scenarios with MTP loss and RL training | | [MinerU2.5](https://huggingface.co/opendatalab/MinerU2.5-2509-1.2B) | 1.2B | Decoupled document parsing VLM with strong text, formula, and table recognition | @@ -44,6 +44,12 @@ To enable GPU acceleration (CUDA), add the feature flag: cargo add oar-ocr-vl --features cuda ``` +HunyuanOCR's custom CUDA kernels compile PTX for the oldest GPU detected by +`nvidia-smi` at build time. For headless, container, or cross-machine builds, +set the target explicitly, for example +`CUDA_COMPUTE_CAP=89 cargo build --features cuda`. The DFlash kernels require +compute capability 8.0 or newer. + ## Usage ### PaddleOCR-VL @@ -191,12 +197,13 @@ cargo run --release --features cuda --example paddleocr_vl -- \ ```bash cargo run --release --features cuda --example hunyuanocr -- \ --model-dir models/HunyuanOCR \ + --dflash-dir models/HunyuanOCR/dflash \ --device cuda \ --prompt "Detect and recognize text in the image, and output the text coordinates in a formatted manner." \ document.jpg ``` -The model repository root contains HunyuanOCR 1.5. The loader detects it automatically; use `--model-dir models/HunyuanOCR/v1.0` for the archived 1.0 checkpoint. +The model repository root contains HunyuanOCR 1.5. The loader detects it automatically; use `--model-dir models/HunyuanOCR/v1.0` for the archived 1.0 checkpoint. `--dflash-dir` enables the official 15-token parallel draft path for 1.5; omit it for ordinary autoregressive decoding. Library callers can use `HunyuanOcr::from_dirs(target_dir, dflash_dir, device)` or `HunyuanOcr::from_dir_with_dflash(model_dir, device)` when the draft is stored in the official `dflash/` subdirectory. ### GLM-OCR (Direct Inference) diff --git a/oar-ocr-vl/build.rs b/oar-ocr-vl/build.rs index c6dd8ed..0821cd0 100644 --- a/oar-ocr-vl/build.rs +++ b/oar-ocr-vl/build.rs @@ -1,8 +1,113 @@ +use std::process::Command; + +const MIN_DFLASH_COMPUTE_CAP: u32 = 80; + +fn parse_compute_cap(value: &str) -> Option<(String, u32)> { + let value = value.trim().to_ascii_lowercase(); + let value = value + .strip_prefix("compute_") + .or_else(|| value.strip_prefix("sm_")) + .unwrap_or(&value) + .replace('.', ""); + let digit_count = value.bytes().take_while(u8::is_ascii_digit).count(); + if digit_count == 0 { + return None; + } + let (digits, suffix) = value.split_at(digit_count); + if !matches!(suffix, "" | "a" | "f") { + return None; + } + let mut base = digits.parse::().ok()?; + if base < 20 { + base *= 10; + } + Some((format!("{base}{suffix}"), base)) +} + +fn detect_local_compute_cap() -> Option { + let output = Command::new("nvidia-smi") + .args(["--query-gpu=compute_cap", "--format=csv,noheader"]) + .output() + .ok()?; + if !output.status.success() { + return None; + } + String::from_utf8_lossy(&output.stdout) + .lines() + .filter_map(|line| parse_compute_cap(line).map(|(_, base)| base)) + // PTX compiled for the oldest GPU reported by nvidia-smi remains + // loadable on newer GPUs in a heterogeneous machine. + .min() +} + +fn cuda_compute_arch() -> String { + if let Ok(value) = std::env::var("CUDA_COMPUTE_CAP") { + let (arch, base) = parse_compute_cap(&value).unwrap_or_else(|| { + panic!( + "invalid CUDA_COMPUTE_CAP={value:?}; expected values such as 89, 8.9, sm_89, or compute_89" + ) + }); + assert!( + base >= MIN_DFLASH_COMPUTE_CAP, + "HunyuanOCR DFlash 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) => { + println!( + "cargo:warning=detected GPU compute capability {base} is below the HunyuanOCR DFlash minimum; compiling forward-compatible compute_{MIN_DFLASH_COMPUTE_CAP} PTX" + ); + format!("compute_{MIN_DFLASH_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)" + ); + format!("compute_{MIN_DFLASH_COMPUTE_CAP}") + } + } +} + fn main() { + println!("cargo:rerun-if-changed=src/hunyuanocr/dynamic_kv.cu"); + println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP"); + println!("cargo:rerun-if-env-changed=NVCC"); let metal_enabled = std::env::var_os("CARGO_FEATURE_METAL").is_some(); + let cuda_enabled = std::env::var_os("CARGO_FEATURE_CUDA").is_some(); let target_os = std::env::var("CARGO_CFG_TARGET_OS").unwrap_or_default(); if metal_enabled && target_os != "macos" { panic!("oar-ocr-vl feature `metal` is only supported on macOS targets"); } + + if cuda_enabled { + let cuda_arch = cuda_compute_arch(); + let nvcc = std::env::var_os("NVCC").unwrap_or_else(|| "nvcc".into()); + let out_dir = std::path::PathBuf::from( + std::env::var_os("OUT_DIR").expect("Cargo always sets OUT_DIR"), + ); + let output = Command::new(&nvcc) + .args(["--ptx", "--std=c++17", "-O3"]) + .arg(format!("--gpu-architecture={cuda_arch}")) + .arg("-o") + .arg(out_dir.join("hunyuan_dynamic_kv.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}", + nvcc + ) + }); + if !output.status.success() { + panic!( + "{:?} failed for HunyuanOCR dynamic KV kernel ({cuda_arch}):\n{}", + nvcc, + String::from_utf8_lossy(&output.stderr) + ); + } + } } diff --git a/oar-ocr-vl/examples/hunyuanocr.rs b/oar-ocr-vl/examples/hunyuanocr.rs index 3d95b26..07bdbd3 100644 --- a/oar-ocr-vl/examples/hunyuanocr.rs +++ b/oar-ocr-vl/examples/hunyuanocr.rs @@ -22,6 +22,7 @@ mod utils; use clap::Parser; use std::path::PathBuf; +use std::time::Duration; use std::time::Instant; use tracing::{error, info}; @@ -37,6 +38,10 @@ struct Args { #[arg(short, long)] model_dir: PathBuf, + /// Optional DFlash draft directory (official checkpoint: /dflash) + #[arg(long)] + dflash_dir: Option, + /// Paths to input images to process #[arg(required = true)] images: Vec, @@ -49,12 +54,26 @@ struct Args { #[arg(long, default_value = "4096")] max_tokens: usize, + /// Override repetition penalty (1.0 matches the official speed benchmark) + #[arg(long)] + repetition_penalty: Option, + /// Instruction prompt (default: text spotting) #[arg( long, default_value = "Detect and recognize text in the image, and output the text coordinates in a formatted manner." )] prompt: String, + + /// Suppress generated text and print aggregate timing/token statistics + #[arg(long)] + 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> { @@ -90,14 +109,33 @@ fn main() -> Result<(), Box> { args.model_dir.display() ); let load_start = Instant::now(); - let model = HunyuanOcr::from_dir(&args.model_dir, device)?; + let mut model = match &args.dflash_dir { + Some(dflash_dir) => { + if !dflash_dir.exists() { + return Err(format!("DFlash directory not found: {}", dflash_dir.display()).into()); + } + HunyuanOcr::from_dirs(&args.model_dir, dflash_dir, device)? + } + None => HunyuanOcr::from_dir(&args.model_dir, device)?, + }; + if let Some(penalty) = args.repetition_penalty { + model.set_repetition_penalty(penalty)?; + } info!( - "HunyuanOCR {} loaded in {:.2}ms", + "HunyuanOCR {} loaded in {:.2}ms{}, repetition penalty {:.3}", model.version(), - load_start.elapsed().as_secs_f64() * 1000.0 + load_start.elapsed().as_secs_f64() * 1000.0, + model + .dflash_num_speculative_tokens() + .map(|n| format!(", DFlash enabled ({n} speculative tokens)")) + .unwrap_or_default(), + model.repetition_penalty(), ); info!("\n=== Processing {} images ===", existing_images.len()); + let mut total_inference = Duration::ZERO; + let mut total_tokens = 0usize; + let mut succeeded = 0usize; for image_path in &existing_images { info!("\nProcessing: {}", image_path.display()); let rgb_img = match load_image(image_path) { @@ -110,20 +148,39 @@ 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) .pop() { - Some(Ok(result)) => { + Some(Ok(tokens)) => { + let elapsed = infer_start.elapsed(); + total_inference += elapsed; + total_tokens += tokens.len(); + succeeded += 1; info!( - " Inference time: {:.2}ms", - infer_start.elapsed().as_secs_f64() * 1000.0 + " Inference time: {:.2}ms, tokens: {}, fingerprint: {:016x}", + elapsed.as_secs_f64() * 1000.0, + tokens.len(), + token_fingerprint(&tokens) ); - println!("{}", result); + if !args.benchmark { + println!("{}", model.decode_tokens(&tokens)?); + } } Some(Err(e)) => error!(" Inference failed: {}", e), None => error!(" No result returned from model"), } } + if succeeded > 0 { + info!( + "Benchmark summary: pages={}, total={:.2}ms, avg={:.2}ms/page, tokens={}, throughput={:.2} tokens/s", + succeeded, + total_inference.as_secs_f64() * 1000.0, + total_inference.as_secs_f64() * 1000.0 / succeeded as f64, + total_tokens, + total_tokens as f64 / total_inference.as_secs_f64() + ); + } + Ok(()) } diff --git a/oar-ocr-vl/src/attention.rs b/oar-ocr-vl/src/attention.rs index 4fccdc3..6ef1044 100644 --- a/oar-ocr-vl/src/attention.rs +++ b/oar-ocr-vl/src/attention.rs @@ -101,6 +101,104 @@ pub fn scaled_dot_product_attention( attn_weights.matmul(v) } +/// Scaled dot-product attention for grouped-query attention without expanding +/// K/V heads. Query heads that share one KV head are folded into the matrix row +/// dimension, preserving the usual head order in the returned tensor. +pub fn scaled_dot_product_attention_gqa( + q: &Tensor, + k: &Tensor, + v: &Tensor, + mask: Option<&Tensor>, + scale: f64, + is_causal: bool, + num_kv_groups: usize, +) -> Result { + if num_kv_groups == 1 { + 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()?; + let (v_batch, v_heads, v_len, v_head_dim) = v.dims4()?; + if batch != k_batch + || batch != v_batch + || num_heads != num_kv_heads * num_kv_groups + || num_kv_heads != v_heads + || kv_len != v_len + || head_dim != k_head_dim + || head_dim != v_head_dim + { + candle_core::bail!( + "invalid GQA shapes q={:?}, k={:?}, v={:?}, groups={num_kv_groups}", + q.dims(), + k.dims(), + v.dims() + ) + } + + let grouped_batch = batch * num_kv_heads; + let grouped_queries = num_kv_groups * query_len; + let grouped_q = q.reshape((grouped_batch, grouped_queries, head_dim))?; + let grouped_k = k + .reshape((grouped_batch, kv_len, head_dim))? + .transpose(1, 2)?; + let mut weights = + (grouped_q.matmul(&grouped_k)? * scale)?.reshape((batch, num_heads, query_len, kv_len))?; + + weights = match mask { + Some(mask) => weights.broadcast_add(mask)?, + None if is_causal => { + let causal = create_causal_mask(query_len, kv_len, weights.dtype(), q.device())?; + weights.broadcast_add(&causal)? + } + None => weights, + }; + + let weight_dtype = weights.dtype(); + let weights = candle_nn::ops::softmax_last_dim(&weights.to_dtype(DType::F32)?)? + .to_dtype(weight_dtype)? + .reshape((grouped_batch, grouped_queries, kv_len))?; + let grouped_v = v.reshape((grouped_batch, kv_len, head_dim))?; + weights + .matmul(&grouped_v)? + .reshape((batch, num_heads, query_len, head_dim)) +} + +#[cfg(any(feature = "cuda", test))] +fn flash_attention_dtype_supported(dtype: DType) -> bool { + matches!(dtype, DType::F16 | DType::BF16) +} + +/// Run CUDA FlashAttention v2 for Q/K/V tensors in `(batch, heads, seq, +/// head_dim)` layout. Returns `None` on non-CUDA devices or for dtypes not +/// supported by the CUDA kernel so callers can retain their portable eager +/// fallback. +pub fn flash_attention( + q: &Tensor, + k: &Tensor, + v: &Tensor, + scale: f64, + causal: bool, +) -> Result> { + #[cfg(feature = "cuda")] + if q.device().is_cuda() + && flash_attention_dtype_supported(q.dtype()) + && k.dtype() == q.dtype() + && v.dtype() == q.dtype() + { + // The CUDA kernel consumes (batch, seq, heads, head_dim) and natively + // supports GQA when K/V have fewer heads than Q. + let q = q.transpose(1, 2)?; + let k = k.transpose(1, 2)?; + let v = v.transpose(1, 2)?; + let output = candle_flash_attn::flash_attn(&q, &k, &v, scale as f32, causal)?; + return Ok(Some(output.transpose(1, 2)?)); + } + + let _ = (q, k, v, scale, causal); + Ok(None) +} + /// Create a causal (lower-triangular) attention mask. /// /// Returns a mask where position i can only attend to positions <= i. @@ -121,18 +219,18 @@ pub fn create_causal_mask( device: &Device, ) -> Result { on_compute_device(device, |compute_device| { - let row_idx = Tensor::arange(0u32, seq_len as u32, compute_device)? - .reshape((seq_len, 1))? - .to_dtype(dtype)?; - let col_idx = Tensor::arange(0u32, kv_len as u32, compute_device)? - .reshape((1, kv_len))? - .to_dtype(dtype)?; + let row_idx = + Tensor::arange(0u32, seq_len as u32, compute_device)?.reshape((seq_len, 1))?; + let col_idx = Tensor::arange(0u32, kv_len as u32, compute_device)?.reshape((1, kv_len))?; - let offset = (kv_len.saturating_sub(seq_len)) as f64; + let offset = kv_len.saturating_sub(seq_len) as u32; // Condition: col <= row + offset - // col - offset <= row - let diff = col_idx.broadcast_sub(&Tensor::new(offset, compute_device)?.to_dtype(dtype)?)?; - let mask_cond = diff.broadcast_le(&row_idx)?; + // Keep this comparison in integer space. BF16 cannot distinguish + // adjacent absolute positions once a document context grows beyond + // 256 tokens, which would let verification queries see future draft + // tokens and invalidate speculative decoding. + let row_limit = row_idx.broadcast_add(&Tensor::new(offset, compute_device)?)?; + let mask_cond = col_idx.broadcast_le(&row_limit)?; let zero = Tensor::new(0f32, compute_device)? .to_dtype(dtype)? @@ -590,6 +688,14 @@ pub fn select_rope_sections( mod tests { use super::*; + #[test] + fn flash_attention_rejects_unsupported_dtypes() { + assert!(flash_attention_dtype_supported(DType::F16)); + assert!(flash_attention_dtype_supported(DType::BF16)); + assert!(!flash_attention_dtype_supported(DType::F32)); + assert!(!flash_attention_dtype_supported(DType::F64)); + } + #[test] fn test_scaled_dot_product_attention() -> Result<()> { let device = Device::Cpu; @@ -612,6 +718,35 @@ mod tests { Ok(()) } + #[test] + fn test_grouped_query_attention_matches_repeated_kv() -> Result<()> { + let device = Device::Cpu; + let q = Tensor::randn(0f32, 1., (1, 4, 3, 8), &device)?; + let k = Tensor::randn(0f32, 1., (1, 2, 5, 8), &device)?; + let v = Tensor::randn(0f32, 1., (1, 2, 5, 8), &device)?; + let mask = create_causal_mask(3, 5, DType::F32, &device)?; + let scale = 1.0 / (8f64).sqrt(); + + let repeated = scaled_dot_product_attention( + &q, + &repeat_kv(&k, 2)?, + &repeat_kv(&v, 2)?, + Some(&mask), + scale, + false, + )?; + let grouped = scaled_dot_product_attention_gqa(&q, &k, &v, Some(&mask), scale, false, 2)?; + let repeated = repeated.flatten_all()?.to_vec1::()?; + let grouped = grouped.flatten_all()?.to_vec1::()?; + assert!( + repeated + .iter() + .zip(grouped) + .all(|(left, right)| (left - right).abs() < 1e-5) + ); + Ok(()) + } + #[test] fn test_causal_mask() -> Result<()> { let device = Device::Cpu; @@ -652,6 +787,29 @@ mod tests { Ok(()) } + #[test] + fn test_bf16_causal_mask_preserves_adjacent_positions_in_long_context() -> Result<()> { + let device = Device::Cpu; + let query_len = 16; + let kv_len = 2048; + let context_len = kv_len - query_len; + let mask = create_causal_mask(query_len, kv_len, DType::BF16, &device)? + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1::()?; + + for row in 0..query_len { + let start = row * kv_len; + let last_visible = context_len + row; + assert_eq!(mask[start + last_visible], 0.0); + if last_visible + 1 < kv_len { + assert!(mask[start + last_visible + 1].is_infinite()); + assert!(mask[start + last_visible + 1].is_sign_negative()); + } + } + Ok(()) + } + #[test] fn test_repeat_kv() -> Result<()> { let device = Device::Cpu; diff --git a/oar-ocr-vl/src/hunyuanocr/dflash.rs b/oar-ocr-vl/src/hunyuanocr/dflash.rs new file mode 100644 index 0000000..581aa02 --- /dev/null +++ b/oar-ocr-vl/src/hunyuanocr/dflash.rs @@ -0,0 +1,1395 @@ +//! DFlash parallel draft model for HunyuanOCR 1.5. +//! +//! The draft consumes intermediate features from the target decoder as cached +//! context K/V. Its queries are one target-produced bonus token followed by a +//! block of mask tokens. All mask positions are predicted in one non-causal +//! pass and are then verified by the target model in one causal pass. + +#[cfg(feature = "cuda")] +use super::dynamic_kv::{ + DynamicPagedKvAppend, FusedAddRmsNormBf16, FusedRmsNormRopeBf16, FusedRopeBf16, + FusedSiluMulBf16, +}; +use crate::attention::{RotaryEmbedding, flash_attention, scaled_dot_product_attention_gqa}; +use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; +use candle_core::{DType, Device, Tensor}; +use candle_nn::{Linear, Module, RmsNorm, VarBuilder, linear_no_bias, rms_norm}; +use oar_ocr_core::core::OCRError; +use serde::Deserialize; +use std::cell::RefCell; +use std::path::Path; + +const ROPE_CACHE_LEN: usize = 16_384; + +fn tensor_err(message: &'static str, error: candle_core::Error) -> OCRError { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + message, + error, + ) +} + +#[derive(Debug, Clone, Deserialize)] +pub struct DFlashTargetConfig { + pub mask_token_id: u32, + pub target_layer_ids: Vec, +} + +/// Configuration stored in `dflash/config.json`. +#[derive(Debug, Clone, Deserialize)] +pub struct DFlashConfig { + pub block_size: usize, + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_attention_heads: usize, + pub num_hidden_layers: usize, + pub num_key_value_heads: usize, + pub head_dim: usize, + pub vocab_size: usize, + pub rms_norm_eps: f64, + pub rope_theta: f64, + pub dflash_config: DFlashTargetConfig, +} + +impl DFlashConfig { + pub fn from_path(path: impl AsRef) -> Result { + crate::utils::load_json_config(path, "HunyuanOCR DFlash", "config.json") + } + + fn validate(&self) -> Result<(), OCRError> { + if self.block_size == 0 + || self.num_hidden_layers == 0 + || self.dflash_config.target_layer_ids.is_empty() + { + return Err(OCRError::ConfigError { + message: + "HunyuanOCR DFlash: block size, layer count, and target layers must be non-zero" + .to_string(), + }); + } + if !self + .num_attention_heads + .is_multiple_of(self.num_key_value_heads) + { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR DFlash: {} attention heads is not divisible by {} KV heads", + self.num_attention_heads, self.num_key_value_heads + ), + }); + } + Ok(()) + } +} + +const CONTEXT_KV_INITIAL_CAPACITY: usize = 16_384; +const CONTEXT_KV_PAGE_SIZE: usize = 32; + +#[derive(Debug)] +struct ContextKv { + storage: Option<(Tensor, Tensor)>, + block_table: Option, + k: Option, + v: Option, + len: usize, + capacity: usize, +} + +impl Default for ContextKv { + fn default() -> Self { + Self { + storage: None, + block_table: None, + k: None, + v: None, + len: 0, + capacity: CONTEXT_KV_INITIAL_CAPACITY, + } + } +} + +impl ContextKv { + fn reset(&mut self) { + self.k = None; + self.v = None; + self.len = 0; + } + + fn ensure_capacity( + &mut self, + template_k: &Tensor, + template_v: &Tensor, + required: usize, + ) -> Result<(), OCRError> { + let (_, heads, _, head_dim) = template_k + .dims4() + .map_err(|e| tensor_err("HunyuanOCR DFlash: context template shape", e))?; + let reusable = self.storage.as_ref().is_some_and(|(storage_k, storage_v)| { + storage_k.dtype() == template_k.dtype() + && storage_v.dtype() == template_v.dtype() + && storage_k.device().same_device(template_k.device()) + && storage_v.device().same_device(template_v.device()) + && storage_k.dims() + == [ + self.capacity / CONTEXT_KV_PAGE_SIZE, + CONTEXT_KV_PAGE_SIZE, + heads, + head_dim, + ] + && storage_v.dims() == storage_k.dims() + }); + if self.storage.is_some() && !reusable { + self.storage = None; + self.block_table = None; + self.reset(); + } + if self.storage.is_none() { + self.capacity = + self.capacity.max(required).div_ceil(CONTEXT_KV_PAGE_SIZE) * CONTEXT_KV_PAGE_SIZE; + let blocks = self.capacity / CONTEXT_KV_PAGE_SIZE; + let shape = (blocks, CONTEXT_KV_PAGE_SIZE, heads, head_dim); + self.storage = Some(( + Tensor::zeros(shape, template_k.dtype(), template_k.device()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: allocate context K", e))?, + Tensor::zeros(shape, template_v.dtype(), template_v.device()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: allocate context V", e))?, + )); + self.block_table = Some( + Tensor::new((0..blocks as u32).collect::>(), template_k.device()) + .and_then(|x| x.reshape((1, blocks))) + .map_err(|e| tensor_err("HunyuanOCR DFlash: create KV block table", e))?, + ); + } else if required > self.capacity { + let new_capacity = (self.capacity * 2) + .max(required) + .div_ceil(CONTEXT_KV_PAGE_SIZE) + * CONTEXT_KV_PAGE_SIZE; + let grow_by = new_capacity - self.capacity; + let (old_k, old_v) = self.storage.as_ref().expect("context storage initialized"); + let old_k = old_k + .reshape((self.capacity, heads, head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: flatten context K", e))?; + let old_v = old_v + .reshape((self.capacity, heads, head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: flatten context V", e))?; + let extra_k = Tensor::zeros( + (grow_by, heads, head_dim), + template_k.dtype(), + template_k.device(), + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: grow context K", e))?; + let extra_v = Tensor::zeros( + (grow_by, heads, head_dim), + template_v.dtype(), + template_v.device(), + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: grow context V", e))?; + let blocks = new_capacity / CONTEXT_KV_PAGE_SIZE; + self.storage = Some(( + Tensor::cat(&[&old_k, &extra_k], 0) + .and_then(|x| x.reshape((blocks, CONTEXT_KV_PAGE_SIZE, heads, head_dim))) + .map_err(|e| tensor_err("HunyuanOCR DFlash: expand context K", e))?, + Tensor::cat(&[&old_v, &extra_v], 0) + .and_then(|x| x.reshape((blocks, CONTEXT_KV_PAGE_SIZE, heads, head_dim))) + .map_err(|e| tensor_err("HunyuanOCR DFlash: expand context V", e))?, + )); + self.block_table = Some( + Tensor::new((0..blocks as u32).collect::>(), template_k.device()) + .and_then(|x| x.reshape((1, blocks))) + .map_err(|e| tensor_err("HunyuanOCR DFlash: grow KV block table", e))?, + ); + self.capacity = new_capacity; + } + Ok(()) + } + + #[cfg(feature = "cuda")] + fn initialize_storage(&mut self, template: &Tensor) -> Result<(), OCRError> { + self.ensure_capacity(template, template, self.capacity) + } + + #[cfg(feature = "cuda")] + fn storage(&self) -> Option<(&Tensor, &Tensor, &Tensor)> { + self.storage + .as_ref() + .zip(self.block_table.as_ref()) + .map(|((k, v), table)| (k, v, table)) + } + + fn append(&mut self, k: &Tensor, v: &Tensor) -> Result<(), OCRError> { + let added = k + .dim(2) + .map_err(|e| tensor_err("HunyuanOCR DFlash: context K length", e))?; + let required = self.len + added; + self.ensure_capacity(k, v, required)?; + // `ensure_capacity` may have discarded incompatible storage and reset + // the logical cache, so do not retain the pre-reset length. + let required = self.len + added; + let (storage_k, storage_v) = self.storage.as_mut().expect("context storage initialized"); + let (_, heads, _, head_dim) = k + .dims4() + .map_err(|e| tensor_err("HunyuanOCR DFlash: append context shape", e))?; + let k = k + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: page context K", e))?; + let v = v + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: page context V", e))?; + let flat_k = storage_k + .reshape((self.capacity, heads, head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: flatten storage K", e))?; + let flat_v = storage_v + .reshape((self.capacity, heads, head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: flatten storage V", e))?; + flat_k + .slice_set(&k, 0, self.len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: append context K", e))?; + flat_v + .slice_set(&v, 0, self.len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: append context V", e))?; + self.len = required; + self.k = Some( + flat_k + .narrow(0, 0, self.len) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.unsqueeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: view context K", e))?, + ); + self.v = Some( + flat_v + .narrow(0, 0, self.len) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.unsqueeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: view context V", e))?, + ); + Ok(()) + } + + /// Place the transient bonus+mask K/V immediately after the accepted + /// context. The logical context length is unchanged, so the next accepted + /// append overwrites this tail instead of copying the whole history. + fn with_queries(&mut self, k: &Tensor, v: &Tensor) -> Result<(Tensor, Tensor), OCRError> { + if self.len == 0 { + return Err(OCRError::InvalidInput { + message: "HunyuanOCR DFlash: context cache is empty".to_string(), + }); + } + let query_len = k + .dim(2) + .map_err(|e| tensor_err("HunyuanOCR DFlash: query K length", e))?; + let required = self.len + query_len; + self.ensure_capacity(k, v, required)?; + let (storage_k, storage_v) = self.storage.as_mut().expect("context storage initialized"); + let (_, heads, _, head_dim) = k + .dims4() + .map_err(|e| tensor_err("HunyuanOCR DFlash: query context shape", e))?; + let k = k + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: page query K", e))?; + let v = v + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: page query V", e))?; + let flat_k = storage_k + .reshape((self.capacity, heads, head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: flatten query storage K", e))?; + let flat_v = storage_v + .reshape((self.capacity, heads, head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: flatten query storage V", e))?; + flat_k + .slice_set(&k, 0, self.len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: write query K", e))?; + flat_v + .slice_set(&v, 0, self.len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: write query V", e))?; + Ok(( + flat_k + .narrow(0, 0, required) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.unsqueeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: view query K", e))?, + flat_v + .narrow(0, 0, required) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.unsqueeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: view query V", e))?, + )) + } +} + +fn apply_rope(x: &Tensor, cos: &Tensor, sin: &Tensor) -> Result { + let a = x + .broadcast_mul(cos) + .map_err(|e| tensor_err("HunyuanOCR DFlash: rope x*cos", e))?; + let b = rotate_half(x)? + .broadcast_mul(sin) + .map_err(|e| tensor_err("HunyuanOCR DFlash: rope rotate_half*sin", e))?; + (a + b).map_err(|e| tensor_err("HunyuanOCR DFlash: rope sum", e)) +} + +#[derive(Debug)] +struct DFlashAttention { + qkv_proj: Linear, + o_proj: Linear, + q_norm: RmsNorm, + k_norm: RmsNorm, + num_heads: usize, + num_kv_heads: usize, + num_kv_groups: usize, + head_dim: usize, + scale: f64, +} + +impl DFlashAttention { + fn load(cfg: &DFlashConfig, vb: VarBuilder) -> Result { + let q_proj = linear_no_bias( + cfg.hidden_size, + cfg.num_attention_heads * cfg.head_dim, + vb.pp("q_proj"), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load q_proj", e))?; + let k_proj = linear_no_bias( + cfg.hidden_size, + cfg.num_key_value_heads * cfg.head_dim, + vb.pp("k_proj"), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load k_proj", e))?; + let v_proj = linear_no_bias( + cfg.hidden_size, + cfg.num_key_value_heads * cfg.head_dim, + vb.pp("v_proj"), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load v_proj", e))?; + let o_proj = linear_no_bias( + cfg.num_attention_heads * cfg.head_dim, + cfg.hidden_size, + vb.pp("o_proj"), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load o_proj", e))?; + let q_norm = rms_norm(cfg.head_dim, cfg.rms_norm_eps, vb.pp("q_norm")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load q_norm", e))?; + let k_norm = rms_norm(cfg.head_dim, cfg.rms_norm_eps, vb.pp("k_norm")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load k_norm", e))?; + let qkv_weight = Tensor::cat(&[q_proj.weight(), k_proj.weight(), v_proj.weight()], 0) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fuse QKV weights", e))?; + Ok(Self { + qkv_proj: Linear::new(qkv_weight, None), + o_proj, + q_norm, + k_norm, + num_heads: cfg.num_attention_heads, + num_kv_heads: cfg.num_key_value_heads, + num_kv_groups: cfg.num_attention_heads / cfg.num_key_value_heads, + head_dim: cfg.head_dim, + scale: (cfg.head_dim as f64).powf(-0.5), + }) + } + + fn shape_projected( + &self, + projected: &Tensor, + heads: usize, + norm: Option<&RmsNorm>, + ) -> Result { + let (batch, seq_len, _) = projected + .dims3() + .map_err(|e| tensor_err("HunyuanOCR DFlash: projected dims", e))?; + let projected = projected + .reshape((batch, seq_len, heads, self.head_dim)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: projection reshape", e))?; + let projected = match norm { + Some(norm) => norm + .forward(&projected) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "Q/K norm", e))?, + None => projected, + }; + projected + .transpose(1, 2) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: projection transpose", e)) + } + + fn qkv_weights(&self) -> Result<(Tensor, Tensor), OCRError> { + let q_width = self.num_heads * self.head_dim; + let kv_width = self.num_kv_heads * self.head_dim; + let k = self + .qkv_proj + .weight() + .narrow(0, q_width, kv_width) + .map_err(|e| tensor_err("HunyuanOCR DFlash: K weight view", e))?; + let v = self + .qkv_proj + .weight() + .narrow(0, q_width + kv_width, kv_width) + .map_err(|e| tensor_err("HunyuanOCR DFlash: V weight view", e))?; + Ok((k, v)) + } + + fn append_projected_context( + &self, + raw_k: &Tensor, + raw_v: &Tensor, + cos: &Tensor, + sin: &Tensor, + cache: &mut ContextKv, + ) -> Result<(), OCRError> { + let k = self.shape_projected(raw_k, self.num_kv_heads, Some(&self.k_norm))?; + let k = apply_rope(&k, cos, sin)?; + let v = self.shape_projected(raw_v, self.num_kv_heads, None)?; + cache.append(&k, &v) + } + + fn project_qkv( + &self, + hidden: &Tensor, + cos: &Tensor, + sin: &Tensor, + cos_sin: Option<&Tensor>, + ) -> Result<(Tensor, Tensor, Tensor), OCRError> { + #[cfg(not(feature = "cuda"))] + let _ = cos_sin; + let q_width = self.num_heads * self.head_dim; + let kv_width = self.num_kv_heads * self.head_dim; + let qkv = self + .qkv_proj + .forward(hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "QKV projection", e))?; + #[cfg(feature = "cuda")] + if qkv.device().is_cuda() + && qkv.dtype() == DType::BF16 + && self.head_dim == 128 + && cos_sin.is_some() + { + let cos_sin = cos_sin.expect("checked above"); + let q = qkv + .apply_op3_no_bwd( + cos_sin, + self.q_norm.weight(), + &FusedRmsNormRopeBf16 { + projection_width: q_width + 2 * kv_width, + projection_offset: 0, + num_heads: self.num_heads, + query_len: hidden.dim(1).map_err(|e| { + tensor_err("HunyuanOCR DFlash: fused Q query length", e) + })?, + head_dim: self.head_dim, + eps: self.q_norm.eps() as f32, + include_v: false, + }, + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused Q RMSNorm RoPE", e))?; + let kv = qkv + .apply_op3_no_bwd( + cos_sin, + self.k_norm.weight(), + &FusedRmsNormRopeBf16 { + projection_width: q_width + 2 * kv_width, + projection_offset: q_width, + num_heads: self.num_kv_heads, + query_len: hidden.dim(1).map_err(|e| { + tensor_err("HunyuanOCR DFlash: fused KV query length", e) + })?, + head_dim: self.head_dim, + eps: self.k_norm.eps() as f32, + include_v: true, + }, + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused KV RMSNorm RoPE", e))?; + let k = kv + .narrow(0, 0, 1) + .and_then(|x| x.squeeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused K view", e))?; + let v = kv + .narrow(0, 1, 1) + .and_then(|x| x.squeeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused V view", e))?; + return Ok((q, k, v)); + } + let q_raw = qkv + .narrow(2, 0, q_width) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: Q projection slice", e))?; + let k_raw = qkv + .narrow(2, q_width, kv_width) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: K projection slice", e))?; + let v_raw = qkv + .narrow(2, q_width + kv_width, kv_width) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: V projection slice", e))?; + let q = self.shape_projected(&q_raw, self.num_heads, Some(&self.q_norm))?; + let q = self.apply_query_rope(&q, cos, sin, cos_sin)?; + let query_k = self.shape_projected(&k_raw, self.num_kv_heads, Some(&self.k_norm))?; + let query_k = self.apply_query_rope(&query_k, cos, sin, cos_sin)?; + let query_v = self.shape_projected(&v_raw, self.num_kv_heads, None)?; + + Ok((q, query_k, query_v)) + } + + fn apply_query_rope( + &self, + input: &Tensor, + cos: &Tensor, + sin: &Tensor, + cos_sin: Option<&Tensor>, + ) -> Result { + #[cfg(not(feature = "cuda"))] + let _ = cos_sin; + #[cfg(feature = "cuda")] + if input.device().is_cuda() + && input.dtype() == DType::BF16 + && cos_sin.is_some() + && self.head_dim.is_multiple_of(2) + { + let output = Tensor::zeros(input.shape(), input.dtype(), input.device()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: allocate fused RoPE output", e))?; + output + .inplace_op3(input, cos_sin.expect("checked above"), &FusedRopeBf16) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused BF16 RoPE", e))?; + return Ok(output); + } + apply_rope(input, cos, sin) + } + + fn attend_projected( + &self, + q: &Tensor, + query_k: &Tensor, + query_v: &Tensor, + cache: &mut ContextKv, + ) -> Result { + let (k, v) = cache.with_queries(query_k, query_v)?; + // Compute GQA without physically repeating the full context K/V cache + // across all query heads. + let output = match flash_attention(q, &k, &v, self.scale, false) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "flash attention", e))? + { + Some(output) => output, + None => scaled_dot_product_attention_gqa( + q, + &k, + &v, + None, + self.scale, + false, + self.num_kv_groups, + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "grouped attention", e))?, + }; + Ok(output) + } + + #[cfg(feature = "cuda")] + fn attend_projected_dynamic( + &self, + q: &Tensor, + query_k: &Tensor, + query_v: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + cache: &ContextKv, + ) -> Result { + let (_, _, query_len, _) = q + .dims4() + .map_err(|e| tensor_err("HunyuanOCR DFlash: dynamic Q dimensions", e))?; + let (cache_k, cache_v, block_table) = + cache.storage().ok_or_else(|| OCRError::ConfigError { + message: "HunyuanOCR DFlash: dynamic context storage is not initialized" + .to_string(), + })?; + let append = DynamicPagedKvAppend { + query_len, + cache_len: CONTEXT_KV_INITIAL_CAPACITY, + }; + cache_k + .inplace_op3(query_k, kv_lengths, &append) + .map_err(|e| tensor_err("HunyuanOCR DFlash: dynamic query K append", e))?; + cache_v + .inplace_op3(query_v, kv_lengths, &append) + .map_err(|e| tensor_err("HunyuanOCR DFlash: dynamic query V append", e))?; + + let q = q + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: dynamic Q layout", e))?; + candle_flash_attn::flash_attn_varlen_paged_windowed( + &q, + cache_k, + cache_v, + query_lengths, + kv_lengths, + block_table, + None, + query_len, + CONTEXT_KV_INITIAL_CAPACITY, + self.scale as f32, + None, + None, + CONTEXT_KV_PAGE_SIZE, + None, + ) + .and_then(|x| x.transpose(0, 1)) + .and_then(|x| x.unsqueeze(0)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "dynamic flash attention", e)) + } + + fn flatten_attention_output(&self, output: &Tensor) -> Result { + let (batch, _, query_len, _) = output + .dims4() + .map_err(|e| tensor_err("HunyuanOCR DFlash: attention output dimensions", e))?; + output + .transpose(1, 2) + .and_then(|x| x.reshape((batch, query_len, self.num_heads * self.head_dim))) + .map_err(|e| tensor_err("HunyuanOCR DFlash: attention output reshape", e)) + } + + fn project_attention_output(&self, output: &Tensor) -> Result { + self.o_proj + .forward(output) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "o_proj", e)) + } +} + +#[derive(Debug)] +struct DFlashMlp { + gate_up_proj: Linear, + down_proj: Linear, + intermediate_size: usize, +} + +impl DFlashMlp { + fn load(cfg: &DFlashConfig, vb: VarBuilder) -> Result { + let gate_proj = linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("gate_proj")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load gate_proj", e))?; + let up_proj = linear_no_bias(cfg.hidden_size, cfg.intermediate_size, vb.pp("up_proj")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load up_proj", e))?; + let down_proj = linear_no_bias(cfg.intermediate_size, cfg.hidden_size, vb.pp("down_proj")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load down_proj", e))?; + let gate_up_weight = Tensor::cat(&[gate_proj.weight(), up_proj.weight()], 0) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fuse gate/up weights", e))?; + Ok(Self { + gate_up_proj: Linear::new(gate_up_weight, None), + down_proj, + intermediate_size: cfg.intermediate_size, + }) + } + + fn forward(&self, input: &Tensor) -> Result { + let gate_up = self + .gate_up_proj + .forward(input) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP gate/up", e))?; + let gate = gate_up + .narrow(2, 0, self.intermediate_size) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP gate", e))?; + let up = gate_up + .narrow(2, self.intermediate_size, self.intermediate_size) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP up", e))?; + #[cfg(feature = "cuda")] + let hidden = if gate.device().is_cuda() + && gate.dtype() == DType::BF16 + && gate.is_contiguous() + && up.is_contiguous() + { + gate.apply_op2_no_bwd(&up, &FusedSiluMulBf16) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused MLP SiLU*up", e))? + } else { + let gate = candle_nn::ops::silu(&gate) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP SiLU", e))?; + (gate * up).map_err(|e| tensor_err("HunyuanOCR DFlash: MLP multiply", e))? + }; + #[cfg(not(feature = "cuda"))] + let hidden = { + let gate = candle_nn::ops::silu(&gate) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP SiLU", e))?; + (gate * up).map_err(|e| tensor_err("HunyuanOCR DFlash: MLP multiply", e))? + }; + self.down_proj + .forward(&hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP down", e)) + } +} + +#[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, + self_attn: DFlashAttention, + post_attention_layernorm: RmsNorm, + mlp: DFlashMlp, +} + +impl DFlashLayer { + fn load(cfg: &DFlashConfig, vb: VarBuilder) -> Result { + let input_layernorm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load input_layernorm", e))?; + let post_attention_layernorm = rms_norm( + cfg.hidden_size, + cfg.rms_norm_eps, + vb.pp("post_attention_layernorm"), + ) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "load post_attention_layernorm", e) + })?; + Ok(Self { + input_layernorm, + self_attn: DFlashAttention::load(cfg, vb.pp("self_attn"))?, + post_attention_layernorm, + mlp: DFlashMlp::load(cfg, vb.pp("mlp"))?, + }) + } + + fn forward( + &self, + hidden: &Tensor, + cos: &Tensor, + sin: &Tensor, + cos_sin: Option<&Tensor>, + cache: &mut ContextKv, + ) -> Result { + let (q, k, v) = self.project_qkv_eager(hidden, cos, sin, cos_sin)?; + let attention = self.self_attn.attend_projected(&q, &k, &v, cache)?; + let attention = self.self_attn.flatten_attention_output(&attention)?; + self.post_attention_eager(hidden, &attention) + } + + fn project_qkv_eager( + &self, + hidden: &Tensor, + cos: &Tensor, + sin: &Tensor, + cos_sin: Option<&Tensor>, + ) -> Result<(Tensor, Tensor, Tensor), OCRError> { + let normed = self + .input_layernorm + .forward(hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "input norm", e))?; + self.self_attn.project_qkv(&normed, cos, sin, cos_sin) + } + + fn post_attention_eager( + &self, + hidden: &Tensor, + attention: &Tensor, + ) -> Result { + let attention = self.self_attn.project_attention_output(attention)?; + #[cfg(feature = "cuda")] + if hidden.device().is_cuda() + && hidden.dtype() == DType::BF16 + && hidden.is_contiguous() + && attention.is_contiguous() + { + let packed = hidden + .apply_op3_no_bwd( + &attention, + self.post_attention_layernorm.weight(), + &FusedAddRmsNormBf16 { + eps: self.post_attention_layernorm.eps() as f32, + }, + ) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "fused residual RMSNorm", e) + })?; + let residual = packed + .narrow(0, 0, 1) + .and_then(|x| x.squeeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused residual view", e))?; + let normed = packed + .narrow(0, 1, 1) + .and_then(|x| x.squeeze(0)) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fused RMSNorm view", e))?; + let mlp = self.mlp.forward(&normed)?; + return (residual + mlp).map_err(|e| tensor_err("HunyuanOCR DFlash: MLP residual", e)); + } + let hidden = (hidden + attention) + .map_err(|e| tensor_err("HunyuanOCR DFlash: attention residual", e))?; + let normed = self + .post_attention_layernorm + .forward(&hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "post-attention norm", e))?; + let mlp = self.mlp.forward(&normed)?; + (hidden + mlp).map_err(|e| tensor_err("HunyuanOCR DFlash: MLP residual", e)) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden: &Tensor, + cos: &Tensor, + sin: &Tensor, + cos_sin: Option<&Tensor>, + query_lengths: &Tensor, + kv_lengths: &Tensor, + cache: &ContextKv, + ) -> Result { + let (q, k, v) = self.project_qkv_eager(hidden, cos, sin, cos_sin)?; + let attention = self.self_attn.attend_projected_dynamic( + &q, + &k, + &v, + query_lengths, + kv_lengths, + cache, + )?; + let attention = self.self_attn.flatten_attention_output(&attention)?; + self.post_attention_eager(hidden, &attention) + } +} + +#[cfg(feature = "cuda")] +struct DFlashCudaGraph { + graph: candle_core::cuda_backend::cudarc::driver::CudaGraph, + query_input: Tensor, + cos_input: Tensor, + sin_input: Tensor, + _query_lengths: Tensor, + kv_lengths: Tensor, + output: Tensor, +} + +#[cfg(feature = "cuda")] +impl std::fmt::Debug for DFlashCudaGraph { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DFlashCudaGraph") + .field("query_input", &self.query_input.shape()) + .field("output", &self.output.shape()) + .finish_non_exhaustive() + } +} + +/// Loaded DFlash draft and its incremental target-context K/V caches. +pub(crate) struct DFlashModel { + #[cfg(feature = "cuda")] + decode_graph: RefCell>, + cfg: DFlashConfig, + fc: Linear, + hidden_norm: RmsNorm, + layers: Vec, + norm: RmsNorm, + rotary: RotaryEmbedding, + rope_cos: Tensor, + rope_sin: Tensor, + context_kv_proj: Linear, + caches: RefCell>, + dtype: DType, + device: Device, +} + +impl DFlashModel { + pub(crate) fn from_dir( + model_dir: impl AsRef, + dtype: DType, + device: &Device, + ) -> Result { + let model_dir = model_dir.as_ref(); + let cfg = DFlashConfig::from_path(model_dir.join("config.json"))?; + cfg.validate()?; + let files = crate::utils::collect_safetensors(model_dir, "HunyuanOCR DFlash")?; + // SAFETY: the checkpoint files remain mapped for the lifetime of the tensors. + let vb = unsafe { + VarBuilder::from_mmaped_safetensors(&files, dtype, device) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load safetensors", e))? + }; + let fc = linear_no_bias( + cfg.hidden_size * cfg.dflash_config.target_layer_ids.len(), + cfg.hidden_size, + vb.pp("fc"), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load fc", e))?; + let hidden_norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("hidden_norm")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load hidden_norm", e))?; + let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm")) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "load norm", e))?; + let mut layers = Vec::with_capacity(cfg.num_hidden_layers); + for index in 0..cfg.num_hidden_layers { + layers.push(DFlashLayer::load(&cfg, vb.pp(format!("layers.{index}")))?); + } + let mut context_weights = Vec::with_capacity(cfg.num_hidden_layers * 2); + for layer in &layers { + let (k, v) = layer.self_attn.qkv_weights()?; + context_weights.push(k); + context_weights.push(v); + } + let context_weight_refs: Vec<&Tensor> = context_weights.iter().collect(); + let context_kv_weight = Tensor::cat(&context_weight_refs, 0) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: fuse context K/V weights", e))?; + let context_kv_proj = Linear::new(context_kv_weight, None); + let rotary = RotaryEmbedding::new_dynamic(cfg.head_dim, cfg.rope_theta, device)?; + let rope_positions = Tensor::arange(0i64, ROPE_CACHE_LEN as i64, device) + .and_then(|x| x.reshape((1, 1, ROPE_CACHE_LEN))) + .map_err(|e| tensor_err("HunyuanOCR DFlash: rope cache positions", e))?; + let (rope_cos, rope_sin) = rotary.forward_multi_axis(&rope_positions, dtype)?; + let caches = (0..cfg.num_hidden_layers) + .map(|_| ContextKv::default()) + .collect(); + let model = Self { + #[cfg(feature = "cuda")] + decode_graph: RefCell::new(None), + cfg, + fc, + hidden_norm, + layers, + norm, + rotary, + rope_cos, + rope_sin, + context_kv_proj, + caches: RefCell::new(caches), + dtype, + device: device.clone(), + }; + // `DynamicPagedKvAppend` (the graph's per-layer append kernel) only + // accepts BF16; capturing under another dtype override would fail + // model load entirely even though the eager DFlash path below + // supports F16/F32. + #[cfg(feature = "cuda")] + if device.is_cuda() + && dtype == DType::BF16 + && std::env::var_os("OAR_HUNYUAN_DISABLE_CUDA_GRAPH").is_none() + { + model.capture_cuda_graph()?; + } + Ok(model) + } + + pub(crate) fn config(&self) -> &DFlashConfig { + &self.cfg + } + + #[cfg(feature = "cuda")] + fn invalidate_cuda_graph(&self) { + // A captured graph owns raw pointers into the fixed-size cache. Once + // that cache must grow, the graph cannot safely be reused, including + // after a later page-level reset. + self.decode_graph.borrow_mut().take(); + } + + fn rope(&self, start: usize, len: usize) -> Result<(Tensor, Tensor), OCRError> { + if start + len <= ROPE_CACHE_LEN { + let cos = self + .rope_cos + .narrow(2, start, len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: rope cos slice", e))?; + let sin = self + .rope_sin + .narrow(2, start, len) + .map_err(|e| tensor_err("HunyuanOCR DFlash: rope sin slice", e))?; + return Ok((cos, sin)); + } + let positions: Vec = (start..start + len).map(|x| x as i64).collect(); + let positions = Tensor::from_vec(positions, (1, 1, len), &self.device) + .map_err(|e| tensor_err("HunyuanOCR DFlash: position tensor", e))?; + self.rotary.forward_multi_axis(&positions, self.dtype) + } + + fn transform_target(&self, aux_hidden: &Tensor) -> Result { + let hidden = self + .fc + .forward(aux_hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "target fc", e))?; + self.hidden_norm + .forward(&hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "target hidden norm", e)) + } + + fn append_projected_context( + &self, + target: &Tensor, + cos: &Tensor, + sin: &Tensor, + caches: &mut [ContextKv], + ) -> Result<(), OCRError> { + let projected = self + .context_kv_proj + .forward(target) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "context K/V", e))?; + let kv_width = self.cfg.num_key_value_heads * self.cfg.head_dim; + for (index, (layer, cache)) in self.layers.iter().zip(caches.iter_mut()).enumerate() { + let offset = index * 2 * kv_width; + let raw_k = projected + .narrow(2, offset, kv_width) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: context K slice", e))?; + let raw_v = projected + .narrow(2, offset + kv_width, kv_width) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: context V slice", e))?; + layer + .self_attn + .append_projected_context(&raw_k, &raw_v, cos, sin, cache)?; + } + Ok(()) + } + + pub(crate) fn reset_context(&self, aux_hidden: &Tensor) -> Result<(), OCRError> { + #[cfg(feature = "cuda")] + let _cuda_htod_cache = match &self.device { + Device::Cuda(device) => Some(device.enable_cuda_graph_htod_cache()), + _ => None, + }; + let len = aux_hidden + .dim(1) + .map_err(|e| tensor_err("HunyuanOCR DFlash: target context length", e))?; + #[cfg(feature = "cuda")] + if len > CONTEXT_KV_INITIAL_CAPACITY { + self.invalidate_cuda_graph(); + } + let target = self.transform_target(aux_hidden)?; + let (cos, sin) = self.rope(0, len)?; + let mut caches = self.caches.borrow_mut(); + for cache in caches.iter_mut() { + cache.reset(); + } + self.append_projected_context(&target, &cos, &sin, &mut caches) + } + + pub(crate) fn append_context(&self, aux_hidden: &Tensor) -> Result<(), OCRError> { + #[cfg(feature = "cuda")] + let _cuda_htod_cache = match &self.device { + Device::Cuda(device) => Some(device.enable_cuda_graph_htod_cache()), + _ => None, + }; + let added = aux_hidden + .dim(1) + .map_err(|e| tensor_err("HunyuanOCR DFlash: appended context length", e))?; + if added == 0 { + return Ok(()); + } + let start = self.context_len(); + #[cfg(feature = "cuda")] + if start.saturating_add(added) > CONTEXT_KV_INITIAL_CAPACITY { + self.invalidate_cuda_graph(); + } + let target = self.transform_target(aux_hidden)?; + let (cos, sin) = self.rope(start, added)?; + let mut caches = self.caches.borrow_mut(); + self.append_projected_context(&target, &cos, &sin, &mut caches) + } + + pub(crate) fn context_len(&self) -> usize { + self.caches.borrow().first().map_or(0, |cache| cache.len) + } + + #[cfg(feature = "cuda")] + fn forward_queries_dynamic( + &self, + query_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let caches = self.caches.borrow(); + let cos_sin = if query_embeds.dtype() == DType::BF16 { + Some( + Tensor::cat(&[cos, sin], 0) + .map_err(|e| tensor_err("HunyuanOCR DFlash: pack dynamic RoPE cos/sin", e))?, + ) + } else { + None + }; + let mut hidden = query_embeds.clone(); + for (layer, cache) in self.layers.iter().zip(caches.iter()) { + hidden = layer.forward_dynamic( + &hidden, + cos, + sin, + cos_sin.as_ref(), + query_lengths, + kv_lengths, + cache, + )?; + } + self.norm + .forward(&hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "dynamic final norm", e)) + } + + #[cfg(feature = "cuda")] + fn capture_cuda_graph(&self) -> 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.device else { + return Ok(()); + }; + let query_len = self.cfg.block_size; + let template = Tensor::zeros( + ( + 1, + self.cfg.num_key_value_heads, + query_len, + self.cfg.head_dim, + ), + self.dtype, + &self.device, + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: dynamic cache template", e))?; + for cache in self.caches.borrow_mut().iter_mut() { + cache.initialize_storage(&template)?; + } + + let query_input = Tensor::zeros( + (1, query_len, self.cfg.hidden_size), + self.dtype, + &self.device, + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: full graph query input", e))?; + let cos_input = Tensor::zeros( + (1, 1, query_len, self.cfg.head_dim), + self.dtype, + &self.device, + ) + .map_err(|e| tensor_err("HunyuanOCR DFlash: full graph cos input", e))?; + let sin_input = Tensor::zeros( + (1, 1, query_len, self.cfg.head_dim), + self.dtype, + &self.device, + ) + .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) + .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(); + + let warm = self.forward_queries_dynamic( + &query_input, + &cos_input, + &sin_input, + &query_lengths, + &kv_lengths, + )?; + sync_dflash_graph_tensor(&warm, "warm full draft graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| dflash_cuda_graph_error("begin full draft graph capture", e))?; + let output = match self.forward_queries_dynamic( + &query_input, + &cos_input, + &sin_input, + &query_lengths, + &kv_lengths, + ) { + 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| dflash_cuda_graph_error("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(&output, "sync full draft graph")?; + self.clear_context(); + *self.decode_graph.borrow_mut() = Some(DFlashCudaGraph { + graph, + query_input, + cos_input, + sin_input, + _query_lengths: query_lengths, + kv_lengths, + output, + }); + Ok(()) + } + + #[cfg(feature = "cuda")] + fn replay_cuda_graph( + &self, + query_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + total_kv_len: usize, + ) -> Result, OCRError> { + if total_kv_len > CONTEXT_KV_INITIAL_CAPACITY { + self.invalidate_cuda_graph(); + return Ok(None); + } + let captured_ref = self.decode_graph.borrow(); + let Some(captured) = captured_ref.as_ref() else { + return Ok(None); + }; + if query_embeds.shape() != captured.query_input.shape() + || cos.shape() != captured.cos_input.shape() + || sin.shape() != captured.sin_input.shape() + { + return Ok(None); + } + captured + .query_input + .slice_set(query_embeds, 0, 0) + .map_err(|e| tensor_err("HunyuanOCR DFlash: copy full graph query", e))?; + captured + .cos_input + .slice_set(cos, 0, 0) + .map_err(|e| tensor_err("HunyuanOCR DFlash: copy full graph cos", e))?; + captured + .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))?; + captured + .graph + .launch() + .map_err(|e| dflash_cuda_graph_error("launch full draft graph", e))?; + Ok(Some(captured.output.clone())) + } + + /// Run the bonus+mask query block. Returns post-norm hidden states for all + /// query positions; the caller samples only rows `1..` (the mask rows). + pub(crate) fn forward_queries(&self, query_embeds: &Tensor) -> Result { + #[cfg(feature = "cuda")] + let _cuda_htod_cache = match &self.device { + Device::Cuda(device) => Some(device.enable_cuda_graph_htod_cache()), + _ => None, + }; + let query_len = query_embeds + .dim(1) + .map_err(|e| tensor_err("HunyuanOCR DFlash: query length", e))?; + let context_len = self.context_len(); + let (cos, sin) = self.rope(context_len, query_len)?; + #[cfg(feature = "cuda")] + if let Some(output) = + self.replay_cuda_graph(query_embeds, &cos, &sin, context_len + query_len)? + { + return Ok(output); + } + #[cfg(feature = "cuda")] + let cos_sin = if query_embeds.device().is_cuda() && query_embeds.dtype() == DType::BF16 { + Some( + Tensor::cat(&[&cos, &sin], 0) + .map_err(|e| tensor_err("HunyuanOCR DFlash: pack query RoPE cos/sin", e))?, + ) + } else { + None + }; + #[cfg(not(feature = "cuda"))] + let cos_sin: Option = None; + let mut caches = self.caches.borrow_mut(); + let mut hidden = query_embeds.clone(); + for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { + hidden = layer.forward(&hidden, &cos, &sin, cos_sin.as_ref(), cache)?; + } + self.norm + .forward(&hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "final norm", e)) + } + + pub(crate) fn clear_context(&self) { + for cache in self.caches.borrow_mut().iter_mut() { + cache.reset(); + } + } +} + +#[cfg(test)] +mod tests { + use super::{ContextKv, DFlashConfig}; + use candle_core::{DType, Device, Tensor}; + + #[test] + fn parses_hunyuan_dflash_config() { + let cfg: DFlashConfig = serde_json::from_str( + r#"{ + "block_size": 16, + "hidden_size": 1024, + "intermediate_size": 3584, + "num_attention_heads": 16, + "num_hidden_layers": 5, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 120818, + "rms_norm_eps": 0.00001, + "rope_theta": 10000.0, + "dflash_config": { + "mask_token_id": 120817, + "target_layer_ids": [1, 8, 15, 22] + } + }"#, + ) + .unwrap(); + assert_eq!(cfg.block_size, 16); + assert_eq!(cfg.dflash_config.target_layer_ids, [1, 8, 15, 22]); + cfg.validate().unwrap(); + } + + #[test] + fn incompatible_context_storage_starts_a_new_logical_cache() { + let device = Device::Cpu; + let mut cache = ContextKv::default(); + let old = Tensor::zeros((1, 1, 3, 4), DType::F32, &device).unwrap(); + cache.append(&old, &old).unwrap(); + + // Changing the number of heads forces replacement storage. + let new = Tensor::ones((1, 2, 2, 4), DType::F32, &device).unwrap(); + cache.append(&new, &new).unwrap(); + + assert_eq!(cache.len, 2); + let k = cache.k.as_ref().unwrap(); + let v = cache.v.as_ref().unwrap(); + assert_eq!(k.dims(), &[1, 2, 2, 4]); + assert_eq!(v.dims(), &[1, 2, 2, 4]); + assert!( + k.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .iter() + .all(|&x| x == 1.0) + ); + } +} diff --git a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu new file mode 100644 index 0000000..cdd006d --- /dev/null +++ b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.cu @@ -0,0 +1,364 @@ +#include +#include +#include + +// Copy one fixed-size verification block into a preallocated KV cache. The +// destination offset is read from a device-side cumulative-length tensor, so +// the same kernel node can be replayed inside a CUDA graph for every decode +// round without patching graph parameters on the host. +extern "C" __global__ void append_kv_f16( + half* cache, + const half* source, + const uint32_t* cumulative_lengths, + uint32_t query_len, + uint32_t num_heads, + uint32_t head_dim, + uint32_t cache_len) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = num_heads * query_len * head_dim; + if (index >= count) { + return; + } + + const uint32_t head_stride = query_len * head_dim; + const uint32_t head = index / head_stride; + const uint32_t within_head = index - head * head_stride; + const uint32_t token = within_head / head_dim; + const uint32_t lane = within_head - token * head_dim; + const uint32_t end = cumulative_lengths[1]; + const uint32_t start = end - query_len; + cache[(head * cache_len + start + token) * head_dim + lane] = source[index]; +} + +extern "C" __global__ void append_kv_bf16( + __nv_bfloat16* cache, + const __nv_bfloat16* source, + const uint32_t* cumulative_lengths, + uint32_t query_len, + uint32_t num_heads, + uint32_t head_dim, + uint32_t cache_len) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = num_heads * query_len * head_dim; + if (index >= count) { + return; + } + + const uint32_t head_stride = query_len * head_dim; + const uint32_t head = index / head_stride; + const uint32_t within_head = index - head * head_stride; + const uint32_t token = within_head / head_dim; + const uint32_t lane = within_head - token * head_dim; + const uint32_t end = cumulative_lengths[1]; + const uint32_t start = end - query_len; + cache[(head * cache_len + start + token) * head_dim + lane] = source[index]; +} + +// Write head-major [head, query, lane] data into the token-major physical +// layout consumed by FlashAttention's paged split-KV kernel. Blocks are stored +// in identity order, so the flattened destination is [token, head, lane]. +extern "C" __global__ void append_paged_kv_bf16( + __nv_bfloat16* cache, + const __nv_bfloat16* source, + const uint32_t* cumulative_lengths, + uint32_t query_len, + uint32_t num_heads, + uint32_t head_dim, + uint32_t cache_len) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = num_heads * query_len * head_dim; + if (index >= count) { + return; + } + + const uint32_t head_stride = query_len * head_dim; + const uint32_t head = index / head_stride; + const uint32_t within_head = index - head * head_stride; + const uint32_t token = within_head / head_dim; + const uint32_t lane = within_head - token * head_dim; + const uint32_t end = cumulative_lengths[1]; + const uint32_t start = end - query_len; + const uint32_t destination_token = start + token; + if (destination_token < cache_len) { + cache[(destination_token * num_heads + head) * head_dim + lane] = source[index]; + } +} + +// Match Candle's separate BF16 SiLU and multiply kernels, including the BF16 +// rounding after every SiLU arithmetic operation and before multiplication. +extern "C" __global__ void silu_mul_bf16( + const __nv_bfloat16* gate, + const __nv_bfloat16* up, + __nv_bfloat16* output, + uint32_t count) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + if (index >= count) { + return; + } + const __nv_bfloat16 one = __float2bfloat16_rn(1.0f); + const __nv_bfloat16 negative = -gate[index]; + const __nv_bfloat16 exponential = hexp(negative); + const __nv_bfloat16 denominator = one + exponential; + const __nv_bfloat16 activated = gate[index] / denominator; + output[index] = activated * up[index]; +} + +// Projected QKV is laid out [token, q+k+v]. Produce the contiguous +// [head, token, lane] Q or K tensor while applying XDRoPE with the exact +// upstream rounding points: BF16 -> F32, two RN multiplies, one RN add, then +// F32 -> BF16. This replaces the generic cast/mul/neg/cat/add/cast sequence. +extern "C" __global__ void xdrope_bf16( + __nv_bfloat16* output, + const __nv_bfloat16* qkv, + const float* cos_sin, + uint32_t projection_width, + uint32_t projection_offset, + uint32_t num_heads, + uint32_t query_len, + uint32_t head_dim) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = num_heads * query_len * head_dim; + if (index >= count) { + return; + } + + const uint32_t head_stride = query_len * head_dim; + const uint32_t head = index / head_stride; + const uint32_t within_head = index - head * head_stride; + const uint32_t token = within_head / head_dim; + const uint32_t lane = within_head - token * head_dim; + const uint32_t half_dim = head_dim / 2; + const uint32_t partner_lane = + lane < half_dim ? lane + half_dim : lane - half_dim; + const uint32_t source_base = + token * projection_width + projection_offset + head * head_dim; + const float x = __bfloat162float(qkv[source_base + lane]); + float rotated = __bfloat162float(qkv[source_base + partner_lane]); + if (lane < half_dim) { + rotated = -rotated; + } + const uint32_t rope_index = token * head_dim + lane; + const uint32_t rope_elements = query_len * head_dim; + const float direct = __fmul_rn(x, cos_sin[rope_index]); + const float cross = + __fmul_rn(rotated, cos_sin[rope_elements + rope_index]); + output[index] = __float2bfloat16_rn(__fadd_rn(direct, cross)); +} + +static __device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + value += __shfl_xor_sync(0xffffffff, value, mask, 32); + } + return value; +} + +// Fuse the target model's XDRoPE, per-head RMSNorm, and BF16-to-F16 +// conversion. Candle uses one 32-thread warp for a 128-wide RMSNorm, with +// each lane accumulating four columns; duplicate that reduction order so the +// verification logits remain bit-for-bit unchanged. For K, optionally emit V +// beside it to also remove the generic slice/transpose/copy/cast chain. +extern "C" __global__ void xdrope_rmsnorm_f16( + const __nv_bfloat16* qkv, + const float* cos_sin, + const __nv_bfloat16* weight, + half* output, + uint32_t projection_width, + uint32_t projection_offset, + uint32_t num_heads, + uint32_t query_len, + uint32_t head_dim, + float eps, + uint32_t include_v) { + const uint32_t row = blockIdx.x; + const uint32_t lane_id = threadIdx.x; + const uint32_t head = row / query_len; + const uint32_t token = row - head * query_len; + const uint32_t half_dim = head_dim / 2; + const uint32_t source_base = + token * projection_width + projection_offset + head * head_dim; + const uint32_t rope_elements = query_len * head_dim; + __nv_bfloat16 values[4]; + float sum = 0.0f; + +#pragma unroll + for (uint32_t item = 0; item < 4; ++item) { + const uint32_t col = lane_id + item * 32; + const uint32_t partner_col = + col < half_dim ? col + half_dim : col - half_dim; + const float x = __bfloat162float(qkv[source_base + col]); + float partner = __bfloat162float(qkv[source_base + partner_col]); + if (col < half_dim) { + partner = -partner; + } + const uint32_t rope_index = token * head_dim + col; + const float direct = __fmul_rn(x, cos_sin[rope_index]); + const float cross = + __fmul_rn(partner, cos_sin[rope_elements + rope_index]); + values[item] = __float2bfloat16_rn(__fadd_rn(direct, cross)); + const float value = __bfloat162float(values[item]); + sum += value * value; + } + sum = warp_sum(sum); + const float scale = rsqrtf(sum / static_cast(head_dim) + eps); + const uint32_t output_base = row * head_dim; + +#pragma unroll + for (uint32_t item = 0; item < 4; ++item) { + const uint32_t col = lane_id + item * 32; + const float normalized = scale * __bfloat162float(values[item]) * + __bfloat162float(weight[col]); + const __nv_bfloat16 rounded = __float2bfloat16_rn(normalized); + output[output_base + col] = __float2half_rn(__bfloat162float(rounded)); + if (include_v) { + const uint32_t count = num_heads * query_len * head_dim; + const uint32_t v_offset = projection_offset + num_heads * head_dim; + output[count + output_base + col] = __float2half_rn(__bfloat162float( + qkv[token * projection_width + v_offset + head * head_dim + col])); + } + } +} + +// Draft attention normalizes Q/K before RoPE. Fuse that ordering while +// matching Candle's 32-thread/128-column RMSNorm and the three BF16 rounding +// points in the existing RoPE path. K can emit V in the same transposed +// [head, token, lane] layout. +extern "C" __global__ void rmsnorm_rope_bf16( + const __nv_bfloat16* qkv, + const __nv_bfloat16* cos_sin, + const __nv_bfloat16* weight, + __nv_bfloat16* output, + uint32_t projection_width, + uint32_t projection_offset, + uint32_t num_heads, + uint32_t query_len, + uint32_t head_dim, + float eps, + uint32_t include_v) { + const uint32_t row = blockIdx.x; + const uint32_t lane_id = threadIdx.x; + const uint32_t head = row / query_len; + const uint32_t token = row - head * query_len; + const uint32_t source_base = + token * projection_width + projection_offset + head * head_dim; + __nv_bfloat16 normalized[4]; + float sum = 0.0f; + +#pragma unroll + for (uint32_t item = 0; item < 4; ++item) { + const uint32_t col = lane_id + item * 32; + const float value = __bfloat162float(qkv[source_base + col]); + sum += value * value; + } + sum = warp_sum(sum); + const float scale = rsqrtf(sum / static_cast(head_dim) + eps); + +#pragma unroll + for (uint32_t item = 0; item < 4; ++item) { + const uint32_t col = lane_id + item * 32; + normalized[item] = __float2bfloat16_rn( + scale * __bfloat162float(qkv[source_base + col]) * + __bfloat162float(weight[col])); + } + + const uint32_t output_base = row * head_dim; + const uint32_t rope_elements = query_len * head_dim; +#pragma unroll + for (uint32_t item = 0; item < 4; ++item) { + const uint32_t col = lane_id + item * 32; + const uint32_t partner_item = col < head_dim / 2 ? item + 2 : item - 2; + float partner = __bfloat162float(normalized[partner_item]); + if (col < head_dim / 2) { + partner = -partner; + } + const uint32_t rope_index = token * head_dim + col; + const __nv_bfloat16 direct = __float2bfloat16_rn( + __bfloat162float(normalized[item]) * + __bfloat162float(cos_sin[rope_index])); + const __nv_bfloat16 cross = __float2bfloat16_rn( + partner * __bfloat162float(cos_sin[rope_elements + rope_index])); + output[output_base + col] = __float2bfloat16_rn( + __bfloat162float(direct) + __bfloat162float(cross)); + if (include_v) { + const uint32_t count = num_heads * query_len * head_dim; + const uint32_t v_offset = projection_offset + num_heads * head_dim; + output[count + output_base + col] = + qkv[token * projection_width + v_offset + head * head_dim + col]; + } + } +} + +// DFlash applies RoPE directly in BF16. Preserve all three BF16 rounding +// points (both products and their sum) while replacing the generic operation +// chain with one kernel. +extern "C" __global__ void rope_bf16( + __nv_bfloat16* output, + const __nv_bfloat16* input, + const __nv_bfloat16* cos_sin, + uint32_t num_heads, + uint32_t query_len, + uint32_t head_dim) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t count = num_heads * query_len * head_dim; + if (index >= count) { + return; + } + + const uint32_t head_stride = query_len * head_dim; + const uint32_t within_head = index % head_stride; + const uint32_t token = within_head / head_dim; + const uint32_t lane = within_head - token * head_dim; + const uint32_t half_dim = head_dim / 2; + const uint32_t partner = + lane < half_dim ? index + half_dim : index - half_dim; + float rotated = __bfloat162float(input[partner]); + if (lane < half_dim) { + rotated = -rotated; + } + const uint32_t rope_index = token * head_dim + lane; + const uint32_t rope_elements = query_len * head_dim; + const float x = __bfloat162float(input[index]); + const __nv_bfloat16 direct = __float2bfloat16_rn( + __fmul_rn(x, __bfloat162float(cos_sin[rope_index]))); + const __nv_bfloat16 cross = __float2bfloat16_rn(__fmul_rn( + rotated, __bfloat162float(cos_sin[rope_elements + rope_index]))); + output[index] = __float2bfloat16_rn(__fadd_rn( + __bfloat162float(direct), __bfloat162float(cross))); +} + +// Match Candle's BF16 badd followed by its F32-accumulating RMSNorm, but write +// both the rounded residual and normalized value in one pass. Output is +// [2, rows, cols]: residual first, normalized second. +extern "C" __global__ void add_rmsnorm_bf16( + const __nv_bfloat16* input, + const __nv_bfloat16* delta, + const __nv_bfloat16* weight, + __nv_bfloat16* output, + uint32_t ncols, + float eps) { + const uint32_t row = blockIdx.x; + const uint32_t col = threadIdx.x; + const uint32_t row_offset = row * ncols; + const uint32_t element_count = gridDim.x * ncols; + + const __nv_bfloat16 residual = __float2bfloat16_rn(__fadd_rn( + __bfloat162float(input[row_offset + col]), + __bfloat162float(delta[row_offset + col]))); + output[row_offset + col] = residual; + const float value = __bfloat162float(residual); + float sum = value * value; + sum = warp_sum(sum); + + __shared__ float warp_sums[32]; + const uint32_t warp = col / 32; + const uint32_t lane = col % 32; + if (lane == 0) { + warp_sums[warp] = sum; + } + __syncthreads(); + sum = warp_sums[lane]; + sum = warp_sum(sum); + const float scale = rsqrtf(sum / static_cast(ncols) + eps); + output[element_count + row_offset + col] = __float2bfloat16_rn( + scale * value * __bfloat162float(weight[col])); +} diff --git a/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs new file mode 100644 index 0000000..0e6e6b2 --- /dev/null +++ b/oar-ocr-vl/src/hunyuanocr/dynamic_kv.rs @@ -0,0 +1,685 @@ +use candle_core::backend::BackendStorage; +use candle_core::{CpuStorage, CustomOp2, CustomOp3, InplaceOp3, Layout, Result, Shape}; + +const PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/hunyuan_dynamic_kv.ptx")); + +pub(super) struct DynamicKvAppend { + pub query_len: usize, + pub cache_len: usize, +} + +pub(super) struct DynamicPagedKvAppend { + pub query_len: usize, + pub cache_len: usize, +} + +pub(super) struct FusedXdRope { + pub projection_width: usize, + pub projection_offset: usize, + pub num_heads: usize, + pub query_len: usize, + pub head_dim: usize, +} + +pub(super) struct FusedXdRopeRmsNormF16 { + pub projection_width: usize, + pub projection_offset: usize, + pub num_heads: usize, + pub query_len: usize, + pub head_dim: usize, + pub eps: f32, + pub include_v: bool, +} + +pub(super) struct FusedRmsNormRopeBf16 { + pub projection_width: usize, + pub projection_offset: usize, + pub num_heads: usize, + pub query_len: usize, + pub head_dim: usize, + pub eps: f32, + pub include_v: bool, +} + +pub(super) struct FusedRopeBf16; + +pub(super) struct FusedAddRmsNormBf16 { + pub eps: f32, +} + +pub(super) struct FusedSiluMulBf16; + +impl CustomOp2 for FusedSiluMulBf16 { + fn name(&self) -> &'static str { + "hunyuan-fused-silu-mul-bf16" + } + + fn cpu_fwd( + &self, + _gate: &CpuStorage, + _gate_layout: &Layout, + _up: &CpuStorage, + _up_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("fused SiLU*up is CUDA-only") + } + + fn cuda_fwd( + &self, + gate: &candle_core::CudaStorage, + gate_layout: &Layout, + up: &candle_core::CudaStorage, + up_layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !gate_layout.is_contiguous() + || !up_layout.is_contiguous() + || gate_layout.shape() != up_layout.shape() + || gate.dtype() != candle_core::DType::BF16 + || up.dtype() != candle_core::DType::BF16 + { + candle_core::bail!("fused SiLU*up requires matching contiguous BF16 tensors") + } + let device = gate.device().clone(); + let gate = gate.as_cuda_slice::()?; + let up = up.as_cuda_slice::()?; + let gate = gate.slice(gate_layout.start_offset()..); + 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 mut builder = function.builder(); + builder.arg(&gate); + builder.arg(&up); + builder.arg(&output); + candle_core::builder_arg!(builder, count as u32); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + gate_layout.shape().clone(), + )) + } +} + +impl CustomOp3 for FusedXdRopeRmsNormF16 { + fn name(&self) -> &'static str { + "hunyuan-fused-xdrope-rmsnorm-f16" + } + + fn cpu_fwd( + &self, + _qkv: &CpuStorage, + _qkv_layout: &Layout, + _cos_sin: &CpuStorage, + _cos_sin_layout: &Layout, + _weight: &CpuStorage, + _weight_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("fused XDRoPE+RMSNorm is CUDA-only") + } + + fn cuda_fwd( + &self, + qkv: &candle_core::CudaStorage, + qkv_layout: &Layout, + cos_sin: &candle_core::CudaStorage, + cos_sin_layout: &Layout, + weight: &candle_core::CudaStorage, + weight_layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !qkv_layout.is_contiguous() + || !cos_sin_layout.is_contiguous() + || !weight_layout.is_contiguous() + { + candle_core::bail!("fused XDRoPE+RMSNorm requires contiguous tensors") + } + let (batch, qkv_len, projection_width) = qkv_layout.shape().dims3()?; + if batch != 1 + || qkv_len != self.query_len + || projection_width != self.projection_width + || self.head_dim != 128 + || self.projection_offset + self.num_heads * self.head_dim > projection_width + || cos_sin_layout.shape().elem_count() != 2 * self.query_len * self.head_dim + || weight_layout.shape().dims1()? != self.head_dim + || (self.include_v + && self.projection_offset + 2 * self.num_heads * self.head_dim > projection_width) + { + candle_core::bail!( + "fused XDRoPE+RMSNorm shape mismatch qkv={:?} cos_sin={:?} weight={:?}", + qkv_layout.shape(), + cos_sin_layout.shape(), + weight_layout.shape() + ) + } + let device = qkv.device().clone(); + let qkv = qkv.as_cuda_slice::()?; + let cos_sin = cos_sin.as_cuda_slice::()?; + let weight = weight.as_cuda_slice::()?; + let qkv = qkv.slice(qkv_layout.start_offset()..); + let cos_sin = cos_sin.slice(cos_sin_layout.start_offset()..); + let weight = weight.slice(weight_layout.start_offset()..); + let count = self.num_heads * self.query_len * self.head_dim; + 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)?; + let mut builder = function.builder(); + builder.arg(&qkv); + builder.arg(&cos_sin); + builder.arg(&weight); + builder.arg(&output); + candle_core::builder_arg!( + builder, + self.projection_width as u32, + self.projection_offset as u32, + self.num_heads as u32, + self.query_len as u32, + self.head_dim as u32, + self.eps, + u32::from(self.include_v) + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: ((self.num_heads * self.query_len) as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + let shape = if self.include_v { + Shape::from_dims(&[2, 1, self.num_heads, self.query_len, self.head_dim]) + } else { + Shape::from_dims(&[1, self.num_heads, self.query_len, self.head_dim]) + }; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + shape, + )) + } +} + +impl CustomOp3 for FusedRmsNormRopeBf16 { + fn name(&self) -> &'static str { + "hunyuan-fused-rmsnorm-rope-bf16" + } + + fn cpu_fwd( + &self, + _qkv: &CpuStorage, + _qkv_layout: &Layout, + _cos_sin: &CpuStorage, + _cos_sin_layout: &Layout, + _weight: &CpuStorage, + _weight_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("fused RMSNorm+RoPE is CUDA-only") + } + + fn cuda_fwd( + &self, + qkv: &candle_core::CudaStorage, + qkv_layout: &Layout, + cos_sin: &candle_core::CudaStorage, + cos_sin_layout: &Layout, + weight: &candle_core::CudaStorage, + weight_layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !qkv_layout.is_contiguous() + || !cos_sin_layout.is_contiguous() + || !weight_layout.is_contiguous() + { + candle_core::bail!("fused RMSNorm+RoPE requires contiguous tensors") + } + let (batch, qkv_len, projection_width) = qkv_layout.shape().dims3()?; + if batch != 1 + || qkv_len != self.query_len + || projection_width != self.projection_width + || self.head_dim != 128 + || self.projection_offset + self.num_heads * self.head_dim > projection_width + || cos_sin_layout.shape().elem_count() != 2 * self.query_len * self.head_dim + || weight_layout.shape().dims1()? != self.head_dim + || (self.include_v + && self.projection_offset + 2 * self.num_heads * self.head_dim > projection_width) + { + candle_core::bail!( + "fused RMSNorm+RoPE shape mismatch qkv={:?} cos_sin={:?} weight={:?}", + qkv_layout.shape(), + cos_sin_layout.shape(), + weight_layout.shape() + ) + } + let device = qkv.device().clone(); + let qkv = qkv.as_cuda_slice::()?; + let cos_sin = cos_sin.as_cuda_slice::()?; + let weight = weight.as_cuda_slice::()?; + let qkv = qkv.slice(qkv_layout.start_offset()..); + let cos_sin = cos_sin.slice(cos_sin_layout.start_offset()..); + let weight = weight.slice(weight_layout.start_offset()..); + let count = self.num_heads * self.query_len * self.head_dim; + 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)?; + let mut builder = function.builder(); + builder.arg(&qkv); + builder.arg(&cos_sin); + builder.arg(&weight); + builder.arg(&output); + candle_core::builder_arg!( + builder, + self.projection_width as u32, + self.projection_offset as u32, + self.num_heads as u32, + self.query_len as u32, + self.head_dim as u32, + self.eps, + u32::from(self.include_v) + ); + unsafe { + builder.launch(LaunchConfig { + grid_dim: ((self.num_heads * self.query_len) as u32, 1, 1), + block_dim: (32, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + let shape = if self.include_v { + Shape::from_dims(&[2, 1, self.num_heads, self.query_len, self.head_dim]) + } else { + Shape::from_dims(&[1, self.num_heads, self.query_len, self.head_dim]) + }; + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + shape, + )) + } +} + +impl CustomOp3 for FusedAddRmsNormBf16 { + fn name(&self) -> &'static str { + "hunyuan-fused-add-rmsnorm-bf16" + } + + fn cpu_fwd( + &self, + _input: &CpuStorage, + _input_layout: &Layout, + _delta: &CpuStorage, + _delta_layout: &Layout, + _weight: &CpuStorage, + _weight_layout: &Layout, + ) -> Result<(CpuStorage, Shape)> { + candle_core::bail!("fused add+RMSNorm is CUDA-only") + } + + fn cuda_fwd( + &self, + input: &candle_core::CudaStorage, + input_layout: &Layout, + delta: &candle_core::CudaStorage, + delta_layout: &Layout, + weight: &candle_core::CudaStorage, + weight_layout: &Layout, + ) -> Result<(candle_core::CudaStorage, Shape)> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !input_layout.is_contiguous() + || !delta_layout.is_contiguous() + || !weight_layout.is_contiguous() + || input_layout.shape() != delta_layout.shape() + { + candle_core::bail!("fused add+RMSNorm requires matching contiguous tensors") + } + let dims = input_layout.shape().dims(); + let ncols = *dims + .last() + .ok_or_else(|| candle_core::Error::Msg("RMSNorm input is scalar".into()))?; + if ncols != 1024 || weight_layout.shape().dims1()? != ncols { + candle_core::bail!( + "fused add+RMSNorm requires hidden width 1024, got input={:?} weight={:?}", + input_layout.shape(), + weight_layout.shape() + ) + } + let element_count = input_layout.shape().elem_count(); + let rows = element_count / ncols; + let device = input.device().clone(); + let input = input.as_cuda_slice::()?; + let delta = delta.as_cuda_slice::()?; + let weight = weight.as_cuda_slice::()?; + let input = input.slice(input_layout.start_offset()..); + let delta = delta.slice(delta_layout.start_offset()..); + 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)?; + let mut builder = function.builder(); + builder.arg(&input); + builder.arg(&delta); + builder.arg(&weight); + builder.arg(&output); + candle_core::builder_arg!(builder, ncols as u32, self.eps); + unsafe { + builder.launch(LaunchConfig { + grid_dim: (rows as u32, 1, 1), + block_dim: (ncols as u32, 1, 1), + shared_mem_bytes: 0, + }) + } + .w()?; + let mut output_shape = Vec::with_capacity(dims.len() + 1); + output_shape.push(2); + output_shape.extend_from_slice(dims); + Ok(( + candle_core::CudaStorage::wrap_cuda_slice(output, device), + Shape::from_dims(&output_shape), + )) + } +} + +impl InplaceOp3 for FusedRopeBf16 { + fn name(&self) -> &'static str { + "hunyuan-fused-rope-bf16" + } + + fn cpu_fwd( + &self, + _output: &mut CpuStorage, + _output_layout: &Layout, + _input: &CpuStorage, + _input_layout: &Layout, + _cos_sin: &CpuStorage, + _cos_sin_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("fused BF16 RoPE is CUDA-only") + } + + fn cuda_fwd( + &self, + output: &mut candle_core::CudaStorage, + output_layout: &Layout, + input: &candle_core::CudaStorage, + input_layout: &Layout, + cos_sin: &candle_core::CudaStorage, + cos_sin_layout: &Layout, + ) -> Result<()> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !output_layout.is_contiguous() + || !input_layout.is_contiguous() + || !cos_sin_layout.is_contiguous() + || output_layout.shape() != input_layout.shape() + { + candle_core::bail!("fused BF16 RoPE requires matching contiguous tensors") + } + let (_, heads, query_len, head_dim) = output_layout.shape().dims4()?; + if !head_dim.is_multiple_of(2) + || cos_sin_layout.shape().elem_count() != 2 * query_len * head_dim + { + candle_core::bail!( + "fused BF16 RoPE shape mismatch input={:?} cos_sin={:?}", + input_layout.shape(), + cos_sin_layout.shape() + ) + } + let device = output.device().clone(); + let output = output.as_cuda_slice_mut::()?; + let input = input.as_cuda_slice::()?; + let cos_sin = cos_sin.as_cuda_slice::()?; + let mut output = output.slice_mut(output_layout.start_offset()..); + 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 mut builder = function.builder(); + builder.arg(&mut output); + builder.arg(&input); + builder.arg(&cos_sin); + candle_core::builder_arg!(builder, heads as u32, query_len as u32, head_dim as u32); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + Ok(()) + } +} + +impl InplaceOp3 for FusedXdRope { + fn name(&self) -> &'static str { + "hunyuan-fused-xdrope" + } + + fn cpu_fwd( + &self, + _output: &mut CpuStorage, + _output_layout: &Layout, + _qkv: &CpuStorage, + _qkv_layout: &Layout, + _cos_sin: &CpuStorage, + _cos_sin_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("fused XDRoPE is CUDA-only") + } + + fn cuda_fwd( + &self, + output: &mut candle_core::CudaStorage, + output_layout: &Layout, + qkv: &candle_core::CudaStorage, + qkv_layout: &Layout, + cos_sin: &candle_core::CudaStorage, + cos_sin_layout: &Layout, + ) -> Result<()> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + if !output_layout.is_contiguous() + || !qkv_layout.is_contiguous() + || !cos_sin_layout.is_contiguous() + { + candle_core::bail!("fused XDRoPE requires contiguous tensors") + } + let (_, heads, query_len, head_dim) = output_layout.shape().dims4()?; + let (_, qkv_len, projection_width) = qkv_layout.shape().dims3()?; + if heads != self.num_heads + || query_len != self.query_len + || qkv_len != query_len + || head_dim != self.head_dim + || projection_width != self.projection_width + || self.projection_offset + heads * head_dim > projection_width + || cos_sin_layout.shape().elem_count() != 2 * query_len * head_dim + { + candle_core::bail!( + "fused XDRoPE shape mismatch output={:?} qkv={:?} cos_sin={:?}", + output_layout.shape(), + qkv_layout.shape(), + cos_sin_layout.shape() + ) + } + let device = output.device().clone(); + let output = output.as_cuda_slice_mut::()?; + let qkv = qkv.as_cuda_slice::()?; + let cos_sin = cos_sin.as_cuda_slice::()?; + let mut output = output.slice_mut(output_layout.start_offset()..); + 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 mut builder = function.builder(); + builder.arg(&mut output); + builder.arg(&qkv); + builder.arg(&cos_sin); + candle_core::builder_arg!( + builder, + projection_width as u32, + self.projection_offset as u32, + heads as u32, + query_len as u32, + head_dim as u32 + ); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + Ok(()) + } +} + +impl InplaceOp3 for DynamicKvAppend { + fn name(&self) -> &'static str { + "hunyuan-dynamic-kv-append" + } + + fn cpu_fwd( + &self, + _cache: &mut CpuStorage, + _cache_layout: &Layout, + _source: &CpuStorage, + _source_layout: &Layout, + _lengths: &CpuStorage, + _lengths_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("dynamic KV append is CUDA-only") + } + + fn cuda_fwd( + &self, + cache: &mut candle_core::CudaStorage, + cache_layout: &Layout, + source: &candle_core::CudaStorage, + source_layout: &Layout, + lengths: &candle_core::CudaStorage, + lengths_layout: &Layout, + ) -> Result<()> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + let (_, num_heads, cache_len, head_dim) = cache_layout.shape().dims4()?; + let (_, source_heads, query_len, source_head_dim) = source_layout.shape().dims4()?; + if num_heads != source_heads + || head_dim != source_head_dim + || query_len != self.query_len + || cache_len != self.cache_len + { + candle_core::bail!( + "dynamic KV shape mismatch cache={:?} source={:?}", + cache_layout.shape(), + source_layout.shape() + ) + } + if lengths_layout.shape().dims1()? != 2 { + candle_core::bail!("dynamic KV cumulative lengths must have shape [2]") + } + let device = cache.device().clone(); + let lengths = lengths.as_cuda_slice::()?; + let lengths = lengths.slice(lengths_layout.start_offset()..); + let count = num_heads * query_len * head_dim; + macro_rules! launch { + ($ty:ty, $function:literal) => {{ + let cache = cache.as_cuda_slice_mut::<$ty>()?; + let source = source.as_cuda_slice::<$ty>()?; + 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)?; + let mut builder = function.builder(); + builder.arg(&mut cache); + builder.arg(&source); + builder.arg(&lengths); + candle_core::builder_arg!( + builder, + query_len as u32, + num_heads as u32, + head_dim as u32, + cache_len as u32 + ); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + }}; + } + match cache.dtype() { + candle_core::DType::F16 => launch!(half::f16, "append_kv_f16"), + candle_core::DType::BF16 => launch!(half::bf16, "append_kv_bf16"), + dtype => candle_core::bail!("dynamic KV append does not support {dtype:?}"), + } + Ok(()) + } +} + +impl InplaceOp3 for DynamicPagedKvAppend { + fn name(&self) -> &'static str { + "hunyuan-dynamic-paged-kv-append" + } + + fn cpu_fwd( + &self, + _cache: &mut CpuStorage, + _cache_layout: &Layout, + _source: &CpuStorage, + _source_layout: &Layout, + _lengths: &CpuStorage, + _lengths_layout: &Layout, + ) -> Result<()> { + candle_core::bail!("dynamic paged KV append is CUDA-only") + } + + fn cuda_fwd( + &self, + cache: &mut candle_core::CudaStorage, + cache_layout: &Layout, + source: &candle_core::CudaStorage, + source_layout: &Layout, + lengths: &candle_core::CudaStorage, + lengths_layout: &Layout, + ) -> Result<()> { + use candle_core::cuda_backend::WrapErr; + use candle_core::cuda_backend::cudarc::driver::{LaunchConfig, PushKernelArg}; + + let (blocks, page_size, num_heads, head_dim) = cache_layout.shape().dims4()?; + let (_, source_heads, query_len, source_head_dim) = source_layout.shape().dims4()?; + if num_heads != source_heads + || head_dim != source_head_dim + || query_len != self.query_len + || blocks * page_size != self.cache_len + { + candle_core::bail!( + "dynamic paged KV shape mismatch cache={:?} source={:?}", + cache_layout.shape(), + source_layout.shape() + ) + } + if lengths_layout.shape().dims1()? != 2 { + candle_core::bail!("dynamic paged KV cumulative lengths must have shape [2]") + } + if cache.dtype() != candle_core::DType::BF16 || source.dtype() != candle_core::DType::BF16 { + candle_core::bail!("dynamic paged KV append requires BF16") + } + let device = cache.device().clone(); + let cache = cache.as_cuda_slice_mut::()?; + let source = source.as_cuda_slice::()?; + let lengths = lengths.as_cuda_slice::()?; + let mut cache = cache.slice_mut(cache_layout.start_offset()..); + let source = source.slice(source_layout.start_offset()..); + 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)?; + let mut builder = function.builder(); + builder.arg(&mut cache); + builder.arg(&source); + builder.arg(&lengths); + candle_core::builder_arg!( + builder, + query_len as u32, + num_heads as u32, + head_dim as u32, + self.cache_len as u32 + ); + unsafe { builder.launch(LaunchConfig::for_num_elems(count as u32)) }.w()?; + Ok(()) + } +} diff --git a/oar-ocr-vl/src/hunyuanocr/llm.rs b/oar-ocr-vl/src/hunyuanocr/llm.rs index f721a20..401a0f6 100644 --- a/oar-ocr-vl/src/hunyuanocr/llm.rs +++ b/oar-ocr-vl/src/hunyuanocr/llm.rs @@ -1,14 +1,27 @@ use super::config::HunyuanOcrConfig; +#[cfg(feature = "cuda")] +use super::dynamic_kv::{ + DynamicKvAppend, FusedAddRmsNormBf16, FusedSiluMulBf16, FusedXdRope, FusedXdRopeRmsNormF16, +}; 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, }; use crate::kv_trim::TrimmableKvCache; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing, rotate_half}; +#[cfg(feature = "cuda")] +use candle_core::Device; use candle_core::Tensor; use candle_nn::Module; use oar_ocr_core::core::OCRError; use std::cell::RefCell; +const DECODE_ROPE_CACHE_LEN: usize = 16_384; + +#[cfg(any(feature = "cuda", test))] +fn target_cuda_graph_dtype_supported(dtype: candle_core::DType) -> bool { + matches!(dtype, candle_core::DType::F16 | candle_core::DType::BF16) +} + /// Apply XDRoPE to `(q, k)` using already-section-mixed F32 `cos`/`sin`. /// /// The section-mix (`select_rope_sections`) and the F32 cast of cos/sin are @@ -87,7 +100,19 @@ fn apply_xdrope_rotary_pos_emb( Ok((q_rot, k_rot)) } -#[derive(Debug, Clone)] +#[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_proj: candle_nn::Linear, up_proj: candle_nn::Linear, @@ -117,14 +142,29 @@ impl HunyuanMlp { .gate_proj .forward(xs) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate_proj", e))?; - let gate = candle_nn::ops::silu(&gate) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp silu", e))?; let up = self .up_proj .forward(xs) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp up_proj", e))?; - let prod = - (&gate * &up).map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate*up", e))?; + #[cfg(feature = "cuda")] + let prod = if gate.device().is_cuda() + && gate.dtype() == candle_core::DType::BF16 + && gate.is_contiguous() + && up.is_contiguous() + { + gate.apply_op2_no_bwd(&up, &FusedSiluMulBf16) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused mlp silu*up", e))? + } else { + let gate = candle_nn::ops::silu(&gate) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp silu", e))?; + (&gate * &up).map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate*up", e))? + }; + #[cfg(not(feature = "cuda"))] + let prod = { + let gate = candle_nn::ops::silu(&gate) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp silu", e))?; + (&gate * &up).map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp gate*up", e))? + }; self.down_proj .forward(&prod) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "mlp down_proj", e)) @@ -133,9 +173,7 @@ impl HunyuanMlp { #[derive(Debug)] struct HunyuanAttention { - q_proj: candle_nn::Linear, - k_proj: candle_nn::Linear, - v_proj: candle_nn::Linear, + qkv_proj: candle_nn::Linear, o_proj: candle_nn::Linear, query_layernorm: Option, key_layernorm: Option, @@ -179,6 +217,10 @@ impl HunyuanAttention { vb.pp("v_proj"), ) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "load v_proj", e))?; + let qkv_weight = Tensor::cat(&[q_proj.weight(), k_proj.weight(), v_proj.weight()], 0) + .and_then(|x| x.contiguous()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fuse qkv weights", e))?; + let qkv_proj = candle_nn::Linear::new(qkv_weight, None); let o_proj = candle_nn::linear_no_bias( cfg.num_attention_heads * cfg.head_dim, cfg.hidden_size, @@ -203,12 +245,10 @@ impl HunyuanAttention { // plus the longest realistic generation. Same growth strategy as // candle_nn::kv_cache::KvCache (Tensor::cat per append). // Trim/gather support allows selective KV retention after generation steps. - let kv_cache = TrimmableKvCache::new(2, 16384); + let kv_cache = TrimmableKvCache::new(2, DECODE_ROPE_CACHE_LEN); Ok(Self { - q_proj, - k_proj, - v_proj, + qkv_proj, o_proj, query_layernorm, key_layernorm, @@ -221,46 +261,173 @@ impl HunyuanAttention { }) } - fn forward( + fn project_qkv( &self, hidden_states: &Tensor, cos: &Tensor, sin: &Tensor, - causal_mask: Option<&Tensor>, - ) -> Result { + cos_sin: Option<&Tensor>, + ) -> Result<(Tensor, Tensor, Tensor), OCRError> { + #[cfg(not(feature = "cuda"))] + let _ = cos_sin; let (b, seq_len, _) = hidden_states .dims3() .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn hidden_states dims3", e))?; - let q = self - .q_proj + let qkv = self + .qkv_proj .forward(hidden_states) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q_proj", e))? - .reshape((b, seq_len, self.num_heads, self.head_dim)) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q reshape", e))? - .transpose(1, 2) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q transpose", e))?; + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn qkv_proj", e))?; + let q_width = self.num_heads * self.head_dim; + let kv_width = self.num_kv_heads * self.head_dim; + #[cfg(feature = "cuda")] + let fused_qkv = + if b == 1 + && self.head_dim == 128 + && qkv.device().is_cuda() + && qkv.dtype() == candle_core::DType::BF16 + && cos.dtype() == candle_core::DType::F32 + && sin.dtype() == candle_core::DType::F32 + && cos_sin.is_some() + { + match (&self.query_layernorm, &self.key_layernorm) { + (Some(query_norm), Some(key_norm)) => { + let cos_sin = cos_sin.expect("checked above"); + let q = qkv + .apply_op3_no_bwd( + cos_sin, + query_norm.weight(), + &FusedXdRopeRmsNormF16 { + projection_width: q_width + 2 * kv_width, + projection_offset: 0, + num_heads: self.num_heads, + query_len: seq_len, + head_dim: self.head_dim, + eps: query_norm.eps() as f32, + include_v: false, + }, + ) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "fused Q XDRoPE RMSNorm", e) + })?; + let kv = qkv + .apply_op3_no_bwd( + cos_sin, + key_norm.weight(), + &FusedXdRopeRmsNormF16 { + projection_width: q_width + 2 * kv_width, + projection_offset: q_width, + num_heads: self.num_kv_heads, + query_len: seq_len, + head_dim: self.head_dim, + eps: key_norm.eps() as f32, + include_v: true, + }, + ) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "fused KV XDRoPE RMSNorm", e) + })?; + let k = kv.narrow(0, 0, 1).and_then(|x| x.squeeze(0)).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "fused K view", e) + })?; + let v = kv.narrow(0, 1, 1).and_then(|x| x.squeeze(0)).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "fused V view", e) + })?; + Some((q, k, v)) + } + _ => None, + } + } else { + None + }; + #[cfg(not(feature = "cuda"))] + let fused_qkv: Option<(Tensor, Tensor, Tensor)> = None; + if let Some(qkv) = fused_qkv { + return Ok(qkv); + } + #[cfg(feature = "cuda")] + let fused_qk = if b == 1 + && self.head_dim.is_multiple_of(2) + && qkv.device().is_cuda() + && qkv.dtype() == candle_core::DType::BF16 + && cos.dtype() == candle_core::DType::F32 + && sin.dtype() == candle_core::DType::F32 + && cos_sin.is_some() + { + let projection_width = q_width + 2 * kv_width; + let cos_sin = cos_sin.expect("checked above"); + let q = Tensor::zeros( + (b, self.num_heads, seq_len, self.head_dim), + qkv.dtype(), + qkv.device(), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "allocate fused XDRoPE Q", e))?; + q.inplace_op3( + &qkv, + cos_sin, + &FusedXdRope { + projection_width, + projection_offset: 0, + num_heads: self.num_heads, + query_len: seq_len, + head_dim: self.head_dim, + }, + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused XDRoPE Q", e))?; + let k = Tensor::zeros( + (b, self.num_kv_heads, seq_len, self.head_dim), + qkv.dtype(), + qkv.device(), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "allocate fused XDRoPE K", e))?; + k.inplace_op3( + &qkv, + cos_sin, + &FusedXdRope { + projection_width, + projection_offset: q_width, + num_heads: self.num_kv_heads, + query_len: seq_len, + head_dim: self.head_dim, + }, + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused XDRoPE K", e))?; + Some((q, k)) + } else { + None + }; + #[cfg(not(feature = "cuda"))] + let fused_qk: Option<(Tensor, Tensor)> = None; - let k = self - .k_proj - .forward(hidden_states) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k_proj", e))? - .reshape((b, seq_len, self.num_kv_heads, self.head_dim)) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k reshape", e))? - .transpose(1, 2) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k transpose", e))?; + let (q, k) = match fused_qk { + Some(qk) => qk, + None => { + let q = qkv + .narrow(2, 0, q_width) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q slice", e))? + .reshape((b, seq_len, self.num_heads, self.head_dim)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q reshape", e))? + .transpose(1, 2) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q transpose", e))?; + let k = qkv + .narrow(2, q_width, kv_width) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k slice", e))? + .reshape((b, seq_len, self.num_kv_heads, self.head_dim)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k reshape", e))? + .transpose(1, 2) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k transpose", e))?; + apply_xdrope_rotary_pos_emb(&q, &k, cos, sin)? + } + }; - let v = self - .v_proj - .forward(hidden_states) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn v_proj", e))? + let v = qkv + .narrow(2, q_width + kv_width, kv_width) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn v slice", e))? .reshape((b, seq_len, self.num_kv_heads, self.head_dim)) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn v reshape", e))? .transpose(1, 2) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn v transpose", e))?; - let (q, k) = apply_xdrope_rotary_pos_emb(&q, &k, cos, sin)?; - // Match upstream HunyuanVL: apply XDRoPE first, then Q/K RMSNorm. // The learned RMSNorm weight is per head dimension, so it does not // commute with the rotary half-dimension mixing. @@ -286,42 +453,98 @@ impl HunyuanAttention { let v = v .contiguous() .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn v contiguous", e))?; + // Keep the model in BF16 but use FP16 inside CUDA attention. The + // FlashAttention BF16 verification path can amplify tiny per-round + // differences into a different greedy branch on long documents. Only + // BF16 gets downcast this way: an explicit `OAR_VL_DTYPE=f32` + // override is meant to raise precision, not get silently discarded. + let attention_dtype = if q.device().is_cuda() && q.dtype() == candle_core::DType::BF16 { + candle_core::DType::F16 + } else { + q.dtype() + }; + let q = q + .to_dtype(attention_dtype) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn q attention dtype", e))?; + let k = k + .to_dtype(attention_dtype) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn k attention dtype", e))?; + let v = v + .to_dtype(attention_dtype) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn v attention dtype", e))?; + + Ok((q, k, v)) + } + + fn attend_projected( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + model_dtype: candle_core::DType, + causal_mask: Option<&Tensor>, + ) -> Result { + let (b, _, seq_len, _) = q + .dims4() + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "projected Q dims", e))?; let (k_all, v_all) = self .kv_cache .borrow_mut() - .append(&k, &v) + .append(k, v) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "append kv_cache", e))?; - let key_states = repeat_kv(&k_all, self.num_kv_groups) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "repeat_kv key", e))?; - let value_states = repeat_kv(&v_all, self.num_kv_groups) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "repeat_kv value", e))?; - - let key_states = key_states - .contiguous() - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn key_states contiguous", e))?; - let value_states = value_states.contiguous().map_err(|e| { - candle_to_ocr_inference("HunyuanOCR", "attn value_states contiguous", e) - })?; - - // Use unified attention implementation (BF16 Q·K, F32 softmax, BF16 A·V). - // The main Hunyuan-specific numerical requirement is above: upstream - // applies XDRoPE in F32 before Q/K RMSNorm. - 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("HunyuanOCR", "scaled_dot_product_attention", e))?; + // Single-document inference has no padding, so FlashAttention can + // replace the explicit causal mask and consume GQA K/V directly. + // Batched/padded and non-CUDA paths retain the portable eager kernel. + let flash_output = if b == 1 { + flash_attention(q, &k_all, &v_all, self.scaling, seq_len > 1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "flash attention", e))? + } else { + None + }; + let attn_output = match flash_output { + Some(output) => output, + None => { + let key_states = repeat_kv(&k_all, self.num_kv_groups) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "repeat_kv key", e))? + .contiguous() + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn key contiguous", e))?; + let value_states = repeat_kv(&v_all, self.num_kv_groups) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "repeat_kv value", e))? + .contiguous() + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "attn value contiguous", e) + })?; + scaled_dot_product_attention( + q, + &key_states, + &value_states, + causal_mask, + self.scaling, + seq_len > 1, + ) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "scaled dot-product attention", e) + })? + } + }; + self.project_attention_output(&attn_output, model_dtype, b, seq_len) + } + fn project_attention_output( + &self, + attn_output: &Tensor, + model_dtype: candle_core::DType, + batch: usize, + seq_len: usize, + ) -> Result { let attn_output = attn_output + .to_dtype(model_dtype) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn output model dtype", e))? .transpose(1, 2) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "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("HunyuanOCR", "attn out reshape", e))?; self.o_proj @@ -329,9 +552,146 @@ impl HunyuanAttention { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "attn o_proj", e)) } + #[cfg(feature = "cuda")] + fn prepare_dynamic_cache(&self, query_len: usize) -> Result<(), OCRError> { + let device = self.qkv_proj.weight().device(); + let template = Tensor::zeros( + (1, self.num_kv_heads, query_len, self.head_dim), + candle_core::DType::F16, + device, + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic KV template", e))?; + self.kv_cache + .borrow_mut() + .initialize_storage(&template) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "initialize dynamic KV", e)) + } + + #[cfg(feature = "cuda")] + fn attend_projected_dynamic( + &self, + q: &Tensor, + k: &Tensor, + v: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + model_dtype: candle_core::DType, + ) -> Result { + let (batch, _, query_len, _) = q + .dims4() + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic Q dims", e))?; + if batch != 1 { + return Err(OCRError::ConfigError { + message: "HunyuanOCR dynamic CUDA-graph attention requires batch size 1" + .to_string(), + }); + } + let (cache_k, cache_v) = + self.kv_cache + .borrow() + .storage() + .ok_or_else(|| OCRError::ConfigError { + message: "HunyuanOCR dynamic KV storage is not initialized".to_string(), + })?; + let append = DynamicKvAppend { + query_len, + cache_len: DECODE_ROPE_CACHE_LEN, + }; + cache_k + .inplace_op3(k, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic key cache append", e))?; + cache_v + .inplace_op3(v, kv_lengths, &append) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic value cache append", e))?; + + let q = q + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic Q layout", e))?; + let cache_k = cache_k + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic K layout", e))?; + let cache_v = cache_v + .squeeze(0) + .and_then(|x| x.transpose(0, 1)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic V layout", e))?; + let output = candle_flash_attn::flash_attn_varlen( + &q, + &cache_k, + &cache_v, + query_lengths, + kv_lengths, + query_len, + DECODE_ROPE_CACHE_LEN, + self.scaling as f32, + true, + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic flash attention", e))? + .transpose(0, 1) + .and_then(|x| x.unsqueeze(0)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic attention layout", e))?; + self.project_attention_output(&output, model_dtype, batch, query_len) + } + fn clear_kv_cache(&self) { self.kv_cache.borrow_mut().reset(); } + + fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { + self.kv_cache + .borrow_mut() + .trim_to(len) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "trim kv_cache", e)) + } + + 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("HunyuanOCR", "set dynamic KV length", e)) + } +} + +#[cfg(feature = "cuda")] +struct TargetDecoderCudaGraph { + // The executable graph owns raw pointers into every tensor below and into + // the decoder weights, so it must be dropped first. + graph: candle_core::cuda_backend::cudarc::driver::CudaGraph, + hidden_input: Tensor, + cos_input: Tensor, + sin_input: Tensor, + _query_lengths: Tensor, + kv_lengths: Tensor, + hidden_output: Tensor, + aux_output: Tensor, + aux_layer_ids: Vec, +} + +#[cfg(feature = "cuda")] +impl std::fmt::Debug for TargetDecoderCudaGraph { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TargetDecoderCudaGraph") + .field("hidden", &self.hidden_input.shape()) + .field("aux_layer_ids", &self.aux_layer_ids) + .finish_non_exhaustive() + } +} + +#[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)] @@ -368,27 +728,88 @@ impl HunyuanDecoderLayer { hidden_states: &Tensor, cos: &Tensor, sin: &Tensor, + cos_sin: Option<&Tensor>, causal_mask: Option<&Tensor>, ) -> Result { - let residual = hidden_states; - let hidden_states = self + let normalized = self .input_layernorm .forward(hidden_states) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm input_layernorm", e))?; - let attn_out = self - .self_attn - .forward(&hidden_states, cos, sin, causal_mask)?; - let hidden_states = (residual + &attn_out) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm attn residual add", e))?; + let (q, k, v) = self.self_attn.project_qkv(&normalized, cos, sin, cos_sin)?; + let attention = + self.self_attn + .attend_projected(&q, &k, &v, hidden_states.dtype(), causal_mask)?; + self.post_attention_eager(hidden_states, &attention) + } + + #[cfg(feature = "cuda")] + fn forward_dynamic( + &self, + hidden_states: &Tensor, + cos: &Tensor, + sin: &Tensor, + cos_sin: Option<&Tensor>, + query_lengths: &Tensor, + kv_lengths: &Tensor, + ) -> Result { + let normalized = self + .input_layernorm + .forward(hidden_states) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm input_layernorm", e))?; + let (q, k, v) = self.self_attn.project_qkv(&normalized, cos, sin, cos_sin)?; + let attention = self.self_attn.attend_projected_dynamic( + &q, + &k, + &v, + query_lengths, + kv_lengths, + hidden_states.dtype(), + )?; + self.post_attention_eager(hidden_states, &attention) + } + fn post_attention_eager( + &self, + hidden_states: &Tensor, + attention: &Tensor, + ) -> Result { + #[cfg(feature = "cuda")] + if hidden_states.device().is_cuda() + && hidden_states.dtype() == candle_core::DType::BF16 + && hidden_states.is_contiguous() + && attention.is_contiguous() + { + let packed = hidden_states + .apply_op3_no_bwd( + attention, + self.post_attention_layernorm.weight(), + &FusedAddRmsNormBf16 { + eps: self.post_attention_layernorm.eps() as f32, + }, + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused residual RMSNorm", e))?; + let residual = packed + .narrow(0, 0, 1) + .and_then(|x| x.squeeze(0)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused residual view", e))?; + let normalized = packed + .narrow(0, 1, 1) + .and_then(|x| x.squeeze(0)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "fused RMSNorm view", e))?; + let mlp_out = self.mlp.forward(&normalized)?; + return (&residual + &mlp_out) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm mlp residual add", e)); + } + let hidden_states = (hidden_states + attention) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm attn residual add", e))?; let residual = hidden_states.clone(); - let hidden_states = self + let normalized = self .post_attention_layernorm .forward(&hidden_states) .map_err(|e| { candle_to_ocr_inference("HunyuanOCR", "llm post_attention_layernorm", e) })?; - let mlp_out = self.mlp.forward(&hidden_states)?; + let mlp_out = self.mlp.forward(&normalized)?; (&residual + &mlp_out) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm mlp residual add", e)) } @@ -396,14 +817,40 @@ impl HunyuanDecoderLayer { fn clear_kv_cache(&self) { self.self_attn.clear_kv_cache(); } + + fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { + self.self_attn.trim_kv_cache(len) + } + + 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) + } +} + +/// Target-model outputs needed by DFlash. `hidden_states` is the regular +/// post-norm decoder output. `aux_hidden_states` concatenates the requested +/// intermediate layer outputs along the hidden dimension, matching vLLM's +/// DFlash target interface. +pub(crate) struct HunyuanLlmOutput { + pub hidden_states: Tensor, + pub aux_hidden_states: Option, } #[derive(Debug)] pub struct HunyuanLlm { + #[cfg(feature = "cuda")] + dflash_decode_graph: RefCell>, embed_tokens: candle_nn::Embedding, layers: Vec, norm: candle_nn::RmsNorm, rotary: RotaryEmbedding, + decode_cos: Tensor, + decode_sin: Tensor, /// XDRoPE section sizes (`config.rope_scaling.xdrope_section`). Used once /// per `forward` to section-mix the rotary `cos`/`sin` tensors before they /// fan out to every layer — see [`apply_xdrope_rotary_pos_emb`]. @@ -432,12 +879,31 @@ impl HunyuanLlm { _ => cfg.rope_theta, }; let rotary = RotaryEmbedding::new_multi_axis(cfg.head_dim, rope_theta, 4, vb.device())?; + let decode_rotary = RotaryEmbedding::new_dynamic(cfg.head_dim, rope_theta, vb.device())?; + let decode_positions = Tensor::arange(0i64, DECODE_ROPE_CACHE_LEN as i64, vb.device()) + .and_then(|x| x.reshape((1, 1, DECODE_ROPE_CACHE_LEN))) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode rope positions", e))?; + let (decode_cos, decode_sin) = + decode_rotary.forward_multi_axis(&decode_positions, vb.dtype())?; + // Decode positions use the same scalar on every XDRoPE axis, so the + // section selection collapses to ordinary 1-D RoPE. Preserve the + // official BF16 quantization point, then hoist the per-round F32 cast. + let decode_cos = decode_cos + .to_dtype(candle_core::DType::F32) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode rope cos f32", e))?; + let decode_sin = decode_sin + .to_dtype(candle_core::DType::F32) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode rope sin f32", e))?; Ok(Self { + #[cfg(feature = "cuda")] + dflash_decode_graph: RefCell::new(None), embed_tokens, layers, norm, rotary, + decode_cos, + decode_sin, xdrope_section: cfg.rope_scaling.xdrope_section.clone(), }) } @@ -458,6 +924,23 @@ impl HunyuanLlm { position_ids: &Tensor, causal_mask: Option<&Tensor>, ) -> Result { + Ok(self + .forward_with_aux(inputs_embeds, position_ids, causal_mask, &[])? + .hidden_states) + } + + /// Forward the target decoder and optionally retain intermediate features. + /// + /// Layer ids are one-based boundary ids: id `1` is the output after the + /// first decoder layer. DFlash checkpoint ids are zero-based layer indices, + /// so callers must add one before passing them here (as vLLM does). + pub(crate) fn forward_with_aux( + &self, + inputs_embeds: &Tensor, + position_ids: &Tensor, + causal_mask: Option<&Tensor>, + aux_layer_ids: &[usize], + ) -> Result { use candle_core::DType; let (cos, sin) = self @@ -476,15 +959,226 @@ impl HunyuanLlm { .to_dtype(DType::F32) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "xdrope sin to_dtype f32", e))?; + self.forward_with_aux_prepared(inputs_embeds, &cos, &sin, causal_mask, aux_layer_ids) + } + + pub(crate) fn forward_with_aux_decode( + &self, + inputs_embeds: &Tensor, + start: usize, + causal_mask: Option<&Tensor>, + aux_layer_ids: &[usize], + ) -> Result { + let len = inputs_embeds + .dim(1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode sequence length", e))?; + if start + len > DECODE_ROPE_CACHE_LEN { + #[cfg(feature = "cuda")] + self.invalidate_dflash_cuda_graph(); + let positions: Vec = (start..start + len) + .flat_map(|position| [position as i64; 4]) + .collect(); + let positions = Tensor::new(positions, inputs_embeds.device()) + .and_then(|t| t.reshape((len, 4))) + .and_then(|t| t.transpose(0, 1)) + .and_then(|t| t.unsqueeze(1)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode positions", e))?; + return self.forward_with_aux(inputs_embeds, &positions, causal_mask, aux_layer_ids); + } + let cos = self + .decode_cos + .narrow(2, start, len) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode rope cos slice", e))?; + let sin = self + .decode_sin + .narrow(2, start, len) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "decode rope sin slice", e))?; + #[cfg(feature = "cuda")] + if let Some(output) = + self.replay_dflash_cuda_graph(inputs_embeds, &cos, &sin, start + len, aux_layer_ids)? + { + return Ok(output); + } + self.forward_with_aux_prepared(inputs_embeds, &cos, &sin, causal_mask, aux_layer_ids) + } + + #[cfg(feature = "cuda")] + fn replay_dflash_cuda_graph( + &self, + inputs_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + kv_len: usize, + aux_layer_ids: &[usize], + ) -> Result, OCRError> { + if kv_len > DECODE_ROPE_CACHE_LEN { + self.invalidate_dflash_cuda_graph(); + return Ok(None); + } + let captured_ref = self.dflash_decode_graph.borrow(); + let Some(captured) = captured_ref.as_ref() else { + return Ok(None); + }; + if inputs_embeds.shape() != captured.hidden_input.shape() + || cos.shape() != captured.cos_input.shape() + || sin.shape() != captured.sin_input.shape() + || aux_layer_ids != captured.aux_layer_ids + { + return Ok(None); + } + captured + .hidden_input + .slice_set(inputs_embeds, 0, 0) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "copy full graph hidden", e))?; + captured + .cos_input + .slice_set(cos, 0, 0) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "copy full graph cos", e))?; + captured + .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 + .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))?; + for layer in &self.layers { + layer.set_kv_cache_len(kv_len)?; + } + Ok(Some(HunyuanLlmOutput { + hidden_states: captured.hidden_output.clone(), + aux_hidden_states: Some(captured.aux_output.clone()), + })) + } + + fn forward_with_aux_prepared( + &self, + inputs_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + causal_mask: Option<&Tensor>, + aux_layer_ids: &[usize], + ) -> Result { + #[cfg(feature = "cuda")] + { + let incoming_len = inputs_embeds.dim(1).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "prepared sequence length", e) + })?; + if self.kv_cache_len().saturating_add(incoming_len) > DECODE_ROPE_CACHE_LEN { + self.invalidate_dflash_cuda_graph(); + } + } + // Decode invokes the same non-contiguous elementwise layouts thousands + // of times. Candle otherwise uploads their dims/strides for every + // kernel. Restrict the built-in HtoD cache to the pure model forward so + // dynamic token tensors created by the caller are never retained. + #[cfg(feature = "cuda")] + let _cuda_htod_cache = match inputs_embeds.device() { + Device::Cuda(device) => Some(device.enable_cuda_graph_htod_cache()), + _ => None, + }; + if aux_layer_ids + .iter() + .any(|&id| id == 0 || id > self.layers.len()) + { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR: DFlash target layer ids must be in 1..={}, got {:?}", + self.layers.len(), + aux_layer_ids + ), + }); + } let mut hidden_states = inputs_embeds.clone(); - for layer in self.layers.iter() { - hidden_states = layer.forward(&hidden_states, &cos, &sin, causal_mask)?; + #[cfg(feature = "cuda")] + let cos_sin = if inputs_embeds.device().is_cuda() + && inputs_embeds.dtype() == candle_core::DType::BF16 + { + Some(Tensor::cat(&[cos, sin], 0).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "pack decoder XDRoPE cos/sin", e) + })?) + } else { + None + }; + #[cfg(not(feature = "cuda"))] + let cos_sin: Option = None; + let mut aux = Vec::with_capacity(aux_layer_ids.len()); + for (index, layer) in self.layers.iter().enumerate() { + hidden_states = + layer.forward(&hidden_states, cos, sin, cos_sin.as_ref(), causal_mask)?; + if aux_layer_ids.contains(&(index + 1)) { + aux.push(hidden_states.clone()); + } } + let aux_hidden_states = if aux.is_empty() { + None + } else { + let refs: Vec<&Tensor> = aux.iter().collect(); + Some(Tensor::cat(&refs, 2).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "concatenate DFlash target features", e) + })?) + }; let hidden_states = self .norm .forward(&hidden_states) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "llm final norm", e))?; - Ok(hidden_states) + Ok(HunyuanLlmOutput { + hidden_states, + aux_hidden_states, + }) + } + + #[cfg(feature = "cuda")] + fn forward_with_aux_dynamic( + &self, + inputs_embeds: &Tensor, + cos: &Tensor, + sin: &Tensor, + query_lengths: &Tensor, + kv_lengths: &Tensor, + aux_layer_ids: &[usize], + ) -> Result { + let mut hidden_states = inputs_embeds.clone(); + let cos_sin = if inputs_embeds.dtype() == candle_core::DType::BF16 { + Some(Tensor::cat(&[cos, sin], 0).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "pack dynamic XDRoPE cos/sin", e) + })?) + } else { + None + }; + let mut aux = Vec::with_capacity(aux_layer_ids.len()); + for (index, layer) in self.layers.iter().enumerate() { + hidden_states = layer.forward_dynamic( + &hidden_states, + cos, + sin, + cos_sin.as_ref(), + query_lengths, + kv_lengths, + )?; + if aux_layer_ids.contains(&(index + 1)) { + aux.push(hidden_states.clone()); + } + } + let refs: Vec<&Tensor> = aux.iter().collect(); + let aux_hidden_states = Tensor::cat(&refs, 2).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "dynamic target auxiliary states", e) + })?; + let hidden_states = self + .norm + .forward(&hidden_states) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "dynamic target final norm", e))?; + Ok(HunyuanLlmOutput { + hidden_states, + aux_hidden_states: Some(aux_hidden_states), + }) } pub fn clear_kv_cache(&self) { @@ -492,4 +1186,199 @@ impl HunyuanLlm { layer.clear_kv_cache(); } } + + pub(crate) fn prepare_dflash_cuda_graphs( + &self, + query_len: usize, + aux_layer_ids: &[usize], + ) -> Result<(), OCRError> { + if std::env::var_os("OAR_HUNYUAN_DISABLE_CUDA_GRAPH").is_some() { + return Ok(()); + } + // Dynamic target attention pins an F16 KV cache for graph replay. + // F32 keeps its higher-precision eager path instead of failing capture + // when the append kernel sees F32 Q/K/V and F16 cache storage. + #[cfg(feature = "cuda")] + if self.decode_cos.device().is_cuda() + && target_cuda_graph_dtype_supported(self.embed_tokens.embeddings().dtype()) + { + let cos = self.decode_cos.narrow(2, 0, query_len).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "graph decode cos template", e) + })?; + let sin = self.decode_sin.narrow(2, 0, query_len).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR", "graph decode sin template", e) + })?; + self.capture_dflash_cuda_graph(query_len, &cos, &sin, aux_layer_ids)?; + } + let _ = query_len; + let _ = aux_layer_ids; + Ok(()) + } + + #[cfg(feature = "cuda")] + fn capture_dflash_cuda_graph( + &self, + query_len: usize, + cos_template: &Tensor, + sin_template: &Tensor, + aux_layer_ids: &[usize], + ) -> Result<(), OCRError> { + use candle_core::cuda_backend::cudarc::driver::sys::{ + CUgraphInstantiate_flags_enum, CUstreamCaptureMode_enum, + }; + + if self.dflash_decode_graph.borrow().is_some() { + return Ok(()); + } + let Device::Cuda(cuda) = self.decode_cos.device() else { + return Ok(()); + }; + if aux_layer_ids.is_empty() { + return Err(OCRError::ConfigError { + message: "HunyuanOCR DFlash CUDA graph requires auxiliary target layers" + .to_string(), + }); + } + for layer in &self.layers { + layer.self_attn.prepare_dynamic_cache(query_len)?; + } + let hidden_size = self + .embed_tokens + .embeddings() + .dim(1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "graph hidden size", e))?; + let dtype = self.embed_tokens.embeddings().dtype(); + let hidden_input = + Tensor::zeros((1, query_len, hidden_size), dtype, self.decode_cos.device()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph hidden input", e))?; + let cos_input = Tensor::zeros( + cos_template.shape(), + cos_template.dtype(), + self.decode_cos.device(), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph cos input", e))?; + let sin_input = Tensor::zeros( + sin_template.shape(), + sin_template.dtype(), + self.decode_sin.device(), + ) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph sin input", e))?; + let query_lengths = Tensor::new(&[0u32, query_len as u32], self.decode_cos.device()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph query lengths", e))?; + // These must not share storage: query length stays fixed at 16 while + // the cumulative KV length is overwritten before every graph replay. + let kv_lengths = Tensor::new(&[0u32, query_len as u32], self.decode_cos.device()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "full graph KV lengths", e))?; + let stream = cuda.cuda_stream(); + let _htod_cache = cuda.enable_cuda_graph_htod_cache(); + + let warm = self.forward_with_aux_dynamic( + &hidden_input, + &cos_input, + &sin_input, + &query_lengths, + &kv_lengths, + aux_layer_ids, + )?; + sync_graph_tensor(&warm.hidden_states, "warm full target decoder graph")?; + + stream + .begin_capture(CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_GLOBAL) + .map_err(|e| cuda_graph_error("begin full target decoder graph capture", e))?; + let output = match self.forward_with_aux_dynamic( + &hidden_input, + &cos_input, + &sin_input, + &query_lengths, + &kv_lengths, + aux_layer_ids, + ) { + 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("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 + .ok_or_else(|| OCRError::ConfigError { + message: "HunyuanOCR full target decoder graph produced no auxiliary states" + .to_string(), + })?; + graph + .launch() + .map_err(|e| cuda_graph_error("warm full target decoder graph", e))?; + sync_graph_tensor(&output.hidden_states, "sync full target decoder graph")?; + self.clear_kv_cache(); + *self.dflash_decode_graph.borrow_mut() = Some(TargetDecoderCudaGraph { + graph, + hidden_input, + cos_input, + sin_input, + _query_lengths: query_lengths, + kv_lengths, + hidden_output: output.hidden_states, + aux_output, + aux_layer_ids: aux_layer_ids.to_vec(), + }); + Ok(()) + } + + #[cfg(feature = "cuda")] + fn invalidate_dflash_cuda_graph(&self) { + // Eager cache growth reallocates the storage whose raw pointers were + // captured by the graph. Dropping it before growth also prevents a + // stale replay after the logical cache is reset for another document. + self.dflash_decode_graph.borrow_mut().take(); + } + + /// Drop any captured target-decoder CUDA graph. + /// + /// Callers must invoke this before feeding a batch shape the graph + /// wasn't captured for (e.g. a multi-image request that disables DFlash + /// and falls back to plain batched prefill): the batched append would + /// otherwise reallocate the per-layer KV storage the graph's raw device + /// pointers still reference, and a later single-image DFlash decode + /// could replay the graph against that freed memory. + pub(crate) fn invalidate_target_cuda_graph(&self) { + #[cfg(feature = "cuda")] + self.invalidate_dflash_cuda_graph(); + } + + pub(crate) fn trim_kv_cache(&self, len: usize) -> Result<(), OCRError> { + for layer in &self.layers { + layer.trim_kv_cache(len)?; + } + Ok(()) + } + + pub(crate) fn kv_cache_len(&self) -> usize { + self.layers + .first() + .map_or(0, HunyuanDecoderLayer::kv_cache_len) + } +} + +#[cfg(test)] +mod tests { + use super::target_cuda_graph_dtype_supported; + use candle_core::DType; + + #[test] + fn target_cuda_graph_skips_unsupported_dtypes() { + assert!(target_cuda_graph_dtype_supported(DType::F16)); + assert!(target_cuda_graph_dtype_supported(DType::BF16)); + assert!(!target_cuda_graph_dtype_supported(DType::F32)); + assert!(!target_cuda_graph_dtype_supported(DType::F64)); + } } diff --git a/oar-ocr-vl/src/hunyuanocr/mod.rs b/oar-ocr-vl/src/hunyuanocr/mod.rs index 1f4aac1..4c96a81 100644 --- a/oar-ocr-vl/src/hunyuanocr/mod.rs +++ b/oar-ocr-vl/src/hunyuanocr/mod.rs @@ -5,6 +5,9 @@ //! Candle. mod config; +mod dflash; +#[cfg(feature = "cuda")] +mod dynamic_kv; mod llm; mod model; mod processing; @@ -14,4 +17,5 @@ pub use config::{ HunyuanOcrConfig, HunyuanOcrImageProcessorConfig, HunyuanOcrRopeScaling, HunyuanOcrVersion, HunyuanOcrVisionConfig, }; +pub use dflash::{DFlashConfig, DFlashTargetConfig}; pub use model::HunyuanOcr; diff --git a/oar-ocr-vl/src/hunyuanocr/model.rs b/oar-ocr-vl/src/hunyuanocr/model.rs index 69f4f2f..4ab9e70 100644 --- a/oar-ocr-vl/src/hunyuanocr/model.rs +++ b/oar-ocr-vl/src/hunyuanocr/model.rs @@ -1,6 +1,7 @@ //! HunyuanOCR model wrapper (HunYuanVLForConditionalGeneration). use super::config::{HunyuanOcrConfig, HunyuanOcrImageProcessorConfig, HunyuanOcrVersion}; +use super::dflash::DFlashModel; use super::llm::HunyuanLlm; use super::processing::{HunyuanOcrImageInputs, preprocess_image}; use super::vision::HunyuanVisionModel; @@ -98,6 +99,7 @@ pub struct HunyuanOcr { image_cfg: HunyuanOcrImageProcessorConfig, tokenizer: Tokenizer, llm: HunyuanLlm, + dflash: Option, vision: HunyuanVisionModel, stop_token_ids: Vec, /// `generation_config.json::repetition_penalty`. HuggingFace's @@ -168,6 +170,7 @@ impl HunyuanOcr { image_cfg, tokenizer, llm, + dflash: None, vision, stop_token_ids, repetition_penalty, @@ -175,6 +178,104 @@ impl HunyuanOcr { }) } + /// Load HunyuanOCR together with an explicitly located DFlash draft. + /// + /// DFlash is supported by HunyuanOCR 1.5 only. The official checkpoint + /// places the draft under `/dflash`. + pub fn from_dirs( + model_dir: impl AsRef, + dflash_dir: impl AsRef, + device: Device, + ) -> Result { + let mut model = Self::from_dir(model_dir, device)?; + if model.version != HunyuanOcrVersion::V1_5 { + return Err(OCRError::ConfigError { + message: "HunyuanOCR: DFlash requires the 1.5 checkpoint".to_string(), + }); + } + let dflash = DFlashModel::from_dir(dflash_dir, model.dtype, &model.device)?; + if dflash.config().hidden_size != model.cfg.hidden_size + || dflash.config().vocab_size != model.cfg.vocab_size + { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR: DFlash target mismatch (hidden/vocab draft={}/{}, target={}/{})", + dflash.config().hidden_size, + dflash.config().vocab_size, + model.cfg.hidden_size, + model.cfg.vocab_size + ), + }); + } + if dflash + .config() + .dflash_config + .target_layer_ids + .iter() + .any(|&id| id >= model.cfg.num_hidden_layers) + { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR: invalid DFlash target layers {:?} for {}-layer target", + dflash.config().dflash_config.target_layer_ids, + model.cfg.num_hidden_layers + ), + }); + } + let target_layers: Vec = dflash + .config() + .dflash_config + .target_layer_ids + .iter() + .map(|id| id + 1) + .collect(); + model + .llm + .prepare_dflash_cuda_graphs(dflash.config().block_size, &target_layers)?; + model.dflash = Some(dflash); + Ok(model) + } + + /// Load the official `/dflash` draft alongside the target. + pub fn from_dir_with_dflash( + model_dir: impl AsRef, + device: Device, + ) -> Result { + let model_dir = model_dir.as_ref(); + Self::from_dirs(model_dir, model_dir.join("dflash"), device) + } + + pub fn dflash_enabled(&self) -> bool { + self.dflash.is_some() + } + + /// Override the checkpoint/reference repetition penalty used by greedy + /// decoding. A value of `1.0` disables the processor, matching the + /// official DFlash speed-benchmark configuration. + pub fn set_repetition_penalty(&mut self, penalty: f64) -> Result<(), OCRError> { + if !penalty.is_finite() || penalty < 1.0 { + return Err(OCRError::ConfigError { + message: format!( + "HunyuanOCR: repetition penalty must be finite and >= 1.0, got {penalty}" + ), + }); + } + self.repetition_penalty = penalty; + Ok(()) + } + + pub fn repetition_penalty(&self) -> f64 { + self.repetition_penalty + } + + /// Number of parallel draft tokens per DFlash step (15 for the official + /// 16-position bonus+mask block). + pub fn dflash_num_speculative_tokens(&self) -> Option { + self.dflash + .as_ref() + .map(|draft| draft.config().block_size.saturating_sub(1)) + } + /// The checkpoint generation detected from `config.json`. pub fn version(&self) -> HunyuanOcrVersion { self.version @@ -417,18 +518,64 @@ impl HunyuanOcr { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "stack pos", e))?; // 4. 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("HunyuanOCR", "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("HunyuanOCR", "create padding", e))?; - let mask = combine_masks(&causal, &padding) - .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "combine masks", e))?; + let mask = if batch_size == 1 && self.device.is_cuda() { + None + } else { + let causal = create_causal_mask(max_seq_len, max_seq_len, self.dtype, &self.device) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "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("HunyuanOCR", "create padding", e))?; + Some( + combine_masks(&causal, &padding) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "combine masks", e))?, + ) + }; // 5. Prefill self.llm.clear_kv_cache(); - let hidden = self - .llm - .forward(&inputs_embeds, &position_ids, Some(&mask))?; + let active_dflash = if batch_size == 1 { + self.dflash.as_ref() + } else { + // Multi-image requests disable DFlash and prefill a batch>1 + // shape, which is incompatible with the batch=1 KV storage a + // captured target CUDA graph points at. Drop the graph before + // that reallocation so a later single-image DFlash decode can't + // replay it against freed memory. + if self.dflash.is_some() { + self.llm.invalidate_target_cuda_graph(); + } + None + }; + let hidden = if let Some(dflash) = active_dflash { + // DFlash stores zero-based target layer indices. HunyuanLlm uses + // one-based post-layer boundaries, matching vLLM's `i + 1` + // conversion before auxiliary-state capture. + let aux_layer_ids: Vec = dflash + .config() + .dflash_config + .target_layer_ids + .iter() + .map(|id| id + 1) + .collect(); + let output = self.llm.forward_with_aux( + &inputs_embeds, + &position_ids, + mask.as_ref(), + &aux_layer_ids, + )?; + let aux = output + .aux_hidden_states + .ok_or_else(|| OCRError::ConfigError { + message: "HunyuanOCR DFlash: target did not return auxiliary hidden states" + .to_string(), + })?; + dflash.reset_context(&aux)?; + output.hidden_states + } else { + self.llm + .forward(&inputs_embeds, &position_ids, mask.as_ref())? + }; // 6. Get initial logits per sample let mut logits_list: Vec = Vec::with_capacity(batch_size); @@ -440,6 +587,17 @@ impl HunyuanOcr { logits_list.push(logits); } + // DFlash currently targets the latency-sensitive single-document path. + // Multi-image calls retain the existing true-batch autoregressive path. + if let Some(dflash) = active_dflash { + let initial_logits = logits_list.pop().ok_or_else(|| OCRError::ConfigError { + message: "HunyuanOCR DFlash: missing target prefill logits".to_string(), + })?; + let tokens = + self.generate_dflash_tokens(dflash, initial_logits, seq_lens[0], max_new_tokens)?; + return Ok(vec![tokens]); + } + // 7. Autoregressive decode let mut generated: Vec> = vec![Vec::new(); batch_size]; let mut finished: Vec = vec![false; batch_size]; @@ -533,6 +691,251 @@ impl HunyuanOcr { Ok(generated) } + fn greedy_token(&self, logits: &Tensor, seen: &[u32]) -> Result { + if self.repetition_penalty > 1.0 && !seen.is_empty() { + argmax_with_repetition_penalty(logits, seen, self.repetition_penalty as f32) + } else { + logits + .argmax(D::Minus1) + .and_then(|t| t.to_scalar::()) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "argmax", e)) + } + } + + /// Greedy DFlash speculative decoding for one document. The target model + /// remains authoritative: accepted draft prefixes plus the first target + /// recovery/bonus token are exactly equivalent to autoregressive greedy + /// decoding, including repetition-penalty processing. + fn generate_dflash_tokens( + &self, + dflash: &DFlashModel, + initial_logits: Tensor, + prompt_len: usize, + max_new_tokens: usize, + ) -> Result, OCRError> { + if max_new_tokens == 0 { + self.llm.clear_kv_cache(); + dflash.clear_context(); + return Ok(Vec::new()); + } + + let first = self.greedy_token(&initial_logits, &[])?; + if self.stop_token_ids.contains(&first) { + self.llm.clear_kv_cache(); + dflash.clear_context(); + return Ok(Vec::new()); + } + + let num_spec = dflash.config().block_size.saturating_sub(1); + if num_spec == 0 { + return Err(OCRError::ConfigError { + message: "HunyuanOCR DFlash: block_size must include at least one mask position" + .to_string(), + }); + } + let mask_id = dflash.config().dflash_config.mask_token_id; + let target_layers: Vec = dflash + .config() + .dflash_config + .target_layer_ids + .iter() + .map(|id| id + 1) + .collect(); + let mut generated = vec![first]; + let mut draft_rounds = 0usize; + let mut accepted_draft_tokens = 0usize; + + while generated.len() < max_new_tokens { + draft_rounds += 1; + let bonus = *generated + .last() + .expect("generated contains the bonus token"); + + // One bonus query followed by parallel mask queries. + let mut query_ids = Vec::with_capacity(num_spec + 1); + query_ids.push(bonus); + query_ids.resize(num_spec + 1, mask_id); + let query_ids = Tensor::new(query_ids, &self.device) + .and_then(|t| t.reshape((1, num_spec + 1))) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "query ids", e))?; + let query_embeds = self.llm.embed(&query_ids)?; + let draft_hidden = dflash.forward_queries(&query_embeds)?; + let draft_rows = draft_hidden + .narrow(1, 1, num_spec) + .and_then(|t| t.squeeze(0)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "draft mask rows", e))?; + let draft_logits = self.logits_from_hidden_batch(&draft_rows)?; + // All draft rows are independent. Keep their sampling as one GPU + // operation; target-side processors remain authoritative during + // verification and therefore preserve the target distribution. + let proposals = draft_logits + .argmax(D::Minus1) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "draft argmax", e))?; + + // The target verifies [bonus, proposal_1, ..., proposal_N] in one + // causal pass. Row i predicts proposal i+1; the final row predicts + // the target-only bonus token used when every proposal is accepted. + let context_len = self.llm.kv_cache_len(); + debug_assert_eq!(context_len, dflash.context_len()); + // Keep proposals on the device while preparing target verification. + // Reading them here would serialize the CPU after the draft pass; + // instead, the target pass is queued immediately and the proposal + // ids are copied back only after target sampling has synchronized + // the stream. + let bonus_id = query_ids + .narrow(1, 0, 1) + .and_then(|t| t.squeeze(0)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "bonus id view", e))?; + let verify_ids = Tensor::cat(&[&bonus_id, &proposals], 0) + .and_then(|t| t.reshape((1, num_spec + 1))) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "verify ids", e))?; + let verify_embeds = self.llm.embed(&verify_ids)?; + let verify_mask = if self.device.is_cuda() { + None + } else { + Some( + create_causal_mask( + num_spec + 1, + context_len + num_spec + 1, + self.dtype, + &self.device, + ) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "verify causal mask", e) + })?, + ) + }; + let output = self.llm.forward_with_aux_decode( + &verify_embeds, + context_len, + verify_mask.as_ref(), + &target_layers, + )?; + let aux = output + .aux_hidden_states + .ok_or_else(|| OCRError::ConfigError { + message: + "HunyuanOCR DFlash: target verification did not return auxiliary states" + .to_string(), + })?; + let target_rows = output.hidden_states.squeeze(0).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "target verify rows", e) + })?; + let target_logits = self.logits_from_hidden_batch(&target_rows)?; + let target_tokens = if self.repetition_penalty <= 1.0 { + Some( + target_logits + .argmax(D::Minus1) + .and_then(|t| t.to_vec1::()) + .map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "batched target argmax", e) + })?, + ) + } else { + None + }; + let proposals = proposals + .to_vec1::() + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "draft proposals", e))?; + + let mut accepted = 0usize; + let mut recovery = None; + let mut seen = if self.repetition_penalty > 1.0 { + generated.clone() + } else { + Vec::new() + }; + for (index, &proposal) in proposals.iter().enumerate() { + let target_token = match &target_tokens { + Some(tokens) => tokens[index], + None => { + let logits = target_logits.i((index, ..)).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "target verify logits", e) + })?; + self.greedy_token(&logits, &seen)? + } + }; + if target_token != proposal { + recovery = Some(target_token); + break; + } + accepted += 1; + seen.push(proposal); + } + accepted_draft_tokens += accepted; + + let (processed_inputs, next_tokens) = if let Some(recovery) = recovery { + let mut tokens = proposals[..accepted].to_vec(); + tokens.push(recovery); + // bonus plus the accepted proposal prefix were actually part + // of the authoritative history before the recovery token. + (1 + accepted, tokens) + } else { + let target_bonus = match &target_tokens { + Some(tokens) => tokens[num_spec], + None => { + let bonus_logits = target_logits.i((num_spec, ..)).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "target bonus logits", e) + })?; + self.greedy_token(&bonus_logits, &seen)? + } + }; + let mut tokens = proposals; + tokens.push(target_bonus); + (num_spec + 1, tokens) + }; + tracing::debug!( + accepted, + processed_inputs, + proposals = ?next_tokens.get(..accepted), + recovery_or_bonus = ?next_tokens.last(), + "HunyuanOCR DFlash verification" + ); + + // Verification appended the whole candidate block. Keep only the + // prefix that precedes the newly emitted recovery/bonus token. + self.llm.trim_kv_cache(context_len + processed_inputs)?; + let accepted_aux = aux.narrow(1, 0, processed_inputs).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "accepted auxiliary states", e) + })?; + dflash.append_context(&accepted_aux)?; + + let mut stop = false; + for token in next_tokens { + if self.stop_token_ids.contains(&token) { + stop = true; + break; + } + if generated.len() == max_new_tokens { + stop = true; + break; + } + generated.push(token); + } + if stop { + break; + } + + debug_assert_eq!(self.llm.kv_cache_len(), dflash.context_len()); + debug_assert!(self.llm.kv_cache_len() >= prompt_len); + } + + if draft_rounds > 0 { + tracing::info!( + rounds = draft_rounds, + accepted = accepted_draft_tokens, + drafted = draft_rounds * num_spec, + acceptance_rate = accepted_draft_tokens as f64 / (draft_rounds * num_spec) as f64, + mean_acceptance_length = 1.0 + accepted_draft_tokens as f64 / draft_rounds as f64, + "HunyuanOCR DFlash metrics" + ); + } + + self.llm.clear_kv_cache(); + dflash.clear_context(); + Ok(generated) + } + pub fn decode_tokens(&self, tokens: &[u32]) -> Result { self.decode_generated_tokens(tokens) } @@ -579,6 +982,20 @@ impl HunyuanOcr { .squeeze(0) .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "squeeze logits", e)) } + + fn logits_from_hidden_batch(&self, hidden: &Tensor) -> Result { + let w = self.llm.token_embedding_weight(); + let wt = w.transpose(0, 1).map_err(|e| { + candle_to_ocr_processing( + oar_ocr_core::core::errors::ProcessingStage::TensorOperation, + "HunyuanOCR: transpose embedding weight failed", + e, + ) + })?; + hidden + .matmul(&wt) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "batched matmul logits", e)) + } } fn build_prompt(instruction: &str, version: HunyuanOcrVersion) -> String { diff --git a/oar-ocr-vl/src/hunyuanocr/vision.rs b/oar-ocr-vl/src/hunyuanocr/vision.rs index 0faab83..5501c8d 100644 --- a/oar-ocr-vl/src/hunyuanocr/vision.rs +++ b/oar-ocr-vl/src/hunyuanocr/vision.rs @@ -1,8 +1,11 @@ use super::config::{HunyuanOcrVersion, HunyuanOcrVisionConfig}; +use crate::attention::flash_attention; use crate::utils::{candle_to_ocr_inference, candle_to_ocr_processing}; use candle_core::{D, DType, Device, IndexOp, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module}; use oar_ocr_core::core::OCRError; +use rayon::prelude::*; +use std::sync::Arc; /// Late vit layers are the cross-implementation drift hotspot: BF16 Q·K /// accumulation drift redirects attention "sink" positions enough to swap the @@ -14,7 +17,9 @@ const VIT_LATE_F32_THRESHOLD: usize = 20; #[derive(Debug, Clone)] struct VisionEmbeddings { patch_embedding: Conv2d, - position_embedding: candle_nn::Embedding, + patch_position_base: Arc<[f32]>, + position_grid: usize, + position_dim: usize, version: HunyuanOcrVersion, } @@ -74,105 +79,81 @@ impl VisionEmbeddings { .map_err(|e| { candle_to_ocr_inference("HunyuanOCR", "load vit position_embedding", e) })?; - - Ok(Self { - patch_embedding, - position_embedding, - version, - }) - } - - fn interpolate_patch_pos( - &self, - height: usize, - width: usize, - device: &Device, - dtype: DType, - ) -> Result { - if height == 0 || width == 0 { - return Err(OCRError::InvalidInput { - message: format!( - "HunyuanOCR: vit interpolate_patch_pos requires height/width > 0, got {height}x{width}" - ), - }); - } - - let pos_w = self.position_embedding.embeddings(); - let (num_positions, dim) = pos_w.dims2().map_err(|e| { + let pos_w = position_embedding.embeddings(); + let (loaded_positions, position_dim) = pos_w.dims2().map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, "HunyuanOCR: vit position_embedding dims2 failed", e, ) })?; - if num_positions < 2 { + if loaded_positions < 2 { return Err(OCRError::ConfigError { message: format!( - "HunyuanOCR: vit position_embedding is too small: {num_positions}" + "HunyuanOCR: vit position_embedding is too small: {loaded_positions}" ), }); } - - let patch_pos = pos_w.i((1..num_positions, ..)).map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "HunyuanOCR: vit slice patch position embedding failed", - e, - ) - })?; - let (num_patch_positions, _) = patch_pos.dims2().map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "HunyuanOCR: vit patch position dims2 failed", - e, - ) - })?; - - let grid = (num_patch_positions as f64).sqrt() as usize; - if grid * grid != num_patch_positions { + let num_patch_positions = loaded_positions - 1; + let position_grid = (num_patch_positions as f64).sqrt() as usize; + if position_grid * position_grid != num_patch_positions { return Err(OCRError::ConfigError { message: format!( "HunyuanOCR: vit patch position grid is not square: num={num_patch_positions}" ), }); } - - let base = patch_pos - .to_dtype(DType::F32) + // Position weights never change. Copy the BF16->F32 base once while + // loading instead of transferring the same 75.5 MB GPU tensor back to + // the CPU for every page before interpolation. + let patch_position_base: Arc<[f32]> = pos_w + .i((1..loaded_positions, ..)) + .and_then(|x| x.to_dtype(DType::F32)) + .and_then(|x| x.flatten_all()) + .and_then(|x| x.to_vec1::()) .map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "HunyuanOCR: vit patch position cast to f32 failed", + "HunyuanOCR: cache vit patch position embedding failed", e, ) })? - .flatten_all() - .map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "HunyuanOCR: vit patch position flatten failed", - e, - ) - })? - .to_vec1::() - .map_err(|e| { - candle_to_ocr_processing( - oar_ocr_core::core::errors::ProcessingStage::TensorOperation, - "HunyuanOCR: vit patch position to_vec failed", - e, - ) - })?; + .into(); + + Ok(Self { + patch_embedding, + patch_position_base, + position_grid, + position_dim, + version, + }) + } + + fn interpolate_patch_pos( + &self, + height: usize, + width: usize, + device: &Device, + dtype: DType, + ) -> Result { + if height == 0 || width == 0 { + return Err(OCRError::InvalidInput { + message: format!( + "HunyuanOCR: vit interpolate_patch_pos requires height/width > 0, got {height}x{width}" + ), + }); + } let out = interpolate_bilinear_align_corners_false( - &base, - grid, - grid, + &self.patch_position_base, + self.position_grid, + self.position_grid, height, width, - dim, + self.position_dim, self.version, ); - Tensor::from_vec(out, (height * width, dim), device) + Tensor::from_vec(out, (height * width, self.position_dim), device) .map_err(|e| { candle_to_ocr_processing( oar_ocr_core::core::errors::ProcessingStage::TensorOperation, @@ -291,6 +272,24 @@ impl VisionAttention { .contiguous() .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "vit attn v contiguous", e))?; + // The official CUDA path uses fused full attention for the vision + // encoder. Besides avoiding the quadratic attention buffer, the + // kernel accumulates the BF16 dot products stably without the late + // F32 eager fallback below. + if let Some(attn_output) = flash_attention(&q, &k, &v, self.scaling, false) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "vit flash attention", e))? + { + let attn_output = attn_output + .transpose(1, 2) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "vit flash transpose", e))? + .reshape((b, seq_len, self.num_heads * self.head_dim)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "vit flash reshape", e))?; + return self + .o_proj + .forward(&attn_output) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "vit attn o_proj", e)); + } + // Chunked attention over the query dimension. Without chunking the // (B, H, N, N) attention matrix at N=4320 (the vit's full patch // sequence) needs ~4 GB just for the BF16 buffer, OOM'ing the 4090. @@ -929,34 +928,28 @@ fn interpolate_bilinear_align_corners_false( } let (inv_scale_y, inv_scale_x) = match version { - // V1.0's Transformers path passed scale_factor=(out + 0.1) / in. HunyuanOcrVersion::V1 => ( in_h as f32 / (out_h as f32 + 0.1), in_w as f32 / (out_w as f32 + 0.1), ), - // V1.5 passes an explicit output size to interpolate(), whose - // align_corners=false source scale is in/out. HunyuanOcrVersion::V1_5 => (in_h as f32 / out_h as f32, in_w as f32 / out_w as f32), }; - for oy in 0..out_h { + let row = |oy: usize, out_row: &mut [f32]| { let fy = ((oy as f32 + 0.5) * inv_scale_y - 0.5).max(0.0); let y0 = fy.floor().max(0.0) as usize; let y1 = (y0 + 1).min(in_h - 1); let wy = fy - y0 as f32; - for ox in 0..out_w { let fx = ((ox as f32 + 0.5) * inv_scale_x - 0.5).max(0.0); let x0 = fx.floor().max(0.0) as usize; let x1 = (x0 + 1).min(in_w - 1); let wx = fx - x0 as f32; - - let out_base = (oy * out_w + ox) * dim; + let out_base = ox * dim; let i00 = (y0 * in_w + x0) * dim; let i01 = (y0 * in_w + x1) * dim; let i10 = (y1 * in_w + x0) * dim; let i11 = (y1 * in_w + x1) * dim; - for c in 0..dim { let v00 = base[i00 + c]; let v01 = base[i01 + c]; @@ -964,9 +957,22 @@ fn interpolate_bilinear_align_corners_false( let v11 = base[i11 + c]; let v0 = v00 + (v01 - v00) * wx; let v1 = v10 + (v11 - v10) * wx; - out[out_base + c] = v0 + (v1 - v0) * wy; + out_row[out_base + c] = v0 + (v1 - v0) * wy; } } + }; + + // Below this row count, Rayon's scheduling overhead outweighs the gain + // from parallelizing; run sequentially instead. + const PAR_ROW_THRESHOLD: usize = 16; + if out_h < PAR_ROW_THRESHOLD { + for (oy, out_row) in out.chunks_mut(out_w * dim).enumerate() { + row(oy, out_row); + } + } else { + out.par_chunks_mut(out_w * dim) + .enumerate() + .for_each(|(oy, out_row)| row(oy, out_row)); } out } diff --git a/oar-ocr-vl/src/kv_trim.rs b/oar-ocr-vl/src/kv_trim.rs index 15c2b0b..3826b1d 100644 --- a/oar-ocr-vl/src/kv_trim.rs +++ b/oar-ocr-vl/src/kv_trim.rs @@ -4,10 +4,9 @@ //! callers roll the cache back to an earlier sequence length or keep an //! arbitrary subset of positions, leaving the rest of the attention path intact. //! -//! Append uses `Tensor::cat` (candle's pre-preallocation behaviour). A -//! preallocation + `slice_set` rewrite was tried for parity with -//! `candle_nn::kv_cache::Cache` but reverted: it didn't improve per-page wall -//! time and introduced subtle K/V value differences. +//! Append uses a fixed-capacity backing tensor and `slice_set`, so speculative +//! verification never copies the accepted history. `kv` contains narrow views +//! over that backing storage and rollback only changes the logical length. use candle_core::{Result, Tensor}; @@ -21,18 +20,22 @@ use candle_core::{Result, Tensor}; pub struct TrimmableKvCache { /// Concatenation axis (typically `2` for the seq dim of `(B, H, T, D)` tensors). cat_dim: usize, - /// Cached `(keys, values)`, each shape `(B, H, cur_len, D)`. `None` until - /// first `append`. K and V are stored together so the "both populated or - /// both empty" invariant is enforced by the type system — earlier - /// revisions used parallel `Option` fields and had to assert this - /// invariant manually. + /// Full-capacity backing storage, allocated lazily once the batch/head + /// dimensions are known and retained across page-level resets. + storage: Option<(Tensor, Tensor)>, + /// Current `(B, H, cur_len, D)` views into `storage`. kv: Option<(Tensor, Tensor)>, cur_len: usize, - /// Configured sequence-length capacity retained for parity with - /// `candle_nn::kv_cache::KvCache::new`. This wrapper does not enforce it. - /// Only read by the `max_seq_len()` accessor. - #[allow(dead_code)] - max_len: usize, + /// Capacity of `storage` along `cat_dim`, zero until the first `append` + /// or `initialize_storage` call. `append` grows this organically (like a + /// `Vec`) instead of jumping straight to `configured_capacity`, so a + /// cache that never needs a fixed CUDA-graph buffer doesn't pay for one. + capacity: usize, + /// Capacity requested via `new()`. Used by `initialize_storage` (CUDA + /// graph decode paths that need a fixed-size buffer pinned before the + /// first append) and the `max_seq_len()` accessor; organic growth in + /// `append` does not preallocate to this size. + configured_capacity: usize, } // `TrimmableKvCache` lives at the crate root so every model's attention path @@ -45,27 +48,80 @@ impl TrimmableKvCache { pub fn new(cat_dim: usize, max_len: usize) -> Self { Self { cat_dim, + storage: None, kv: None, cur_len: 0, - max_len, + capacity: 0, + configured_capacity: max_len, } } - /// Append `(k_new, v_new)` and return the concatenated `(K_all, V_all)` - /// tensors the attention path consumes. Uses cat-based growth (see the - /// module note on why preallocation was not adopted). + /// Append `(k_new, v_new)` into preallocated storage and return current + /// logical K/V views. pub fn append(&mut self, k_new: &Tensor, v_new: &Tensor) -> Result<(Tensor, Tensor)> { let new_len = k_new.dim(self.cat_dim)?; - let (k_all, v_all) = match self.kv.as_ref() { - None => (k_new.clone(), v_new.clone()), - Some((k_old, v_old)) => { - let k = Tensor::cat(&[k_old, k_new], self.cat_dim)?.contiguous()?; - let v = Tensor::cat(&[v_old, v_new], self.cat_dim)?.contiguous()?; - (k, v) - } - }; + let reusable = self.storage.as_ref().is_some_and(|(storage_k, storage_v)| { + storage_k.dtype() == k_new.dtype() + && storage_v.dtype() == v_new.dtype() + && storage_k.device().same_device(k_new.device()) + && storage_v.device().same_device(v_new.device()) + && storage_k.dims().len() == k_new.dims().len() + && storage_k + .dims() + .iter() + .zip(k_new.dims()) + .enumerate() + .all(|(dim, (stored, new))| dim == self.cat_dim || stored == new) + && storage_v + .dims() + .iter() + .zip(v_new.dims()) + .enumerate() + .all(|(dim, (stored, new))| dim == self.cat_dim || stored == new) + }); + if self.storage.is_some() && !reusable { + self.storage = None; + self.kv = None; + self.cur_len = 0; + } + // Compatibility changes start a new logical cache. Compute this only + // after the reset above; otherwise the old length becomes a zero-filled + // prefix in the replacement storage. + let required = self.cur_len + new_len; + if self.storage.is_none() { + // Grow from what's actually needed rather than jumping straight + // to `configured_capacity`: most callers never trigger a CUDA + // graph (which instead pins a fixed buffer up front via + // `initialize_storage`), so a short document should not reserve + // e.g. 16K tokens of K/V per layer on its first token. + let mut shape = k_new.dims().to_vec(); + shape[self.cat_dim] = required; + self.capacity = required; + self.storage = Some(( + Tensor::zeros(shape.as_slice(), k_new.dtype(), k_new.device())?, + Tensor::zeros(shape.as_slice(), v_new.dtype(), v_new.device())?, + )); + } else if required > self.capacity { + let grow_by = self.capacity.max(new_len); + let mut shape = k_new.dims().to_vec(); + shape[self.cat_dim] = grow_by; + let (old_k, old_v) = self.storage.as_ref().expect("storage initialized"); + let extra_k = Tensor::zeros(shape.as_slice(), k_new.dtype(), k_new.device())?; + let extra_v = Tensor::zeros(shape.as_slice(), v_new.dtype(), v_new.device())?; + self.storage = Some(( + Tensor::cat(&[old_k, &extra_k], self.cat_dim)?.contiguous()?, + Tensor::cat(&[old_v, &extra_v], self.cat_dim)?.contiguous()?, + )); + self.capacity += grow_by; + } + + let (storage_k, storage_v) = self.storage.as_mut().expect("storage initialized"); + storage_k.slice_set(k_new, self.cat_dim, self.cur_len)?; + storage_v.slice_set(v_new, self.cat_dim, self.cur_len)?; + self.cur_len = required; + let k_all = storage_k.narrow(self.cat_dim, 0, self.cur_len)?; + let v_all = storage_v.narrow(self.cat_dim, 0, self.cur_len)?; self.kv = Some((k_all.clone(), v_all.clone())); - self.cur_len += new_len; Ok((k_all, v_all)) } @@ -80,13 +136,18 @@ impl TrimmableKvCache { } // `cur_len > 0` implies `kv.is_some()` by the invariant maintained in // `append` / `reset`. - let Some((k_old, v_old)) = self.kv.as_ref() else { + if self.kv.is_none() { return Err(candle_core::Error::Msg( "TrimmableKvCache::trim_to: cache empty but cur_len > 0".into(), )); - }; - let k = k_old.narrow(self.cat_dim, 0, len)?.contiguous()?; - let v = v_old.narrow(self.cat_dim, 0, len)?.contiguous()?; + } + let (storage_k, storage_v) = self.storage.as_ref().ok_or_else(|| { + candle_core::Error::Msg( + "TrimmableKvCache::trim_to: storage empty but cur_len > 0".into(), + ) + })?; + let k = storage_k.narrow(self.cat_dim, 0, len)?; + let v = storage_v.narrow(self.cat_dim, 0, len)?; self.kv = Some((k, v)); self.cur_len = len; Ok(()) @@ -125,17 +186,64 @@ impl TrimmableKvCache { let idx_t = Tensor::new(indices, device)?; let new_k = k.index_select(&idx_t, self.cat_dim)?.contiguous()?; let new_v = v.index_select(&idx_t, self.cat_dim)?.contiguous()?; - self.kv = Some((new_k, new_v)); - self.cur_len = indices.len(); - Ok(()) + self.cur_len = 0; + self.kv = None; + self.append(&new_k, &new_v).map(|_| ()) } pub fn current_seq_len(&self) -> usize { self.cur_len } + /// Ensure fixed-capacity storage exists without retaining a logical token. + /// This is used by CUDA-graph decode paths whose device kernel writes at a + /// runtime offset while the host only tracks the resulting logical length. + pub fn initialize_storage(&mut self, template: &Tensor) -> Result<()> { + if self.storage.is_none() { + let mut shape = template.dims().to_vec(); + shape[self.cat_dim] = self.configured_capacity; + self.capacity = self.configured_capacity; + self.storage = Some(( + Tensor::zeros(shape.as_slice(), template.dtype(), template.device())?, + Tensor::zeros(shape.as_slice(), template.dtype(), template.device())?, + )); + } + Ok(()) + } + + /// Return the fixed backing tensors used by dynamic CUDA-graph appends. + pub fn storage(&self) -> Option<(Tensor, Tensor)> { + self.storage.as_ref().map(|(k, v)| (k.clone(), v.clone())) + } + + /// Update only the logical length after a device-side graph append. + pub fn set_current_len(&mut self, len: usize) -> Result<()> { + if len > self.capacity { + return Err(candle_core::Error::Msg(format!( + "TrimmableKvCache::set_current_len: {len} exceeds capacity {}", + self.capacity + ))); + } + if len == 0 { + self.kv = None; + self.cur_len = 0; + return Ok(()); + } + let (storage_k, storage_v) = self.storage.as_ref().ok_or_else(|| { + candle_core::Error::Msg( + "TrimmableKvCache::set_current_len: storage is not initialized".into(), + ) + })?; + self.kv = Some(( + storage_k.narrow(self.cat_dim, 0, len)?, + storage_v.narrow(self.cat_dim, 0, len)?, + )); + self.cur_len = len; + Ok(()) + } + pub fn max_seq_len(&self) -> usize { - self.max_len + self.configured_capacity } pub fn reset(&mut self) { @@ -223,6 +331,24 @@ mod tests { Ok(()) } + #[test] + fn incompatible_storage_starts_a_new_logical_cache() -> Result<()> { + let mut c = TrimmableKvCache::new(2, 64); + let old = Tensor::zeros((1, 1, 3, 4), DType::F32, &dev())?; + c.append(&old, &old)?; + + // Changing the number of heads makes the retained storage + // incompatible with the new sequence. + let new = Tensor::ones((1, 2, 2, 4), DType::F32, &dev())?; + let (k, v) = c.append(&new, &new)?; + + assert_eq!(c.current_seq_len(), 2); + assert_eq!(k.dims(), &[1, 2, 2, 4]); + assert_eq!(v.dims(), &[1, 2, 2, 4]); + assert!(k.flatten_all()?.to_vec1::()?.iter().all(|&x| x == 1.0)); + Ok(()) + } + #[test] fn trim_then_append_concats_correctly() -> Result<()> { let mut c = TrimmableKvCache::new(2, 64); diff --git a/oar-ocr-vl/src/lib.rs b/oar-ocr-vl/src/lib.rs index b21f909..47b675b 100644 --- a/oar-ocr-vl/src/lib.rs +++ b/oar-ocr-vl/src/lib.rs @@ -48,7 +48,7 @@ pub use paddleocr_vl::{ }; pub use glmocr::GlmOcr; -pub use hunyuanocr::{HunyuanOcr, HunyuanOcrVersion}; +pub use hunyuanocr::{DFlashConfig, DFlashTargetConfig, HunyuanOcr, HunyuanOcrVersion}; pub use mineru::MinerU; pub use mineru_diffusion::{DiffusionGenerationConfig, MinerUDiffusion};