Optimize Higgs Audio v3 inference

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(

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
 

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<int32_t> & values);
 void read_tensor_f32_into(const ggml_tensor * tensor, std::vector<float> & values);
 std::vector<float> read_tensor_f32(const ggml_tensor * tensor);
+void read_tensor_f16_into(const ggml_tensor * tensor, std::vector<float> & values);
+std::vector<float> read_tensor_f16(const ggml_tensor * tensor);
 void read_tensor_i32_into(const ggml_tensor * tensor, std::vector<int32_t> & values);
 std::vector<int32_t> read_tensor_i32(const ggml_tensor * tensor);
 

include/engine/framework/modules/attention/qwen_causal_decoder.h

@@ -79,6 +79,11 @@ std::vector<int32_t> qwen_position_ids(int64_t steps, int64_t offset = 0);
 
 std::vector<ggml_fp16_t> qwen_causal_prefill_mask_values(int64_t batch_size, int64_t steps);
 
+std::vector<ggml_fp16_t> 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,

include/engine/framework/modules/attention/qwen_decoder.h

@@ -84,6 +84,7 @@ struct QwenDecoderLayerConfig {
 struct QwenMLPWeights {
     LinearWeights gate_proj;
     LinearWeights up_proj;
+    std::optional<LinearWeights> gate_up_proj;
     LinearWeights down_proj;
 };
 

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(

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<ReferencePrefixCache> reference_prefix_cache_;
+    bool reference_kv_ready_ = false;
     std::optional<HiggsCudaSamplingPolicy> cuda_sampling_policy_;
     std::unique_ptr<HiggsARKVCache> ar_kv_cache_;
     std::unique_ptr<HiggsARPrefillGraph> prefill_graph_;

src/framework/core/backend.cpp

@@ -527,10 +527,22 @@ void read_tensor_f32_into(const ggml_tensor * tensor, std::vector<float> & value
     read_tensor_typed_into<float>(tensor, GGML_TYPE_F32, values);
 }
 
-std::vector<float> read_tensor_f32(const ggml_tensor * tensor) {
-    return read_tensor_typed<float>(tensor, GGML_TYPE_F32);
-}
-
+std::vector<float> read_tensor_f32(const ggml_tensor * tensor) {
+    return read_tensor_typed<float>(tensor, GGML_TYPE_F32);
+}
+
+void read_tensor_f16_into(const ggml_tensor * tensor, std::vector<float> & values) {
+    const auto fp16_values = read_tensor_typed<ggml_fp16_t>(tensor, GGML_TYPE_F16);
+    values.resize(fp16_values.size());
+    ggml_fp16_to_fp32_row(fp16_values.data(), values.data(), static_cast<int64_t>(values.size()));
+}
+
+std::vector<float> read_tensor_f16(const ggml_tensor * tensor) {
+    std::vector<float> values;
+    read_tensor_f16_into(tensor, values);
+    return values;
+}
+
 void read_tensor_i32_into(const ggml_tensor * tensor, std::vector<int32_t> & values) {
     read_tensor_typed_into<int32_t>(tensor, GGML_TYPE_I32, values);
 }

src/framework/modules/attention/qwen_causal_decoder.cpp

@@ -175,6 +175,36 @@ std::vector<ggml_fp16_t> qwen_causal_prefill_mask_values(int64_t batch_size, int
     return out;
 }
 
+std::vector<ggml_fp16_t> 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<ggml_fp16_t> one(static_cast<size_t>(query_steps * key_steps), masked);
+    for (int64_t row = 0; row < query_steps; ++row) {
+        const size_t row_offset = static_cast<size_t>(row * key_steps);
+        std::fill_n(
+            one.begin() + static_cast<std::ptrdiff_t>(row_offset),
+            prefix_steps + row + 1,
+            visible);
+    }
+    std::vector<ggml_fp16_t> out;
+    out.reserve(static_cast<size_t>(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,

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<core::TensorValue> 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);

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);
 }