Add HunyuanOCR 1.5 DFlash decoding

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)
 

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: <model-dir>/dflash)
+    #[arg(long)]
+    dflash_dir: Option<PathBuf>,
+
     /// Paths to input images to process
     #[arg(required = true)]
     images: Vec<PathBuf>,
@@ -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<dyn std::error::Error>> {
@@ -90,14 +105,29 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         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<dyn std::error::Error>> {
 
         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(())
 }

oar-ocr-vl/src/attention.rs

@@ -121,18 +121,18 @@ pub fn create_causal_mask(
     device: &Device,
 ) -> Result<Tensor> {
     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::<f32>()?;
+
+        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;