feat(vl): accelerate HunyuanOCR 1.5 with DFlash#158
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the DFlash parallel draft model for HunyuanOCR 1.5, introducing custom CUDA kernels for fused operations, CUDA graph execution, and preallocated backing storage for the KV cache to optimize speculative decoding. The review feedback highlights several critical issues: a missing bounds check in replay_cuda_graph that could trigger a GPU Use-After-Free upon cache reallocation, incorrect non-causal attention in the fallback eager path when causal_mask is absent, and state-desynchronization bugs in both kv_trim.rs and dflash.rs where resetting non-reusable storage fails to clear the logical cache lengths.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e2973d376c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 338950f867
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces speculative decoding support for HunyuanOCR 1.5 using the DFlash parallel draft model, alongside several performance optimizations. Key changes include the addition of the dflash module, custom CUDA kernels for fused operations (such as fused SiLU-multiplication, XDRoPE, and RMSNorm), and CUDA graph capture/replay for accelerated decoding. Additionally, TrimmableKvCache has been refactored to use preallocated backing storage to avoid redundant copies, and the vision encoder has been optimized with Rayon-parallelized bilinear interpolation and cached position embeddings. Feedback on these changes suggests addressing a potential capacity growth issue in TrimmableKvCache and considering a sequential fallback for Rayon parallelization on small images to avoid scheduling overhead.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Avoid TrimmableKvCache preallocating its full configured capacity for every caller (only CUDA-graph paths need a fixed-size buffer up front, via initialize_storage); invalidate the target CUDA graph before a multi-image request reallocates its batch=1 KV storage out from under it; gate DFlash graph capture to BF16 since the append kernel rejects other dtypes; stop silently discarding an OAR_VL_DTYPE=f32 override in CUDA attention; and fall back to sequential bilinear interpolation for small position-embedding grids.
GreatV
force-pushed
the
agent/hunyuanocr-dflash
branch
from
0a9217d to
15f0ec4
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 15f0ec4d64
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 604de2b1f1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| assert!( | ||
| base >= MIN_DFLASH_COMPUTE_CAP, | ||
| "HunyuanOCR DFlash CUDA kernels require compute capability 8.0 or newer; got CUDA_COMPUTE_CAP={value:?}" | ||
| ); |
There was a problem hiding this comment.
When CUDA_COMPUTE_CAP is set for a headless/cross-machine CUDA build targeting a pre-Ampere GPU (for example 75 for Turing), this assertion aborts the entire oar-ocr-vl CUDA build even though the crate still has supported non-DFlash paths and select_dtype already 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 👍 / 👎.
| 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)?; |
There was a problem hiding this comment.
On CUDA GPUs that Candle can otherwise run in F16 but FlashAttention-2 does not support (for example Turing cards after select_dtype falls back from BF16), this helper still calls the FlashAttention kernel solely because the tensor dtype is F16/BF16. That makes ordinary Hunyuan vision/AR inference return a kernel error instead of using the eager attention path; gate this branch on a supported compute capability or convert unsupported-kernel failures into Ok(None).
Useful? React with 👍 / 👎.
1 participant
Summary
dflash/checkpoint layoutWhy
The initial native implementation spent substantial time on per-layer kernel launches, host/device synchronization, and contiguous draft-context attention. Profiling the official Python/vLLM implementation showed that its draft path relies on paged split-KV attention and extensive fusion. This change brings those execution strategies to the Candle implementation while keeping target-model verification authoritative.
Performance
RTX 4090, OmniDocBench v1.5 first 30 pages,
max_tokens=8000, natural EOS, concurrency 1:The official run benefits from an exact first-image multimodal-cache warm-up. Excluding that page, the remaining runtime gap is about 9.2%. Representative medium and long-page token fingerprints remained stable across the optimizations; the long-page runtime is 7.71s versus 7.60s for the official implementation.
Validation
cargo fmt --all -- --checkcargo check -p oar-ocr-vl --libcargo check -p oar-ocr-vl --lib --features cudacargo test -p oar-ocr-vl --lib hunyuanocr(5 passed)cargo test -p oar-ocr-vl --lib kv_trim(12 passed)git diff --check