diff --git a/CMakeLists.txt b/CMakeLists.txt index a5abb3ba6..bf9092e22 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -353,6 +353,7 @@ add_library(engine_core OBJECT src/framework/assets/torch_bin.cpp src/framework/core/module.cpp src/framework/core/backend.cpp + src/framework/core/attention_fallback.cpp src/framework/core/deferred_tensor_writer.cpp src/framework/core/execution_context.cpp src/framework/core/host_memory.cpp @@ -2213,6 +2214,14 @@ if (ENGINE_BUILD_TESTS) COMMAND backend_device_resolution_test ) + add_engine_unittest(attention_fallback_test tests/unittests/test_attention_fallback.cpp) + target_include_directories(attention_fallback_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) + + add_test( + NAME attention_fallback_test + COMMAND attention_fallback_test + ) + add_executable(streaming_audio_input_test tests/unittests/test_streaming_audio_input.cpp app/streaming/pcm_source.cpp diff --git a/docs/models/breeze_tts.md b/docs/models/breeze_tts.md index 9aca60d01..584713bb8 100644 --- a/docs/models/breeze_tts.md +++ b/docs/models/breeze_tts.md @@ -67,6 +67,7 @@ audiocpp_cli \ | `--request-option top_p=` | `0..1` | `1.0` | Top-p sampling limit. | | `--request-option seed=` | integer >= 0 | `0` | Generation seed. | | `--session-option breeze_tts.reference_cache_slots=` | integer >= 0 | `1` | Prepared reference-audio cache slots. | +| `--session-option breeze_tts.attention=` | `auto`, `flash`, `eager` | `auto` | Attention kernel. `auto` uses flash except on Volta/Turing GPUs (e.g. V100), where it falls back to eager to avoid missing MMA kernels. | | `--session-option weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0`, `q4_0`, `q4_k` | `native` | Weight storage type; quantized types convert at load time from the BF16 package. | Quantized weight storage is the largest measured speedup and applies to CUDA @@ -76,4 +77,3 @@ and HIP alike: `q8_0` cut the fixed 100-token regression case from RTF ~1.5 to voice-design regression cases. Counter to intuition, fp32 is the one configuration known to be *worse* for this model (mispronunciations and runaway repetition), because the model is trained and tuned in bf16. - diff --git a/docs/tts.md b/docs/tts.md index 4b2910fdf..87af27914 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -559,6 +559,7 @@ python3 tools/model_manager_v2.py install --models-root models higgs_audio_tts_4 | `--top-k` | integer | `30` | AR top-k sampling limit. The narrower default is less prone to premature EOC than the Python client's `50`. | | `--top-p` | float | `0.8` | AR nucleus sampling limit. The Python client's unfiltered equivalent is `1.0`. | | `--repetition-penalty` | float | `1.1` | Accepted for Python API compatibility; Higgs audio-code sampling does not consume it. | +| `--session-option higgs_audio_tts.attention=` | `auto`, `flash`, `eager` | `auto` | Attention kernel. `auto` uses flash except on Volta/Turing GPUs (e.g. V100), where it falls back to eager to avoid missing MMA kernels. | ## Fish Audio S2 Pro diff --git a/include/engine/framework/core/attention_fallback.h b/include/engine/framework/core/attention_fallback.h new file mode 100644 index 000000000..0a238aadf --- /dev/null +++ b/include/engine/framework/core/attention_fallback.h @@ -0,0 +1,44 @@ +#pragma once + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include + +namespace engine::core { + +// Session-option vocabulary shared by families that lower attention with +// ggml_flash_attn_ext: "auto" (default), "flash", or "eager". +// +// Background: ggml-cuda only instantiates the MMA/wmma flash-attention kernels +// for compute capability >= 8.0 (Ampere). On older GPUs such as Volta/sm70 +// (e.g. Tesla V100) a graph containing GGML_OP_FLASH_ATTN_EXT fails at compute +// time with "no device code compatible with CUDA arch 700", even when ggml was +// built with that arch enabled. The eager (explicit matmul + softmax) lowering +// computes the same operation with generic ops and runs everywhere (output +// logits may differ at ulp level, as with any kernel change). +enum class AttentionPreference { + Auto, + Flash, + Eager, +}; + +// Parses a ".attention" session-option value. Throws std::runtime_error +// naming option_name on invalid input. +AttentionPreference parse_attention_preference(const std::string & value, const char * option_name); + +// Resolves whether flash attention may be used for the given backend and head +// dimension. Flash forces true, Eager forces false, Auto gates on the CUDA +// device compute capability (Volta/Turing resolve to eager; see the .cpp for +// why supports_op cannot be used). A null or non-CUDA backend, or a device +// query failure, preserves historical behavior (true). +// +// Adopting in other families (currently wired for higgs_audio_tts and +// breeze_tts only): resolve once per runtime with the model's head_dim and +// the family's ".attention" session option, then switch the +// QwenDecoder prefill/static modes (or SDPA/GQA lowerings) between flash and +// their ManualRepeat/Explicit equivalents based on the result. +bool resolve_flash_attention(ggml_backend_t backend, int64_t head_dim, AttentionPreference preference); + +} // namespace engine::core diff --git a/include/engine/framework/modules/transformers/qwen_decoder.h b/include/engine/framework/modules/transformers/qwen_decoder.h index eaa80c12e..fa1c5ac99 100644 --- a/include/engine/framework/modules/transformers/qwen_decoder.h +++ b/include/engine/framework/modules/transformers/qwen_decoder.h @@ -77,6 +77,9 @@ struct QwenDecoderAttentionPolicy { QwenDecoderAttentionMode static_mode = QwenDecoderAttentionMode::FlashGrouped; QwenDecoderPrefixAttentionMode prefix_mode = QwenDecoderPrefixAttentionMode::Exact; int64_t grouped_query_min_steps = 0; + // False routes flash branches through repeat-KV + matmul/softmax for GPUs + // without a flash kernel (e.g. CUDA sm70). True preserves historical behavior. + bool allow_flash_attention = true; }; struct QwenDecoderStaticCachePolicy { diff --git a/include/engine/models/breeze_tts/generator.h b/include/engine/models/breeze_tts/generator.h index 7b3a95ff7..7156ea6b5 100644 --- a/include/engine/models/breeze_tts/generator.h +++ b/include/engine/models/breeze_tts/generator.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/runtime/session.h" #include "engine/models/breeze_tts/assets.h" #include "engine/models/breeze_tts/speech_decoder.h" @@ -35,7 +36,8 @@ class BreezeGeneratorRuntime { engine::core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - engine::assets::TensorStorageType storage_type); + engine::assets::TensorStorageType storage_type, + engine::core::AttentionPreference attention_preference = engine::core::AttentionPreference::Auto); ~BreezeGeneratorRuntime(); engine::runtime::AudioBuffer generate(const BreezeGenerationRequest & request); diff --git a/include/engine/models/breeze_tts/speech_decoder.h b/include/engine/models/breeze_tts/speech_decoder.h index bb27bc33d..a5c269c6d 100644 --- a/include/engine/models/breeze_tts/speech_decoder.h +++ b/include/engine/models/breeze_tts/speech_decoder.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/runtime/session.h" @@ -36,7 +37,8 @@ class BreezeSpeechDecoderRuntime { size_t graph_arena_bytes, size_t constant_context_bytes, engine::assets::TensorStorageType linear_weight_storage_type, - engine::assets::TensorStorageType conv_weight_storage_type); + engine::assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto); ~BreezeSpeechDecoderRuntime(); runtime::AudioBuffer decode(const BreezeSpeechCodes & codec_codes) const; @@ -50,6 +52,7 @@ class BreezeSpeechDecoderRuntime { core::ExecutionContext * execution_context_ = nullptr; std::shared_ptr weights_; size_t graph_arena_bytes_ = 0; + bool allow_flash_attention_ = true; std::unique_ptr constants_; mutable std::unique_ptr graph_; // Always present to keep this public class layout identical when the private diff --git a/include/engine/models/breeze_tts/speech_encoder.h b/include/engine/models/breeze_tts/speech_encoder.h index f2dbae6c4..818a8cb2a 100644 --- a/include/engine/models/breeze_tts/speech_encoder.h +++ b/include/engine/models/breeze_tts/speech_encoder.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/assets/tensor_source.h" #include "engine/framework/runtime/session.h" @@ -35,7 +36,8 @@ class BreezeSpeechEncoderRuntime { core::ExecutionContext & execution_context, size_t graph_arena_bytes, engine::assets::TensorStorageType linear_weight_storage_type, - engine::assets::TensorStorageType conv_weight_storage_type); + engine::assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto); ~BreezeSpeechEncoderRuntime(); BreezeSpeechCodes encode(const runtime::AudioBuffer & audio) const; @@ -46,6 +48,7 @@ class BreezeSpeechEncoderRuntime { std::shared_ptr weights_; core::ExecutionContext * execution_context_ = nullptr; size_t graph_arena_bytes_ = 0; + bool allow_flash_attention_ = true; std::unique_ptr constants_; mutable std::unique_ptr conv_graph_; mutable std::unique_ptr transformer_graph_; diff --git a/include/engine/models/higgs_audio_tts/ar.h b/include/engine/models/higgs_audio_tts/ar.h index 77549e596..d91818f81 100644 --- a/include/engine/models/higgs_audio_tts/ar.h +++ b/include/engine/models/higgs_audio_tts/ar.h @@ -1,6 +1,7 @@ #pragma once #include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/core/execution_context.h" #include "engine/framework/core/module.h" #include "engine/framework/modules/transformers/qwen_decoder.h" @@ -44,7 +45,8 @@ class HiggsARRuntime { std::shared_ptr assets, core::ExecutionContext & execution, size_t weight_context_bytes, - assets::TensorStorageType weight_storage_type); + assets::TensorStorageType weight_storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto); const HiggsAssets & assets() const noexcept; const HiggsARWeights & weights() const noexcept; @@ -52,6 +54,7 @@ class HiggsARRuntime { core::BackendType backend_type() const noexcept; int device() const noexcept; int threads() const noexcept; + bool allow_flash_attention() const noexcept; private: std::shared_ptr assets_; @@ -59,6 +62,7 @@ class HiggsARRuntime { core::BackendType backend_type_ = core::BackendType::Cpu; int device_ = 0; int threads_ = 1; + bool allow_flash_attention_ = true; std::shared_ptr weights_; }; diff --git a/model_specs/breeze_tts.json b/model_specs/breeze_tts.json index 95a0a2cc5..d94e62cbf 100644 --- a/model_specs/breeze_tts.json +++ b/model_specs/breeze_tts.json @@ -161,6 +161,14 @@ "required": false, "min": 0, "default": 1 + }, + { + "name": "attention", + "type": "enum", + "description": "Attention lowering; auto probes the backend and falls back to eager on GPUs without a flash kernel (e.g. sm70); default auto.", + "required": false, + "values": ["auto", "flash", "eager"], + "default": "auto" } ], "load": [] diff --git a/src/framework/core/attention_fallback.cpp b/src/framework/core/attention_fallback.cpp new file mode 100644 index 000000000..e53d69c3b --- /dev/null +++ b/src/framework/core/attention_fallback.cpp @@ -0,0 +1,121 @@ +#include "engine/framework/core/attention_fallback.h" + +#include +#include +#include +#include +#include +#include + +#include "ggml.h" +#include "ggml-backend.h" + +#ifdef GGML_USE_CUDA +// CUDA driver API, declared manually so this translation unit needs neither +// the CUDA headers on its include path nor any CMake changes. The driver +// library is already linked transitively through ggml-cuda. Attribute ids +// CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR=75 / MINOR=76 are stable ABI. +extern "C" { +typedef int kCcProbeCuDevice; +typedef int kCcProbeCuResult; +kCcProbeCuResult cuDeviceGet(kCcProbeCuDevice * device, int ordinal); +kCcProbeCuResult cuDeviceGetAttribute(int * value, int attrib, kCcProbeCuDevice device); +} +#endif // GGML_USE_CUDA + +namespace engine::core { +namespace { + +std::string to_lower(std::string value) { + std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { + return static_cast(std::tolower(c)); + }); + return value; +} + +AttentionPreference parse_preference_value(const std::string & value, const char * option_name) { + const std::string lowered = to_lower(value); + if (lowered == "auto") { + return AttentionPreference::Auto; + } + if (lowered == "flash" || lowered == "on" || lowered == "1") { + return AttentionPreference::Flash; + } + if (lowered == "eager" || lowered == "off" || lowered == "0") { + return AttentionPreference::Eager; + } + throw std::runtime_error( + std::string(option_name) + " must be 'auto', 'flash', or 'eager' (got '" + value + "')"); +} + +// Auto-resolution for the CUDA flash-attention path. +// +// ggml_backend_supports_op() cannot be used here: on Volta it returns true +// (the MMA kernel is "selected") yet large prefill shapes die at launch with +// "flash_attn_ext_f16 has no device code compatible with CUDA arch 700". +// Instead, gate on compute capability, mirroring the kernel guards in +// ggml-cuda: flash below 700 (only generic TILE/VEC kernels exist) and at or +// above 800 (MMA fully instantiated); eager on 700-800, where large shapes +// select the MMA kernel with no usable device code. Unknown backends and +// query failures fail OPEN to preserve current behavior. +bool cuda_device_wants_eager(ggml_backend_t backend) { +#ifdef GGML_USE_CUDA + if (backend == nullptr) { + return false; + } + ggml_backend_dev_t device = ggml_backend_get_device(backend); + if (device == nullptr) { + return false; + } + if (ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_GPU) { + return false; + } + const char * name = ggml_backend_dev_name(device); + if (name == nullptr || std::strncmp(name, "CUDA", 4) != 0) { + return false; // HIP / Vulkan / Metal / CPU: unchanged behavior. + } + char * end = nullptr; + const long ordinal = std::strtol(name + 4, &end, 10); + if (end == name + 4 || ordinal < 0) { + return false; + } + kCcProbeCuDevice cu_device = -1; + if (cuDeviceGet(&cu_device, static_cast(ordinal)) != 0) { + return false; + } + int major = 0; + int minor = 0; + if (cuDeviceGetAttribute(&major, 75 /* COMPUTE_CAPABILITY_MAJOR */, cu_device) != 0) { + return false; + } + if (cuDeviceGetAttribute(&minor, 76 /* COMPUTE_CAPABILITY_MINOR */, cu_device) != 0) { + return false; + } + const int cc = major * 100 + minor * 10; + return cc >= 700 && cc < 800; +#else + (void) backend; + return false; +#endif // GGML_USE_CUDA +} + +} // namespace + +AttentionPreference parse_attention_preference(const std::string & value, const char * option_name) { + return parse_preference_value(value, option_name != nullptr ? option_name : "attention"); +} + +bool resolve_flash_attention(ggml_backend_t backend, int64_t head_dim, AttentionPreference preference) { + (void) head_dim; + switch (preference) { + case AttentionPreference::Flash: + return true; + case AttentionPreference::Eager: + return false; + case AttentionPreference::Auto: + break; + } + return !cuda_device_wants_eager(backend); +} + +} // namespace engine::core diff --git a/src/framework/modules/transformers/qwen_decoder.cpp b/src/framework/modules/transformers/qwen_decoder.cpp index d7031270b..b40fc98cf 100644 --- a/src/framework/modules/transformers/qwen_decoder.cpp +++ b/src/framework/modules/transformers/qwen_decoder.cpp @@ -255,6 +255,10 @@ struct QKVProjections { core::TensorValue v; }; +bool flash_branches_allowed(const QwenDecoderLayerConfig & config) { + return config.runtime.attention.allow_flash_attention; +} + QKVProjections build_qkv_projections( core::ModuleBuildContext & ctx, const core::TensorValue & input, @@ -544,13 +548,18 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( v = core::ensure_backend_addressable_layout(ctx, v); auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); + const bool allow_flash = flash_branches_allowed(config_); const bool use_prefix_flash = + allow_flash && prefix_key.has_value() && config_.runtime.attention.prefix_mode == QwenDecoderPrefixAttentionMode::FlashWithPrefix && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV; core::TensorValue all_k = k; core::TensorValue all_v = v; - if (use_prefix_flash) { + // Cached prefix KV may be stored in a different dtype than the current + // K/V (e.g. Higgs reference state); cast before concat on every path. + // (The eager branch previously skipped this and died in ggml_concat.) + if (prefix_key.has_value()) { auto attention_prefix_key = prefix_key; auto attention_prefix_value = prefix_value; if (attention_prefix_key->type != k.type) { @@ -567,12 +576,9 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( } all_k = ConcatModule({1}).build(ctx, *attention_prefix_key, k); all_v = ConcatModule({1}).build(ctx, *attention_prefix_value, v); - } else if (prefix_key.has_value()) { - all_k = ConcatModule({1}).build(ctx, *prefix_key, k); - all_v = ConcatModule({1}).build(ctx, *prefix_value, v); } core::TensorValue context; - if (!prefix_key.has_value() && attention_mask.has_value() && + if (allow_flash && !prefix_key.has_value() && attention_mask.has_value() && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); @@ -585,10 +591,10 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( dim, *attention_mask, config_.attention_precision); - } else if (attention_mask.has_value() && - ((!prefix_key.has_value() && - config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) || - use_prefix_flash)) { + } else if (allow_flash && attention_mask.has_value() && + ((!prefix_key.has_value() && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) || + use_prefix_flash)) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); auto v_heads = TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v); @@ -738,6 +744,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( auto k_heads = TransposeModule({{0, 2, 1, 3}, attention_key_cache.shape.rank}).build(ctx, attention_key_cache); auto v_heads = TransposeModule({{0, 2, 1, 3}, attention_value_cache.shape.rank}).build(ctx, attention_value_cache); core::TensorValue context; + const bool allow_flash = flash_branches_allowed(config_); const bool use_grouped_query = config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery && config_.runtime.attention.grouped_query_min_steps > 0 && @@ -755,7 +762,8 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( config_.num_attention_heads, config_.num_key_value_heads, attention_mask); - } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || + } else if (!allow_flash || + config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); @@ -898,6 +906,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat auto k_heads = TransposeModule({{0, 2, 1, 3}, attention_key_cache.shape.rank}).build(ctx, attention_key_cache); auto v_heads = TransposeModule({{0, 2, 1, 3}, attention_value_cache.shape.rank}).build(ctx, attention_value_cache); core::TensorValue context; + const bool allow_flash = flash_branches_allowed(config_); const bool use_grouped_query = config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery && config_.runtime.attention.grouped_query_min_steps > 0 && @@ -915,7 +924,8 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail_bat config_.num_attention_heads, config_.num_key_value_heads, attention_mask); - } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || + } else if (!allow_flash || + config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeat || config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); diff --git a/src/models/breeze_tts/generator.cpp b/src/models/breeze_tts/generator.cpp index 28458f122..54b5e360b 100644 --- a/src/models/breeze_tts/generator.cpp +++ b/src/models/breeze_tts/generator.cpp @@ -82,7 +82,8 @@ modules::QwenDecoderActivationCastPolicy breeze_bf16_activation_policy(core::Bac modules::QwenCausalDecodeRuntimeConfig backbone_config( const BreezeTTSConfig & config, core::BackendType backend_type, - size_t graph_arena_bytes) { + size_t graph_arena_bytes, + bool allow_flash_attention = true) { modules::QwenCausalDecodeRuntimeConfig out; out.trace_name = "breeze_tts.backbone"; out.prefill_graph_arena_bytes = graph_arena_bytes; @@ -101,8 +102,15 @@ modules::QwenCausalDecodeRuntimeConfig backbone_config( out.decoder.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.decoder.stack.attention_precision = GGML_PREC_F32; out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; - out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + // Eager graph for GPUs without a flash kernel (e.g. sm70). + out.decoder.stack.runtime.attention.allow_flash_attention = allow_flash_attention; + if (allow_flash_attention) { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + } else { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + } out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || @@ -129,7 +137,8 @@ modules::QwenCausalDecodeRuntimeConfig backbone_config( modules::QwenCausalDecodeRuntimeConfig depth_config( const BreezeTTSConfig & config, core::BackendType backend_type, - size_t graph_arena_bytes) { + size_t graph_arena_bytes, + bool allow_flash_attention = true) { modules::QwenCausalDecodeRuntimeConfig out; out.trace_name = "breeze_tts.depth_decoder"; out.prefill_graph_arena_bytes = graph_arena_bytes; @@ -148,8 +157,14 @@ modules::QwenCausalDecodeRuntimeConfig depth_config( out.decoder.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.decoder.stack.attention_precision = GGML_PREC_F32; out.decoder.stack.projection_precision = GGML_PREC_DEFAULT; - out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.allow_flash_attention = allow_flash_attention; + if (allow_flash_attention) { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + } else { + out.decoder.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.decoder.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + } out.decoder.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; out.decoder.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; if (backend_type == core::BackendType::Cuda || backend_type == core::BackendType::Hip || @@ -720,7 +735,8 @@ struct BreezeGeneratorRuntime::Impl { core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - assets::TensorStorageType storage_type) + assets::TensorStorageType storage_type, + core::AttentionPreference attention_preference = core::AttentionPreference::Auto) : assets(std::move(assets)), execution(execution), tokenizer(this->assets), @@ -735,8 +751,14 @@ struct BreezeGeneratorRuntime::Impl { throw std::runtime_error("BreezeTTS generator requires assets"); } const auto & config = this->assets->config; - backbone_runtime_config = backbone_config(config, execution.backend_type(), graph_arena_bytes); - depth_runtime_config = depth_config(config, execution.backend_type(), graph_arena_bytes); + const bool allow_backbone_flash = core::resolve_flash_attention( + execution.backend(), config.head_dim, attention_preference); + const bool allow_depth_flash = core::resolve_flash_attention( + execution.backend(), config.depth_head_dim, attention_preference); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_backbone_flash", allow_backbone_flash); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_depth_flash", allow_depth_flash); + backbone_runtime_config = backbone_config(config, execution.backend_type(), graph_arena_bytes, allow_backbone_flash); + depth_runtime_config = depth_config(config, execution.backend_type(), graph_arena_bytes, allow_depth_flash); weights = load_weights(*this->assets, execution, weight_context_bytes, storage_type, backbone_runtime_config); backbone_cond = std::make_unique(execution, backbone_runtime_config, weights->backbone); backbone_uncond = std::make_unique(execution, backbone_runtime_config, weights->backbone); @@ -754,14 +776,16 @@ struct BreezeGeneratorRuntime::Impl { execution, graph_arena_bytes, storage_type, - storage_type); + storage_type, + attention_preference); speech_decoder = std::make_unique( this->assets, execution, graph_arena_bytes, weight_context_bytes, storage_type, - storage_type); + storage_type, + attention_preference); depth_first_embed_staging_.assign(static_cast(config.depth_hidden_size), 0.0F); depth_projected_pair_staging_.assign(static_cast(2 * config.depth_hidden_size), 0.0F); depth_prefill_staging_.assign(static_cast(4 * config.depth_hidden_size), 0.0F); @@ -1158,8 +1182,10 @@ BreezeGeneratorRuntime::BreezeGeneratorRuntime( engine::core::ExecutionContext & execution, size_t graph_arena_bytes, size_t weight_context_bytes, - engine::assets::TensorStorageType storage_type) - : impl_(std::make_unique(std::move(assets), execution, graph_arena_bytes, weight_context_bytes, storage_type)) {} + engine::assets::TensorStorageType storage_type, + engine::core::AttentionPreference attention_preference) + : impl_(std::make_unique( + std::move(assets), execution, graph_arena_bytes, weight_context_bytes, storage_type, attention_preference)) {} BreezeGeneratorRuntime::~BreezeGeneratorRuntime() = default; diff --git a/src/models/breeze_tts/session.cpp b/src/models/breeze_tts/session.cpp index 10f7e5e86..25338dd7b 100644 --- a/src/models/breeze_tts/session.cpp +++ b/src/models/breeze_tts/session.cpp @@ -1,5 +1,6 @@ #include "engine/models/breeze_tts/session.h" +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/runtime/options.h" #include "engine/framework/runtime/spec_backed_model.h" @@ -11,6 +12,7 @@ #include #include #include +#include #include namespace engine::models::breeze_tts { @@ -47,6 +49,23 @@ std::vector split_request(const runtime::TaskRequest & req return runtime::chunk_text_request(request, text_chunk_size, text_chunk_mode); } +core::AttentionPreference attention_preference_from_options(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"breeze_tts.attention"})) { + return core::parse_attention_preference(*value, "breeze_tts.attention"); + } + return core::AttentionPreference::Auto; +} + +void trace_attention_preference(core::AttentionPreference preference) { + const char * name = "auto"; + if (preference == core::AttentionPreference::Flash) { + name = "flash"; + } else if (preference == core::AttentionPreference::Eager) { + name = "eager"; + } + engine::debug::trace_log_scalar("breeze_tts.attention.preference", std::string_view(name)); +} + std::size_t reference_cache_slots_from_options(const runtime::SessionOptions & options) { const int64_t slots = runtime::parse_i64_option( options.options, @@ -124,12 +143,15 @@ BreezeTTSSession::BreezeTTSSession( options.options, {"weight_context_mb"}, 2048ull * 1024ull * 1024ull); + const auto attention_preference = attention_preference_from_options(options); + trace_attention_preference(attention_preference); generator_ = std::make_unique( assets_, execution_context(), graph_arena_bytes, weight_context_bytes, - storage_type); + storage_type, + attention_preference); } BreezeTTSSession::~BreezeTTSSession() = default; diff --git a/src/models/breeze_tts/speech_decoder.cpp b/src/models/breeze_tts/speech_decoder.cpp index bd1c220af..1edf0c1c4 100644 --- a/src/models/breeze_tts/speech_decoder.cpp +++ b/src/models/breeze_tts/speech_decoder.cpp @@ -651,7 +651,8 @@ core::TensorValue attention( ggml_tensor * positions, const core::TensorValue & attention_mask, const modules::AttentionWeights & weights, - const DecoderConfig & config) { + const DecoderConfig & config, + modules::ScaledDotProductAttentionLowering lowering = modules::ScaledDotProductAttentionLowering::Flash) { const int64_t kv_repeat = config.num_heads / config.num_kv_heads; auto q_value = modules::LinearModule(binding::linear_config(config.hidden_size, config.num_heads * config.head_dim, false)) .build(build_ctx, input, {weights.q_weight, weights.q_bias}); @@ -717,7 +718,7 @@ core::TensorValue attention( } auto context = modules::ScaledDotProductAttentionModule({ config.head_dim, - modules::ScaledDotProductAttentionLowering::Flash, + lowering, GGML_PREC_F32, modules::AttentionCausality::NonCausal, }).build( @@ -804,10 +805,12 @@ class BreezeSpeechDecoderGraph { int64_t code_frames, core::ExecutionContext & execution_context, core::ConstantTensorCache & constants, - size_t graph_arena_bytes) + size_t graph_arena_bytes, + bool allow_flash_attention = true) : weights_(std::move(weights)), code_frames_(code_frames), backend_(execution_context.backend()), + allow_flash_attention_(allow_flash_attention), compute_threads_(std::max(1, execution_context.config().threads)) { if (weights_ == nullptr) { throw std::runtime_error("Breeze speech decoder graph requires weights"); @@ -863,7 +866,16 @@ class BreezeSpeechDecoderGraph { mask_, core::TensorShape::from_dims({1, 1, code_frames_, code_frames_}), GGML_TYPE_F16); - auto attn_out = attention(ctx_.get(), build_ctx, attn_in, positions_, attention_mask, layer.attention, config); + auto attn_out = attention( + ctx_.get(), + build_ctx, + attn_in, + positions_, + attention_mask, + layer.attention, + config, + allow_flash_attention_ ? modules::ScaledDotProductAttentionLowering::Flash + : modules::ScaledDotProductAttentionLowering::Explicit); attn_out = modules::LayerScaleModule{}.build( build_ctx, attn_out, @@ -1022,6 +1034,7 @@ class BreezeSpeechDecoderGraph { int64_t code_frames_ = 0; int64_t waveform_frames_ = 0; ggml_backend_t backend_ = nullptr; + bool allow_flash_attention_ = true; int compute_threads_ = 1; std::unique_ptr ctx_; ggml_tensor * codes_ = nullptr; @@ -1040,7 +1053,8 @@ BreezeSpeechDecoderRuntime::BreezeSpeechDecoderRuntime( size_t graph_arena_bytes, size_t constant_context_bytes, assets::TensorStorageType linear_weight_storage_type, - assets::TensorStorageType conv_weight_storage_type) + assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference) : assets_(std::move(assets)), execution_context_(&execution_context), graph_arena_bytes_(graph_arena_bytes) { @@ -1053,6 +1067,9 @@ BreezeSpeechDecoderRuntime::BreezeSpeechDecoderRuntime( execution_context_->backend_type(), linear_weight_storage_type, conv_weight_storage_type); + allow_flash_attention_ = core::resolve_flash_attention( + execution_context_->backend(), weights_->config.head_dim, attention_preference); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_decoder_flash", allow_flash_attention_); constants_ = std::make_unique( execution_context_->backend(), std::max(1, execution_context_->config().threads), @@ -1107,7 +1124,8 @@ runtime::AudioBuffer BreezeSpeechDecoderRuntime::decode(const BreezeSpeechCodes chunk_frames, *execution_context_, *constants_, - graph_arena_bytes_); + graph_arena_bytes_, + allow_flash_attention_); graph = std::move(replacement); } auto decoded = graph->run(chunk.data(), chunk.size()); diff --git a/src/models/breeze_tts/speech_encoder.cpp b/src/models/breeze_tts/speech_encoder.cpp index b1d59b400..d9453ce9f 100644 --- a/src/models/breeze_tts/speech_encoder.cpp +++ b/src/models/breeze_tts/speech_encoder.cpp @@ -178,7 +178,8 @@ core::TensorValue mimi_self_attention( const core::TensorValue & input, const core::TensorValue & positions, const TransformerLayerWeights & weights, - const std::optional & attention_mask) { + const std::optional & attention_mask, + modules::ScaledDotProductAttentionLowering lowering = modules::ScaledDotProductAttentionLowering::Flash) { constexpr int64_t kHeads = 8; constexpr int64_t kHeadDim = 64; auto q = modules::LinearModule(binding::linear_config(kHiddenSize, kHiddenSize, false)) @@ -206,7 +207,7 @@ core::TensorValue mimi_self_attention( auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v); auto context = modules::ScaledDotProductAttentionModule({ kHeadDim, - modules::ScaledDotProductAttentionLowering::Flash, + lowering, GGML_PREC_F32, modules::AttentionCausality::Causal, }).build(ctx, q_heads, k_heads, v_heads, attention_mask); @@ -221,12 +222,13 @@ core::TensorValue transformer_block( const core::TensorValue & input, const core::TensorValue & positions, const TransformerLayerWeights & weights, - const std::optional & attention_mask) { + const std::optional & attention_mask, + modules::ScaledDotProductAttentionLowering lowering = modules::ScaledDotProductAttentionLowering::Flash) { const modules::LayerNormModule norm({kHiddenSize, 1.0e-5F, true, true}); auto x = norm.build(ctx, input, weights.norm1); auto attn_out = modules::LayerScaleModule{}.build( ctx, - mimi_self_attention(ctx, x, positions, weights, attention_mask), + mimi_self_attention(ctx, x, positions, weights, attention_mask, lowering), weights.scale1); x = modules::AddModule{}.build(ctx, input, attn_out); auto y = norm.build(ctx, x, weights.norm2); @@ -545,8 +547,10 @@ class BreezeSpeechEncoderTransformerGraph { int64_t frames, core::ExecutionContext & execution_context, core::ConstantTensorCache & constants, - size_t graph_arena_bytes) + size_t graph_arena_bytes, + bool allow_flash_attention = true) : weights_(std::move(weights)), + allow_flash_attention_(allow_flash_attention), frames_(frames), backend_(execution_context.backend()), compute_threads_(std::max(1, execution_context.config().threads)) { @@ -590,7 +594,14 @@ class BreezeSpeechEncoderTransformerGraph { core::TensorShape::from_dims({1, 1, frames_, frames_}), GGML_TYPE_F16); for (const auto & layer : weights_->transformer_layers) { - seq = transformer_block(build_ctx, seq, positions_value, layer, attention_mask); + seq = transformer_block( + build_ctx, + seq, + positions_value, + layer, + attention_mask, + allow_flash_attention_ ? modules::ScaledDotProductAttentionLowering::Flash + : modules::ScaledDotProductAttentionLowering::Explicit); } auto x = modules::TransposeModule({{0, 2, 1, 3}, seq.shape.rank}).build(build_ctx, seq); x = core::ensure_backend_addressable_layout(build_ctx, x); @@ -676,6 +687,7 @@ class BreezeSpeechEncoderTransformerGraph { } std::shared_ptr weights_; + bool allow_flash_attention_ = true; int64_t frames_ = 0; int64_t output_frames_ = 0; std::unique_ptr ctx_; @@ -697,7 +709,8 @@ BreezeSpeechEncoderRuntime::BreezeSpeechEncoderRuntime( core::ExecutionContext & execution_context, size_t graph_arena_bytes, assets::TensorStorageType linear_weight_storage_type, - assets::TensorStorageType conv_weight_storage_type) + assets::TensorStorageType conv_weight_storage_type, + core::AttentionPreference attention_preference) : assets_(std::move(assets)), execution_context_(&execution_context), graph_arena_bytes_(graph_arena_bytes) { @@ -710,6 +723,9 @@ BreezeSpeechEncoderRuntime::BreezeSpeechEncoderRuntime( execution_context_->backend_type(), linear_weight_storage_type, conv_weight_storage_type); + // Mimi encoder self-attention head dim (kHeadDim in mimi_self_attention). + allow_flash_attention_ = core::resolve_flash_attention(execution_context_->backend(), 64, attention_preference); + engine::debug::trace_log_scalar("breeze_tts.attention.allow_encoder_flash", allow_flash_attention_); constants_ = std::make_unique( execution_context_->backend(), std::max(1, execution_context_->config().threads), @@ -754,7 +770,8 @@ BreezeSpeechCodes BreezeSpeechEncoderRuntime::encode(const runtime::AudioBuffer graph_frames, *execution_context_, *constants_, - graph_arena_bytes_); + graph_arena_bytes_, + allow_flash_attention_); } std::vector features(static_cast(kHiddenSize * graph_frames)); diff --git a/src/models/higgs_audio_tts/ar.cpp b/src/models/higgs_audio_tts/ar.cpp index e7646dda9..cd4166f0c 100644 --- a/src/models/higgs_audio_tts/ar.cpp +++ b/src/models/higgs_audio_tts/ar.cpp @@ -39,7 +39,9 @@ struct GgmlContextDeleter { } }; -modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConfig & config) { +modules::QwenDecoderStackConfig make_higgs_qwen_stack_config( + const HiggsTextConfig & config, + bool allow_flash_attention = true) { modules::QwenDecoderStackConfig out; out.hidden_size = config.hidden_size; out.num_attention_heads = config.num_attention_heads; @@ -53,9 +55,17 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf out.projection_precision = GGML_PREC_DEFAULT; out.qkv_layout = modules::QwenDecoderQKVLayout::Separate; out.use_qk_norm = true; - out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; - out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::FlashWithPrefix; + // Eager graph for GPUs without a flash kernel (e.g. sm70). + out.runtime.attention.allow_flash_attention = allow_flash_attention; + if (allow_flash_attention) { + out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::FlashWithPrefix; + } else { + out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::Exact; + } out.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; out.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; out.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; @@ -64,8 +74,8 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf class HiggsQwenDecoderComponent { public: - HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv) - : stack_config_(make_higgs_qwen_stack_config(config)), + HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv, bool allow_flash_attention = true) + : stack_config_(make_higgs_qwen_stack_config(config, allow_flash_attention)), layer_config_(modules::qwen_decoder_layer_config_from_stack(stack_config_)), layer_module_([&] { layer_config_.qkv_layout = packed_qkv @@ -376,12 +386,18 @@ HiggsARRuntime::HiggsARRuntime( std::shared_ptr assets, core::ExecutionContext & execution, size_t weight_context_bytes, - assets::TensorStorageType weight_storage_type) + assets::TensorStorageType weight_storage_type, + core::AttentionPreference attention_preference) : assets_(std::move(assets)), backend_(execution.backend()), backend_type_(execution.backend_type()), device_(execution.config().device), - threads_(std::max(1, execution.config().threads)) { + threads_(std::max(1, execution.config().threads)), + allow_flash_attention_(core::resolve_flash_attention( + execution.backend(), + this->assets_->config.text.head_dim, + attention_preference)) { + engine::debug::trace_log_scalar("higgs_audio_tts.attention.allow_flash", allow_flash_attention_); if (assets_ == nullptr) { throw std::runtime_error("Higgs TTS AR runtime requires assets"); } @@ -404,6 +420,10 @@ ggml_backend_t HiggsARRuntime::backend() const noexcept { return backend_; } +bool HiggsARRuntime::allow_flash_attention() const noexcept { + return allow_flash_attention_; +} + core::BackendType HiggsARRuntime::backend_type() const noexcept { return backend_type_; } @@ -600,7 +620,8 @@ struct HiggsARDecodeGraph::Impl { GGML_TYPE_F16); graph = ggml_new_graph_custom(ctx.get(), 65536, false); - const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); + const HiggsQwenDecoderComponent decoder( + config.text, tensor_weights.packed_qkv, runtime->allow_flash_attention()); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { auto out = decoder.build_decode_layer( build_ctx, @@ -845,7 +866,8 @@ struct HiggsARPrefillGraph::Impl { graph = ggml_new_graph_custom(ctx.get(), 262144, false); keys.reserve(tensor_weights.decoder.layers.size()); values.reserve(tensor_weights.decoder.layers.size()); - const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); + const HiggsQwenDecoderComponent decoder( + config.text, tensor_weights.packed_qkv, runtime->allow_flash_attention()); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { std::optional prefix_key; std::optional prefix_value; @@ -1034,7 +1056,8 @@ struct HiggsARPrefillGraph::Impl { attention_mask, core::TensorShape::from_dims({1, 1, steps, steps}), GGML_TYPE_F16); - const HiggsQwenDecoderComponent decoder(config.text, runtime.weights().packed_qkv); + const HiggsQwenDecoderComponent decoder( + config.text, runtime.weights().packed_qkv, runtime.allow_flash_attention()); auto out = decoder.build_prefill_layer( build_ctx, x, diff --git a/src/models/higgs_audio_tts/loader.cpp b/src/models/higgs_audio_tts/loader.cpp index efc600c72..5d77f9ccc 100644 --- a/src/models/higgs_audio_tts/loader.cpp +++ b/src/models/higgs_audio_tts/loader.cpp @@ -56,6 +56,7 @@ runtime::ModelCliInterface cli(const HiggsAssets &) { {"higgs_audio_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, {"higgs_audio_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, {"higgs_audio_tts.reference_cache_slots", "n", "Encoded reference-audio cache slots; default 1."}, + {"higgs_audio_tts.attention", "auto|flash|eager", "Attention lowering; auto probes the backend and falls back to eager on GPUs without a flash kernel (e.g. sm70); default auto."}, }; return out; } diff --git a/src/models/higgs_audio_tts/session.cpp b/src/models/higgs_audio_tts/session.cpp index 0a5d47790..ca3110b38 100644 --- a/src/models/higgs_audio_tts/session.cpp +++ b/src/models/higgs_audio_tts/session.cpp @@ -1,5 +1,6 @@ #include "engine/models/higgs_audio_tts/session.h" +#include "engine/framework/core/attention_fallback.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/debug/trace.h" #include "engine/framework/runtime/options.h" @@ -10,6 +11,7 @@ #include #include #include +#include #include namespace engine::models::higgs_audio_tts { @@ -50,6 +52,23 @@ uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { return hash; } +core::AttentionPreference resolve_attention_preference(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"higgs_audio_tts.attention"})) { + return core::parse_attention_preference(*value, "higgs_audio_tts.attention"); + } + return core::AttentionPreference::Auto; +} + +void trace_attention_preference(core::AttentionPreference preference) { + const char * name = "auto"; + if (preference == core::AttentionPreference::Flash) { + name = "flash"; + } else if (preference == core::AttentionPreference::Eager) { + name = "eager"; + } + debug::trace_log_scalar("higgs_audio_tts.attention.preference", std::string_view(name)); +} + std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { const int64_t slots = runtime::parse_i64_option( options.options, @@ -169,6 +188,7 @@ HiggsTTSSession::HiggsTTSSession( key != "higgs_audio_tts.codec_decode_graph_arena_mb" && key != "higgs_audio_tts.codec_encode_graph_arena_mb" && key != "higgs_audio_tts.reference_cache_slots" && + key != "higgs_audio_tts.attention" && key != "higgs_audio_tts.weight_type" && key != "higgs_audio_tts.ar_weight_type" && key != "higgs_audio_tts.codec_weight_type") { @@ -176,11 +196,14 @@ HiggsTTSSession::HiggsTTSSession( } } + const auto attention_preference = resolve_attention_preference(options); + trace_attention_preference(attention_preference); ar_ = std::make_shared( assets_, execution_context(), ar_weight_context_bytes_, - ar_weight_storage_type_); + ar_weight_storage_type_, + attention_preference); codec_ = std::make_shared( assets_, execution_context(), diff --git a/tests/unittests/test_attention_fallback.cpp b/tests/unittests/test_attention_fallback.cpp new file mode 100644 index 000000000..d3a35f2d0 --- /dev/null +++ b/tests/unittests/test_attention_fallback.cpp @@ -0,0 +1,68 @@ +#include "engine/framework/core/attention_fallback.h" + +#include "test_assert.h" + +#include +#include +#include + +namespace { + +using engine::core::AttentionPreference; +using engine::test::require; +using engine::test::require_eq; + +void test_parse_attention_preference() { + require_eq( + static_cast(engine::core::parse_attention_preference("auto", "attention")), + static_cast(AttentionPreference::Auto), + "parse auto"); + require_eq( + static_cast(engine::core::parse_attention_preference("flash", "attention")), + static_cast(AttentionPreference::Flash), + "parse flash"); + require_eq( + static_cast(engine::core::parse_attention_preference("eager", "attention")), + static_cast(AttentionPreference::Eager), + "parse eager"); + bool threw = false; + try { + engine::core::parse_attention_preference("sometimes", "breeze_tts.attention"); + } catch (const std::runtime_error & error) { + threw = true; + require( + std::string(error.what()).find("breeze_tts.attention") != std::string::npos, + "parse error names the option"); + } + require(threw, "parse invalid must throw"); +} + +void test_resolve_flash_attention() { + require( + engine::core::resolve_flash_attention(nullptr, 128, AttentionPreference::Flash), + "explicit flash resolves true"); + require( + !engine::core::resolve_flash_attention(nullptr, 128, AttentionPreference::Eager), + "explicit eager resolves false"); + // Null backend preserves historical behavior regardless of head_dim. + require( + engine::core::resolve_flash_attention(nullptr, 128, AttentionPreference::Auto), + "auto with null backend preserves flash"); + require( + engine::core::resolve_flash_attention(nullptr, -1, AttentionPreference::Auto), + "auto with bad head_dim preserves flash"); +} + +} // namespace + +int main() { + try { + test_parse_attention_preference(); + test_resolve_flash_attention(); + } catch (const std::exception & error) { + std::cerr << "attention_fallback_test failed: " << error.what() << '\n'; + return 1; + } + std::cout << "attention_fallback_test passed\n"; + return 0; +}