diff --git a/oar-ocr-vl/README.md b/oar-ocr-vl/README.md index 749ebdb..c060458 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 | @@ -191,12 +191,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/examples/hunyuanocr.rs b/oar-ocr-vl/examples/hunyuanocr.rs index 3d95b26..91e7f20 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, @@ -55,6 +60,16 @@ struct Args { 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 +105,29 @@ fn main() -> Result<(), Box> { args.model_dir.display() ); let load_start = Instant::now(); - let model = HunyuanOcr::from_dir(&args.model_dir, device)?; + let 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)?, + }; info!( - "HunyuanOCR {} loaded in {:.2}ms", + "HunyuanOCR {} loaded in {:.2}ms{}", 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() ); 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 +140,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..8db85a1 100644 --- a/oar-ocr-vl/src/attention.rs +++ b/oar-ocr-vl/src/attention.rs @@ -121,18 +121,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)? @@ -652,6 +652,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..8876a3d --- /dev/null +++ b/oar-ocr-vl/src/hunyuanocr/dflash.rs @@ -0,0 +1,533 @@ +//! 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. + +use crate::attention::{RotaryEmbedding, repeat_kv, scaled_dot_product_attention}; +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; + +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(()) + } +} + +#[derive(Debug, Default)] +struct ContextKv { + k: Option, + v: Option, + len: usize, +} + +impl ContextKv { + fn reset(&mut self) { + self.k = None; + self.v = None; + self.len = 0; + } + + 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 (new_k, new_v) = match (&self.k, &self.v) { + (Some(old_k), Some(old_v)) => ( + Tensor::cat(&[old_k, k], 2) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: append context K", e))?, + Tensor::cat(&[old_v, v], 2) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: append context V", e))?, + ), + _ => (k.clone(), v.clone()), + }; + self.k = Some(new_k); + self.v = Some(new_v); + self.len += added; + Ok(()) + } +} + +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 { + q_proj: Linear, + k_proj: Linear, + v_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))?; + Ok(Self { + q_proj, + k_proj, + v_proj, + 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 project( + &self, + projection: &Linear, + input: &Tensor, + heads: usize, + norm: Option<&RmsNorm>, + ) -> Result { + let (batch, seq_len, _) = input + .dims3() + .map_err(|e| tensor_err("HunyuanOCR DFlash: projection input dims", e))?; + let projected = projection + .forward(input) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "attention projection", e))? + .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 append_context( + &self, + target_hidden: &Tensor, + cos: &Tensor, + sin: &Tensor, + cache: &mut ContextKv, + ) -> Result<(), OCRError> { + let k = self.project( + &self.k_proj, + target_hidden, + self.num_kv_heads, + Some(&self.k_norm), + )?; + let k = apply_rope(&k, cos, sin)?; + let v = self.project(&self.v_proj, target_hidden, self.num_kv_heads, None)?; + cache.append(&k, &v) + } + + fn forward( + &self, + hidden: &Tensor, + cos: &Tensor, + sin: &Tensor, + cache: &ContextKv, + ) -> Result { + let q = self.project(&self.q_proj, hidden, self.num_heads, Some(&self.q_norm))?; + let q = apply_rope(&q, cos, sin)?; + let query_k = self.project(&self.k_proj, hidden, self.num_kv_heads, Some(&self.k_norm))?; + let query_k = apply_rope(&query_k, cos, sin)?; + let query_v = self.project(&self.v_proj, hidden, self.num_kv_heads, None)?; + + let (Some(context_k), Some(context_v)) = (&cache.k, &cache.v) else { + return Err(OCRError::InvalidInput { + message: "HunyuanOCR DFlash: context cache is empty".to_string(), + }); + }; + let k = Tensor::cat(&[context_k, &query_k], 2) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: concatenate attention K", e))?; + let v = Tensor::cat(&[context_v, &query_v], 2) + .and_then(|x| x.contiguous()) + .map_err(|e| tensor_err("HunyuanOCR DFlash: concatenate attention V", e))?; + let k = repeat_kv(&k, self.num_kv_groups) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "repeat K", e))?; + let v = repeat_kv(&v, self.num_kv_groups) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "repeat V", e))?; + let output = scaled_dot_product_attention(&q, &k, &v, None, self.scale, false) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "attention", e))?; + let (batch, _, query_len, _) = output + .dims4() + .map_err(|e| tensor_err("HunyuanOCR DFlash: attention output dims", e))?; + let output = 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))?; + self.o_proj + .forward(&output) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "o_proj", e)) + } +} + +#[derive(Debug)] +struct DFlashMlp { + gate_proj: Linear, + up_proj: Linear, + down_proj: Linear, +} + +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))?; + Ok(Self { + gate_proj, + up_proj, + down_proj, + }) + } + + fn forward(&self, input: &Tensor) -> Result { + let gate = self + .gate_proj + .forward(input) + .and_then(|x| candle_nn::ops::silu(&x)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP gate", e))?; + let up = self + .up_proj + .forward(input) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "MLP up", e))?; + let hidden = (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)) + } +} + +#[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, + cache: &ContextKv, + ) -> Result { + let normed = self + .input_layernorm + .forward(hidden) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "input norm", e))?; + let attention = self.self_attn.forward(&normed, cos, sin, cache)?; + 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)) + } +} + +/// Loaded DFlash draft and its incremental target-context K/V caches. +pub(crate) struct DFlashModel { + cfg: DFlashConfig, + fc: Linear, + hidden_norm: RmsNorm, + layers: Vec, + norm: RmsNorm, + rotary: RotaryEmbedding, + 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 rotary = RotaryEmbedding::new_dynamic(cfg.head_dim, cfg.rope_theta, device)?; + let caches = (0..cfg.num_hidden_layers) + .map(|_| ContextKv::default()) + .collect(); + Ok(Self { + cfg, + fc, + hidden_norm, + layers, + norm, + rotary, + caches: RefCell::new(caches), + dtype, + device: device.clone(), + }) + } + + pub(crate) fn config(&self) -> &DFlashConfig { + &self.cfg + } + + fn rope(&self, start: usize, len: usize) -> Result<(Tensor, Tensor), OCRError> { + 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)) + } + + pub(crate) fn reset_context(&self, aux_hidden: &Tensor) -> Result<(), OCRError> { + let len = aux_hidden + .dim(1) + .map_err(|e| tensor_err("HunyuanOCR DFlash: target context length", e))?; + 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(); + } + for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { + layer.self_attn.append_context(&target, &cos, &sin, cache)?; + } + Ok(()) + } + + pub(crate) fn append_context(&self, aux_hidden: &Tensor) -> Result<(), OCRError> { + 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(); + let target = self.transform_target(aux_hidden)?; + let (cos, sin) = self.rope(start, added)?; + let mut caches = self.caches.borrow_mut(); + for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { + layer.self_attn.append_context(&target, &cos, &sin, cache)?; + } + Ok(()) + } + + pub(crate) fn context_len(&self) -> usize { + self.caches.borrow().first().map_or(0, |cache| cache.len) + } + + /// 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 { + 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)?; + let caches = self.caches.borrow(); + let mut hidden = query_embeds.clone(); + for (layer, cache) in self.layers.iter().zip(caches.iter()) { + hidden = layer.forward(&hidden, &cos, &sin, 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::DFlashConfig; + + #[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(); + } +} diff --git a/oar-ocr-vl/src/hunyuanocr/llm.rs b/oar-ocr-vl/src/hunyuanocr/llm.rs index f721a20..388dca3 100644 --- a/oar-ocr-vl/src/hunyuanocr/llm.rs +++ b/oar-ocr-vl/src/hunyuanocr/llm.rs @@ -332,6 +332,17 @@ impl HunyuanAttention { 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() + } } #[derive(Debug)] @@ -396,6 +407,23 @@ 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() + } +} + +/// 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)] @@ -458,8 +486,38 @@ 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. This is the convention used by DFlash/vLLM and by + /// the checkpoint's `dflash_config.target_layer_ids`. + 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; + 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 (cos, sin) = self .rotary .forward_multi_axis(position_ids, inputs_embeds.dtype())?; @@ -477,14 +535,29 @@ impl HunyuanLlm { .map_err(|e| candle_to_ocr_inference("HunyuanOCR", "xdrope sin to_dtype f32", e))?; let mut hidden_states = inputs_embeds.clone(); - for layer in self.layers.iter() { + 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, 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, + }) } pub fn clear_kv_cache(&self) { @@ -492,4 +565,17 @@ impl HunyuanLlm { layer.clear_kv_cache(); } } + + 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) + } } diff --git a/oar-ocr-vl/src/hunyuanocr/mod.rs b/oar-ocr-vl/src/hunyuanocr/mod.rs index 1f4aac1..9c19ae2 100644 --- a/oar-ocr-vl/src/hunyuanocr/mod.rs +++ b/oar-ocr-vl/src/hunyuanocr/mod.rs @@ -5,6 +5,7 @@ //! Candle. mod config; +mod dflash; mod llm; mod model; mod processing; @@ -14,4 +15,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..ac63368 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,75 @@ 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 == 0 || 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 + ), + }); + } + 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() + } + + /// 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 @@ -426,9 +498,30 @@ impl HunyuanOcr { // 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 { + None + }; + let hidden = if let Some(dflash) = active_dflash { + let output = self.llm.forward_with_aux( + &inputs_embeds, + &position_ids, + Some(&mask), + &dflash.config().dflash_config.target_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, Some(&mask))? + }; // 6. Get initial logits per sample let mut logits_list: Vec = Vec::with_capacity(batch_size); @@ -440,6 +533,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 +637,196 @@ 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 = &dflash.config().dflash_config.target_layer_ids; + let mut generated = vec![first]; + + while generated.len() < max_new_tokens { + 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) + .and_then(|t| t.to_vec1::()) + .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()); + let mut verify_ids = Vec::with_capacity(num_spec + 1); + verify_ids.push(bonus); + verify_ids.extend_from_slice(&proposals); + let verify_ids = Tensor::new(verify_ids, &self.device) + .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 positions: Vec = (context_len..context_len + num_spec + 1) + .flat_map(|position| [position as i64; 4]) + .collect(); + let positions = Tensor::new(positions, &self.device) + .and_then(|t| t.reshape((num_spec + 1, 4))) + .and_then(|t| t.transpose(0, 1)) + .and_then(|t| t.unsqueeze(1)) + .map_err(|e| candle_to_ocr_inference("HunyuanOCR DFlash", "verify positions", e))?; + let mask = 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( + &verify_embeds, + &positions, + Some(&mask), + 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 mut accepted = 0usize; + let mut recovery = None; + let mut seen = generated.clone(); + for (index, &proposal) in proposals.iter().enumerate() { + let logits = target_logits.i((index, ..)).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "target verify logits", e) + })?; + let target_token = self.greedy_token(&logits, &seen)?; + if target_token != proposal { + recovery = Some(target_token); + break; + } + accepted += 1; + seen.push(proposal); + } + + 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 bonus_logits = target_logits.i((num_spec, ..)).map_err(|e| { + candle_to_ocr_inference("HunyuanOCR DFlash", "target bonus logits", e) + })?; + let target_bonus = 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); + } + + 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 +873,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/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};