diff --git a/src/arch/canary/canary.h b/src/arch/canary/canary.h index 7b0e9e7e..fc25fb73 100644 --- a/src/arch/canary/canary.h +++ b/src/arch/canary/canary.h @@ -135,7 +135,8 @@ struct CanarySession final : public transcribe_session { // Per-stage flash-attention controls. ggml's Metal flash kernel lacks a // path for some encoder rel-pos MHSA dk values (96 for 180m-flash, 128 - // otherwise), so encoder flash defaults off on Metal; decoder defaults on. + // otherwise), and CUDA/ROCm can't take the per-head rel-pos mask, so + // encoder flash defaults off on those backends; decoder defaults on. bool encoder_use_flash = true; bool decoder_use_flash = true; diff --git a/src/arch/canary/model.cpp b/src/arch/canary/model.cpp index cdaf3ac6..245e7038 100644 --- a/src/arch/canary/model.cpp +++ b/src/arch/canary/model.cpp @@ -538,7 +538,7 @@ transcribe_status init_context(transcribe_model * model, auto * cm = static_cast(model); const bool is_metal = (cm->plan.primary_kind == transcribe::BackendKind::Metal); - cc->encoder_use_flash = !is_metal; + cc->encoder_use_flash = !is_metal && transcribe::flash::per_head_mask_supported(cm->plan.primary_kind); cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); diff --git a/src/arch/cohere/cohere.h b/src/arch/cohere/cohere.h index 09f1449e..356eb671 100644 --- a/src/arch/cohere/cohere.h +++ b/src/arch/cohere/cohere.h @@ -170,7 +170,8 @@ struct CohereSession final : public transcribe_session { std::vector enc_host; // Flash-attn is per-stage: encoder dk=160 has no Metal flash_attn_ext - // kernel (default OFF on Metal; manual path ties anyway); decoder + // kernel (default OFF on Metal; manual path ties anyway), and CUDA/ROCm + // can't take its per-head rel-pos mask (default OFF there too); decoder // dk=128 works everywhere (default ON). TRANSCRIBE_NO_FLASH / // TRANSCRIBE_FORCE_FLASH apply to both stages at once. bool encoder_use_flash = true; diff --git a/src/arch/cohere/model.cpp b/src/arch/cohere/model.cpp index e7f9e054..ed332152 100644 --- a/src/arch/cohere/model.cpp +++ b/src/arch/cohere/model.cpp @@ -635,11 +635,12 @@ transcribe_status init_context(transcribe_model * model, cc->n_ctx = transcribe_session_params_n_ctx(params); // Flash-attention policy — see CohereSession (cohere.h). Encoder - // auto-disables on Metal; decoder defaults on; env overrides apply globally. + // auto-disables on Metal, CUDA and ROCm; decoder defaults on; env + // overrides apply globally. auto * cm = static_cast(model); const bool is_metal = (cm->plan.primary_kind == transcribe::BackendKind::Metal); - cc->encoder_use_flash = !is_metal; + cc->encoder_use_flash = !is_metal && transcribe::flash::per_head_mask_supported(cm->plan.primary_kind); cc->decoder_use_flash = true; transcribe::flash::apply_env_overrides(cc->encoder_use_flash, cc->decoder_use_flash); diff --git a/src/arch/parakeet/encoder.cpp b/src/arch/parakeet/encoder.cpp index ba0e6253..ab2d659f 100644 --- a/src/arch/parakeet/encoder.cpp +++ b/src/arch/parakeet/encoder.cpp @@ -25,7 +25,6 @@ #include "conformer/conformer.h" #include "ggml.h" #include "transcribe-debug.h" -#include "transcribe-flash-policy.h" #include "transcribe-log.h" #include "weights.h" @@ -316,6 +315,7 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, const ParakeetHParams & hp, int n_mel_frames, ggml_type kv_type, + bool use_flash, const char * backend_name, const BufferedStreamMaskOverride * buf_mask, int n_batch, @@ -589,16 +589,11 @@ EncoderBuild build_encoder_graph(ggml_context * ctx, } conf::BlockParams bparams{}; - bparams.d_model = hp.enc_d_model; - bparams.n_head = hp.enc_n_heads; - bparams.conv_kernel = hp.enc_conv_kernel; - bparams.kv_type = kv_type; - // Flash attention is the default. TRANSCRIBE_NO_FLASH selects the F32 - // mul_mat + soft_max path used by CPU tensor validation; flash casts - // the relative-position mask to F16. Parakeet has no attention decoder. - bool enc_use_flash = true, dec_use_flash = false; - transcribe::flash::apply_env_overrides(enc_use_flash, dec_use_flash); - bparams.use_flash = enc_use_flash; + bparams.d_model = hp.enc_d_model; + bparams.n_head = hp.enc_n_heads; + bparams.conv_kernel = hp.enc_conv_kernel; + bparams.kv_type = kv_type; + bparams.use_flash = use_flash; bparams.policy = policy; bparams.att_context_left = band_covers_all ? -1 : hp.enc_att_context_left; bparams.att_context_right = band_covers_all ? -1 : hp.enc_att_context_right; @@ -775,6 +770,7 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * ctx, int drop_extra_pre_encoded, StreamingEncoderCacheIO & cache_io, ggml_type kv_type, + bool use_flash, const char * backend_name, bool spk_supervision) { conf::ConvPolicy policy{}; @@ -916,17 +912,11 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * ctx, } conf::BlockParams bparams{}; - bparams.d_model = hp.enc_d_model; - bparams.n_head = hp.enc_n_heads; - bparams.conv_kernel = hp.enc_conv_kernel; - bparams.kv_type = kv_type; - // Flash-attention opt-out (TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH). - // On CPU at streaming geometry the flash kernel is dramatically slower - // (~33 ms/layer vs ~1 ms on a Cortex-A55). Parakeet has no attention - // decoder, so the decoder flag is a throwaway. - bool enc_use_flash = true, dec_use_flash = false; - transcribe::flash::apply_env_overrides(enc_use_flash, dec_use_flash); - bparams.use_flash = enc_use_flash; + bparams.d_model = hp.enc_d_model; + bparams.n_head = hp.enc_n_heads; + bparams.conv_kernel = hp.enc_conv_kernel; + bparams.kv_type = kv_type; + bparams.use_flash = use_flash; bparams.policy = policy; // No att_context_left/right: ChunkedLimited derives the band entirely // from attn_chunked_mask (build_conformer_block reads them only on the diff --git a/src/arch/parakeet/encoder.h b/src/arch/parakeet/encoder.h index e77c8c46..79f9470c 100644 --- a/src/arch/parakeet/encoder.h +++ b/src/arch/parakeet/encoder.h @@ -193,6 +193,8 @@ struct EncoderBuild { // count) the struct has nullptr fields and a diagnostic is logged. // kv_type: GGML type for K/V activations in flash attention; // GGML_TYPE_COUNT means "auto" (f16 for quantized weights, f32 for f32). +// use_flash: flash attention in the conformer blocks; callers pass their +// model's encoder_use_flash (resolved at load). // backend_name: primary backend name ("MTL0", "Vulkan0", "CPU") for // auto-detecting conv strategy; env vars override. // @@ -221,6 +223,7 @@ EncoderBuild build_encoder_graph(ggml_context * compute_ctx, const ParakeetHParams & hp, int n_mel_frames, ggml_type kv_type = GGML_TYPE_COUNT, + bool use_flash = true, const char * backend_name = "", const BufferedStreamMaskOverride * buf_mask = nullptr, int n_batch = 1, @@ -334,6 +337,7 @@ EncoderBuild build_encoder_graph_streaming(ggml_context * compute_ctx int drop_extra_pre_encoded, StreamingEncoderCacheIO & cache_io, ggml_type kv_type = GGML_TYPE_COUNT, + bool use_flash = true, const char * backend_name = "", // Multitalker streaming pass: expose the // per-chunk speaker/background supervision diff --git a/src/arch/parakeet/model.cpp b/src/arch/parakeet/model.cpp index a877781e..87890547 100644 --- a/src/arch/parakeet/model.cpp +++ b/src/arch/parakeet/model.cpp @@ -23,6 +23,7 @@ #include "transcribe-batch-util.h" #include "transcribe-debug.h" #include "transcribe-env.h" +#include "transcribe-flash-policy.h" #include "transcribe-load-common.h" #include "transcribe-loader.h" #include "transcribe-log.h" @@ -536,6 +537,10 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par m->backend = ggml_backend_name(m->plan.primary); m->primary_backend = m->plan.primary; + bool dec_use_flash = false; // no attention decoder + m->encoder_use_flash = transcribe::flash::per_head_mask_supported(m->plan.primary_kind); + transcribe::flash::apply_env_overrides(m->encoder_use_flash, dec_use_flash); + if (const transcribe_status st = transcribe::load_common::load_weights( loader.path(), gguf_data, m->ctx_meta, m->plan, params, "parakeet", m->weight_storage); st != TRANSCRIBE_OK) { @@ -1138,9 +1143,10 @@ transcribe_status run_one_shot_inner(ParakeetSession * pc, } const int mt_keep_frames = static_cast(mt_keep.size()); const bool kestrel_masks = pm->hparams.kestrel_length_masking; - EncoderBuild eb = build_encoder_graph(pc->compute_ctx, pm->weights, pm->hparams, mel_n_frames, resolved_kv, - pm->backend.c_str(), /*buf_mask=*/nullptr, /*n_batch=*/1, - /*batch_var_len=*/kestrel_masks, spk_supervision, mt_keep_frames); + EncoderBuild eb = + build_encoder_graph(pc->compute_ctx, pm->weights, pm->hparams, mel_n_frames, resolved_kv, pm->encoder_use_flash, + pm->backend.c_str(), /*buf_mask=*/nullptr, /*n_batch=*/1, + /*batch_var_len=*/kestrel_masks, spk_supervision, mt_keep_frames); if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; // build_encoder_graph logged } @@ -1852,7 +1858,7 @@ static transcribe_status run_batch_encode(ParakeetSession * } EncoderBuild eb = build_encoder_graph(pc->compute_ctx, pm->weights, pm->hparams, T_max, resolved_kv, - pm->backend.c_str(), /*buf_mask=*/nullptr, + pm->encoder_use_flash, pm->backend.c_str(), /*buf_mask=*/nullptr, /*n_batch=*/n, /*batch_var_len=*/var_len); if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; @@ -2569,7 +2575,7 @@ transcribe_status emit_streaming_chunk(ParakeetSession * pc, } eb = build_encoder_graph_streaming(pc->compute_ctx, pm->weights, hp, n_mel_chunk_frames, drop_extra_pre_encoded, - cache_io, resolved_kv, pm->backend.c_str(), + cache_io, resolved_kv, pm->encoder_use_flash, pm->backend.c_str(), /*spk_supervision=*/mt_supervised); if (eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; @@ -2986,7 +2992,7 @@ transcribe_status emit_buffered_chunk(ParakeetSession * pc, buf_mask.valid_frames = effective_T; EncoderBuild eb = build_encoder_graph(pc->compute_ctx, pm->weights, pm->hparams, mel_n_frames, resolved_kv, - pm->backend.c_str(), &buf_mask); + pm->encoder_use_flash, pm->backend.c_str(), &buf_mask); if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; } diff --git a/src/arch/parakeet/multitalker.cpp b/src/arch/parakeet/multitalker.cpp index 0c5b1812..d55c00a6 100644 --- a/src/arch/parakeet/multitalker.cpp +++ b/src/arch/parakeet/multitalker.cpp @@ -411,8 +411,8 @@ transcribe_status run_multitalker(ParakeetSession * pc, const int64_t t_diar_start = ggml_time_us(); transcribe::debug::push_name_prefix("mt.diar"); const transcribe_status dst = sf::run_diar_streaming_core( - pc->diar_scratch, diar.hp, diar.conformer_hp, diar.conformer, diar.weights, pm->backend.c_str(), pc->sched, - pc->n_threads, P, diar_mel.data(), dm_mels, dm_frames, pc); + pc->diar_scratch, diar.hp, diar.conformer_hp, diar.conformer, diar.weights, pm->backend.c_str(), + pm->encoder_use_flash, pc->sched, pc->n_threads, P, diar_mel.data(), dm_mels, dm_frames, pc); transcribe::debug::pop_name_prefix(); if (dst != TRANSCRIBE_OK) { return dst; diff --git a/src/arch/parakeet/parakeet.h b/src/arch/parakeet/parakeet.h index 30ae2743..e2791d7c 100644 --- a/src/arch/parakeet/parakeet.h +++ b/src/arch/parakeet/parakeet.h @@ -108,6 +108,15 @@ struct ParakeetModel final : public transcribe_model { transcribe::BackendPlan plan; transcribe::load_common::WeightStorage weight_storage; + // Encoder flash attention, resolved at load() from the primary backend + // (off on CUDA/ROCm, see transcribe::flash::per_head_mask_supported) + // and the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH overrides. + // NO_FLASH selects the F32 mul_mat + soft_max path used by CPU tensor + // validation (flash casts the relative-position mask to F16). On CPU at + // streaming geometry the flash kernel is much slower (~33 ms/layer vs + // ~1 ms on a Cortex-A55). + bool encoder_use_flash = true; + // Fused BN parameters live in a separate ggml context + buffer, // computed at load time from the raw BN tensors. Freed in dtor. ggml_context * bn_fused_ctx = nullptr; diff --git a/src/arch/sortformer/model.cpp b/src/arch/sortformer/model.cpp index 09eb4455..b1488fd8 100644 --- a/src/arch/sortformer/model.cpp +++ b/src/arch/sortformer/model.cpp @@ -463,6 +463,7 @@ StreamInferBuild build_stream_infer_graph(ggml_context * ctx, const pk::ParakeetWeights & conformer, const SortformerWeights & w, const char * backend, + bool use_flash, int T_concat) { StreamInferBuild b{}; const int ed = chp.enc_d_model; @@ -479,13 +480,11 @@ StreamInferBuild build_stream_infer_graph(ggml_context * ctx, ggml_set_input(pos_emb_in); conf::BlockParams bp{}; - bp.d_model = ed; - bp.n_head = chp.enc_n_heads; - bp.conv_kernel = chp.enc_conv_kernel; - bp.kv_type = GGML_TYPE_F32; - bool enc_use_flash = true, dec_use_flash = false; - transcribe::flash::apply_env_overrides(enc_use_flash, dec_use_flash); - bp.use_flash = enc_use_flash; + bp.d_model = ed; + bp.n_head = chp.enc_n_heads; + bp.conv_kernel = chp.enc_conv_kernel; + bp.kv_type = GGML_TYPE_F32; + bp.use_flash = use_flash; bp.policy = sf_conv_policy(backend); bp.att_context_left = -1; // full (offline) attention over the concat bp.att_context_right = -1; @@ -588,6 +587,10 @@ transcribe_status load(Loader & loader, const transcribe_model_load_params * par m->backend = ggml_backend_name(m->plan.primary); m->primary_backend = m->plan.primary; + bool dec_use_flash = false; // no attention decoder + m->encoder_use_flash = transcribe::flash::per_head_mask_supported(m->plan.primary_kind); + transcribe::flash::apply_env_overrides(m->encoder_use_flash, dec_use_flash); + if (const transcribe_status st = transcribe::load_common::load_weights( loader.path(), gguf_data, m->ctx_meta, m->plan, params, "sortformer", m->weight_storage); st != TRANSCRIBE_OK) { @@ -670,8 +673,8 @@ static transcribe_status run_offline_forward(SortformerSession * pc, SortformerM } ggml_context * ctx = pc->compute_ctx; - pk::EncoderBuild eb = - pk::build_encoder_graph(ctx, pm->conformer, pm->conformer_hp, mel_n_frames, GGML_TYPE_F32, pm->backend.c_str()); + pk::EncoderBuild eb = pk::build_encoder_graph(ctx, pm->conformer, pm->conformer_hp, mel_n_frames, GGML_TYPE_F32, + pm->encoder_use_flash, pm->backend.c_str()); if (eb.mel_in == nullptr || eb.out == nullptr || eb.graph == nullptr) { return TRANSCRIBE_ERR_GGUF; } @@ -740,6 +743,7 @@ transcribe_status run_diar_streaming_core(DiarStreamScratch & sc, const pk::ParakeetWeights & conformer, const SortformerWeights & w, const char * backend, + bool use_flash, ggml_backend_sched_t sched, int n_threads, const SortformerStreamParams & P, @@ -872,7 +876,7 @@ transcribe_status run_diar_streaming_core(DiarStreamScratch & sc, sc.concat_host.begin() + static_cast(S + F) * ed); // ---- Graph B: xscale + 17 blocks + proj + 18 tf + head ---- - StreamInferBuild B = build_stream_infer_graph(ctx, hp, chp, conformer, w, backend, T_concat); + StreamInferBuild B = build_stream_infer_graph(ctx, hp, chp, conformer, w, backend, use_flash, T_concat); if (B.graph == nullptr) { transcribe::debug::pop_name_prefix(); return TRANSCRIBE_ERR_GGUF; @@ -923,9 +927,9 @@ static transcribe_status run_streaming(SortformerSession * pc, } const int64_t t_stream_start = ggml_time_us(); - const transcribe_status st = run_diar_streaming_core(pc->scratch, pm->hparams, pm->conformer_hp, pm->conformer, - pm->weights, pm->backend.c_str(), pc->sched, pc->n_threads, P, - pc->mel_buf.data(), mel_n_mels, mel_n_frames, pc); + const transcribe_status st = run_diar_streaming_core( + pc->scratch, pm->hparams, pm->conformer_hp, pm->conformer, pm->weights, pm->backend.c_str(), + pm->encoder_use_flash, pc->sched, pc->n_threads, P, pc->mel_buf.data(), mel_n_mels, mel_n_frames, pc); if (st != TRANSCRIBE_OK) { return st; } diff --git a/src/arch/sortformer/sortformer.h b/src/arch/sortformer/sortformer.h index a9cb6a20..7852f6d8 100644 --- a/src/arch/sortformer/sortformer.h +++ b/src/arch/sortformer/sortformer.h @@ -57,6 +57,10 @@ struct SortformerModel final : public transcribe_model { transcribe::BackendPlan plan; transcribe::load_common::WeightStorage weight_storage; + // Conformer flash attention, resolved at load() like + // ParakeetModel::encoder_use_flash. + bool encoder_use_flash = true; + // Fused conformer BN params (computed at load from the raw BN tensors). ggml_context * bn_fused_ctx = nullptr; ggml_backend_buffer_t bn_fused_buffer = nullptr; @@ -239,6 +243,7 @@ transcribe_status run_diar_streaming_core(DiarStreamScratch & const transcribe::parakeet::ParakeetWeights & conformer, const SortformerWeights & w, const char * backend, + bool use_flash, ggml_backend_sched_t sched, int n_threads, const SortformerStreamParams & P, diff --git a/src/transcribe-flash-policy.cpp b/src/transcribe-flash-policy.cpp index 0cb0ab0b..8cd0d2ac 100644 --- a/src/transcribe-flash-policy.cpp +++ b/src/transcribe-flash-policy.cpp @@ -4,10 +4,15 @@ #include "transcribe-flash-policy.h" +#include "transcribe-backend.h" #include "transcribe-env.h" namespace transcribe::flash { +bool per_head_mask_supported(BackendKind kind) { + return kind != BackendKind::Cuda && kind != BackendKind::Rocm; +} + void apply_env_overrides(bool & encoder_use_flash, bool & decoder_use_flash) { if (transcribe::env::flag("TRANSCRIBE_NO_FLASH")) { encoder_use_flash = false; diff --git a/src/transcribe-flash-policy.h b/src/transcribe-flash-policy.h index b0db8b96..7011d805 100644 --- a/src/transcribe-flash-policy.h +++ b/src/transcribe-flash-policy.h @@ -4,14 +4,27 @@ // // Flash-attention support in ggml is per-backend and per-head-dim, so the // default on/off decision is per-family (it depends on head dim and the -// backend it was loaded on). What lives here is the one cross-family -// concern: honoring the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH global -// debug overrides. +// backend it was loaded on). What lives here are the cross-family concerns: +// backend limits that every family with the same attention shape hits, and +// honoring the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH global debug +// overrides. #pragma once +namespace transcribe { +enum class BackendKind; // transcribe-backend.h +} + namespace transcribe::flash { +// Whether `kind`'s flash kernels take a mask with one slice per head, which +// is how the conformer's relative-position attention passes its position +// bias. ggml-cuda's kernels (ROCm builds the same ones) only take a mask +// shared by all heads, so on CUDA/ROCm the scheduler runs FLASH_ATTN_EXT on +// the CPU, with a device round trip in every block; the manual path stays on +// the GPU. +bool per_head_mask_supported(BackendKind kind); + // Apply the TRANSCRIBE_NO_FLASH / TRANSCRIBE_FORCE_FLASH environment // overrides to a pair of encoder/decoder flash flags. Either override forces // both flags in the same direction; if both are set, FORCE wins. The flags