perf(vl): accelerate GLM-OCR, MinerU, and PaddleOCR-VL

oar-ocr-vl/build.rs

@@ -1,6 +1,6 @@
 use std::process::Command;
 
-const MIN_DFLASH_COMPUTE_CAP: u32 = 80;
+const MIN_CUDA_COMPUTE_CAP: u32 = 80;
 
 fn parse_compute_cap(value: &str) -> Option<(String, u32)> {
     let value = value.trim().to_ascii_lowercase();
@@ -48,25 +48,25 @@ fn cuda_compute_arch() -> String {
             )
         });
         assert!(
-            base >= MIN_DFLASH_COMPUTE_CAP,
-            "HunyuanOCR DFlash CUDA kernels require compute capability 8.0 or newer; got CUDA_COMPUTE_CAP={value:?}"
+            base >= MIN_CUDA_COMPUTE_CAP,
+            "oar-ocr-vl CUDA kernels require compute capability 8.0 or newer; got CUDA_COMPUTE_CAP={value:?}"
         );
         return format!("compute_{arch}");
     }
 
     match detect_local_compute_cap() {
-        Some(base) if base >= MIN_DFLASH_COMPUTE_CAP => format!("compute_{base}"),
+        Some(base) if base >= MIN_CUDA_COMPUTE_CAP => format!("compute_{base}"),
         Some(base) => {
             println!(
-                "cargo:warning=detected GPU compute capability {base} is below the HunyuanOCR DFlash minimum; compiling forward-compatible compute_{MIN_DFLASH_COMPUTE_CAP} PTX"
+                "cargo:warning=detected GPU compute capability {base} is below the oar-ocr-vl CUDA kernel minimum; compiling forward-compatible compute_{MIN_CUDA_COMPUTE_CAP} PTX"
             );
-            format!("compute_{MIN_DFLASH_COMPUTE_CAP}")
+            format!("compute_{MIN_CUDA_COMPUTE_CAP}")
         }
         None => {
             println!(
-                "cargo:warning=could not detect a CUDA GPU; compiling compute_{MIN_DFLASH_COMPUTE_CAP} PTX (set CUDA_COMPUTE_CAP to override for cross/headless builds)"
+                "cargo:warning=could not detect a CUDA GPU; compiling compute_{MIN_CUDA_COMPUTE_CAP} PTX (set CUDA_COMPUTE_CAP to override for cross/headless builds)"
             );
-            format!("compute_{MIN_DFLASH_COMPUTE_CAP}")
+            format!("compute_{MIN_CUDA_COMPUTE_CAP}")
         }
     }
 }
