From 9ba3c4b82e3cc0f9cc1f45baaf672d34e40386bb Mon Sep 17 00:00:00 2001 From: mirek190 Date: Sun, 19 Jul 2026 10:44:14 +0100 Subject: [PATCH] Optimize Higgs Audio v3 inference - pack QKV and gate/up projections and use fused SwiGLU in the shared Qwen decoder path - use grouped FlashAttention, F16 bucketed KV caches, direct set-rows updates, and CUDA-friendly RoPE/view/set-rows graphs - retain and reuse cloned-reference KV prefixes so repeated server requests prefill only their text suffix - add graph and F16 KV correctness coverage plus fixed-seed request-level CUDA benchmark/parity tooling - document the optimized Higgs runtime controls and reproducible validation flow --- CMakeLists.txt | 11 + docs/tts.md | 6 +- include/engine/framework/core/backend.h | 2 + .../modules/attention/qwen_causal_decoder.h | 5 + .../modules/attention/qwen_decoder.h | 1 + include/engine/models/higgs_tts/ar.h | 1 + include/engine/models/higgs_tts/generator.h | 1 + src/framework/core/backend.cpp | 20 +- .../modules/attention/qwen_causal_decoder.cpp | 30 ++ .../modules/attention/qwen_decoder.cpp | 122 +++-- .../modules/optimizations/fast_kv_modules.cpp | 31 +- src/framework/runtime/kv_cache.cpp | 38 +- src/models/higgs_tts/ar.cpp | 148 +++-- src/models/higgs_tts/generator.cpp | 61 ++- tests/higgs_tts/.gitignore | 2 + tests/higgs_tts/README.md | 38 ++ tests/higgs_tts/compare_warmbench_results.py | 160 ++++++ .../higgs_tts/higgs_tts_cuda_bench_cases.json | 35 ++ .../higgs_tts/higgs_tts_cuda_mixed_cases.json | 44 ++ .../higgs_tts/higgs_tts_cuda_perf_cases.json | 57 ++ tests/higgs_tts/higgs_tts_warm_bench.cpp | 51 +- tests/higgs_tts/run_cuda_performance.ps1 | 69 +++ .../test_qwen_decoder_packed_projections.cpp | 504 ++++++++++++++++++ 23 files changed, 1321 insertions(+), 116 deletions(-) create mode 100644 tests/higgs_tts/.gitignore create mode 100644 tests/higgs_tts/README.md create mode 100644 tests/higgs_tts/compare_warmbench_results.py create mode 100644 tests/higgs_tts/higgs_tts_cuda_bench_cases.json create mode 100644 tests/higgs_tts/higgs_tts_cuda_mixed_cases.json create mode 100644 tests/higgs_tts/higgs_tts_cuda_perf_cases.json create mode 100644 tests/higgs_tts/run_cuda_performance.ps1 create mode 100644 tests/unittests/test_qwen_decoder_packed_projections.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index cd0b9a6b..18dbec27 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -709,6 +709,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(chatterbox_warm_bench tests/chatterbox/chatterbox_warm_bench.cpp) add_engine_warmbench(citrinet_asr_warm_bench tests/citrinet_asr/citrinet_asr_warm_bench.cpp) add_engine_warmbench(higgs_audio_stt_warm_bench tests/higgs_audio_stt/higgs_audio_stt_warm_bench.cpp) + add_engine_warmbench(higgs_tts_warm_bench tests/higgs_tts/higgs_tts_warm_bench.cpp) add_engine_warmbench(hviske_asr_warm_bench tests/hviske_asr/hviske_asr_warm_bench.cpp) add_engine_warmbench(index_tts2_warm_bench tests/index_tts2/index_tts2_warm_bench.cpp) add_engine_warmbench(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) @@ -857,6 +858,16 @@ if (ENGINE_BUILD_TESTS) COMMAND encoder_module_test ) + add_engine_unittest( + qwen_decoder_packed_projection_test + tests/unittests/test_qwen_decoder_packed_projections.cpp + ) + + add_test( + NAME qwen_decoder_packed_projection_test + COMMAND qwen_decoder_packed_projection_test + ) + add_engine_unittest(conv_transpose_fast_path_test tests/unittests/test_conv_transpose_fast_path.cpp) add_test( diff --git a/docs/tts.md b/docs/tts.md index d0ff8eee..6a3a6768 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -368,9 +368,9 @@ audiocpp_cli --task tts --family higgs_tts --model models/higgs-audio-v3-tts-4b | `--text-chunk-size` | integer chars | `512` | Long-form chunk size. | | `--max-tokens` | integer | `1024` | Maximum generated AR tokens per chunk. | | `--temperature` | float | `0.8` | AR sampling temperature. | -| `--top-k` | integer | `30` | AR top-k sampling limit. | -| `--top-p` | float | `0.8` | AR nucleus sampling limit. | -| `--repetition-penalty` | float | `1.1` | AR repetition penalty. | +| `--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. | ## IndexTTS2 diff --git a/include/engine/framework/core/backend.h b/include/engine/framework/core/backend.h index ac8c0277..4597e357 100644 --- a/include/engine/framework/core/backend.h +++ b/include/engine/framework/core/backend.h @@ -74,6 +74,8 @@ void write_tensor_i32(const TensorValue & tensor, const int32_t * values, size_t void write_tensor_i32(const TensorValue & tensor, const std::vector & values); void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_f32(const ggml_tensor * tensor); +void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values); +std::vector read_tensor_f16(const ggml_tensor * tensor); void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_i32(const ggml_tensor * tensor); diff --git a/include/engine/framework/modules/attention/qwen_causal_decoder.h b/include/engine/framework/modules/attention/qwen_causal_decoder.h index ffc51353..0ee5ebe7 100644 --- a/include/engine/framework/modules/attention/qwen_causal_decoder.h +++ b/include/engine/framework/modules/attention/qwen_causal_decoder.h @@ -79,6 +79,11 @@ std::vector qwen_position_ids(int64_t steps, int64_t offset = 0); std::vector qwen_causal_prefill_mask_values(int64_t batch_size, int64_t steps); +std::vector qwen_causal_suffix_mask_values( + int64_t batch_size, + int64_t query_steps, + int64_t prefix_steps); + void write_qwen_causal_prefill_mask( ggml_tensor * tensor, int64_t batch_size, diff --git a/include/engine/framework/modules/attention/qwen_decoder.h b/include/engine/framework/modules/attention/qwen_decoder.h index 968d6d60..0cfaf6bc 100644 --- a/include/engine/framework/modules/attention/qwen_decoder.h +++ b/include/engine/framework/modules/attention/qwen_decoder.h @@ -84,6 +84,7 @@ struct QwenDecoderLayerConfig { struct QwenMLPWeights { LinearWeights gate_proj; LinearWeights up_proj; + std::optional gate_up_proj; LinearWeights down_proj; }; diff --git a/include/engine/models/higgs_tts/ar.h b/include/engine/models/higgs_tts/ar.h index 73270860..fd47ba3f 100644 --- a/include/engine/models/higgs_tts/ar.h +++ b/include/engine/models/higgs_tts/ar.h @@ -28,6 +28,7 @@ struct HiggsARWeights { core::TensorValue modality_embedding; HiggsQwenDecoderStackWeights decoder; core::TensorValue norm; + bool packed_qkv = false; }; HiggsARWeights load_higgs_ar_weights( diff --git a/include/engine/models/higgs_tts/generator.h b/include/engine/models/higgs_tts/generator.h index 2d3733cf..4eb3e761 100644 --- a/include/engine/models/higgs_tts/generator.h +++ b/include/engine/models/higgs_tts/generator.h @@ -67,6 +67,7 @@ class HiggsGenerator { HiggsTextTokenizer tokenizer_; size_t ar_decode_graph_arena_bytes_ = 0; std::optional reference_prefix_cache_; + bool reference_kv_ready_ = false; std::optional cuda_sampling_policy_; std::unique_ptr ar_kv_cache_; std::unique_ptr prefill_graph_; diff --git a/src/framework/core/backend.cpp b/src/framework/core/backend.cpp index b34cdea3..255d59b4 100644 --- a/src/framework/core/backend.cpp +++ b/src/framework/core/backend.cpp @@ -527,10 +527,22 @@ void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & value read_tensor_typed_into(tensor, GGML_TYPE_F32, values); } -std::vector read_tensor_f32(const ggml_tensor * tensor) { - return read_tensor_typed(tensor, GGML_TYPE_F32); -} - +std::vector read_tensor_f32(const ggml_tensor * tensor) { + return read_tensor_typed(tensor, GGML_TYPE_F32); +} + +void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values) { + const auto fp16_values = read_tensor_typed(tensor, GGML_TYPE_F16); + values.resize(fp16_values.size()); + ggml_fp16_to_fp32_row(fp16_values.data(), values.data(), static_cast(values.size())); +} + +std::vector read_tensor_f16(const ggml_tensor * tensor) { + std::vector values; + read_tensor_f16_into(tensor, values); + return values; +} + void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values) { read_tensor_typed_into(tensor, GGML_TYPE_I32, values); } diff --git a/src/framework/modules/attention/qwen_causal_decoder.cpp b/src/framework/modules/attention/qwen_causal_decoder.cpp index 229243c1..63c5fcd8 100644 --- a/src/framework/modules/attention/qwen_causal_decoder.cpp +++ b/src/framework/modules/attention/qwen_causal_decoder.cpp @@ -175,6 +175,36 @@ std::vector qwen_causal_prefill_mask_values(int64_t batch_size, int return out; } +std::vector qwen_causal_suffix_mask_values( + int64_t batch_size, + int64_t query_steps, + int64_t prefix_steps) { + if (batch_size <= 0) { + throw std::runtime_error("qwen_causal_suffix_mask_values requires positive batch size"); + } + validate_steps(query_steps, "qwen_causal_suffix_mask_values"); + if (prefix_steps < 0) { + throw std::runtime_error("qwen_causal_suffix_mask_values requires non-negative prefix steps"); + } + const int64_t key_steps = prefix_steps + query_steps; + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::vector one(static_cast(query_steps * key_steps), masked); + for (int64_t row = 0; row < query_steps; ++row) { + const size_t row_offset = static_cast(row * key_steps); + std::fill_n( + one.begin() + static_cast(row_offset), + prefix_steps + row + 1, + visible); + } + std::vector out; + out.reserve(static_cast(batch_size) * one.size()); + for (int64_t batch = 0; batch < batch_size; ++batch) { + out.insert(out.end(), one.begin(), one.end()); + } + return out; +} + void write_qwen_causal_prefill_mask( ggml_tensor * tensor, int64_t batch_size, diff --git a/src/framework/modules/attention/qwen_decoder.cpp b/src/framework/modules/attention/qwen_decoder.cpp index 886f1c03..56b7c601 100644 --- a/src/framework/modules/attention/qwen_decoder.cpp +++ b/src/framework/modules/attention/qwen_decoder.cpp @@ -332,35 +332,75 @@ core::TensorValue build_mlp( const core::TensorValue & input, const QwenDecoderLayerConfig & config, const QwenMLPWeights & weights) { - auto gate = LinearModule( - { - config.hidden_size, - config.intermediate_size, - weights.gate_proj.bias.has_value(), - config.projection_precision, - }) - .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { - gate = activation_cast(ctx, gate, config.activation_cast); - } - gate = SiluModule{}.build(ctx, gate); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { - gate = activation_cast(ctx, gate, config.activation_cast); - } - auto up = LinearModule( - { - config.hidden_size, - config.intermediate_size, - weights.up_proj.bias.has_value(), - config.projection_precision, - }) - .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { - up = activation_cast(ctx, up, config.activation_cast); - } - auto gated = MulModule{}.build(ctx, gate, up); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { - gated = activation_cast(ctx, gated, config.activation_cast); + core::TensorValue gate; + core::TensorValue up; + std::optional packed_gate_up; + if (weights.gate_up_proj.has_value()) { + auto gate_up = LinearModule( + { + config.hidden_size, + config.intermediate_size * 2, + weights.gate_up_proj->bias.has_value(), + config.projection_precision, + }) + .build( + ctx, + input, + require_linear(*weights.gate_up_proj, false, "QwenMLPWeights.gate_up_proj")); + packed_gate_up = gate_up; + gate = SliceModule({2, 0, config.intermediate_size}).build(ctx, gate_up); + up = SliceModule({2, config.intermediate_size, config.intermediate_size}).build(ctx, gate_up); + } else { + gate = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.gate_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); + up = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.up_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); + } + const bool can_use_fused_swiglu = + !config.activation_cast.enabled || + (!config.activation_cast.after_mlp_projection && + !config.activation_cast.after_mlp_silu && + !config.activation_cast.after_mlp_mul); + core::TensorValue gated; + if (can_use_fused_swiglu && packed_gate_up.has_value()) { + gated = core::wrap_tensor( + ggml_swiglu(ctx.ggml, packed_gate_up->tensor), + core::TensorShape::from_dims({ + input.shape.dims[0], + input.shape.dims[1], + config.intermediate_size, + }), + packed_gate_up->type); + } else if (can_use_fused_swiglu) { + gated = core::wrap_tensor( + ggml_swiglu_split(ctx.ggml, gate.tensor, up.tensor), + gate.shape, + gate.type); + } else { + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + gate = activation_cast(ctx, gate, config.activation_cast); + up = activation_cast(ctx, up, config.activation_cast); + } + gate = SiluModule{}.build(ctx, gate); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + gated = MulModule{}.build(ctx, gate, up); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { + gated = activation_cast(ctx, gated, config.activation_cast); + } } auto down = LinearModule( { @@ -441,8 +481,22 @@ 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); - auto all_k = prefix_key.has_value() ? ConcatModule({1}).build(ctx, *prefix_key, k) : k; - auto all_v = prefix_value.has_value() ? ConcatModule({1}).build(ctx, *prefix_value, v) : v; + auto attention_prefix_key = prefix_key; + auto attention_prefix_value = prefix_value; + if (attention_prefix_key.has_value() && attention_prefix_key->type != k.type) { + attention_prefix_key = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_key->tensor, k.type), + attention_prefix_key->shape, + k.type); + } + if (attention_prefix_value.has_value() && attention_prefix_value->type != v.type) { + attention_prefix_value = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_value->tensor, v.type), + attention_prefix_value->shape, + v.type); + } + auto all_k = attention_prefix_key.has_value() ? ConcatModule({1}).build(ctx, *attention_prefix_key, k) : k; + auto all_v = attention_prefix_value.has_value() ? ConcatModule({1}).build(ctx, *attention_prefix_value, v) : v; core::TensorValue context; if (!prefix_key.has_value() && attention_mask.has_value() && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { @@ -457,8 +511,10 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( dim, *attention_mask, config_.attention_precision); - } else if (!prefix_key.has_value() && attention_mask.has_value() && - config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) { + } else if (attention_mask.has_value() && + (config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped || + (prefix_key.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); auto v_heads = TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v); diff --git a/src/framework/modules/optimizations/fast_kv_modules.cpp b/src/framework/modules/optimizations/fast_kv_modules.cpp index ce9e08eb..64389763 100644 --- a/src/framework/modules/optimizations/fast_kv_modules.cpp +++ b/src/framework/modules/optimizations/fast_kv_modules.cpp @@ -51,8 +51,8 @@ core::TensorValue FastKVSetRowsModule::build( if (row_index.shape.rank != 1 || (row_index.shape.dims[0] != 1 && row_index.shape.dims[0] != batch)) { throw std::runtime_error("FastKVSetRowsModule row_index must have shape {1} or {batch}"); } - if (cache.type != GGML_TYPE_F32 || row.type != GGML_TYPE_F32) { - throw std::runtime_error("FastKVSetRowsModule requires f32 cache and row tensors"); + if ((cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16) || row.type != GGML_TYPE_F32) { + throw std::runtime_error("FastKVSetRowsModule requires an f32/f16 cache and an f32 row tensor"); } if (row_index.type != GGML_TYPE_I32 && row_index.type != GGML_TYPE_I64) { throw std::runtime_error("FastKVSetRowsModule requires i32 or i64 row_index tensor"); @@ -69,16 +69,39 @@ core::TensorValue FastKVSetRowsModule::build( } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({1, row_elems})); + auto flat_row = core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + 1, + contiguous_row.tensor->nb[2], + 0), + core::TensorShape::from_dims({1, row_elems}), + row.type); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); + // ggml_set_rows src[2] is only a legacy dependency anchor for the + // destination. Point it at the underlying cache so the metadata-only + // flatten does not interrupt CUDA's ROPE -> VIEW -> SET_ROWS fusion. + updated->src[2] = cache.tensor; auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({batch * steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({batch, row_elems})); + auto flat_row = core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + batch, + contiguous_row.tensor->nb[3], + 0), + core::TensorShape::from_dims({batch, row_elems}), + row.type); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); + updated->src[2] = cache.tensor; auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 24f9925c..17e236fa 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -9,6 +9,30 @@ namespace engine::runtime { +namespace { + +void write_cache_tensor(const core::TensorValue & tensor, const std::vector & values) { + if (tensor.type == GGML_TYPE_F32) { + core::write_tensor_f32(tensor, values); + } else if (tensor.type == GGML_TYPE_F16) { + core::write_tensor_f16(tensor, values); + } else { + throw std::runtime_error("TransformerKVCache supports only f32 and f16 cache tensors"); + } +} + +std::vector read_cache_tensor(const core::TensorValue & tensor) { + if (tensor.type == GGML_TYPE_F32) { + return core::read_tensor_f32(tensor.tensor); + } + if (tensor.type == GGML_TYPE_F16) { + return core::read_tensor_f16(tensor.tensor); + } + throw std::runtime_error("TransformerKVCache supports only f32 and f16 cache tensors"); +} + +} // namespace + TransformerKVCache::TransformerKVCache( int64_t cache_steps, int64_t step_elems, @@ -68,8 +92,8 @@ void TransformerKVCache::import_state(const TransformerKVState & state) { std::copy(source.key.begin(), source.key.end(), cache.import_key_scratch.begin()); std::copy(source.value.begin(), source.value.end(), cache.import_value_scratch.begin()); } - core::write_tensor_f32(cache.key_tensor, cache.import_key_scratch); - core::write_tensor_f32(cache.value_tensor, cache.import_value_scratch); + write_cache_tensor(cache.key_tensor, cache.import_key_scratch); + write_cache_tensor(cache.value_tensor, cache.import_value_scratch); } } } @@ -85,8 +109,8 @@ TransformerKVState TransformerKVCache::export_state() const { if (keep_elems == 0) { continue; } - const auto key_values = core::read_tensor_f32(layers_[layer].key_tensor.tensor); - const auto value_values = core::read_tensor_f32(layers_[layer].value_tensor.tensor); + const auto key_values = read_cache_tensor(layers_[layer].key_tensor); + const auto value_values = read_cache_tensor(layers_[layer].value_tensor); out.key.assign(key_values.begin(), key_values.begin() + static_cast(keep_elems)); out.value.assign(value_values.begin(), value_values.begin() + static_cast(keep_elems)); } @@ -141,11 +165,11 @@ void TransformerKVCache::trace_log_state(const std::string & name, int64_t num_h return; } const size_t keep_elems = static_cast(valid_steps_ * step_elems_); - const auto first_key = core::read_tensor_f32(layers_.front().key_tensor.tensor); + const auto first_key = read_cache_tensor(layers_.front().key_tensor); std::vector first_key_keep(first_key.begin(), first_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer0.key", {1, valid_steps_, num_heads, head_dim}, first_key_keep); if (layers_.size() > 1) { - const auto last_key = core::read_tensor_f32(layers_.back().key_tensor.tensor); + const auto last_key = read_cache_tensor(layers_.back().key_tensor); std::vector last_key_keep(last_key.begin(), last_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer_last.key", {1, valid_steps_, num_heads, head_dim}, last_key_keep); } @@ -175,7 +199,7 @@ core::TensorValue view_transformer_kv_cache_steps( cache.tensor->nb[3], static_cast(start) * cache.tensor->nb[2]), core::TensorShape::from_dims({1, steps, heads, head_dim}), - GGML_TYPE_F32); + cache.type); } } // namespace engine::runtime diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_tts/ar.cpp index 6ffe2c4c..922d4830 100644 --- a/src/models/higgs_tts/ar.cpp +++ b/src/models/higgs_tts/ar.cpp @@ -62,18 +62,25 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf class HiggsQwenDecoderComponent { public: - explicit HiggsQwenDecoderComponent(const HiggsTextConfig & config) + HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv) : stack_config_(make_higgs_qwen_stack_config(config)), layer_config_(modules::qwen_decoder_layer_config_from_stack(stack_config_)), - layer_module_(layer_config_) {} + layer_module_([&] { + layer_config_.qkv_layout = packed_qkv + ? modules::QwenDecoderQKVLayout::PackedQKV + : modules::QwenDecoderQKVLayout::Separate; + return layer_config_; + }()) {} modules::QwenDecoderLayerOutputs build_prefill_layer( core::ModuleBuildContext & ctx, const core::TensorValue & input, const core::TensorValue & positions, const modules::QwenDecoderLayerWeights & weights, - const core::TensorValue & attention_mask) const { - return layer_module_.build(ctx, input, positions, weights, std::nullopt, std::nullopt, attention_mask); + const core::TensorValue & attention_mask, + const std::optional & prefix_key = std::nullopt, + const std::optional & prefix_value = std::nullopt) const { + return layer_module_.build(ctx, input, positions, weights, prefix_key, prefix_value, attention_mask); } modules::QwenDecoderLayerOutputs build_decode_layer( @@ -144,23 +151,33 @@ modules::QwenDecoderLayerWeights load_layer_weights( store.load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}), std::nullopt, }; - // Keep Q/K/V separate here so Higgs exercises the framework Qwen decoder path. - // A packed-QKV fast path can be evaluated later as a framework-level optimization. - weights.self_attention.q_weight = store.load_tensor( - source, - prefix + ".self_attn.q_proj.weight", - storage_type, - {q_out, config.hidden_size}); - weights.self_attention.k_weight = store.load_tensor( - source, - prefix + ".self_attn.k_proj.weight", - storage_type, - {kv_out, config.hidden_size}); - weights.self_attention.v_weight = store.load_tensor( - source, - prefix + ".self_attn.v_proj.weight", - storage_type, - {kv_out, config.hidden_size}); + { + const auto q = source.require_tensor( + prefix + ".self_attn.q_proj.weight", + storage_type, + {q_out, config.hidden_size}); + const auto k = source.require_tensor( + prefix + ".self_attn.k_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + const auto v = source.require_tensor( + prefix + ".self_attn.v_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + if (q.type != k.type || q.type != v.type) { + throw std::runtime_error("Higgs TTS packed QKV weights require matching storage types"); + } + std::vector packed; + packed.reserve(q.bytes.size() + k.bytes.size() + v.bytes.size()); + packed.insert(packed.end(), q.bytes.begin(), q.bytes.end()); + packed.insert(packed.end(), k.bytes.begin(), k.bytes.end()); + packed.insert(packed.end(), v.bytes.begin(), v.bytes.end()); + weights.self_attention.qkv_weight = store.make_tensor( + core::TensorShape::from_dims({q_out + 2 * kv_out, config.hidden_size}), + q.type, + packed.data(), + packed.size()); + } weights.self_attention.out_weight = store.load_tensor( source, prefix + ".self_attn.o_proj.weight", @@ -178,22 +195,31 @@ modules::QwenDecoderLayerWeights load_layer_weights( store.load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size}), std::nullopt, }; - weights.mlp.gate_proj = { - store.load_tensor( - source, + { + const auto gate = source.require_tensor( prefix + ".mlp.gate_proj.weight", storage_type, - {config.intermediate_size, config.hidden_size}), - std::nullopt, - }; - weights.mlp.up_proj = { - store.load_tensor( - source, + {config.intermediate_size, config.hidden_size}); + const auto up = source.require_tensor( prefix + ".mlp.up_proj.weight", storage_type, - {config.intermediate_size, config.hidden_size}), - std::nullopt, - }; + {config.intermediate_size, config.hidden_size}); + if (gate.type != up.type) { + throw std::runtime_error("Higgs TTS packed gate/up weights require matching storage types"); + } + std::vector packed; + packed.reserve(gate.bytes.size() + up.bytes.size()); + packed.insert(packed.end(), gate.bytes.begin(), gate.bytes.end()); + packed.insert(packed.end(), up.bytes.begin(), up.bytes.end()); + weights.mlp.gate_up_proj = modules::LinearWeights{ + store.make_tensor( + core::TensorShape::from_dims({config.intermediate_size * 2, config.hidden_size}), + gate.type, + packed.data(), + packed.size()), + std::nullopt, + }; + } weights.mlp.down_proj = { store.load_tensor( source, @@ -331,6 +357,7 @@ HiggsARWeights load_higgs_ar_weights( {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}); weights.decoder = load_decoder_weights(*weights.store, source, config.text, weight_storage_type); weights.norm = weights.store->load_f32_tensor(source, "body.norm.weight", {config.text.hidden_size}); + weights.packed_qkv = true; weights.store->upload(); return weights; } @@ -414,11 +441,11 @@ struct HiggsARKVCache::Impl { for (size_t layer = 0; layer < tensor_weights.decoder.layers.size(); ++layer) { key_tensors.push_back(core::make_tensor( build_ctx, - GGML_TYPE_F32, + GGML_TYPE_F16, core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); value_tensors.push_back(core::make_tensor( build_ctx, - GGML_TYPE_F32, + GGML_TYPE_F16, core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); } cache = runtime::TransformerKVCache( @@ -568,7 +595,7 @@ struct HiggsARDecodeGraph::Impl { GGML_TYPE_F16); graph = ggml_new_graph_custom(ctx.get(), 65536, false); - const HiggsQwenDecoderComponent decoder(config.text); + const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { auto out = decoder.build_decode_layer( build_ctx, @@ -761,7 +788,7 @@ struct HiggsARPrefillGraph::Impl { start_step(input_start_step), run_steps(input_prompt_steps - input_start_step), prefill_cache_steps(input_prompt_steps), - layerwise(input_prompt_steps >= kLayerwisePrefillMinSteps), + layerwise(input_prompt_steps >= kLayerwisePrefillMinSteps && input_start_step == 0), graph_arena_bytes(graph_arena_bytes) { if (runtime == nullptr) { throw std::runtime_error("Higgs TTS AR prefill graph requires runtime"); @@ -772,8 +799,8 @@ struct HiggsARPrefillGraph::Impl { if (start_step < 0 || start_step >= prompt_steps) { throw std::runtime_error("Higgs TTS AR prefill graph start step is outside the prompt"); } - if (start_step != 0) { - throw std::runtime_error("Higgs TTS AR prefill graph requires full prompt prefill"); + if (start_step > 0 && target_cache == nullptr) { + throw std::runtime_error("Higgs TTS AR suffix prefill requires a target KV cache"); } if (layerwise) { engine::debug::timing_log_scalar("higgs_tts.ar.prefill.graph.build_ms", 0.0); @@ -813,28 +840,48 @@ 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); + const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { + std::optional prefix_key; + std::optional prefix_value; + if (start_step > 0) { + prefix_key = higgs_cache_view( + build_ctx, + target_cache->key_tensor(layer_index), + 0, + start_step, + config.text.num_key_value_heads, + config.text.head_dim); + prefix_value = higgs_cache_view( + build_ctx, + target_cache->value_tensor(layer_index), + 0, + start_step, + config.text.num_key_value_heads, + config.text.head_dim); + } auto out = decoder.build_prefill_layer( build_ctx, x, positions_value, tensor_weights.decoder.layers[layer_index], - attention_mask_value); + attention_mask_value, + prefix_key, + prefix_value); x = out.output; if (target_cache != nullptr) { auto key_dest = higgs_cache_view( build_ctx, target_cache->key_tensor(layer_index), - 0, - prompt_steps, + start_step, + run_steps, config.text.num_key_value_heads, config.text.head_dim); auto value_dest = higgs_cache_view( build_ctx, target_cache->value_tensor(layer_index), - 0, - prompt_steps, + start_step, + run_steps, config.text.num_key_value_heads, config.text.head_dim); ggml_build_forward_expand(graph, ggml_cpy(ctx.get(), out.key.tensor, key_dest.tensor)); @@ -861,7 +908,7 @@ struct HiggsARPrefillGraph::Impl { text_gate_values.assign(static_cast(run_steps), 0.0F); code_gate_values.assign(static_cast(run_steps), 0.0F); positions_values = modules::qwen_position_ids(run_steps, start_step); - attention_mask_values = modules::qwen_causal_prefill_mask_values(1, run_steps); + attention_mask_values = modules::qwen_causal_suffix_mask_values(1, run_steps, start_step); engine::debug::timing_log_scalar( "higgs_tts.ar.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); @@ -982,7 +1029,7 @@ struct HiggsARPrefillGraph::Impl { attention_mask, core::TensorShape::from_dims({1, 1, steps, steps}), GGML_TYPE_F16); - const HiggsQwenDecoderComponent decoder(config.text); + const HiggsQwenDecoderComponent decoder(config.text, runtime.weights().packed_qkv); auto out = decoder.build_prefill_layer( build_ctx, x, @@ -1157,6 +1204,11 @@ struct HiggsARPrefillGraph::Impl { if (candidate_start_step != start_step) { throw std::runtime_error("Higgs TTS AR prefill graph start step mismatch"); } + if (start_step > 0 && + (target_cache == nullptr || target_cache->valid_steps() < start_step || + target_cache->current_end() != start_step)) { + throw std::runtime_error("Higgs TTS AR suffix prefill requires the retained prefix in KV cache"); + } if (layerwise) { return run_layerwise(input); } @@ -1201,7 +1253,7 @@ struct HiggsARPrefillGraph::Impl { 0, out.output.codebook_logits.size() * sizeof(float)); if (target_cache != nullptr) { - target_cache->advance_after_direct_append(prompt_steps); + target_cache->advance_after_direct_append(run_steps); out.wrote_cache = true; out.kv_state.current_end = prompt_steps; return out; diff --git a/src/models/higgs_tts/generator.cpp b/src/models/higgs_tts/generator.cpp index 51a4dd3b..058fa34d 100644 --- a/src/models/higgs_tts/generator.cpp +++ b/src/models/higgs_tts/generator.cpp @@ -20,7 +20,21 @@ namespace { using Clock = std::chrono::steady_clock; -constexpr int64_t kInitialGeneratedCacheSteps = 512; +constexpr int64_t kInitialGeneratedCacheSteps = 128; +constexpr int64_t kMinimumCacheBucketSteps = 128; + +int64_t bucketed_initial_cache_steps(int64_t prompt_steps, int64_t max_tokens) { + const int64_t maximum = prompt_steps + max_tokens; + const int64_t required = prompt_steps + std::min(max_tokens, kInitialGeneratedCacheSteps); + int64_t bucket = kMinimumCacheBucketSteps; + while (bucket < required && bucket <= maximum / 2) { + bucket *= 2; + } + if (bucket < required) { + bucket = required; + } + return std::min(bucket, maximum); +} void validate_generation_options(const HiggsGenerationOptions & options) { if (options.max_tokens <= 0) { @@ -205,9 +219,21 @@ void HiggsGenerator::prepare(const HiggsGenerationRequest & request) { cache.prefix_tokens.assign(prepared.prompt.token_ids.begin(), prepared.prompt.token_ids.begin() + static_cast(prepared.prefix_steps)); + const bool same_reference = + reference_prefix_cache_.has_value() && + reference_prefix_cache_->reference_text == cache.reference_text && + reference_prefix_cache_->reference_codes == cache.reference_codes && + reference_prefix_cache_->reference_frames == cache.reference_frames && + reference_prefix_cache_->reference_codebooks == cache.reference_codebooks && + reference_prefix_cache_->prefix_steps == cache.prefix_steps && + reference_prefix_cache_->prefix_tokens == cache.prefix_tokens; reference_prefix_cache_ = std::move(cache); + if (!same_reference) { + reference_kv_ready_ = false; + } } else { reference_prefix_cache_.reset(); + reference_kv_ready_ = false; } } @@ -328,20 +354,38 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_cache_hit", reference_cache_hit); engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_steps", prepared.prefix_steps); const int64_t max_cache_steps = prompt_steps + request.options.max_tokens; - const int64_t initial_cache_steps = - prompt_steps + std::min(request.options.max_tokens, kInitialGeneratedCacheSteps); - if (ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps)) { + const int64_t initial_cache_steps = bucketed_initial_cache_steps(prompt_steps, request.options.max_tokens); + const bool cache_rebuild = + ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps) || + ar_kv_cache_->cache_steps() != initial_cache_steps; + if (cache_rebuild) { decode_graph_.reset(); ar_kv_cache_ = std::make_unique(ar_, initial_cache_steps); + reference_kv_ready_ = false; + } + const bool reference_kv_cache_hit = + reference_cache_hit && reference_kv_ready_ && + ar_kv_cache_->valid_steps() >= prepared.prefix_steps; + const int64_t prefill_start_step = reference_kv_cache_hit ? prepared.prefix_steps : 0; + if (reference_kv_cache_hit) { + ar_kv_cache_->retain_prefix(prefill_start_step); + } else { + ar_kv_cache_->reset(); } - ar_kv_cache_->reset(); - if (prefill_graph_ == nullptr || !prefill_graph_->matches(*ar_, prompt_steps, 0)) { + engine::debug::trace_log_scalar("higgs_tts.generator.reference_kv_cache_hit", reference_kv_cache_hit); + engine::debug::trace_log_scalar("higgs_tts.generator.prefill_start_step", prefill_start_step); + engine::debug::trace_log_scalar("higgs_tts.generator.prefill_run_steps", prompt_steps - prefill_start_step); + engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_steps", initial_cache_steps); + engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_rebuild", cache_rebuild); + if (prefill_graph_ == nullptr || + !prefill_graph_->matches(*ar_, prompt_steps, prefill_start_step)) { prefill_graph_.reset(); prefill_graph_ = std::make_unique( - ar_, prompt_steps, 0, ar_kv_cache_.get(), ar_decode_graph_arena_bytes_); + ar_, prompt_steps, prefill_start_step, ar_kv_cache_.get(), ar_decode_graph_arena_bytes_); } - auto prefill_output = prefill_graph_->run(prepared.ar_input, 0); + auto prefill_output = prefill_graph_->run(prepared.ar_input, prefill_start_step); prefill_graph_.reset(); + reference_kv_ready_ = reference_cache_hit; if (decode_graph_ == nullptr || !decode_graph_->can_run(*ar_, ar_kv_cache_->cache_steps())) { decode_graph_ = std::make_unique( @@ -420,6 +464,7 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re decode_graph_.reset(); ar_kv_cache_ = std::make_unique(ar_, grown_cache_steps); ar_kv_cache_->import_state(kv_state); + engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_grown_steps", grown_cache_steps); decode_graph_ = std::make_unique( ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); decode_graph_->begin_decode_run(); diff --git a/tests/higgs_tts/.gitignore b/tests/higgs_tts/.gitignore new file mode 100644 index 00000000..b45f5596 --- /dev/null +++ b/tests/higgs_tts/.gitignore @@ -0,0 +1,2 @@ +results/ +__pycache__/ diff --git a/tests/higgs_tts/README.md b/tests/higgs_tts/README.md new file mode 100644 index 00000000..393a1e63 --- /dev/null +++ b/tests/higgs_tts/README.md @@ -0,0 +1,38 @@ +# Higgs Audio v3 TTS tests + +The focused framework unit test validates that packed QKV/gate-up projections +match the separate projections, suffix causal masks are correct, F16 KV writes +preserve their values, and the decode graph exposes the intended CUDA paths: +grouped FlashAttention, packed SwiGLU, direct KV updates, and the +`ROPE -> VIEW -> SET_ROWS` fusion pattern. + +```powershell +cmake --build build/windows-cuda-release --config Release --target qwen_decoder_packed_projection_test higgs_tts_warm_bench -j 8 +ctest --test-dir build/windows-cuda-release -C Release -R qwen_decoder_packed_projection_test --output-on-failure +``` + +Run the fixed-seed, five-request CUDA benchmark and save every generated WAV: + +```powershell +tests/higgs_tts/run_cuda_performance.ps1 ` + -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` + -Label candidate +``` + +Compare a candidate run with a prior result directory request by request: + +```powershell +tests/higgs_tts/run_cuda_performance.ps1 ` + -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` + -Label candidate ` + -Baseline tests/higgs_tts/results/baseline +``` + +The comparison reports frame counts, wall time, RTF, speedup, waveform cosine, +and 80-band log-mel cosine. Result WAVs, logs, and JSON reports are written below +`tests/higgs_tts/results/`, which is intentionally ignored by Git. +The comparison helper requires Python 3 with NumPy. + +Add `-RequireSameFrames` when comparing paths that are expected to be +deterministic and numerically identical. Sampled or mixed-precision paths still +report their frame drift and similarity metrics without hiding the results. diff --git a/tests/higgs_tts/compare_warmbench_results.py b/tests/higgs_tts/compare_warmbench_results.py new file mode 100644 index 00000000..3facf3a5 --- /dev/null +++ b/tests/higgs_tts/compare_warmbench_results.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +"""Compare Higgs warmbench runs request by request. + +Each result directory is expected to contain timing.log and audio/audio_N.wav, +as emitted by higgs_tts_warm_bench. The report includes exact frame counts, +wall time, RTF, speedup, waveform cosine, and log-mel cosine per request. +""" + +from __future__ import annotations + +import argparse +import json +import math +import wave +from pathlib import Path + +import numpy as np + + +def read_wav(path: Path) -> tuple[int, np.ndarray]: + with wave.open(str(path), "rb") as wav: + if wav.getsampwidth() != 2: + raise ValueError(f"{path}: expected PCM16 WAV") + channels = wav.getnchannels() + sample_rate = wav.getframerate() + samples = np.frombuffer(wav.readframes(wav.getnframes()), dtype=" 1: + samples = samples.reshape(-1, channels).mean(axis=1) + return sample_rate, samples / 32768.0 + + +def cosine(a: np.ndarray, b: np.ndarray) -> float: + count = min(a.size, b.size) + if count == 0: + return math.nan + a = a[:count].astype(np.float64, copy=False) + b = b[:count].astype(np.float64, copy=False) + denom = np.linalg.norm(a) * np.linalg.norm(b) + return float(np.dot(a, b) / denom) if denom > 0.0 else math.nan + + +def hz_to_mel(hz: np.ndarray | float) -> np.ndarray | float: + return 2595.0 * np.log10(1.0 + np.asarray(hz) / 700.0) + + +def mel_to_hz(mel: np.ndarray | float) -> np.ndarray | float: + return 700.0 * (np.power(10.0, np.asarray(mel) / 2595.0) - 1.0) + + +def log_mel(samples: np.ndarray, sample_rate: int, n_fft: int = 1024, hop: int = 256, bands: int = 80) -> np.ndarray: + if samples.size < n_fft: + samples = np.pad(samples, (0, n_fft - samples.size)) + frame_count = 1 + (samples.size - n_fft) // hop + shape = (frame_count, n_fft) + strides = (samples.strides[0] * hop, samples.strides[0]) + frames = np.lib.stride_tricks.as_strided(samples, shape=shape, strides=strides) + spectrum = np.abs(np.fft.rfft(frames * np.hanning(n_fft), axis=1)) ** 2 + + mel_points = np.linspace(hz_to_mel(0.0), hz_to_mel(sample_rate / 2.0), bands + 2) + bins = np.floor((n_fft + 1) * mel_to_hz(mel_points) / sample_rate).astype(np.int64) + bins = np.clip(bins, 0, spectrum.shape[1] - 1) + filters = np.zeros((bands, spectrum.shape[1]), dtype=np.float64) + for band in range(bands): + left, center, right = bins[band : band + 3] + if center > left: + filters[band, left:center] = np.arange(center - left) / (center - left) + if right > center: + filters[band, center:right] = np.arange(right - center, 0, -1) / (right - center) + return np.log(np.maximum(spectrum @ filters.T, 1.0e-10)).astype(np.float32) + + +def timings(path: Path) -> dict[int, float]: + result: dict[int, float] = {} + for line in (path / "timing.log").read_text(encoding="utf-8").splitlines(): + prefix = "higgs_tts.cpp.request_" + if not line.startswith(prefix) or ".wall_ms=" not in line: + continue + index_text, value = line[len(prefix) :].split(".wall_ms=", 1) + result[int(index_text)] = float(value) + return result + + +def audio_files(path: Path) -> dict[int, Path]: + result: dict[int, Path] = {} + for wav_path in (path / "audio").glob("audio_*.wav"): + result[int(wav_path.stem.removeprefix("audio_"))] = wav_path + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--baseline", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--output", type=Path) + parser.add_argument("--require-same-frames", action="store_true") + parser.add_argument("--min-wav-cosine", type=float) + parser.add_argument("--min-logmel-cosine", type=float) + args = parser.parse_args() + + baseline_times = timings(args.baseline) + candidate_times = timings(args.candidate) + baseline_audio = audio_files(args.baseline) + candidate_audio = audio_files(args.candidate) + indices = sorted(set(baseline_times) & set(candidate_times) & set(baseline_audio) & set(candidate_audio)) + if not indices: + raise ValueError("no matching warmbench requests were found") + + failed = False + requests = [] + for index in indices: + baseline_rate, baseline_samples = read_wav(baseline_audio[index]) + candidate_rate, candidate_samples = read_wav(candidate_audio[index]) + if baseline_rate != candidate_rate: + raise ValueError(f"request {index}: sample rates differ") + same_frames = baseline_samples.size == candidate_samples.size + wav_cosine = cosine(baseline_samples, candidate_samples) + baseline_mel = log_mel(baseline_samples, baseline_rate) + candidate_mel = log_mel(candidate_samples, candidate_rate) + mel_frames = min(baseline_mel.shape[0], candidate_mel.shape[0]) + logmel_cosine = cosine(baseline_mel[:mel_frames].reshape(-1), candidate_mel[:mel_frames].reshape(-1)) + baseline_duration_sec = baseline_samples.size / baseline_rate + candidate_duration_sec = candidate_samples.size / candidate_rate + baseline_ms = baseline_times[index] + candidate_ms = candidate_times[index] + baseline_rtf = baseline_ms / 1000.0 / baseline_duration_sec + candidate_rtf = candidate_ms / 1000.0 / candidate_duration_sec + request = { + "request_index": index, + "baseline_frames": int(baseline_samples.size), + "candidate_frames": int(candidate_samples.size), + "same_frames": same_frames, + "baseline_wall_ms": baseline_ms, + "candidate_wall_ms": candidate_ms, + "wall_speedup": baseline_ms / candidate_ms, + "baseline_rtf": baseline_rtf, + "candidate_rtf": candidate_rtf, + "rtf_speedup": baseline_rtf / candidate_rtf, + "wav_cosine": wav_cosine, + "logmel_cosine": logmel_cosine, + } + requests.append(request) + failed = failed or (args.require_same_frames and not same_frames) + failed = failed or (args.min_wav_cosine is not None and wav_cosine < args.min_wav_cosine) + failed = failed or (args.min_logmel_cosine is not None and logmel_cosine < args.min_logmel_cosine) + print( + f"request={index} frames={baseline_samples.size}/{candidate_samples.size} " + f"wall_ms={baseline_ms:.3f}/{candidate_ms:.3f} " + f"rtf={baseline_rtf:.4f}/{candidate_rtf:.4f} speedup={baseline_rtf / candidate_rtf:.3f}x " + f"wav_cos={wav_cosine:.8f} mel_cos={logmel_cosine:.8f}" + ) + + payload = {"baseline": str(args.baseline), "candidate": str(args.candidate), "requests": requests} + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/higgs_tts/higgs_tts_cuda_bench_cases.json b/tests/higgs_tts/higgs_tts_cuda_bench_cases.json new file mode 100644 index 00000000..cc4ff314 --- /dev/null +++ b/tests/higgs_tts/higgs_tts_cuda_bench_cases.json @@ -0,0 +1,35 @@ +[ + { + "id": "clone_prefix_first", + "text": "Hello. This is the first retained prefix test.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 1234 + }, + { + "id": "clone_prefix_second", + "text": "The second request should reuse the cloned voice prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 2234 + }, + { + "id": "clone_prefix_third", + "text": "A third short sentence checks stable repeated generation.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 3234 + } +] diff --git a/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json b/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json new file mode 100644 index 00000000..345b8f3c --- /dev/null +++ b/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json @@ -0,0 +1,44 @@ +[ + { + "id": "clone_first", + "text": "This cloned request prepares the reusable reference prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6134 + }, + { + "id": "clone_reuse", + "text": "This cloned request reuses the same reference prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6234 + }, + { + "id": "unconditioned", + "text": "This request intentionally generates an unconditioned random voice.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6334 + }, + { + "id": "clone_after_unconditioned", + "text": "The cloned voice is rebuilt safely after the unconditioned request.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6434 + } +] diff --git a/tests/higgs_tts/higgs_tts_cuda_perf_cases.json b/tests/higgs_tts/higgs_tts_cuda_perf_cases.json new file mode 100644 index 00000000..83403b80 --- /dev/null +++ b/tests/higgs_tts/higgs_tts_cuda_perf_cases.json @@ -0,0 +1,57 @@ +[ + { + "id": "control_room_update", + "text": "The control room confirmed the overnight checks and the field team can restart the survey.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 1234 + }, + { + "id": "lab_briefing", + "text": "The lab briefing is ready. Please confirm the calibration notes before the afternoon check-in.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 2234 + }, + { + "id": "dispatch_note", + "text": "Dispatch logged the revised route. The north access road is clear and the receiver test can begin after lunch.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 3234 + }, + { + "id": "short_status", + "text": "All systems are ready for the next test.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 4234 + }, + { + "id": "weather_report", + "text": "Light rain is expected this evening, but tomorrow morning should remain calm and clear.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 5234 + } +] diff --git a/tests/higgs_tts/higgs_tts_warm_bench.cpp b/tests/higgs_tts/higgs_tts_warm_bench.cpp index 665e3881..73c400b9 100644 --- a/tests/higgs_tts/higgs_tts_warm_bench.cpp +++ b/tests/higgs_tts/higgs_tts_warm_bench.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -106,13 +107,14 @@ void set_optional_option( } } -engine::runtime::AudioBuffer read_reference_audio(const engine::io::json::Value & object) { +std::optional read_reference_audio( + const engine::io::json::Value & object) { auto reference_path = optional_string(object, "reference_audio"); if (reference_path.empty()) { reference_path = optional_string(object, "voice_ref"); } if (reference_path.empty()) { - throw std::runtime_error("Higgs TTS warmbench request missing field: reference_audio"); + return std::nullopt; } const auto wav = engine::audio::read_wav_f32(resolve_path(reference_path)); return engine::runtime::AudioBuffer{wav.sample_rate, wav.channels, wav.samples}; @@ -121,10 +123,12 @@ engine::runtime::AudioBuffer read_reference_audio(const engine::io::json::Value engine::runtime::TaskRequest make_request(const engine::io::json::Value & object) { engine::runtime::TaskRequest request; request.text_input = engine::runtime::Transcript{required_string(object, "text"), ""}; - request.voice = engine::runtime::VoiceCondition{}; - request.voice->speaker = engine::runtime::VoiceReference{}; - request.voice->speaker->audio = read_reference_audio(object); - request.options["reference_text"] = required_string(object, "reference_text"); + if (auto reference_audio = read_reference_audio(object); reference_audio.has_value()) { + request.voice = engine::runtime::VoiceCondition{}; + request.voice->speaker = engine::runtime::VoiceReference{}; + request.voice->speaker->audio = std::move(*reference_audio); + request.options["reference_text"] = required_string(object, "reference_text"); + } set_optional_option(request, object, "max_tokens", "max_tokens"); set_optional_option(request, object, "temperature", "temperature"); set_optional_option(request, object, "top_p", "top_p"); @@ -205,10 +209,18 @@ engine::io::json::Value step_json( if (!audio_path.empty()) { stem.emplace("audio", string(audio_path.string())); } + const auto & audio = *result.audio_output; + const double frames = static_cast( + audio.samples.size() / static_cast(std::max(1, audio.channels))); + const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; + const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; return engine::io::json::Value::make_object({ {"request_index", number(static_cast(request_index))}, {"stems", engine::io::json::Value::make_array({engine::io::json::Value::make_object(std::move(stem))})}, - {"metrics", engine::io::json::Value::make_object({{"wall_ms", number(wall_ms)}})}, + {"metrics", engine::io::json::Value::make_object({ + {"wall_ms", number(wall_ms)}, + {"rtf", number(rtf)}, + })}, }); } @@ -238,7 +250,19 @@ int main(int argc, char ** argv) { const int threads = int_arg(argc, argv, "--threads", 8); const int warmup = int_arg(argc, argv, "--warmup", 0); const int iterations = int_arg(argc, argv, "--iterations", 1); - const std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); + std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); + const std::filesystem::path request_sequence_file = + arg_value(argc, argv, "--request-sequence-file", ""); + if (request_sequence_json.empty() && !request_sequence_file.empty()) { + std::ifstream input(request_sequence_file, std::ios::binary); + if (!input) { + throw std::runtime_error( + "failed to open Higgs TTS request sequence: " + request_sequence_file.string()); + } + request_sequence_json.assign( + std::istreambuf_iterator(input), + std::istreambuf_iterator()); + } const std::filesystem::path output_dir = arg_value(argc, argv, "--output-dir", ""); const std::filesystem::path timing_path = arg_value(argc, argv, "--timing-file", "/tmp/higgs_tts_warm_bench_timing.log"); @@ -308,7 +332,16 @@ int main(int argc, char ** argv) { } timing_lines.push_back( "higgs_tts.cpp.request_" + std::to_string(request_index) + ".wall_ms=" + std::to_string(wall_ms)); - std::cout << "higgs_tts.cpp.wall_ms=" << wall_ms << "\n"; + const auto & audio = *last_result.audio_output; + const double frames = static_cast( + audio.samples.size() / static_cast(std::max(1, audio.channels))); + const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; + const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; + timing_lines.push_back( + "higgs_tts.cpp.request_" + std::to_string(request_index) + ".rtf=" + std::to_string(rtf)); + std::cout << "higgs_tts.cpp.request=" << request_index + << " wall_ms=" << wall_ms + << " rtf=" << rtf << "\n"; steps.push_back(step_json(last_result, static_cast(request_index), wall_ms, audio_path)); } diff --git a/tests/higgs_tts/run_cuda_performance.ps1 b/tests/higgs_tts/run_cuda_performance.ps1 new file mode 100644 index 00000000..456a3549 --- /dev/null +++ b/tests/higgs_tts/run_cuda_performance.ps1 @@ -0,0 +1,69 @@ +param( + [Parameter(Mandatory = $true)] + [string]$Model, + [string]$BuildDir = "build/windows-cuda-release", + [string]$Label = (Get-Date -Format "yyyyMMdd-HHmmss"), + [string]$Baseline = "", + [switch]$RequireSameFrames, + [int]$Device = 0, + [int]$Threads = 8, + [int]$Warmup = 1, + [int]$Iterations = 1 +) + +$ErrorActionPreference = "Stop" +$RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "../..")).Path +$ModelPath = (Resolve-Path $Model).Path +$Bench = Join-Path $RepoRoot "$BuildDir/bin/higgs_tts_warm_bench.exe" +$Cases = Join-Path $PSScriptRoot "higgs_tts_cuda_perf_cases.json" +$ResultDir = Join-Path $PSScriptRoot "results/$Label" + +if (-not (Test-Path -LiteralPath $Bench)) { + throw "Warmbench binary does not exist: $Bench" +} + +New-Item -ItemType Directory -Force (Join-Path $ResultDir "audio") | Out-Null +$PreviousErrorActionPreference = $ErrorActionPreference +$ErrorActionPreference = "Continue" +& $Bench ` + --model $ModelPath ` + --backend cuda ` + --device $Device ` + --threads $Threads ` + --warmup $Warmup ` + --iterations $Iterations ` + --request-sequence-file $Cases ` + --output-dir (Join-Path $ResultDir "audio") ` + --timing-file (Join-Path $ResultDir "timing.log") 2>&1 | + ForEach-Object { + if ($_ -is [System.Management.Automation.ErrorRecord]) { + $_.Exception.Message + } else { + $_ + } + } | + Tee-Object -FilePath (Join-Path $ResultDir "console.log") +$BenchExitCode = $LASTEXITCODE +$ErrorActionPreference = $PreviousErrorActionPreference +if ($BenchExitCode -ne 0) { + exit $BenchExitCode +} + +if ($Baseline) { + $BaselinePath = (Resolve-Path $Baseline).Path + $CompareArgs = @( + (Join-Path $PSScriptRoot "compare_warmbench_results.py"), + "--baseline", $BaselinePath, + "--candidate", $ResultDir, + "--output", (Join-Path $ResultDir "comparison.json") + ) + if ($RequireSameFrames) { + $CompareArgs += "--require-same-frames" + } + python @CompareArgs + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } +} + +Write-Host "Higgs CUDA performance artifacts: $ResultDir" diff --git a/tests/unittests/test_qwen_decoder_packed_projections.cpp b/tests/unittests/test_qwen_decoder_packed_projections.cpp new file mode 100644 index 00000000..c34c4535 --- /dev/null +++ b/tests/unittests/test_qwen_decoder_packed_projections.cpp @@ -0,0 +1,504 @@ +#include "engine/framework/core/backend.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr size_t kGraphBytes = 16 * 1024 * 1024; +constexpr size_t kGraphNodes = 4096; + +std::vector patterned(size_t count, float phase, float scale) { + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + const float x = static_cast(i); + values[i] = scale * (std::sin(phase + 0.19f * x) + 0.35f * std::cos(phase + 0.07f * x)); + } + return values; +} + +void require_allclose( + const std::vector & actual, + const std::vector & expected, + float tolerance, + const std::string & label) { + if (actual.size() != expected.size()) { + throw std::runtime_error(label + " size mismatch"); + } + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + if (diff > tolerance) { + std::ostringstream message; + message << label << " mismatch at " << i << ": expected " << expected[i] + << ", got " << actual[i] << ", diff=" << diff; + throw std::runtime_error(message.str()); + } + } +} + +struct LayerResult { + std::vector output; + std::vector key; + std::vector value; +}; + +LayerResult run_layer(bool packed) { + constexpr int64_t batch = 1; + constexpr int64_t steps = 3; + constexpr int64_t hidden = 8; + constexpr int64_t heads = 2; + constexpr int64_t kv_heads = 1; + constexpr int64_t head_dim = 4; + constexpr int64_t intermediate = 12; + constexpr int64_t q_out = heads * head_dim; + constexpr int64_t kv_out = kv_heads * head_dim; + + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "qwen_packed_projection_test", engine::core::BackendType::Cpu}; + auto make_f32 = [&](std::initializer_list dims) { + return engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims(dims)); + }; + + auto input = make_f32({batch, steps, hidden}); + auto positions = engine::core::make_tensor( + ctx, + GGML_TYPE_I32, + engine::core::TensorShape::from_dims({steps})); + + engine::modules::QwenDecoderLayerWeights weights; + weights.input_norm = {make_f32({hidden}), std::nullopt}; + weights.post_norm = {make_f32({hidden}), std::nullopt}; + weights.self_attention.out_weight = make_f32({hidden, hidden}); + weights.mlp.down_proj = {make_f32({hidden, intermediate}), std::nullopt}; + + const auto q_values = patterned(static_cast(q_out * hidden), 0.1f, 0.12f); + const auto k_values = patterned(static_cast(kv_out * hidden), 0.5f, 0.10f); + const auto v_values = patterned(static_cast(kv_out * hidden), 0.9f, 0.08f); + const auto gate_values = patterned(static_cast(intermediate * hidden), 1.3f, 0.11f); + const auto up_values = patterned(static_cast(intermediate * hidden), 1.7f, 0.09f); + + if (packed) { + weights.self_attention.qkv_weight = make_f32({q_out + 2 * kv_out, hidden}); + weights.mlp.gate_up_proj = engine::modules::LinearWeights{ + make_f32({intermediate * 2, hidden}), + std::nullopt, + }; + } else { + weights.self_attention.q_weight = make_f32({q_out, hidden}); + weights.self_attention.k_weight = make_f32({kv_out, hidden}); + weights.self_attention.v_weight = make_f32({kv_out, hidden}); + weights.mlp.gate_proj = {make_f32({intermediate, hidden}), std::nullopt}; + weights.mlp.up_proj = {make_f32({intermediate, hidden}), std::nullopt}; + } + + engine::modules::QwenDecoderLayerConfig config; + config.hidden_size = hidden; + config.num_attention_heads = heads; + config.num_key_value_heads = kv_heads; + config.head_dim = head_dim; + config.intermediate_size = intermediate; + config.rms_norm_eps = 1e-5f; + config.qkv_layout = packed + ? engine::modules::QwenDecoderQKVLayout::PackedQKV + : engine::modules::QwenDecoderQKVLayout::Separate; + config.use_qk_norm = false; + config.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::ManualRepeat; + + const auto outputs = engine::modules::QwenDecoderLayerModule(config).build( + ctx, + input, + positions, + weights); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, outputs.output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate test tensors"); + } + + engine::core::write_tensor_f32(input, patterned(static_cast(batch * steps * hidden), 2.1f, 0.20f)); + engine::core::write_tensor_i32(positions, {0, 1, 2}); + engine::core::write_tensor_f32(*weights.input_norm.weight, patterned(hidden, 0.3f, 0.7f)); + engine::core::write_tensor_f32(*weights.post_norm.weight, patterned(hidden, 0.7f, 0.8f)); + engine::core::write_tensor_f32( + weights.self_attention.out_weight, + patterned(static_cast(hidden * hidden), 1.1f, 0.10f)); + engine::core::write_tensor_f32( + weights.mlp.down_proj.weight, + patterned(static_cast(hidden * intermediate), 1.9f, 0.10f)); + + if (packed) { + std::vector qkv_values; + qkv_values.reserve(q_values.size() + k_values.size() + v_values.size()); + qkv_values.insert(qkv_values.end(), q_values.begin(), q_values.end()); + qkv_values.insert(qkv_values.end(), k_values.begin(), k_values.end()); + qkv_values.insert(qkv_values.end(), v_values.begin(), v_values.end()); + engine::core::write_tensor_f32(*weights.self_attention.qkv_weight, qkv_values); + + std::vector gate_up_values; + gate_up_values.reserve(gate_values.size() + up_values.size()); + gate_up_values.insert(gate_up_values.end(), gate_values.begin(), gate_values.end()); + gate_up_values.insert(gate_up_values.end(), up_values.begin(), up_values.end()); + engine::core::write_tensor_f32(weights.mlp.gate_up_proj->weight, gate_up_values); + } else { + engine::core::write_tensor_f32(weights.self_attention.q_weight, q_values); + engine::core::write_tensor_f32(weights.self_attention.k_weight, k_values); + engine::core::write_tensor_f32(weights.self_attention.v_weight, v_values); + engine::core::write_tensor_f32(weights.mlp.gate_proj.weight, gate_values); + engine::core::write_tensor_f32(weights.mlp.up_proj.weight, up_values); + } + + ggml_backend_graph_compute(backend, graph); + LayerResult result; + engine::core::read_tensor_f32_into(outputs.output.tensor, result.output); + engine::core::read_tensor_f32_into(outputs.key.tensor, result.key); + engine::core::read_tensor_f32_into(outputs.value.tensor, result.value); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + return result; + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +void test_packed_qkv_and_gate_up_match_separate_projections() { + const auto separate = run_layer(false); + const auto packed = run_layer(true); + require_allclose(packed.output, separate.output, 2.0e-5f, "decoder output"); + require_allclose(packed.key, separate.key, 2.0e-5f, "decoder key"); + require_allclose(packed.value, separate.value, 2.0e-5f, "decoder value"); +} + +void test_suffix_causal_mask() { + const auto values = engine::modules::qwen_causal_suffix_mask_values(2, 3, 2); + if (values.size() != 30) { + throw std::runtime_error("suffix causal mask size mismatch"); + } + const std::vector expected{ + true, true, true, false, false, + true, true, true, true, false, + true, true, true, true, true, + }; + for (int batch = 0; batch < 2; ++batch) { + for (size_t i = 0; i < expected.size(); ++i) { + const float actual = ggml_fp16_to_fp32(values[static_cast(batch) * expected.size() + i]); + if ((expected[i] && actual != 0.0F) || (!expected[i] && !std::isinf(actual))) { + throw std::runtime_error("suffix causal mask visibility mismatch"); + } + } + } +} + +void test_f16_kv_set_rows() { + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_test", engine::core::BackendType::Cpu}; + const auto cache = engine::core::make_tensor( + ctx, + GGML_TYPE_F16, + engine::core::TensorShape::from_dims({1, 3, 1, 2})); + const auto row = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({1, 1, 1, 2})); + const auto row_index = engine::core::make_tensor( + ctx, + GGML_TYPE_I64, + engine::core::TensorShape::from_dims({1})); + const auto output = engine::modules::FastKVSetRowsModule{}.build(ctx, cache, row, row_index); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate f16 KV test tensors"); + } + + engine::core::write_tensor_f16(cache, std::vector(6, 0.0F)); + engine::core::write_tensor_f32(row, {1.25F, -2.5F}); + const int64_t index = 1; + ggml_backend_tensor_set(row_index.tensor, &index, 0, sizeof(index)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("f16 KV set-rows graph compute failed"); + } + + const auto values = engine::core::read_tensor_f16(output.tensor); + require_allclose(values, {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F}, 1.0e-3F, "f16 KV cache"); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +void test_f16_kv_set_rows_batched() { + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_batched_test", engine::core::BackendType::Cpu}; + const auto cache = engine::core::make_tensor( + ctx, + GGML_TYPE_F16, + engine::core::TensorShape::from_dims({2, 3, 1, 2})); + const auto rows = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({2, 1, 1, 2})); + const auto row_indices = engine::core::make_tensor( + ctx, + GGML_TYPE_I64, + engine::core::TensorShape::from_dims({2})); + const auto output = engine::modules::FastKVSetRowsModule{}.build(ctx, cache, rows, row_indices); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate batched F16 KV test tensors"); + } + + engine::core::write_tensor_f16(cache, std::vector(12, 0.0F)); + engine::core::write_tensor_f32(rows, {1.25F, -2.5F, 3.5F, -4.5F}); + const std::vector indices{1, 4}; + ggml_backend_tensor_set(row_indices.tensor, indices.data(), 0, indices.size() * sizeof(int64_t)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("batched F16 KV set-rows graph compute failed"); + } + + const auto values = engine::core::read_tensor_f16(output.tensor); + require_allclose( + values, + {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F, + 0.0F, 0.0F, 3.5F, -4.5F, 0.0F, 0.0F}, + 1.0e-3F, + "batched f16 KV cache"); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +int count_graph_op(ggml_cgraph * graph, ggml_op op) { + int count = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + count += node != nullptr && node->op == op ? 1 : 0; + } + return count; +} + +bool graph_contains_sequence(ggml_cgraph * graph, std::initializer_list ops) { + if (ops.size() == 0 || static_cast(ops.size()) > ggml_graph_n_nodes(graph)) { + return false; + } + for (int start = 0; start + static_cast(ops.size()) <= ggml_graph_n_nodes(graph); ++start) { + bool matches = true; + int offset = 0; + for (const ggml_op op : ops) { + const ggml_tensor * node = ggml_graph_node(graph, start + offset++); + matches = matches && node != nullptr && node->op == op; + } + if (matches) { + return true; + } + } + return false; +} + +std::string graph_ops(ggml_cgraph * graph) { + std::ostringstream out; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + if (i != 0) { + out << ','; + } + out << (node != nullptr ? ggml_op_name(node->op) : "null"); + } + return out.str(); +} + +void test_higgs_decode_graph_exposes_cuda_fast_paths() { + constexpr int64_t hidden = 8; + constexpr int64_t heads = 2; + constexpr int64_t kv_heads = 1; + constexpr int64_t head_dim = 4; + constexpr int64_t intermediate = 12; + constexpr int64_t cache_steps = 8; + constexpr int64_t qkv_out = heads * head_dim + 2 * kv_heads * head_dim; + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + throw std::runtime_error("failed to initialize Higgs decode graph test context"); + } + + try { + engine::core::ModuleBuildContext ctx{ggml, "higgs_decode_fast_path_test", engine::core::BackendType::Cuda}; + auto make_tensor = [&](ggml_type type, std::initializer_list dims) { + return engine::core::make_tensor(ctx, type, engine::core::TensorShape::from_dims(dims)); + }; + + const auto input = make_tensor(GGML_TYPE_F32, {1, 1, hidden}); + const auto positions = make_tensor(GGML_TYPE_I32, {1}); + const auto cache_key = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim}); + const auto cache_value = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim}); + const auto cache_slot = make_tensor(GGML_TYPE_I64, {1}); + const auto attention_mask = make_tensor(GGML_TYPE_F16, {1, 1, 1, cache_steps}); + + engine::modules::QwenDecoderLayerWeights weights; + weights.input_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt}; + weights.q_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt}; + weights.k_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt}; + weights.post_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt}; + weights.self_attention.qkv_weight = make_tensor(GGML_TYPE_F32, {qkv_out, hidden}); + weights.self_attention.out_weight = make_tensor(GGML_TYPE_F32, {hidden, hidden}); + weights.mlp.gate_up_proj = engine::modules::LinearWeights{ + make_tensor(GGML_TYPE_F32, {intermediate * 2, hidden}), + std::nullopt, + }; + weights.mlp.down_proj = { + make_tensor(GGML_TYPE_F32, {hidden, intermediate}), + std::nullopt, + }; + + engine::modules::QwenDecoderLayerConfig config; + config.hidden_size = hidden; + config.num_attention_heads = heads; + config.num_key_value_heads = kv_heads; + config.head_dim = head_dim; + config.intermediate_size = intermediate; + config.qkv_layout = engine::modules::QwenDecoderQKVLayout::PackedQKV; + config.use_qk_norm = true; + config.runtime.attention.static_mode = + engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + config.runtime.static_cache.update_mode = + engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + const auto outputs = engine::modules::QwenDecoderLayerModule(config).build_with_static_cache_tail( + ctx, + graph, + input, + positions, + weights, + cache_key, + cache_value, + cache_slot, + attention_mask); + ggml_build_forward_expand(graph, outputs.output.tensor); + + if (count_graph_op(graph, GGML_OP_FLASH_ATTN_EXT) != 1) { + throw std::runtime_error("Higgs decode graph must contain one grouped FlashAttention op"); + } + if (count_graph_op(graph, GGML_OP_SET_ROWS) != 2) { + throw std::runtime_error("Higgs decode graph must update both F16 KV caches with set-rows"); + } + if (count_graph_op(graph, GGML_OP_GLU) != 1) { + throw std::runtime_error("Higgs decode graph must contain one packed SwiGLU op"); + } + if (count_graph_op(graph, GGML_OP_REPEAT) != 0) { + throw std::runtime_error("grouped FlashAttention must not materialize repeated KV heads"); + } + if (!graph_contains_sequence(graph, {GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS})) { + throw std::runtime_error( + "Higgs key-cache update must expose CUDA RoPE/view/set-rows fusion; graph=" + + graph_ops(graph)); + } + + ggml_free(ggml); + } catch (...) { + ggml_free(ggml); + throw; + } +} + +} // namespace + +int main() { + try { + test_packed_qkv_and_gate_up_match_separate_projections(); + test_suffix_causal_mask(); + test_f16_kv_set_rows(); + test_f16_kv_set_rows_batched(); + test_higgs_decode_graph_exposes_cuda_fast_paths(); + std::cout << "qwen_decoder_packed_projection_test: ok\n"; + return 0; + } catch (const std::exception & ex) { + std::cerr << "qwen_decoder_packed_projection_test: failed: " << ex.what() << "\n"; + return 1; + } +}