GreatV · GreatV · Jul 14, 2026 · Jul 13, 2026 · Jul 14, 2026 · Jul 14, 2026
diff --git 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
@@ -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/build.rs b/oar-ocr-vl/build.rs
@@ -1,8 +1,34 @@
 fn main() {
+    println!("cargo:rerun-if-changed=src/hunyuanocr/dynamic_kv.cu");
     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 out_dir = std::path::PathBuf::from(
+            std::env::var_os("OUT_DIR").expect("Cargo always sets OUT_DIR"),
+        );
+        let output = std::process::Command::new("nvcc")
+            .args([
+                "--ptx",
+                "--std=c++17",
+                "-O3",
+                "--gpu-architecture=compute_80",
+                "-o",
+            ])
+            .arg(out_dir.join("hunyuan_dynamic_kv.ptx"))
+            .arg("src/hunyuanocr/dynamic_kv.cu")
+            .output()
+            .expect("failed to invoke nvcc for HunyuanOCR dynamic KV kernel");
+        if !output.status.success() {
+            panic!(
+                "nvcc failed for HunyuanOCR dynamic KV kernel:\n{}",
+                String::from_utf8_lossy(&output.stderr)
+            );
+        }
+    }
 }
diff --git 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: <model-dir>/dflash)
+    #[arg(long)]
+    dflash_dir: Option<PathBuf>,
+
     /// Paths to input images to process
     #[arg(required = true)]
     images: Vec<PathBuf>,
@@ -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<f64>,
+
     /// 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<dyn std::error::Error>> {
@@ -90,14 +109,33 @@ 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 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<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(())
 }
diff --git a/oar-ocr-vl/src/attention.rs b/oar-ocr-vl/src/attention.rs
@@ -101,6 +101,94 @@ 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<Tensor> {
+    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))
+}
+
+/// Run CUDA FlashAttention v2 for Q/K/V tensors in `(batch, heads, seq,
+/// head_dim)` layout. Returns `None` on non-CUDA devices so callers can retain
+/// their portable eager fallback.
+pub fn flash_attention(
+    q: &Tensor,
+    k: &Tensor,
+    v: &Tensor,
+    scale: f64,
+    causal: bool,
+) -> Result<Option<Tensor>> {
+    #[cfg(feature = "cuda")]
+    if q.device().is_cuda() {
+        // 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 +209,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)?
@@ -612,6 +700,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::<f32>()?;
+        let grouped = grouped.flatten_all()?.to_vec1::<f32>()?;
+        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 +769,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;