| 1091 | } |
| 1092 | |
| 1093 | HeartMuLaDecoderPrefillOutput run(const std::vector<float> & embeddings, int64_t codebook_index) { |
| 1094 | allocate_workspace(); |
| 1095 | const auto & full_config = runtime_->assets().mula_config; |
| 1096 | const auto & config = full_config.decoder; |
| 1097 | const int64_t head_dim = require_head_dim(config); |
| 1098 | if (static_cast<int64_t>(embeddings.size()) != batch_size_ * steps_ * full_config.backbone.embed_dim) { |
| 1099 | throw std::runtime_error("HeartMuLa decoder prefill embedding payload size mismatch"); |
| 1100 | } |
| 1101 | if (codebook_index < 0 || codebook_index >= full_config.audio_num_codebooks - 1) { |
| 1102 | throw std::runtime_error("HeartMuLa decoder prefill codebook index is out of range"); |
| 1103 | } |
| 1104 | ggml_backend_tensor_set(input_, embeddings.data(), 0, embeddings.size() * sizeof(float)); |
| 1105 | const int32_t codebook = static_cast<int32_t>(codebook_index); |
| 1106 | ggml_backend_tensor_set(codebook_index_, &codebook, 0, sizeof(int32_t)); |
| 1107 | core::set_backend_threads(runtime_->backend(), runtime_->threads()); |
| 1108 | const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); |
| 1109 | ggml_backend_synchronize(runtime_->backend()); |
| 1110 | if (status != GGML_STATUS_SUCCESS) { |
| 1111 | throw std::runtime_error("HeartMuLa decoder prefill graph compute failed"); |
| 1112 | } |
| 1113 | HeartMuLaDecoderPrefillOutput out; |
| 1114 | out.result.logits.vocab_size = full_config.audio_vocab_size; |
| 1115 | out.result.logits.values.resize(static_cast<size_t>(batch_size_ * full_config.audio_vocab_size)); |
| 1116 | ggml_backend_tensor_get( |
| 1117 | logits_output_, |
| 1118 | out.result.logits.values.data(), |
| 1119 | 0, |
| 1120 | out.result.logits.values.size() * sizeof(float)); |
| 1121 | out.result.last_hidden.dims = config.embed_dim; |
| 1122 | out.result.last_hidden.values.resize(static_cast<size_t>(batch_size_ * config.embed_dim)); |
| 1123 | ggml_backend_tensor_get( |
| 1124 | hidden_output_, |
| 1125 | out.result.last_hidden.values.data(), |
| 1126 | 0, |
| 1127 | out.result.last_hidden.values.size() * sizeof(float)); |
| 1128 | out.state.current_end = steps_; |
| 1129 | out.state.layers.resize(keys_.size()); |
| 1130 | const size_t layer_values = static_cast<size_t>(steps_ * layer_step_elems_); |
| 1131 | for (size_t layer = 0; layer < keys_.size(); ++layer) { |
| 1132 | auto & state_layer = out.state.layers[layer]; |
| 1133 | state_layer.valid_steps = steps_; |
| 1134 | std::vector<float> key_values(layer_values); |
| 1135 | std::vector<float> value_values(layer_values); |
| 1136 | ggml_backend_tensor_get(keys_[layer], key_values.data(), 0, key_values.size() * sizeof(float)); |
| 1137 | ggml_backend_tensor_get(values_[layer], value_values.data(), 0, value_values.size() * sizeof(float)); |
| 1138 | state_layer.key = flatten_prefill_cache_by_step( |
| 1139 | key_values, |
| 1140 | batch_size_, |
| 1141 | steps_, |
| 1142 | config.num_kv_heads, |
| 1143 | head_dim); |
| 1144 | state_layer.value = flatten_prefill_cache_by_step( |
| 1145 | value_values, |
| 1146 | batch_size_, |
| 1147 | steps_, |
| 1148 | config.num_kv_heads, |
| 1149 | head_dim); |
| 1150 | } |
no test coverage detected