Skip to content

flash: keep conformer encoder attention on the GPU on CUDA and ROCm - #233

Open
AliOsm wants to merge 1 commit into
handy-computer:mainfrom
AliOsm:encoder-flash-per-head-mask
Open

AliOsm wants to merge 1 commit into
handy-computer:mainfrom
AliOsm:encoder-flash-per-head-mask

Conversation

@AliOsm

@AliOsm AliOsm commented Oct 26, 2026

Copy link
Copy Markdown

Summary

On CUDA, the relative-position attention in the conformer encoders runs on the CPU. The conformer passes the position bias to ggml_flash_attn_ext as a mask with one slice per head (mask->ne[2] == n_head), and ggml-cuda's flash kernels only take a mask shared by all heads (fattn.cu: if (mask && mask->ne[2] != 1) return BEST_FATTN_KERNEL_NONE;). So the scheduler moves every FLASH_ATTN_EXT node to the CPU, with a round trip through host memory in every block. ROCm builds the same kernels, so it should behave the same way.

This PR adds transcribe::flash::per_head_mask_supported(BackendKind) to the shared flash policy, false on CUDA and ROCm, and uses it for the encoder flash default of the four families that build this attention: Cohere, Canary, Parakeet and Sortformer. On those two backends the encoders use the existing manual attention path (mul_mat + soft_max_ext + mul_mat), which stays on the GPU. Other backends keep their current defaults, and TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH work as before.

Scope

  • src/transcribe-flash-policy: the new helper. The header forward-declares BackendKind, so it doesn't pull ggml-backend.h into more files.
  • cohere, canary: encoder_use_flash at session init also checks the helper, next to the existing Metal check.
  • parakeet, sortformer: these hard-coded enc_use_flash = true (and read the env overrides) inside each graph builder, in three places. They now resolve the flag once at load from plan.primary_kind and the overrides, the same way Cohere and Canary do it per session, and pass it to build_encoder_graph, build_encoder_graph_streaming, build_stream_infer_graph and run_diar_streaming_core. One small behavior change: for these two families the env overrides are now read at load, not at every graph build.
  • The shared conformer code is not touched, so this is independent of conformer: make the GLU sigmoid input contiguous so CUDA keeps it on the GPU #228. The only CPU nodes left in these encoders after this PR are the SIGMOID ones that conformer: make the GLU sigmoid input contiguous so CUDA keeps it on the GPU #228 moves.

AI Assistance

Yes. As in my other PRs, I used an AI assistant to help build the code and the wording of this description. I reviewed the changes and ran the measurements below. I found this while profiling my project, which uses transcribe.cpp with Cohere on CUDA.

Validation

GTX 1660 SUPER (6 GB, sm_75), CUDA 12.9, Linux. Release builds of main (8facde6) and this branch. The machine was shared and busy (load average 15–50 during the timings), so timings are medians of runs in ABBA order. Main's flash fallback runs on the CPU, so its times grow with CPU load; the branch's encoder only uses the GPU.

Node placement with GGML_SCHED_DEBUG=2 (transcribe-cli --backend cuda):

                                          FLASH_ATTN on CPU    graph splits
                                          main -> branch       main -> branch
parakeet-tdt-0.6b-v3 Q8_0                 24 -> 0               97 -> 49
cohere-transcribe-arabic-07-2026 Q8_0     48 -> 0              196 -> 100
canary-180m-flash Q8_0                    17 -> 0               72 -> 38
diar_streaming_sortformer_4spk-v2.1 Q8_0  17 -> 0               70 -> 36

transcribe-bench --backend cuda --warmup 2 --iters 10, median of 20 iterations per build:

                                     encode ms
parakeet Q8_0, dots.wav (35.3 s)     1471 -> 440 (3.3x)
cohere Q8_0, ar-long.wav (25.7 s)    3049 -> 994 (3.1x)
canary Q8_0, jfk.wav (11.0 s)         251 ->  55 (4.5x)

Total time went from 3282 to 1277 ms for Cohere and from 311 to 120 ms for Canary. Parakeet's total is mostly its TDT decode on the CPU, which this PR doesn't change and which was too noisy under the load to compare.

On a quieter run (load around 12), parakeet's encoder on dots.wav took 784–1205 ms on main and 400–438 ms with the manual path (main with TRANSCRIBE_NO_FLASH=1, the same path). On 600 s of audio it took 12.2–13.7 s and 6.7–6.8 s.

Sortformer, transcribe-cli --repeat 5 on 120 s of audio: 31.5 and 32.0 s on main, 7.7 and 7.9 s on the branch.

GPU memory barely changes. Peak VRAM for parakeet (sampled with nvidia-smi) with the manual path was 1084, 2172, 3474 and 4800 MiB at 60 s, 306 s, 600 s and 900 s of audio, 8 MiB more than the flash path at each length. I measured this on main with TRANSCRIBE_NO_FLASH=1, which is the path this PR selects for parakeet on CUDA.

Accuracy on CUDA:

  • cohere: 500 human-transcribed Arabic clips (FLEURS, Common Voice, SADA, Casablanca and Quran, 79 minutes) with transcribe-cli --batch --batch-size 8. WER went from 21.31% to 21.23% and CER from 9.69% to 9.63%. The paired bootstrap 95% interval for the WER change is -0.37 to +0.22 points, so I don't see an accuracy change. 24 clips changed (5 better, 3 worse, 16 with the same error count). The run took 238 s on the branch; main's output is from my earlier run of 4d7c1cc (349.6 s) for conformer: make the GLU sigmoid input contiguous so CUDA keeps it on the GPU #228. The only change between 4d7c1cc and 8facde6 is in Whisper.
  • parakeet: identical transcripts on jfk, dots, dots-full, death, cj-swimming-drop and fleurs-en.
  • canary-180m-flash: identical transcripts on jfk, dots, cj-swimming-drop and fleurs-en.
  • sortformer: identical segments on sortformer-2spk-mix and multitalker-2spk-mix.

scripts/ci/clang-format.sh --check passes, and ctest passes (61 tests, including the diarize CLI smoke with TRANSCRIBE_SORTFORMER_GGUF set to the Q8_0 model).

Not tested: ROCm, and the larger Canary and Parakeet variants, which take the same code path. Metal, Vulkan and CPU are unchanged. My GPU is Turing; on Ampere and newer, ggml turns on TF32 for cuBLAS, so the F32 attention matmuls may run in TF32 there.

@AliOsm
AliOsm requested a review from cjpais as a code owner October 26, 2026 08:53
@AliOsm

AliOsm commented Oct 26, 2026

Copy link
Copy Markdown
Author

This is the last PR in the PRs series I found while using transcribe.cpp in my project.

@cjpais

cjpais commented Oct 26, 2026

Copy link
Copy Markdown
Contributor

Thank you. I appreciate it. I will pull it in when I have a chance. I'm just finishing some bigger tasks with handy right now, and after those are merged, I will come back and review these.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants