Skip to content

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

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.

Sign up for GitHub

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

Open
AliOsm wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion src/arch/canary/canary.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
2 changes: 1 addition & 1 deletion src/arch/canary/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -538,7 +538,7 @@ transcribe_status init_context(transcribe_model * model,
auto * cm = static_cast<CanaryModel *>(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);
Expand Down
3 changes: 2 additions & 1 deletion src/arch/cohere/cohere.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,7 +170,8 @@ struct CohereSession final : public transcribe_session {
std::vector<float> 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;
Expand Down
5 changes: 3 additions & 2 deletions src/arch/cohere/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<CohereModel *>(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);
Expand Down
34 changes: 12 additions & 22 deletions src/arch/parakeet/encoder.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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{};
Expand Down Expand Up @@ -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
Expand Down
4 changes: 4 additions & 0 deletions src/arch/parakeet/encoder.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
//
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
18 changes: 12 additions & 6 deletions src/arch/parakeet/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -1138,9 +1143,10 @@ transcribe_status run_one_shot_inner(ParakeetSession * pc,
}
const int mt_keep_frames = static_cast<int>(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
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
}
Expand Down
4 changes: 2 additions & 2 deletions src/arch/parakeet/multitalker.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
9 changes: 9 additions & 0 deletions src/arch/parakeet/parakeet.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
30 changes: 17 additions & 13 deletions src/arch/sortformer/model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -872,7 +876,7 @@ transcribe_status run_diar_streaming_core(DiarStreamScratch & sc,
sc.concat_host.begin() + static_cast<size_t>(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;
Expand Down Expand Up @@ -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;
}
Expand Down
5 changes: 5 additions & 0 deletions src/arch/sortformer/sortformer.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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,
Expand Down
5 changes: 5 additions & 0 deletions src/transcribe-flash-policy.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Loading