-
Notifications
You must be signed in to change notification settings - Fork 21
feat(vl): accelerate HunyuanOCR 1.5 with DFlash #158
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit
Hold shift + click to select a range
e2973d3
feat(vl): accelerate HunyuanOCR with DFlash
GreatV 338950f
fix(vl): address HunyuanOCR DFlash review feedback
GreatV 15f0ec4
fix(vl): address second round of HunyuanOCR DFlash review feedback
GreatV 604de2b
fix(vl): skip target CUDA graph for F32
GreatV File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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::<u32>().ok()?; | ||
| if base < 20 { | ||
| base *= 10; | ||
| } | ||
| Some((format!("{base}{suffix}"), base)) | ||
| } | ||
|
|
||
| fn detect_local_compute_cap() -> Option<u32> { | ||
| 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) | ||
| ); | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
When
CUDA_COMPUTE_CAPis set for a headless/cross-machine CUDA build targeting a pre-Ampere GPU (for example75for Turing), this assertion aborts the entireoar-ocr-vlCUDA build even though the crate still has supported non-DFlash paths andselect_dtypealready falls back away from BF16 on older devices. The autodetect branch below merely warns and emits compute_80 PTX, so explicit older targets should not be fatal unless the user actually tries to load the DFlash-only kernels; otherwise users building PaddleOCR/GLM/Hunyuan AR for older CUDA GPUs are blocked at build time.Useful? React with 👍 / 👎.