@@ -93,18 +93,18 @@ fn main() {
             .args(["--ptx", "--std=c++17", "-O3"])
             .arg(format!("--gpu-architecture={cuda_arch}"))
             .arg("-o")
-            .arg(out_dir.join("hunyuan_dynamic_kv.ptx"))
+            .arg(out_dir.join("oar_vl_kernels.ptx"))
             .arg("src/hunyuanocr/dynamic_kv.cu")
             .output()
             .unwrap_or_else(|error| {
                 panic!(
-                    "failed to invoke {:?} for HunyuanOCR dynamic KV kernel; install the CUDA toolkit or set NVCC to the compiler path: {error}",
+                    "failed to invoke {:?} for oar-ocr-vl CUDA kernels; install the CUDA toolkit or set NVCC to the compiler path: {error}",
                     nvcc
                 )
             });
         if !output.status.success() {
             panic!(
-                "{:?} failed for HunyuanOCR dynamic KV kernel ({cuda_arch}):\n{}",
+                "{:?} failed for oar-ocr-vl CUDA kernels ({cuda_arch}):\n{}",
                 nvcc,
                 String::from_utf8_lossy(&output.stderr)
             );

oar-ocr-vl/examples/glmocr.rs

@@ -27,6 +27,7 @@ use tracing::{error, info};
 use oar_ocr_core::utils::load_image;
 use oar_ocr_vl::GlmOcr;
 use oar_ocr_vl::utils::parse_device;
+use utils::token_fingerprint;
 
 #[derive(Parser)]
 #[command(name = "glmocr")]
@@ -102,16 +103,18 @@ fn main() -> Result<(), Box<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)
             .into_iter()
             .next()
         {
-            Some(Ok(result)) => {
+            Some(Ok(tokens)) => {
                 info!(
-                    "  Inference time: {:.2}ms",
-                    infer_start.elapsed().as_secs_f64() * 1000.0
+                    "  Inference time: {:.2}ms, tokens: {}, fingerprint: {:016x}",
+                    infer_start.elapsed().as_secs_f64() * 1000.0,
+                    tokens.len(),
+                    token_fingerprint(&tokens)
                 );
-                println!("{}", result);
+                println!("{}", model.decode_tokens(&tokens)?);
             }
             Some(Err(e)) => error!("  Inference failed: {}", e),
             None => error!("  No result returned from model"),

oar-ocr-vl/examples/hunyuanocr.rs

@@ -29,6 +29,7 @@ use tracing::{error, info};
 use oar_ocr_core::utils::load_image;
 use oar_ocr_vl::HunyuanOcr;
 use oar_ocr_vl::utils::parse_device;
+use utils::token_fingerprint;
 
 #[derive(Parser)]
 #[command(name = "hunyuanocr")]
@@ -70,12 +71,6 @@ struct Args {
     benchmark: bool,
 }
 
-fn token_fingerprint(tokens: &[u32]) -> u64 {
-    tokens.iter().fold(0xcbf29ce484222325_u64, |hash, token| {
-        (hash ^ u64::from(*token)).wrapping_mul(0x100000001b3)
-    })
-}
-
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     utils::init_tracing();
     let args = Args::parse();

oar-ocr-vl/examples/mineru.rs

@@ -34,6 +34,7 @@ use oar_ocr_vl::utils::{convert_otsl_to_html, truncate_repetitive_content};
 use utils::mineru_layout::{
     ContentBlock, LAYOUT_IMAGE_SIZE, LAYOUT_PROMPT, parse_layout_output, prepare_for_extract,
 };
+use utils::token_fingerprint;
 
 #[derive(Parser)]
 #[command(name = "mineru")]
@@ -156,11 +157,17 @@ fn two_step_extract(
         LAYOUT_IMAGE_SIZE,
         imageops::FilterType::CatmullRom,
     );
-    let layout = model
-        .generate(&[layout_image], &[LAYOUT_PROMPT], max_tokens)
+    let layout_tokens = model
+        .generate_tokens(&[layout_image], &[LAYOUT_PROMPT], max_tokens)
         .into_iter()
         .next()
         .ok_or("Layout detection returned no result")??;
+    info!(
+        "  Layout tokens: {}, fingerprint: {:016x}",
+        layout_tokens.len(),
+        token_fingerprint(&layout_tokens)
+    );
+    let layout = model.decode_tokens(&layout_tokens)?;
 
     if dump_layout {
         info!("Layout raw output:\n{}", layout);
@@ -181,11 +188,18 @@ fn two_step_extract(
     for (i, (block_image, prompt)) in block_images.into_iter().zip(prompts.iter()).enumerate() {
         let idx = indices[i];
         let output = model
-            .generate(&[block_image], &[prompt], max_tokens)
+            .generate_tokens(&[block_image], &[prompt], max_tokens)
             .into_iter()
             .next();
         match output {
-            Some(Ok(content)) => {
+            Some(Ok(tokens)) => {
+                info!(
+                    "  Block {} tokens: {}, fingerprint: {:016x}",
+                    idx,
+                    tokens.len(),
+                    token_fingerprint(&tokens)
+                );
+                let content = model.decode_tokens(&tokens)?;
                 let cleaned = truncate_repetitive_content(&content, 10, 10, 10);
                 let content = if blocks[idx].block_type == "table" {
                     convert_otsl_to_html(&cleaned)

oar-ocr-vl/examples/paddleocr_vl.rs

@@ -60,6 +60,7 @@ use tracing::{error, info};
 use oar_ocr_core::utils::load_image;
 use oar_ocr_vl::utils::parse_device;
 use oar_ocr_vl::{PaddleOcrVl, PaddleOcrVlTask};
+use utils::token_fingerprint;
 
 /// Command-line arguments for the PaddleOCR-VL example
 #[derive(Parser)]
@@ -174,17 +175,19 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         // Run inference
         let infer_start = Instant::now();
         match model
-            .generate(&[rgb_img], &[task], args.max_tokens)
+            .generate_tokens(&[rgb_img], &[task], args.max_tokens)
             .pop()
             .unwrap()
         {
-            Ok(result) => {
+            Ok(tokens) => {
                 let infer_duration = infer_start.elapsed();
                 info!(
-                    "  Inference time: {:.2}ms",
-                    infer_duration.as_secs_f64() * 1000.0
+                    "  Inference time: {:.2}ms, tokens: {}, fingerprint: {:016x}",
+                    infer_duration.as_secs_f64() * 1000.0,
+                    tokens.len(),
+                    token_fingerprint(&tokens)
                 );
-                println!("{}", result);
+                println!("{}", model.decode_tokens(&tokens, task)?.1);
             }
             Err(e) => {
                 error!("  Inference failed: {}", e);

oar-ocr-vl/examples/utils/mod.rs

@@ -20,6 +20,14 @@ pub fn init_tracing() {
         .init();
 }
 
+/// Stable FNV-1a fingerprint used to compare generated token streams.
+#[allow(dead_code)]
+pub fn token_fingerprint(tokens: &[u32]) -> u64 {
+    tokens.iter().fold(0xcbf29ce484222325_u64, |hash, token| {
+        (hash ^ u64::from(*token)).wrapping_mul(0x100000001b3)
+    })
+}
+
 #[allow(dead_code)]
 pub fn print_preview(tag: &str, text: &str) {
     let preview: String = text.chars().take(400).collect();

oar-ocr-vl/src/attention.rs

@@ -21,6 +21,18 @@
 use crate::utils::candle_to_ocr_processing;
 use candle_core::{D, DType, Device, IndexOp, Result, Tensor};
 use oar_ocr_core::core::errors::OCRError;
+use std::sync::OnceLock;
+
+fn grouped_query_attention_disabled() -> bool {
+    static DISABLED: OnceLock<bool> = OnceLock::new();
+    *DISABLED.get_or_init(|| std::env::var_os("OAR_VL_DISABLE_GQA").is_some())
+}
+
+#[cfg(feature = "cuda")]
+fn flash_attention_disabled() -> bool {
+    static DISABLED: OnceLock<bool> = OnceLock::new();
+    *DISABLED.get_or_init(|| std::env::var_os("OAR_VL_DISABLE_FLASH_ATTN").is_some())
+}
 
 /// Helper function to handle Metal device computation.
 ///
@@ -116,6 +128,11 @@ pub fn scaled_dot_product_attention_gqa(
     if num_kv_groups == 1 {
         return scaled_dot_product_attention(q, k, v, mask, scale, is_causal);
     }
+    if grouped_query_attention_disabled() {
+        let k = repeat_kv(k, num_kv_groups)?;
+        let v = repeat_kv(v, num_kv_groups)?;
+        return scaled_dot_product_attention(q, &k, &v, mask, scale, is_causal);
+    }
 
     let (batch, num_heads, query_len, head_dim) = q.dims4()?;
     let (k_batch, num_kv_heads, kv_len, k_head_dim) = k.dims4()?;
@@ -181,7 +198,8 @@ pub fn flash_attention(
     causal: bool,
 ) -> Result<Option<Tensor>> {
     #[cfg(feature = "cuda")]
-    if q.device().is_cuda()
+    if !flash_attention_disabled()
+        && q.device().is_cuda()
         && flash_attention_dtype_supported(q.dtype())
         && k.dtype() == q.dtype()
         && v.dtype() == q.dtype()