From bddd0f8447975d9ede049374b396d433e79697ba Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Mon, 17 Aug 2026 07:21:54 +0530 Subject: [PATCH 01/19] Enable frame stacking to support magpie tts v2607 Signed-off-by: Anand Joseph --- app/serve.cpp | 12 +- app/synthesize.cpp | 7 +- conversion/tts.py | 31 ++++- docs/tts/models.md | 16 +++ server/http/http_server.cpp | 18 +++ server/http/http_server.h | 2 + src/tts/magpietts/decoder.cpp | 115 +++++++++++------- src/tts/magpietts/lt.cpp | 28 +++-- src/tts/magpietts/magpietts.cpp | 67 +++++----- src/tts/magpietts/model.cpp | 81 ++++++++---- src/tts/magpietts/model.h | 43 +++++++ tests/cpp/tts/CMakeLists.txt | 3 + .../cpp/tts/test_magpietts_frame_stacking.cpp | 37 ++++++ 13 files changed, 346 insertions(+), 114 deletions(-) create mode 100644 tests/cpp/tts/test_magpietts_frame_stacking.cpp diff --git a/app/serve.cpp b/app/serve.cpp index 9552c1b..1f31665 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -498,8 +498,13 @@ run_server(int argc, char** argv) { tts_config.runtime.codec_cpu = gpu < 0; } else { tts_config.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cuda; - tts_config.runtime.sampling_backend = - nemo_speech::tts::MagpieBackendPreference::Cuda; + // CUDA decoding does not imply CUDA local-transformer sampling. Keep sampling + // on CPU unless the TTS-specific setting explicitly requests CUDA. + if (tts_config.runtime.sampling_backend != + nemo_speech::tts::MagpieBackendPreference::Cuda) { + tts_config.runtime.sampling_backend = + nemo_speech::tts::MagpieBackendPreference::Cpu; + } } } tts_config.runtime.magpie_model = magpie_path; @@ -513,6 +518,9 @@ run_server(int argc, char** argv) { config.default_voice_name = tts_config.default_voice_name; engines.load_tts(std::move(config)); } + // The HTTP server owns request handling, so carry the existing TTS + // benchmark switch across from the TTS server configuration. + server_config.tts_benchmark = tts_config.benchmark; #endif if (!engines.ready()) throw std::runtime_error( diff --git a/app/synthesize.cpp b/app/synthesize.cpp index 5c7388a..4f6becc 100644 --- a/app/synthesize.cpp +++ b/app/synthesize.cpp @@ -211,7 +211,12 @@ command_synthesize(int argc, char** argv) { parsed.runtime.codec_cpu = true; } else if (cuda_device) { parsed.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cuda; - parsed.runtime.sampling_backend = nemo_speech::tts::MagpieBackendPreference::Cuda; + // Preserve CPU sampling by default; CUDA sampling is opt-in via + // --tts.sampling-backend cuda. + if (parsed.runtime.sampling_backend != + nemo_speech::tts::MagpieBackendPreference::Cuda) { + parsed.runtime.sampling_backend = nemo_speech::tts::MagpieBackendPreference::Cpu; + } parsed.runtime.codec_cpu = false; } else { parsed.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cpu; diff --git a/conversion/tts.py b/conversion/tts.py index bee7b17..ce0d3b1 100644 --- a/conversion/tts.py +++ b/conversion/tts.py @@ -93,10 +93,29 @@ def add_metadata( decoder = cfg["decoder"] lt_hidden = int(cfg.get("local_transformer_hidden_dim", 256)) frame_stacking = int(cfg.get("frame_stacking_factor", 1)) - n_codebooks = int( + n_stacked_codebooks = int( len([k for k in sd if k.startswith("audio_embeddings.") and k.endswith(".weight")]) ) - n_codebooks //= frame_stacking + if frame_stacking < 1 or n_stacked_codebooks == 0 or n_stacked_codebooks % frame_stacking: + raise ValueError( + "audio embedding count must be a positive multiple of frame_stacking_factor: " + f"embeddings={n_stacked_codebooks} frame_stacking_factor={frame_stacking}" + ) + n_codebooks = n_stacked_codebooks // frame_stacking + expected_logits = n_stacked_codebooks * audio_vocab + if int(sd["final_proj.weight"].shape[0]) != expected_logits: + raise ValueError( + "final_proj rows do not match stacked audio layout: " + f"rows={sd['final_proj.weight'].shape[0]} expected={expected_logits}" + ) + n_lt_heads = len( + [k for k in sd if k.startswith("local_transformer_out_projections.") and k.endswith(".weight")] + ) + if n_lt_heads != n_stacked_codebooks: + raise ValueError( + "local transformer output heads do not match audio embeddings: " + f"heads={n_lt_heads} embeddings={n_stacked_codebooks}" + ) inf = cfg.get("inference_parameters", {}) @@ -108,6 +127,7 @@ def add_metadata( "codec_model": cfg.get("codecmodel_path"), "text_vocab_size": text_vocab, "audio_codebooks": n_codebooks, + "stacked_audio_codebooks": n_stacked_codebooks, "audio_codebook_size": codebook_size, "audio_vocab_size": audio_vocab, "frame_stacking_factor": frame_stacking, @@ -135,6 +155,7 @@ def add_metadata( add_i32(writer, "magpietts.text_vocab_size", text_vocab) add_i32(writer, "magpietts.audio_codebooks", n_codebooks) + add_i32(writer, "magpietts.stacked_audio_codebooks", n_stacked_codebooks) add_i32(writer, "magpietts.audio_codebook_size", codebook_size) add_i32(writer, "magpietts.audio_vocab_size", audio_vocab) add_i32(writer, "magpietts.audio_bos_id", codebook_size + 0) @@ -308,9 +329,9 @@ def convert( if cfg.get("target") != "nemo.collections.tts.models.magpietts.MagpieTTSModel": raise RuntimeError(f"unsupported target: {cfg.get('target')}") - if cfg.get("model_type") != "decoder_ce" or not cfg.get( - "has_baked_context_embedding", False - ): + # v2602 records this in the config while v2607 retains the baked + # tensor but omits the legacy config flag. + if cfg.get("model_type") != "decoder_ce" or "baked_context_embedding.weight" not in sd: raise RuntimeError( "this GGML example expects MagpieTTS decoder_ce with baked context embeddings" ) diff --git a/docs/tts/models.md b/docs/tts/models.md index 19e11be..07a17ad 100644 --- a/docs/tts/models.md +++ b/docs/tts/models.md @@ -25,6 +25,22 @@ tar -xf models/magpie-tts/magpie_tts_multilingual_357m.nemo \ -C models/magpie-tts/extracted ``` +MagpieTTS v2607 uses factor-2 frame stacking and must currently be converted +locally before use: + +```bash +hf download nvidia/magpie_tts_multilingual_357m \ + magpie_tts_multilingual_357m.nemo \ + --revision v2607 --local-dir models/magpie-tts-v2607 +python3 convert_model.py models/magpie-tts-v2607/magpie_tts_multilingual_357m.nemo \ + --outfile models/magpie-tts-v2607/magpie_tts_multilingual_357m.v2607.f16.gguf +mkdir -p models/magpie-tts-v2607/extracted +tar -xf models/magpie-tts-v2607/magpie_tts_multilingual_357m.nemo \ + -C models/magpie-tts-v2607/extracted +``` + +Both v2602 (factor 1) and v2607 (factor 2) use the same NanoCodec decoder. + **Tokenizer.** MagpieTTS's tokenizer assets live *inside* the `.nemo` archive - they are not part of the GGUF. Extract the `.nemo` and pass that directory to the server as `--tts.tokenizer-model-dir` (here diff --git a/server/http/http_server.cpp b/server/http/http_server.cpp index a962c22..7d7da2e 100644 --- a/server/http/http_server.cpp +++ b/server/http/http_server.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -644,6 +645,23 @@ struct Server::Impl { pcm += chunk; return true; }); + if (this->config.tts_benchmark) { + const auto& stats = result.stats; + std::cerr << "[nemo_http][tts_benchmark]" + << " text_chars=" << result.metadata.original_text.size() + << " frames=" << stats.generated_frames + << " audio_s=" << stats.audio_s + << " decoder_ttft_ms=" << stats.decoder_ttft_ms + << " decoder_itl_avg_ms=" << stats.decoder_itl_avg_ms + << " decoder_itl_p95_ms=" << stats.decoder_itl_p95_ms + << " decoder_itl_p99_ms=" << stats.decoder_itl_p99_ms + << " codec_ttfa_ms=" << stats.codec_ttfa_ms + << " codec_icl_avg_ms=" << stats.codec_icl_avg_ms + << " codec_icl_p95_ms=" << stats.codec_icl_p95_ms + << " codec_icl_p99_ms=" << stats.codec_icl_p99_ms + << " e2e_ttfa_ms=" << stats.e2e_ttfa_ms + << " e2e_rtfx=" << stats.e2e_rtfx << '\n'; + } if (format == "pcm") response.set_content(std::move(pcm), "audio/pcm"); else diff --git a/server/http/http_server.h b/server/http/http_server.h index 2c04bfb..6d8776b 100644 --- a/server/http/http_server.h +++ b/server/http/http_server.h @@ -19,6 +19,8 @@ struct ServerConfig { size_t max_upload_bytes = 512ULL * 1024ULL * 1024ULL; bool access_log = false; bool json_logs = false; + // Emit per-request TTS runtime timings to stderr. + bool tts_benchmark = false; std::string tls_certificate; std::string tls_private_key; std::string api_key; diff --git a/src/tts/magpietts/decoder.cpp b/src/tts/magpietts/decoder.cpp index 7e71a07..2f7a38f 100644 --- a/src/tts/magpietts/decoder.cpp +++ b/src/tts/magpietts/decoder.cpp @@ -502,11 +502,40 @@ build_audio_embedding( ggml_context* ctx, const magpietts_model& model, const std::vector& audio_tok_inputs) { ggml_tensor* sum = nullptr; - for (int c = 0; c < model.hparams.audio_codebooks; ++c) { + for (int c = 0; c < model.hparams.stacked_audio_codebooks(); ++c) { ggml_tensor* emb = ggml_get_rows(ctx, model.audio_embeddings[c], audio_tok_inputs[c]); sum = sum ? ggml_add(ctx, sum, emb) : emb; } - return ggml_scale(ctx, sum, 1.0f / (float)model.hparams.audio_codebooks); + return ggml_scale(ctx, sum, 1.0f / (float)model.hparams.stacked_audio_codebooks()); +} + +static bool +stack_audio_codes( + const std::vector>& audio_codes, const magpietts_hparams& h, + std::vector>& stacked) { + if ((int)audio_codes.size() != h.audio_codebooks || audio_codes.empty() || + audio_codes[0].empty() || + (int)audio_codes[0].size() % h.frame_stacking_factor != 0) { + return false; + } + const int raw_len = (int)audio_codes[0].size(); + for (const auto& codes : audio_codes) { + if ((int)codes.size() != raw_len) { + return false; + } + } + const int stacked_len = raw_len / h.frame_stacking_factor; + stacked.assign((size_t)h.stacked_audio_codebooks(), std::vector(stacked_len)); + for (int lane = 0; lane < h.frame_stacking_factor; ++lane) { + for (int codebook = 0; codebook < h.audio_codebooks; ++codebook) { + auto& dst = stacked[(size_t)(codebook + lane * h.audio_codebooks)]; + const auto& src = audio_codes[(size_t)codebook]; + for (int pos = 0; pos < stacked_len; ++pos) { + dst[(size_t)pos] = src[(size_t)(pos * h.frame_stacking_factor + lane)]; + } + } + } + return true; } static bool @@ -519,17 +548,18 @@ decoder_eval_impl( const ggml_nvtx::range nvtx_range( conditional ? "magpietts_decoder_eval_cond" : "magpietts_decoder_eval_uncond"); const auto& h = model.hparams; - const int audio_len = audio_codes.empty() ? 0 : (int)audio_codes[0].size(); - if (audio_len <= 0) { + std::vector> stacked_audio; + if (!stack_audio_codes(audio_codes, h, stacked_audio)) { fprintf(stderr, "decoder_eval requires at least one audio token\n"); return false; } + const int audio_len = (int)stacked_audio[0].size(); const int total_len = h.baked_context_length + audio_len; ggml_context* ctx = new_graph_context(); ggml_cgraph* gf = ggml_new_graph_custom(ctx, MAGPIETTS_MAX_NODES, false); - std::vector audio_tok_inputs(h.audio_codebooks); + std::vector audio_tok_inputs(h.stacked_audio_codebooks()); std::vector>> i32_inputs; std::vector>> f32_inputs; @@ -538,12 +568,12 @@ decoder_eval_impl( ggml_set_input(speaker_in); i32_inputs.push_back({"magpietts_decoder_speaker", {speaker}}); - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { const std::string name = "magpietts_decoder_audio_tokens_" + std::to_string(c); audio_tok_inputs[c] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, audio_len); ggml_set_name(audio_tok_inputs[c], name.c_str()); ggml_set_input(audio_tok_inputs[c]); - i32_inputs.push_back({name, audio_codes[c]}); + i32_inputs.push_back({name, stacked_audio[c]}); } ggml_tensor* pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, total_len); @@ -623,14 +653,14 @@ decoder_eval_impl( return false; } - const size_t logits_last_size = (size_t)h.audio_codebooks * h.audio_vocab_size; + const size_t logits_last_size = (size_t)h.stacked_audio_codebooks() * h.audio_vocab_size; const size_t logits_off_floats = logits_last_size * (total_len - 1); if (hidden_out && hidden_last) { ggml_backend_tensor_copy(hidden_last, hidden_out->tensor); } if (cuda_sample) { const bool sampled = MagpieCodebookSampler::runCuda( - model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, h.audio_codebooks, + model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, h.stacked_audio_codebooks(), 0); ggml_gallocr_free(allocr); ggml_free(ctx); @@ -666,17 +696,18 @@ decoder_eval_pair_impl( MagpiePinnedHostScratch& output_staging, const magpietts_decoder_attention* attention) { const ggml_nvtx::range nvtx_range("magpietts_decoder_eval_pair"); const auto& h = model.hparams; - const int audio_len = audio_codes.empty() ? 0 : (int)audio_codes[0].size(); - if (audio_len <= 0) { + std::vector> stacked_audio; + if (!stack_audio_codes(audio_codes, h, stacked_audio)) { fprintf(stderr, "decoder_eval_pair requires at least one audio token\n"); return false; } + const int audio_len = (int)stacked_audio[0].size(); const int total_len = h.baked_context_length + audio_len; ggml_context* ctx = new_graph_context(); ggml_cgraph* gf = ggml_new_graph_custom(ctx, MAGPIETTS_MAX_NODES, false); - std::vector audio_tok_inputs(h.audio_codebooks); + std::vector audio_tok_inputs(h.stacked_audio_codebooks()); std::vector>> i32_inputs; std::vector>> f32_inputs; @@ -685,12 +716,12 @@ decoder_eval_pair_impl( ggml_set_input(speaker_in); i32_inputs.push_back({"magpietts_decoder_speaker", {speaker}}); - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { const std::string name = "magpietts_decoder_audio_tokens_" + std::to_string(c); audio_tok_inputs[c] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, audio_len); ggml_set_name(audio_tok_inputs[c], name.c_str()); ggml_set_input(audio_tok_inputs[c]); - i32_inputs.push_back({name, audio_codes[c]}); + i32_inputs.push_back({name, stacked_audio[c]}); } ggml_tensor* pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, total_len); @@ -786,7 +817,7 @@ decoder_eval_pair_impl( return false; } - const size_t logits_last_size = (size_t)h.audio_codebooks * h.audio_vocab_size; + const size_t logits_last_size = (size_t)h.stacked_audio_codebooks() * h.audio_vocab_size; const size_t logits_off_floats = logits_last_size * (total_len - 1); if (cond_hidden_out && cond_hidden_last) { ggml_backend_tensor_copy(cond_hidden_last, cond_hidden_out->tensor); @@ -797,7 +828,7 @@ decoder_eval_pair_impl( if (cuda_sample) { const bool sampled = MagpieCodebookSampler::runCuda( model.backend, h, cuda_sample, logits_cond, logits_uncond, logits_off_floats, - h.audio_codebooks, 0); + h.stacked_audio_codebooks(), 0); ggml_gallocr_free(allocr); ggml_free(ctx); return sampled; @@ -842,11 +873,12 @@ decoder_eval_cached_impl( conditional ? "magpietts_decoder_eval_cached_cond" : "magpietts_decoder_eval_cached_uncond"); const auto& h = model.hparams; - const int audio_len = audio_codes.empty() ? 0 : (int)audio_codes[0].size(); - if (audio_len <= 0) { + std::vector> stacked_audio; + if (!stack_audio_codes(audio_codes, h, stacked_audio)) { fprintf(stderr, "decoder_eval_cached requires at least one audio token\n"); return false; } + const int audio_len = (int)stacked_audio[0].size(); if (h.dec_kernel != 1) { return decoder_eval_impl( model, text_cond, text_len, audio_codes, speaker, conditional, threads, result, @@ -878,7 +910,7 @@ decoder_eval_cached_impl( ggml_context* ctx = new_graph_context(); ggml_cgraph* gf = ggml_new_graph_custom(ctx, MAGPIETTS_MAX_NODES, false); - std::vector audio_tok_inputs(h.audio_codebooks); + std::vector audio_tok_inputs(h.stacked_audio_codebooks()); std::vector>> i32_inputs; std::vector>> f32_inputs; @@ -895,22 +927,22 @@ decoder_eval_cached_impl( ctx_emb = ggml_scale(ctx, ctx_emb, 0.0f); } - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { const std::string name = "magpietts_decoder_audio_tokens_" + std::to_string(c); audio_tok_inputs[c] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_audio_in); ggml_set_name(audio_tok_inputs[c], name.c_str()); ggml_set_input(audio_tok_inputs[c]); - i32_inputs.push_back({name, audio_codes[c]}); + i32_inputs.push_back({name, stacked_audio[c]}); } ggml_tensor* audio_emb = build_audio_embedding(ctx, model, audio_tok_inputs); dec_in = ggml_concat(ctx, ctx_emb, audio_emb, 1); } else { - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { const std::string name = "magpietts_decoder_audio_tokens_" + std::to_string(c); audio_tok_inputs[c] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_audio_in); ggml_set_name(audio_tok_inputs[c], name.c_str()); ggml_set_input(audio_tok_inputs[c]); - i32_inputs.push_back({name, {audio_codes[c].back()}}); + i32_inputs.push_back({name, {stacked_audio[c].back()}}); } dec_in = build_audio_embedding(ctx, model, audio_tok_inputs); } @@ -983,14 +1015,14 @@ decoder_eval_cached_impl( return false; } - const size_t logits_last_size = (size_t)h.audio_codebooks * h.audio_vocab_size; + const size_t logits_last_size = (size_t)h.stacked_audio_codebooks() * h.audio_vocab_size; const size_t logits_off_floats = logits_last_size * (n_graph_tokens - 1); if (hidden_out && hidden_last) { ggml_backend_tensor_copy(hidden_last, hidden_out->tensor); } if (cuda_sample) { const bool sampled = MagpieCodebookSampler::runCuda( - model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, h.audio_codebooks, + model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, h.stacked_audio_codebooks(), 0); ggml_gallocr_free(allocr); ggml_free(ctx); @@ -1033,11 +1065,12 @@ decoder_eval_cached_pair_impl( MagpiePinnedHostScratch& output_staging, const magpietts_decoder_attention* attention) { const ggml_nvtx::range nvtx_range("magpietts_decoder_eval_cached_pair"); const auto& h = model.hparams; - const int audio_len = audio_codes.empty() ? 0 : (int)audio_codes[0].size(); - if (audio_len <= 0) { + std::vector> stacked_audio; + if (!stack_audio_codes(audio_codes, h, stacked_audio)) { fprintf(stderr, "decoder_eval_cached_pair requires at least one audio token\n"); return false; } + const int audio_len = (int)stacked_audio[0].size(); if (h.dec_kernel != 1) { return decoder_eval_pair_impl( model, text_cond, text_len, audio_codes, speaker, threads, cond_result, uncond_result, @@ -1090,7 +1123,7 @@ decoder_eval_cached_pair_impl( ggml_context* ctx = new_graph_context(); ggml_cgraph* gf = ggml_new_graph_custom(ctx, MAGPIETTS_MAX_NODES, false); - std::vector audio_tok_inputs(h.audio_codebooks); + std::vector audio_tok_inputs(h.stacked_audio_codebooks()); std::vector>> i32_inputs; std::vector>> f32_inputs; @@ -1102,12 +1135,12 @@ decoder_eval_cached_pair_impl( ggml_set_input(speaker_in); i32_inputs.push_back({"magpietts_decoder_speaker", {speaker}}); - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { const std::string name = "magpietts_decoder_audio_tokens_" + std::to_string(c); audio_tok_inputs[c] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_audio_in); ggml_set_name(audio_tok_inputs[c], name.c_str()); ggml_set_input(audio_tok_inputs[c]); - i32_inputs.push_back({name, audio_codes[c]}); + i32_inputs.push_back({name, stacked_audio[c]}); } ggml_tensor* ctx_flat = ggml_get_rows(ctx, model.baked_context, speaker_in); @@ -1118,12 +1151,12 @@ decoder_eval_cached_pair_impl( dec_in_cond = ggml_concat(ctx, ctx_emb_cond, audio_emb, 1); dec_in_uncond = ggml_concat(ctx, ctx_emb_uncond, audio_emb, 1); } else { - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { const std::string name = "magpietts_decoder_audio_tokens_" + std::to_string(c); audio_tok_inputs[c] = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_audio_in); ggml_set_name(audio_tok_inputs[c], name.c_str()); ggml_set_input(audio_tok_inputs[c]); - i32_inputs.push_back({name, {audio_codes[c].back()}}); + i32_inputs.push_back({name, {stacked_audio[c].back()}}); } ggml_tensor* audio_emb = build_audio_embedding(ctx, model, audio_tok_inputs); dec_in_cond = audio_emb; @@ -1219,7 +1252,7 @@ decoder_eval_cached_pair_impl( return false; } - const size_t logits_last_size = (size_t)h.audio_codebooks * h.audio_vocab_size; + const size_t logits_last_size = (size_t)h.stacked_audio_codebooks() * h.audio_vocab_size; const size_t logits_off_floats = logits_last_size * (n_graph_tokens - 1); if (cond_hidden_out && cond_hidden_last) { ggml_backend_tensor_copy(cond_hidden_last, cond_hidden_out->tensor); @@ -1230,7 +1263,7 @@ decoder_eval_cached_pair_impl( if (cuda_sample) { const bool sampled = MagpieCodebookSampler::runCuda( model.backend, h, cuda_sample, logits_cond, logits_uncond, logits_off_floats, - h.audio_codebooks, 0); + h.stacked_audio_codebooks(), 0); ggml_gallocr_free(allocr); ggml_free(ctx); if (sampled) { @@ -1374,14 +1407,14 @@ MagpieCodebookSampler::sampleParallel( const magpietts_hparams& h, bool use_cfg, float cfg_scale, float temperature, int top_k, bool forbid_audio_eos, std::mt19937& rng, std::vector* argmax_codes) { const ggml_nvtx::range nvtx_range("magpietts_sample_parallel_codebooks"); - std::vector codes(h.audio_codebooks); + std::vector codes(h.stacked_audio_codebooks()); if (argmax_codes) { - argmax_codes->assign(h.audio_codebooks, 0); + argmax_codes->assign(h.stacked_audio_codebooks(), 0); } std::vector> futures; - futures.reserve(h.audio_codebooks); - for (int c = 0; c < h.audio_codebooks; ++c) { + futures.reserve(h.stacked_audio_codebooks()); + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { futures.push_back(std::async(std::launch::async, [&, c]() { std::vector logits = slice_codebook_logits(cond_logits, h, c); if (use_cfg) { @@ -1395,11 +1428,11 @@ MagpieCodebookSampler::sampleParallel( })); } - std::vector prepared((size_t)h.audio_codebooks); - for (int c = 0; c < h.audio_codebooks; ++c) { + std::vector prepared((size_t)h.stacked_audio_codebooks()); + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { prepared[(size_t)c] = futures[(size_t)c].get(); } - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { codes[c] = sample_from_prepared_codebook(prepared[(size_t)c], temperature, rng); if (argmax_codes) { (*argmax_codes)[c] = prepared[(size_t)c].greedy; diff --git a/src/tts/magpietts/lt.cpp b/src/tts/magpietts/lt.cpp index 292768b..b2e6bc0 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -253,8 +253,8 @@ magpietts_model_init_local_transformer_copy( std::vector> copies; if (!magpietts_local_add_tensor_vector( dst.ctx, src.audio_embeddings, dst.audio_embeddings, copies, fp32) || - !magpietts_local_add_tensor(dst.ctx, src.lt_in_w, &dst.lt_in_w, copies, fp32) || - !magpietts_local_add_tensor(dst.ctx, src.lt_in_b, &dst.lt_in_b, copies, fp32) || + (src.lt_in_w && !magpietts_local_add_tensor(dst.ctx, src.lt_in_w, &dst.lt_in_w, copies, fp32)) || + (src.lt_in_b && !magpietts_local_add_tensor(dst.ctx, src.lt_in_b, &dst.lt_in_b, copies, fp32)) || !magpietts_local_add_tensor_vector(dst.ctx, src.lt_out_w, dst.lt_out_w, copies, fp32) || !magpietts_local_add_tensor_vector(dst.ctx, src.lt_out_b, dst.lt_out_b, copies, fp32) || !magpietts_local_copy_transformer_layout(dst.ctx, src.local, dst.local, copies, fp32)) { @@ -462,7 +462,7 @@ local_transformer_graph_init( graph.reset(); const auto& h = model.hparams; - if (codebook_idx < 0 || codebook_idx >= h.audio_codebooks) { + if (codebook_idx < 0 || codebook_idx >= h.stacked_audio_codebooks()) { fprintf(stderr, "invalid local-transformer codebook index: %d\n", codebook_idx); return false; } @@ -509,12 +509,16 @@ local_transformer_graph_init( graph.dec_cond, pair ? "magpietts_local_transformer_dec_cond" : "magpietts_local_transformer_dec_last"); ggml_set_input(graph.dec_cond); - cur_cond = linear(graph.ctx, model.lt_in_w, graph.dec_cond, model.lt_in_b); + cur_cond = model.lt_in_w + ? linear(graph.ctx, model.lt_in_w, graph.dec_cond, model.lt_in_b) + : graph.dec_cond; if (pair) { graph.dec_uncond = ggml_new_tensor_2d(graph.ctx, GGML_TYPE_F32, h.n_embd, 1); ggml_set_name(graph.dec_uncond, "magpietts_local_transformer_dec_uncond"); ggml_set_input(graph.dec_uncond); - cur_uncond = linear(graph.ctx, model.lt_in_w, graph.dec_uncond, model.lt_in_b); + cur_uncond = model.lt_in_w + ? linear(graph.ctx, model.lt_in_w, graph.dec_uncond, model.lt_in_b) + : graph.dec_uncond; } } else { const std::string name = "magpietts_local_transformer_prev_code"; @@ -523,7 +527,7 @@ local_transformer_graph_init( ggml_set_input(graph.prev_token); ggml_tensor* emb = ggml_get_rows( graph.ctx, model.audio_embeddings[codebook_idx - 1], graph.prev_token); - cur_cond = linear(graph.ctx, model.lt_in_w, emb, model.lt_in_b); + cur_cond = model.lt_in_w ? linear(graph.ctx, model.lt_in_w, emb, model.lt_in_b) : emb; if (pair) { cur_uncond = cur_cond; } @@ -843,7 +847,7 @@ sample_local_codebooks_impl( codes.clear(); argmax_codes.clear(); std::vector prev; - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { std::vector logits; if (use_cfg) { std::vector uncond; @@ -864,7 +868,7 @@ sample_local_codebooks_impl( logits, h, temperature, top_k, rng, forbid_audio_eos); const int greedy = MagpieCodebookSampler::argmaxFromLogits(logits, h, forbid_audio_eos); dump_local_codebook_logits(logit_dump, c, sampled, greedy, logits); - const int emitted = forced_codes && (int)forced_codes->size() == h.audio_codebooks + const int emitted = forced_codes && (int)forced_codes->size() == h.stacked_audio_codebooks() ? (*forced_codes)[c] : sampled; codes.push_back(emitted); @@ -889,7 +893,7 @@ sample_local_codebooks_cuda_impl( } codes.clear(); argmax_codes.clear(); - for (int c = 0; c < h.audio_codebooks; ++c) { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { magpietts_cuda_sample_request cuda_sample; #if defined(MAGPIETTS_CUDA_SAMPLING) cuda_sample.sampler = cuda_sampler; @@ -909,11 +913,11 @@ sample_local_codebooks_cuda_impl( return false; } } - codes.assign((size_t)h.audio_codebooks, 0); - argmax_codes.assign((size_t)h.audio_codebooks, 0); + codes.assign((size_t)h.stacked_audio_codebooks(), 0); + argmax_codes.assign((size_t)h.stacked_audio_codebooks(), 0); char error[256] = {}; if (!magpietts_cuda_copy_sampled_codebooks( - cuda_sampler, h.audio_codebooks, codes.data(), argmax_codes.data(), error, + cuda_sampler, h.stacked_audio_codebooks(), codes.data(), argmax_codes.data(), error, sizeof(error))) { fprintf( stderr, "CUDA local-transformer sampled-code host copy failed: %s\n", diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index 43cccb8..33073f3 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1135,7 +1135,7 @@ stream_magpie_to_audio( metrics.begin(); outputs.metrics = &metrics; - if (!workspace.beginRequest(params.threads, use_cuda_sampling, h.audio_codebooks)) { + if (!workspace.beginRequest(params.threads, use_cuda_sampling, h.stacked_audio_codebooks())) { return false; } LocalCodebookSampler* local_sampler = nullptr; @@ -1250,7 +1250,7 @@ stream_magpie_to_audio( cond_cross_kv.clear(); std::vector> audio_codes(h.audio_codebooks); for (int c = 0; c < h.audio_codebooks; ++c) { - audio_codes[c].push_back(h.audio_bos_id); + audio_codes[c].assign((size_t)h.frame_stacking_factor, h.audio_bos_id); } attention_prior.beginChunk( h, left_offset, text_len, (int)current_tokens.size(), first_text_chunk); @@ -1302,7 +1302,9 @@ stream_magpie_to_audio( int near_end_frames = 0; bool suppress_nonfinal_codec_output = false; int suppressed_nonfinal_frames = 0; - for (int step = 0; step < h.max_decoder_steps; ++step) { + const int max_decoder_positions = + (h.max_decoder_steps + h.frame_stacking_factor - 1) / h.frame_stacking_factor; + for (int step = 0; step < max_decoder_positions; ++step) { const ggml_nvtx::range nvtx_step("magpietts_stream_generation_step"); if (codec_worker.is_failed()) { codec_worker.join(); @@ -1311,12 +1313,12 @@ stream_magpie_to_audio( if (params.verbose && step % 10 == 0) { fprintf( stderr, "%s generating codec frame %d/%d for text chunk %zu/%zu\n", label, - step, h.max_decoder_steps, chunk_index + 1, token_chunks.size()); + step, max_decoder_positions, chunk_index + 1, token_chunks.size()); } decoder_result cond; decoder_result uncond; - const bool forbid_eos = step < h.min_generated_frames; + const bool forbid_eos = step * h.frame_stacking_factor < h.min_generated_frames; std::vector next_codes; std::vector argmax_codes; std::vector alignment_scores; @@ -1415,8 +1417,8 @@ stream_magpie_to_audio( next_codes = std::move(cuda_sample.codes); argmax_codes = std::move(cuda_sample.argmax_codes); } - if ((int)next_codes.size() != h.audio_codebooks || - (int)argmax_codes.size() != h.audio_codebooks) { + if ((int)next_codes.size() != h.stacked_audio_codebooks() || + (int)argmax_codes.size() != h.stacked_audio_codebooks()) { fprintf( stderr, "CUDA sampler returned an unexpected number of codebooks\n"); return cancel_worker(); @@ -1518,8 +1520,15 @@ stream_magpie_to_audio( } } - const bool has_eos = !forbid_eos && MagpieCodebookSampler::hasEos( - next_codes, argmax_codes, h.audio_eos_id); + std::vector> codec_frames; + if (!magpietts_unstack_codes(next_codes, h, codec_frames)) { + fprintf(stderr, "sampled an invalid stacked MagpieTTS frame\n"); + return cancel_worker(); + } + const int eos_lane = + forbid_eos ? -1 : magpietts_first_eos_lane(next_codes, argmax_codes, h); + const bool has_eos = eos_lane >= 0; + bool terminate_after_frame = false; if (has_eos) { ggml_nvtx::mark("magpietts_stream_eos"); if (params.verbose) { @@ -1528,7 +1537,7 @@ stream_magpie_to_audio( step, chunk_index + 1, token_chunks.size()); } if (final_text_chunk || reached_chunk_end || !can_catch_up_nonfinal) { - break; + terminate_after_frame = true; } if (!suppress_nonfinal_codec_output) { suppress_nonfinal_codec_output = true; @@ -1544,31 +1553,30 @@ stream_magpie_to_audio( } } - const bool emit_frame = !suppress_nonfinal_codec_output; - if (emit_frame) { - if (!code_writer.write_frame(next_codes)) { - fprintf(stderr, "failed to write streamed codec frame\n"); - return cancel_worker(); - } - } else { - ++suppressed_nonfinal_frames; - } - bool first_frame = false; metrics.record_decoder_frame(ggml_time_us(), first_frame); for (int c = 0; c < h.audio_codebooks; ++c) { - audio_codes[c].push_back(next_codes[c]); + for (int lane = 0; lane < h.frame_stacking_factor; ++lane) { + audio_codes[c].push_back(next_codes[c + lane * h.audio_codebooks]); + } } - ++decoder_frames_generated; - if (emit_frame) { - ++frames_generated; - ++chunk_frames_generated; - if (!codec_worker.write_frame(next_codes)) { - codec_worker.join(); - return false; + decoder_frames_generated += h.frame_stacking_factor; + const int frames_to_emit = has_eos ? eos_lane : h.frame_stacking_factor; + if (!suppress_nonfinal_codec_output) { + for (int lane = 0; lane < frames_to_emit; ++lane) { + const auto& frame = codec_frames[(size_t)lane]; + if (!code_writer.write_frame(frame) || !codec_worker.write_frame(frame)) { + fprintf(stderr, "failed to write streamed codec frame\n"); + codec_worker.join(); + return false; + } + ++frames_generated; + ++chunk_frames_generated; } + } else { + suppressed_nonfinal_frames += frames_to_emit; } if (start_suppressing_after_frame) { suppress_nonfinal_codec_output = true; @@ -1593,6 +1601,9 @@ stream_magpie_to_audio( } break; } + if (terminate_after_frame) { + break; + } } if (longform_active && !final_text_chunk) { diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index e7ad6a4..5bf55e3 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -721,6 +721,8 @@ magpietts_model_load_impl( h.mask_token_id = gguf_i32(model.gguf, "magpietts.mask_token_id", h.mask_token_id); h.frame_stacking_factor = gguf_i32(model.gguf, "magpietts.frame_stacking_factor", h.frame_stacking_factor); + const int32_t stored_stacked_codebooks = gguf_i32( + model.gguf, "magpietts.stacked_audio_codebooks", h.stacked_audio_codebooks()); h.n_embd = gguf_i32(model.gguf, "magpietts.embedding_dim", h.n_embd); h.n_ffn = gguf_i32(model.gguf, "magpietts.ffn_dim", h.n_ffn); h.n_ctx = gguf_i32(model.gguf, "magpietts.context_length", h.n_ctx); @@ -799,12 +801,19 @@ magpietts_model_load_impl( return false; } - if (h.frame_stacking_factor != 1) { + if (h.frame_stacking_factor < 1) { fprintf( - stderr, "unsupported frame_stacking_factor=%d; this example currently supports 1\n", + stderr, "invalid frame_stacking_factor=%d\n", h.frame_stacking_factor); return false; } + if (h.audio_codebooks < 1 || stored_stacked_codebooks != h.stacked_audio_codebooks()) { + fprintf( + stderr, + "invalid stacked audio layout: codebooks=%d frame_stacking_factor=%d stored_slots=%d\n", + h.audio_codebooks, h.frame_stacking_factor, stored_stacked_codebooks); + return false; + } model.cuda_unified_memory = force_cpu ? false : magpietts_configure_unified_memory(uma_mode); ggml_backend_load_all(); @@ -882,11 +891,22 @@ magpietts_model_load_impl( model.baked_context = require_tensor(model, "baked_context_embedding.weight"); model.final_proj_w = require_tensor(model, "final_proj.weight"); model.final_proj_b = require_tensor(model, "final_proj.bias"); - model.lt_in_w = require_tensor(model, "local_transformer_in_projection.weight"); - model.lt_in_b = require_tensor(model, "local_transformer_in_projection.bias"); + // v2602 projects decoder/audio embeddings into the local-transformer + // dimension. In v2607 those dimensions are both 768 and the checkpoint + // intentionally omits the projection (identity input). + model.lt_in_w = ggml_get_tensor(model.ctx, "local_transformer_in_projection.weight"); + model.lt_in_b = ggml_get_tensor(model.ctx, "local_transformer_in_projection.bias"); + if ((model.lt_in_w == nullptr) != (model.lt_in_b == nullptr)) { + fprintf(stderr, "local-transformer input projection must include both weight and bias\n"); + return false; + } + if (model.lt_in_w == nullptr && h.n_embd != h.lt_hidden) { + fprintf(stderr, "local-transformer input projection is missing for incompatible dimensions\n"); + return false; + } - model.audio_embeddings.resize(h.audio_codebooks); - for (int i = 0; i < h.audio_codebooks; ++i) { + model.audio_embeddings.resize(h.stacked_audio_codebooks()); + for (int i = 0; i < h.stacked_audio_codebooks(); ++i) { model.audio_embeddings[i] = require_tensor(model, "audio_embeddings." + std::to_string(i) + ".weight"); } @@ -911,9 +931,9 @@ magpietts_model_load_impl( } #endif - model.lt_out_w.resize(h.audio_codebooks); - model.lt_out_b.resize(h.audio_codebooks); - for (int i = 0; i < h.audio_codebooks; ++i) { + model.lt_out_w.resize(h.stacked_audio_codebooks()); + model.lt_out_b.resize(h.stacked_audio_codebooks()); + for (int i = 0; i < h.stacked_audio_codebooks(); ++i) { model.lt_out_w[i] = require_tensor( model, "local_transformer_out_projections." + std::to_string(i) + ".weight"); model.lt_out_b[i] = require_tensor( @@ -922,10 +942,10 @@ magpietts_model_load_impl( fprintf( stderr, - "loaded MagpieTTS GGUF: text_vocab=%d audio_codebooks=%d audio_vocab=%d speakers=%d " + "loaded MagpieTTS GGUF: text_vocab=%d audio_codebooks=%d stacked_slots=%d audio_vocab=%d speakers=%d " "attention_prior=%s epsilon=%.4g lookahead=%d start_step=%d advance_threshold=%d " "decay_threshold=%d estimate_layers=%s apply_layers=%s\n", - h.text_vocab_size, h.audio_codebooks, h.audio_vocab_size, h.baked_speakers, + h.text_vocab_size, h.audio_codebooks, h.stacked_audio_codebooks(), h.audio_vocab_size, h.baked_speakers, h.apply_attention_prior ? "on" : "off", h.attention_prior_epsilon, h.attention_prior_lookahead_window, h.start_prior_after_n_audio_steps, h.attention_prior_advance_threshold, h.attention_prior_decay_threshold, @@ -1905,7 +1925,7 @@ MagpieCodeGenerator::generate( std::vector> audio_codes(h.audio_codebooks); for (int c = 0; c < h.audio_codebooks; ++c) { - audio_codes[c].push_back(h.audio_bos_id); + audio_codes[c].assign((size_t)h.frame_stacking_factor, h.audio_bos_id); } std::vector> generated_frames; @@ -1942,7 +1962,7 @@ MagpieCodeGenerator::generate( } } #if defined(MAGPIETTS_CUDA_SAMPLING) - cuda_sampler.reset(magpietts_cuda_sampler_create(h.audio_codebooks)); + cuda_sampler.reset(magpietts_cuda_sampler_create(h.stacked_audio_codebooks())); if (!cuda_sampler) { fprintf(stderr, "failed to create CUDA sampler\n"); return false; @@ -1958,14 +1978,16 @@ MagpieCodeGenerator::generate( const int64_t generation_start_us = ggml_time_us(); { const ggml_nvtx::range nvtx_loop("magpietts_generate_loop"); - for (int step = 0; step < h.max_decoder_steps; ++step) { + const int max_decoder_positions = + (h.max_decoder_steps + h.frame_stacking_factor - 1) / h.frame_stacking_factor; + for (int step = 0; step < max_decoder_positions; ++step) { const ggml_nvtx::range nvtx_step("magpietts_generate_step"); const int64_t frame_start_us = ggml_time_us(); if (step % 10 == 0) { - fprintf(stderr, "%s decoding frame %d/%d\n", label, step, h.max_decoder_steps); + fprintf(stderr, "%s decoding frame %d/%d\n", label, step, max_decoder_positions); } - const bool forbid_eos = step < h.min_generated_frames; + const bool forbid_eos = step * h.frame_stacking_factor < h.min_generated_frames; std::vector next_codes; std::vector argmax_codes; std::vector alignment_scores; @@ -2062,8 +2084,8 @@ MagpieCodeGenerator::generate( next_codes = std::move(cuda_sample.codes); argmax_codes = std::move(cuda_sample.argmax_codes); } - if ((int)next_codes.size() != h.audio_codebooks || - (int)argmax_codes.size() != h.audio_codebooks) { + if ((int)next_codes.size() != h.stacked_audio_codebooks() || + (int)argmax_codes.size() != h.stacked_audio_codebooks()) { fprintf(stderr, "CUDA sampler returned an unexpected number of codebooks\n"); return false; } @@ -2115,18 +2137,27 @@ MagpieCodeGenerator::generate( attention_prior.update(h, step, (int)tokens.size(), alignment_scores); } - if (!forbid_eos && - MagpieCodebookSampler::hasEos(next_codes, argmax_codes, h.audio_eos_id)) { + std::vector> codec_frames; + if (!magpietts_unstack_codes(next_codes, h, codec_frames)) { + fprintf(stderr, "sampled an invalid stacked MagpieTTS frame\n"); + return false; + } + const int eos_lane = + forbid_eos ? -1 : magpietts_first_eos_lane(next_codes, argmax_codes, h); + for (int c = 0; c < h.audio_codebooks; ++c) { + for (int lane = 0; lane < h.frame_stacking_factor; ++lane) { + audio_codes[c].push_back(next_codes[c + lane * h.audio_codebooks]); + } + } + for (int lane = 0; lane < (eos_lane >= 0 ? eos_lane : h.frame_stacking_factor); ++lane) { + generated_frames.push_back(codec_frames[(size_t)lane]); + } + if (eos_lane >= 0) { ggml_nvtx::mark("magpietts_eos"); fprintf(stderr, "%s EOS detected at frame %d\n", label, step); break; } - generated_frames.push_back(next_codes); - for (int c = 0; c < h.audio_codebooks; ++c) { - audio_codes[c].push_back(next_codes[c]); - } - bool first_frame = false; const int64_t frame_done_us = ggml_time_us(); const double inter_ms = metrics.record_frame(frame_done_us, first_frame); diff --git a/src/tts/magpietts/model.h b/src/tts/magpietts/model.h index c123527..b5511a9 100644 --- a/src/tts/magpietts/model.h +++ b/src/tts/magpietts/model.h @@ -50,6 +50,13 @@ struct magpietts_hparams { int32_t mask_token_id = 2020; int32_t frame_stacking_factor = 1; + // Number of model prediction slots in one decoder position. v2602 has + // one slot per codec codebook, while v2607 predicts two consecutive codec + // frames at once. + int32_t stacked_audio_codebooks() const { + return audio_codebooks * frame_stacking_factor; + } + int32_t n_embd = 768; int32_t n_ffn = 3072; int32_t n_ctx = 2048; @@ -87,6 +94,42 @@ struct magpietts_hparams { std::vector apply_prior_to_layers; }; +inline bool +magpietts_unstack_codes( + const std::vector& stacked_codes, const magpietts_hparams& h, + std::vector>& frames) { + if ((int)stacked_codes.size() != h.stacked_audio_codebooks()) { + return false; + } + frames.assign((size_t)h.frame_stacking_factor, std::vector(h.audio_codebooks)); + for (int lane = 0; lane < h.frame_stacking_factor; ++lane) { + for (int codebook = 0; codebook < h.audio_codebooks; ++codebook) { + frames[(size_t)lane][(size_t)codebook] = + stacked_codes[(size_t)(codebook + lane * h.audio_codebooks)]; + } + } + return true; +} + +inline int +magpietts_first_eos_lane( + const std::vector& sampled, const std::vector& greedy, + const magpietts_hparams& h) { + if ((int)sampled.size() != h.stacked_audio_codebooks() || + (int)greedy.size() != h.stacked_audio_codebooks()) { + return -1; + } + for (int lane = 0; lane < h.frame_stacking_factor; ++lane) { + for (int codebook = 0; codebook < h.audio_codebooks; ++codebook) { + const size_t index = (size_t)(codebook + lane * h.audio_codebooks); + if (sampled[index] == h.audio_eos_id || greedy[index] == h.audio_eos_id) { + return lane; + } + } + } + return -1; +} + struct magpietts_layer { ggml_tensor* norm_self = nullptr; ggml_tensor* self_qkv = nullptr; diff --git a/tests/cpp/tts/CMakeLists.txt b/tests/cpp/tts/CMakeLists.txt index 8f62b6b..c06c038 100644 --- a/tests/cpp/tts/CMakeLists.txt +++ b/tests/cpp/tts/CMakeLists.txt @@ -8,6 +8,9 @@ target_link_libraries(test_magpietts_file PRIVATE nemo_speech_tts_magpietts) add_executable(test_magpietts_attention_prior test_magpietts_attention_prior.cpp) target_link_libraries(test_magpietts_attention_prior PRIVATE nemo_speech_tts_magpietts) +add_executable(test_magpietts_frame_stacking test_magpietts_frame_stacking.cpp) +target_link_libraries(test_magpietts_frame_stacking PRIVATE nemo_speech_tts_magpietts) + add_executable(test_magpietts_asr test_magpietts_asr.cpp) target_link_libraries(test_magpietts_asr PRIVATE tests_cpp_wer diff --git a/tests/cpp/tts/test_magpietts_frame_stacking.cpp b/tests/cpp/tts/test_magpietts_frame_stacking.cpp new file mode 100644 index 0000000..d094651 --- /dev/null +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -0,0 +1,37 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include + +#include "tts/magpietts/model.h" + +namespace tts = nemo_speech::tts; + +int +main() { + tts::magpietts_hparams h; + h.audio_codebooks = 2; + h.frame_stacking_factor = 2; + h.audio_eos_id = 99; + + const std::vector stacked = {10, 11, 20, 21}; + std::vector> frames; + if (!tts::magpietts_unstack_codes(stacked, h, frames) || frames.size() != 2 || + frames[0] != std::vector({10, 11}) || + frames[1] != std::vector({20, 21})) { + std::fprintf(stderr, "stacked-frame reconstruction failed\n"); + return 1; + } + + std::vector greedy = stacked; + greedy[3] = h.audio_eos_id; + if (tts::magpietts_first_eos_lane(stacked, greedy, h) != 1) { + std::fprintf(stderr, "EOS lane detection failed\n"); + return 1; + } + if (tts::magpietts_first_eos_lane(stacked, stacked, h) != -1) { + std::fprintf(stderr, "unexpected EOS lane\n"); + return 1; + } + return 0; +} From e80bcd8ceeb31137ae93c354339bb0ac47bb4a1d Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Mon, 17 Aug 2026 08:34:18 +0530 Subject: [PATCH 02/19] Add preempt option for tts Signed-off-by: Anand Joseph --- app/serve.cpp | 4 + server/http/http_server.cpp | 92 ++++++++++++++++++++-- server/http/http_server.h | 2 + tests/cli/cli_contract_test.py | 6 +- tests/cpp/CMakeLists.txt | 18 +++++ tests/integration/http_conformance_test.py | 40 ++++++++++ 6 files changed, 156 insertions(+), 6 deletions(-) diff --git a/app/serve.cpp b/app/serve.cpp index 1f31665..0cefccf 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -225,6 +225,9 @@ run_server(int argc, char** argv) { "tts.enabled", [&](const std::string& value) { tts_enabled = parse_enablement("tts.enabled", value); }, "Enable TTS: auto, true, or false"); + parser.Register( + "tts.preempt", &server_config.preempt_tts, + "Cancel older HTTP TTS synthesis when a newer request arrives"); #endif auto value = [&](int& index, const std::string& option) { if (++index >= argc) @@ -619,6 +622,7 @@ print_serve_help(const char* program) { " --read-timeout SEC Socket read timeout (default: 30)\n" " --write-timeout SEC Socket write timeout (default: 30)\n" " --access-log Log completed HTTP requests\n" + " --tts.preempt Cancel older HTTP TTS synthesis for the newest request\n" " --log-format text|json Access-log format (--json also selects JSON)\n" " --api-key KEY Require Authorization: Bearer KEY\n" " --cors-origin ORIGIN Allow one cross-origin browser origin\n" diff --git a/server/http/http_server.cpp b/server/http/http_server.cpp index 7d7da2e..0dafecc 100644 --- a/server/http/http_server.cpp +++ b/server/http/http_server.cpp @@ -10,10 +10,12 @@ #include #include #include +#include #include #include #include #include +#include #include #include #include @@ -318,11 +320,61 @@ transcript_response( } // namespace +// Magpie owns one mutable streaming workspace, so synthesis is serialized by the +// runtime. This coordinator sits in front of it when preemption is requested: +// a newer request makes every older request ineligible to start (or continue) +// and waits for the active request to release the runtime. +class TtsPreemptionCoordinator { + public: + uint64_t claim() { + std::unique_lock lock(mutex_); + const uint64_t generation = ++newest_generation_; + ready_.wait(lock, [&] { return !active_ || generation != newest_generation_; }); + if (generation != newest_generation_) + return 0; + active_ = true; + return generation; + } + + bool superseded(uint64_t generation) const { + std::lock_guard lock(mutex_); + return generation != newest_generation_; + } + + void release() { + std::lock_guard lock(mutex_); + active_ = false; + ready_.notify_all(); + } + + private: + mutable std::mutex mutex_; + std::condition_variable ready_; + uint64_t newest_generation_ = 0; + bool active_ = false; +}; + +class TtsPreemptionLease { + public: + explicit TtsPreemptionLease(TtsPreemptionCoordinator* coordinator) : coordinator_(coordinator) {} + ~TtsPreemptionLease() { + if (coordinator_) + coordinator_->release(); + } + + TtsPreemptionLease(const TtsPreemptionLease&) = delete; + TtsPreemptionLease& operator=(const TtsPreemptionLease&) = delete; + + private: + TtsPreemptionCoordinator* coordinator_; +}; + struct Server::Impl { EngineRegistry& models; ServerConfig config; std::unique_ptr server; std::atomic request_id{1}; + TtsPreemptionCoordinator tts_preemption; Impl(EngineRegistry& engines, ServerConfig config) : models(engines), config(std::move(config)) { @@ -639,12 +691,42 @@ struct Server::Impl { const std::string format = body.string_or("response_format", "wav"); if (format != "wav" && format != "pcm") throw std::invalid_argument("response_format must be wav or pcm"); + + uint64_t generation = 0; + if (this->config.preempt_tts) { + generation = this->tts_preemption.claim(); + if (generation == 0) { + fail(response, 409, "TTS synthesis was canceled by a newer request"); + return; + } + } + TtsPreemptionLease lease( + this->config.preempt_tts ? &this->tts_preemption : nullptr); std::string pcm; - const auto result = - synthesizer->synthesize(synthesis, [&](const auto&, const std::string& chunk) { - pcm += chunk; - return true; - }); + tts::SynthesisResult result; + try { + result = synthesizer->synthesize( + synthesis, [&](const auto&, const std::string& chunk) { + if (this->config.preempt_tts && + this->tts_preemption.superseded(generation)) { + return false; + } + pcm += chunk; + return true; + }); + } + catch (const std::exception&) { + if (this->config.preempt_tts && + this->tts_preemption.superseded(generation)) { + fail(response, 409, "TTS synthesis was canceled by a newer request"); + return; + } + throw; + } + if (this->config.preempt_tts && this->tts_preemption.superseded(generation)) { + fail(response, 409, "TTS synthesis was canceled by a newer request"); + return; + } if (this->config.tts_benchmark) { const auto& stats = result.stats; std::cerr << "[nemo_http][tts_benchmark]" diff --git a/server/http/http_server.h b/server/http/http_server.h index 6d8776b..94ea243 100644 --- a/server/http/http_server.h +++ b/server/http/http_server.h @@ -21,6 +21,8 @@ struct ServerConfig { bool json_logs = false; // Emit per-request TTS runtime timings to stderr. bool tts_benchmark = false; + // Keep only the latest HTTP TTS request; newer requests cancel older synthesis. + bool preempt_tts = false; std::string tls_certificate; std::string tls_private_key; std::string api_key; diff --git a/tests/cli/cli_contract_test.py b/tests/cli/cli_contract_test.py index 3c7b322..46ddf95 100644 --- a/tests/cli/cli_contract_test.py +++ b/tests/cli/cli_contract_test.py @@ -74,6 +74,7 @@ def main() -> None: serve_help = run(binary, "serve", "--help") if "--cors-origin" in serve_help.stdout: + assert "--tts.preempt" in serve_help.stdout with socket.socket() as stalled_server: stalled_server.bind(("127.0.0.1", 0)) stalled_server.listen() @@ -141,7 +142,10 @@ def stall_response() -> None: missing_paths.append(str(path)) if len(missing_paths) >= 2: error = expect_json_error( - run(binary, "--json", "serve", *missing_arguments, "--no-warmup"), + run( + binary, "--json", "serve", "--tts.preempt", *missing_arguments, + "--no-warmup" + ), 3, "missing_model", ) diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index e951ea6..b0b362a 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -58,6 +58,16 @@ if(TARGET nemo_speech_cli) --binary $ --asr-model $ENV{NEMO_SPEECH_TEST_ASR_MODEL} --audio $ENV{NEMO_SPEECH_TEST_AUDIO}) + set(nemo_speech_http_tts_conformance OFF) + if(NEMO_SPEECH_BUILD_TTS AND DEFINED ENV{NEMO_SPEECH_TEST_TTS_MODEL} AND + DEFINED ENV{NEMO_SPEECH_TEST_CODEC_MODEL} AND + DEFINED ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR}) + set(nemo_speech_http_tts_conformance ON) + list(APPEND http_conformance_args + --tts-model $ENV{NEMO_SPEECH_TEST_TTS_MODEL} + --codec-model $ENV{NEMO_SPEECH_TEST_CODEC_MODEL} + --tokenizer-dir $ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR}) + endif() find_program(NEMO_SPEECH_NODE_EXECUTABLE NAMES node nodejs) if(NEMO_SPEECH_NODE_EXECUTABLE AND DEFINED ENV{NEMO_SPEECH_TEST_OPENAI_JS_PACKAGE}) @@ -73,6 +83,14 @@ if(TARGET nemo_speech_cli) ${CMAKE_SOURCE_DIR}/tests/integration/http_conformance_test.py ${http_conformance_args}) set_tests_properties(http_openai_conformance PROPERTIES TIMEOUT 600) + if(nemo_speech_http_tts_conformance) + add_test( + NAME http_tts_preemption + COMMAND ${Python3_EXECUTABLE} + ${CMAKE_SOURCE_DIR}/tests/integration/http_conformance_test.py + ${http_conformance_args} --tts-preempt) + set_tests_properties(http_tts_preemption PROPERTIES TIMEOUT 600) + endif() endif() endif() endif() diff --git a/tests/integration/http_conformance_test.py b/tests/integration/http_conformance_test.py index 4d40d52..912be6e 100644 --- a/tests/integration/http_conformance_test.py +++ b/tests/integration/http_conformance_test.py @@ -13,6 +13,7 @@ import subprocess import sys import tempfile +import threading import time import wave from pathlib import Path @@ -41,6 +42,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--openai-js-package") parser.add_argument("--openai-js-example") parser.add_argument("--timeout", type=int, default=180) + parser.add_argument("--tts-preempt", action="store_true") return parser.parse_args() @@ -128,6 +130,8 @@ def main() -> None: ): if value: command.extend((option, value)) + if args.tts_preempt: + command.append("--tts.preempt") server_log = tempfile.TemporaryFile(mode="w+", encoding="utf-8") process = subprocess.Popen( @@ -417,6 +421,42 @@ def main() -> None: json={"input": "Invalid sample rate.", "voice": voice, "sample_rate": 96001}, ) require(response.status_code == 400, "invalid TTS sample rate was accepted") + if args.tts_preempt: + first_started = threading.Event() + + def synthesize_long_text() -> tuple[int, bytes]: + first_started.set() + with httpx.Client(timeout=args.timeout, trust_env=False) as concurrent_client: + response = concurrent_client.post( + f"{base}/v1/audio/speech", + headers=headers, + json={"input": "NeMo-Speech.cpp runs locally. " * 100, "voice": voice}, + ) + return response.status_code, response.content + + first_result: list[tuple[int, bytes]] = [] + + def collect_first() -> None: + first_result.append(synthesize_long_text()) + + first = threading.Thread(target=collect_first) + first.start() + require(first_started.wait(5), "long TTS request did not start") + time.sleep(0.5) + with httpx.Client(timeout=args.timeout, trust_env=False) as concurrent_client: + latest = concurrent_client.post( + f"{base}/v1/audio/speech", + headers=headers, + json={"input": "Newest synthesis request.", "voice": voice}, + ) + first.join(args.timeout) + require(not first.is_alive(), "preempted TTS request did not finish") + require( + first_result and first_result[0][0] == 409, + "older TTS request was not canceled", + ) + require(latest.status_code == 200, f"newest TTS request failed: {latest.text}") + require(latest.content.startswith(b"RIFF"), "newest TTS response is not WAV") if args.nmt_model: with Path(args.audio).open("rb") as audio: response = client.post( From 82ed85122eb00facef61157a9e6129552af6065c Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Mon, 17 Aug 2026 12:09:19 +0530 Subject: [PATCH 03/19] Add support for AR, KO, PT Signed-off-by: Anand Joseph --- docs/tts/configuration.md | 5 +- src/services/grpc_tts.cc | 2 +- src/tts/magpietts/config.cpp | 3 + src/tts/synthesizer.cpp | 3 +- src/tts/tokenizer/tokenizer.cpp | 71 ++++++++++++---- src/tts/tokenizer/tokenizer.h | 5 +- src/tts/tokenizer/tokenizer_impl.cpp | 84 +++++++++++++++++++ tests/cpp/tts/test_tokenizer_single_chars.cpp | 22 +++++ 8 files changed, 175 insertions(+), 20 deletions(-) diff --git a/docs/tts/configuration.md b/docs/tts/configuration.md index 2496706..97ccb67 100644 --- a/docs/tts/configuration.md +++ b/docs/tts/configuration.md @@ -42,7 +42,8 @@ when present. The older split layout (`classify/tokenize_and_classify.far`, The optional `riva_server` adapter implements `RivaSpeechSynthesis.Synthesize`, `SynthesizeOnline`, and `GetRivaSynthesisConfig`. It takes plain text in `SynthesizeSpeechRequest.text`, supports native Magpie tokenizers for `en`, -`es`, `de`, `fr`, `it`, `vi`, `zh`, `hi`, and `ja`, and returns `LINEAR_PCM` +`es`, `de`, `fr`, `it`, `vi`, `zh`, `hi`, `ja`, `ar-AE`, `ar-SA`, `ar-MSA`, +`ko`, and `pt-BR`, and returns `LINEAR_PCM` s16le at the NanoCodec sample rate. Japanese and Mandarin are included when `NEMO_SPEECH_TTS_WITH_JA` and `NEMO_SPEECH_TTS_WITH_ZH`, respectively, are enabled at build time; both default to `OFF`. Native tokenizers are cached by @@ -127,7 +128,7 @@ All keys nest under `tts.`. Defaults shown; CLI alias listed where one exists. | `tts.magpie-model` | - | - | MagpieTTS GGUF token generator (required) | | `tts.codec-model` | - | - | NanoCodec decoder GGUF (required) | | `tts.tokenizer-model-dir` | - | - | extracted Magpie `.nemo` dir (required) | -| `tts.tokenizer.sentence-limit.` | - | per language (`en` 45 ... `ja` 40) | sentence-chunking threshold in words (characters for `zh`/`ja`); subkeys `en`, `es`, `fr`, `vi`, `it`, `de`, `zh`, `hi`, `ja` | +| `tts.tokenizer.sentence-limit.` | - | per language (`en` 45 ... `ja` 40) | sentence-chunking threshold in words (characters for `zh`/`ja`); subkeys `en`, `es`, `fr`, `vi`, `it`, `de`, `zh`, `hi`, `ja`, `ar`, `ko`, `pt` | | `tts.tn-model-dir` | - | - | enables Sparrowhawk TN with this grammar dir; requires `NEMO_SPEECH_WITH_NORM=ON` | | `tts.language-code` | - | `en-US` | default Riva language code | | `tts.voice-name` | - | - | default voice name or speaker index | diff --git a/src/services/grpc_tts.cc b/src/services/grpc_tts.cc index 975d4c1..033abe7 100644 --- a/src/services/grpc_tts.cc +++ b/src/services/grpc_tts.cc @@ -322,7 +322,7 @@ GrpcTtsService::GetRivaSynthesisConfig( grpc::StatusCode::NOT_FOUND, "unknown TTS model '" + req->model_name() + "'"); } - const std::vector language_codes = tts::supported_language_codes(); + const std::vector language_codes = synthesizer_->supported_language_codes(); const std::vector speaker_names = synthesizer_->speaker_names(); const std::vector dotted_voices = dotted_voice_names(synthesizer_->model_name(), speaker_names); diff --git a/src/tts/magpietts/config.cpp b/src/tts/magpietts/config.cpp index 637ea02..ef84e32 100644 --- a/src/tts/magpietts/config.cpp +++ b/src/tts/magpietts/config.cpp @@ -22,6 +22,9 @@ MagpieTokenizerSentenceLimits::Register(common::ParameterParser& p) { p.Register("zh", &zh, "Mandarin sentence-chunking character threshold"); p.Register("hi", &hi, "Hindi sentence-chunking word threshold"); p.Register("ja", &ja, "Japanese sentence-chunking character threshold"); + p.Register("ar", &ar, "Arabic sentence-chunking word threshold"); + p.Register("ko", &ko, "Korean sentence-chunking word threshold"); + p.Register("pt", &pt, "Portuguese sentence-chunking word threshold"); } void diff --git a/src/tts/synthesizer.cpp b/src/tts/synthesizer.cpp index 349ec37..5f33e7b 100644 --- a/src/tts/synthesizer.cpp +++ b/src/tts/synthesizer.cpp @@ -304,7 +304,8 @@ Synthesizer::speaker_names() const { std::vector Synthesizer::supported_language_codes() const { - return MagpieNativeTokenizer::supported_language_codes(); + return impl_->tokenizer ? impl_->tokenizer->supported_language_codes() + : std::vector{}; } const std::string& diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index 05d5623..5a7270a 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -49,6 +50,12 @@ sentence_limit_for_language( return limits.hi; if (language == "ja") return limits.ja; + if (language == "ar" || language.rfind("ar-", 0) == 0) + return limits.ar; + if (language == "ko") + return limits.ko; + if (language == "pt" || language == "pt-br") + return limits.pt; return limits.en; } @@ -143,6 +150,9 @@ class MagpieNativeTokenizer::Impl { if (model_dir_.empty()) { throw std::invalid_argument("tokenizer model directory is required"); } + std::ifstream config_file(fs::path(model_dir_) / "model_config.yaml"); + std::string contents((std::istreambuf_iterator(config_file)), {}); + v2607_ = contents.find("portuguese_Brazilian_phoneme") != std::string::npos; } MagpieTokenizationResult tokenize( @@ -156,6 +166,7 @@ class MagpieNativeTokenizer::Impl { p.model = model_dir_; p.text = text; p.language = MagpieNativeTokenizer::normalize_language_code(language_code); + p.v2607 = v2607_; p.chunk_text_transform = chunk_text_transform; const std::string tokenizer_name = tokenizer_for_language(p.language); @@ -170,6 +181,25 @@ class MagpieNativeTokenizer::Impl { p.sentence_chunking = should_tokenize_by_sentence(text, p.language, config_.sentence_limit); tokenizer_result native = tokenize_native(p); + if (v2607_) { + int offset_delta = 0; + if (p.language == "es" || p.language == "de" || p.language == "zh" || + p.language == "ja") + offset_delta = 384; + else if (p.language == "fr") + offset_delta = 1188; + else if (p.language == "it" || p.language == "vi") + offset_delta = 997; + else if (p.language == "hi") + offset_delta = 111; + for (auto& chunk : native.chunks) + for (int& token : chunk.tokens) { + if (token == native.eos_id) + token = 3358; + else + token += offset_delta; + } + } MagpieTokenizationResult out; out.language = native.language; out.tokenizer_name = native.tokenizer_name; @@ -188,6 +218,22 @@ class MagpieNativeTokenizer::Impl { return out; } + std::vector supported_language_codes() const { + std::vector languages = {"en-US", "es-ES", "de-DE", "fr-FR", + "it-IT", "vi-VN", "hi-IN"}; + if (v2607_) { + languages.insert( + languages.end(), {"ar-AE", "ar-SA", "ar-MSA", "ko-KR", "pt-BR"}); + } +#ifdef NEMO_SPEECH_TTS_WITH_ZH + languages.emplace_back("zh-CN"); +#endif +#ifdef NEMO_SPEECH_TTS_WITH_JA + languages.emplace_back("ja-JP"); +#endif + return languages; + } + private: tokenizer_result tokenize_native(const params& p) const { if (p.language == "en" || p.language == "es" || p.language == "de") { @@ -252,6 +298,7 @@ class MagpieNativeTokenizer::Impl { std::string model_dir_; MagpieTokenizerConfig config_; + bool v2607_ = false; mutable std::mutex cache_mutex_; mutable std::map> ipa_cache_; #ifdef NEMO_SPEECH_TTS_WITH_ZH @@ -299,27 +346,20 @@ MagpieNativeTokenizer::normalize_language_code(const std::string& language_code) return "en"; } std::replace(lang.begin(), lang.end(), '_', '-'); - const size_t dash = lang.find('-'); - if (dash != std::string::npos) { - lang.resize(dash); - } std::transform(lang.begin(), lang.end(), lang.begin(), [](unsigned char c) { return (char)std::tolower(c); }); + if (lang == "ar-ae" || lang == "ar-sa" || lang == "ar-msa" || lang == "pt-br") + return lang; + const size_t dash = lang.find('-'); + if (dash != std::string::npos) + lang.resize(dash); return lang; } std::vector -MagpieNativeTokenizer::supported_language_codes() { - std::vector languages = {"en-US", "es-ES", "de-DE", "fr-FR", - "it-IT", "vi-VN", "hi-IN"}; -#ifdef NEMO_SPEECH_TTS_WITH_ZH - languages.emplace_back("zh-CN"); -#endif -#ifdef NEMO_SPEECH_TTS_WITH_JA - languages.emplace_back("ja-JP"); -#endif - return languages; +MagpieNativeTokenizer::supported_language_codes() const { + return impl_->supported_language_codes(); } std::string @@ -337,7 +377,8 @@ ensure_terminal_punctuation(const std::string& text, const std::string& language terminal = "。"; } else if ( language == "en" || language == "es" || language == "fr" || language == "de" || - language == "it" || language == "vi") { + language == "it" || language == "vi" || language == "pt-br" || language == "ko" || + language.rfind("ar-", 0) == 0) { terminal = "."; } else { return text; diff --git a/src/tts/tokenizer/tokenizer.h b/src/tts/tokenizer/tokenizer.h index 6dedc50..6b2b366 100644 --- a/src/tts/tokenizer/tokenizer.h +++ b/src/tts/tokenizer/tokenizer.h @@ -26,6 +26,9 @@ struct MagpieTokenizerSentenceLimits { int zh = 100; int hi = 40; int ja = 40; + int ar = 45; + int ko = 45; + int pt = 45; void Register(common::ParameterParser& parser); }; @@ -79,7 +82,7 @@ class MagpieNativeTokenizer { const PositionedChunkTextTransform& chunk_text_transform) const; static std::string normalize_language_code(const std::string& language_code); - static std::vector supported_language_codes(); + std::vector supported_language_codes() const; private: class Impl; diff --git a/src/tts/tokenizer/tokenizer_impl.cpp b/src/tts/tokenizer/tokenizer_impl.cpp index 7948fd6..0e4084e 100644 --- a/src/tts/tokenizer/tokenizer_impl.cpp +++ b/src/tts/tokenizer/tokenizer_impl.cpp @@ -53,6 +53,7 @@ struct params { fs::path model; std::string text; std::string language = "en"; + bool v2607 = false; bool sentence_chunking = true; std::function chunk_text_transform; }; @@ -864,6 +865,16 @@ tokenizer_for_language(const std::string& lang) { return "italian_phoneme"; if (lang == "vi") return "vietnamese_phoneme"; + if (lang == "pt-br") + return "portuguese_Brazilian_phoneme"; + if (lang == "ko") + return "korean_chartokenizer"; + if (lang == "ar-ae") + return "arabic_AE_chartokenizer"; + if (lang == "ar-sa") + return "arabic_SA_chartokenizer"; + if (lang == "ar-msa") + return "arabic_MSA_chartokenizer"; #ifdef NEMO_SPEECH_TTS_WITH_ZH if (lang == "zh") return "mandarin_phoneme"; @@ -934,6 +945,10 @@ ipa_config_for_language(const std::string& lang) { true, true}; } + if (lang == "pt-br") { + return {"portuguese_Brazilian_phoneme", 1017, "pt_br_prondict", "", "pt-BR", "upper", + "#", true, true}; + } throw std::runtime_error("no native IPA tokenizer for language " + lang); } @@ -976,6 +991,9 @@ exact_ipa_tokens(const std::string& tokenizer_name) { "“", "„", "‹", "›", " ", "", "", }; } + if (tokenizer_name == "portuguese_Brazilian_phoneme") { + return {"!", "\"", "#A", "#B", "#C", "#D", "#E", "#F", "#G", "#H", "#I", "#J", "#K", "#L", "#M", "#N", "#O", "#P", "#Q", "#R", "#S", "#T", "#U", "#V", "#W", "#X", "#Y", "#Z", "#À", "#Á", "#Â", "#Ã", "#Ç", "#É", "#Ê", "#Í", "#Ó", "#Ô", "#Õ", "#Ú", "#Ü", "'", "(", ")", ",", "-", ".", "/", ":", ";", "?", "[", "]", "a", "b", "d", "e", "f", "h", "i", "j", "k", "l", "m", "n", "o", "p", "r", "s", "t", "u", "v", "w", "x", "y", "z", "{", "}", "ð", "õ", "ĩ", "ŋ", "ũ", "ɐ", "ɑ", "ɒ", "ɔ", "ə", "ɛ", "ɜ", "ɡ", "ɪ", "ɲ", "ɹ", "ɾ", "ʁ", "ʃ", "ʊ", "ʌ", "ʎ", "ʒ", "ʲ", "ˈ", "ˌ", "ː", "̃", "θ", "ẽ", " ", "", ""}; + } return {}; } @@ -1421,6 +1439,53 @@ run_hindi_native(const params& p) { return result; } +static std::vector +arabic_tokens() { + std::vector tokens = {" ", "ء", "آ", "أ", "إ", "ؤ", "ئ", "ا", "ب", "ة", "ت", "ث", "ج", "ح", "خ", "د", "ذ", "ر", "ز", "س", "ش", "ص", "ض", "ط", "ظ", "ع", "غ", "ف", "ق", "ك", "ل", "م", "ن", "ه", "و", "ى", "ي", "ً", "ٌ", "ٍ", "َ", "ُ", "ِ", "ّ", "ٰ", "ْ"}; + // v2607 pins NeMo ArabicCharsTokenizer charset_version=1, whose mixed-case + // Arabic character set deliberately contains this second copy. + const std::vector arabic_chars(tokens.begin() + 1, tokens.end()); + tokens.insert(tokens.end(), arabic_chars.begin(), arabic_chars.end()); + for (char c = 'a'; c <= 'z'; ++c) tokens.emplace_back(1, c); + for (char c = 'A'; c <= 'Z'; ++c) tokens.emplace_back(1, c); + tokens.push_back("'"); + const auto punct = default_punct(); + tokens.insert(tokens.end(), punct.begin(), punct.end()); + tokens.insert(tokens.end(), {"،", "؛", "؟", "", ""}); + return tokens; +} + +static tokenizer_result +run_arabic_native(const params& p, int offset, const std::string& tokenizer_name) { + const auto tokens = arabic_tokens(); + std::unordered_map ids; + for (size_t i = 0; i < tokens.size(); ++i) ids[tokens[i]] = (int)i; + tokenizer_result result; + result.language = p.language; + result.tokenizer_name = tokenizer_name; + const int pad_id = token_id_for_symbol(tokens, offset, ""); + for (const std::string& sentence : tokenizer_input_units(p, p.text)) { + chunk ch; + ch.text = sentence; + std::vector symbols; + for (const auto& symbol : split_utf8(sentence)) { + if (symbol == " ") { + if (!symbols.empty() && symbols.back() != symbol) symbols.push_back(symbol); + } else if (ids.count(symbol) != 0) { + symbols.push_back(symbol); + } + } + while (!symbols.empty() && symbols.back() == " ") symbols.pop_back(); + symbols.insert(symbols.begin(), " "); + symbols.push_back(" "); + for (const auto& symbol : symbols) ch.tokens.push_back(offset + ids[symbol]); + ch.tokens.push_back(result.eos_id); + pad_short_text_chunk_before_eos(ch, result.eos_id, pad_id); + result.chunks.push_back(std::move(ch)); + } + return result; +} + #ifdef NEMO_SPEECH_TTS_WITH_ZH static tokenizer_result run_mandarin_native(const params& p, const mandarin_tokenizer& tok) { @@ -1596,6 +1661,10 @@ supports_native(const params& p) { if (p.language == "hi") { return true; } + if ((p.language == "pt-br" || p.language == "ko" || p.language == "ar-ae" || + p.language == "ar-sa" || p.language == "ar-msa") && p.v2607) { + return p.language != "pt-br" || fs::is_directory(p.model); + } #ifdef NEMO_SPEECH_TTS_WITH_JA if (p.language == "ja") { return !find_openjtalk_dictionary_dir(p.model).empty(); @@ -1627,6 +1696,21 @@ run_native(const params& p) { if (p.language == "hi") { return run_hindi_native(p); } + if (p.language == "pt-br") { + return run_ipa_native(p); + } + if (p.language == "ko") { + return run_byt5_native(p, 2973, "korean_chartokenizer"); + } + if (p.language == "ar-ae") { + return run_arabic_native(p, 1329, "arabic_AE_chartokenizer"); + } + if (p.language == "ar-sa") { + return run_arabic_native(p, 1493, "arabic_SA_chartokenizer"); + } + if (p.language == "ar-msa") { + return run_arabic_native(p, 1657, "arabic_MSA_chartokenizer"); + } #ifdef NEMO_SPEECH_TTS_WITH_JA if (p.language == "ja") { return run_japanese_native(p); diff --git a/tests/cpp/tts/test_tokenizer_single_chars.cpp b/tests/cpp/tts/test_tokenizer_single_chars.cpp index a04a03b..072503b 100644 --- a/tests/cpp/tts/test_tokenizer_single_chars.cpp +++ b/tests/cpp/tts/test_tokenizer_single_chars.cpp @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 #include #include +#include #include #include #include @@ -130,6 +131,27 @@ main(int argc, char** argv) { } tts::MagpieNativeTokenizer tokenizer(argv[1]); + std::ifstream tokenizer_config(std::string(argv[1]) + "/model_config.yaml"); + const std::string config_text( + (std::istreambuf_iterator(tokenizer_config)), std::istreambuf_iterator()); + const bool v2607 = config_text.find("portuguese_Brazilian_phoneme") != std::string::npos; + if (v2607) { + bool v2607_ok = true; + v2607_ok &= check_tokens(tokenizer, "A", "en-US", {90, 94, 94, 94, 88, 3358}); + v2607_ok &= check_tokens(tokenizer, "مرحبا", "ar-AE", {1329, 1405, 1391, 1387, 1382, + 1381, 1329, 3358}); + v2607_ok &= check_tokens(tokenizer, "مرحبا", "ar-SA", {1493, 1569, 1555, 1551, 1546, + 1545, 1493, 3358}); + v2607_ok &= check_tokens(tokenizer, "مرحبا", "ar-MSA", {1657, 1733, 1719, 1715, 1710, + 1709, 1657, 3358}); + v2607_ok &= check_tokens(tokenizer, "안녕하세요.", "ko-KR", {3212, 3125, 3112, 3211, 3109, + 3125, 3213, 3125, 3128, 3212, + 3108, 3160, 3212, 3130, 3124, + 3022, 2974, 3358}); + v2607_ok &= check_tokens(tokenizer, "Olá!", "pt-BR", {1125, 1082, 1079, 1119, 1070, + 1017, 1125, 3358}); + return v2607_ok ? 0 : 1; + } bool ok = true; ok &= check_tokens(tokenizer, "A", "en-US", {90, 94, 94, 53, 84, 2361}); From ad6aba2044591ec473e9d058c3f5e82890fc8684 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:45:20 +0000 Subject: [PATCH 04/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- app/synthesize.cpp | 3 +- conversion/tts.py | 6 +++- server/http/http_server.cpp | 6 ++-- src/tts/magpietts/decoder.cpp | 11 +++---- src/tts/magpietts/lt.cpp | 6 ++-- src/tts/magpietts/model.cpp | 21 ++++++------ src/tts/magpietts/model.h | 4 +-- src/tts/tokenizer/tokenizer.cpp | 3 +- src/tts/tokenizer/tokenizer_impl.cpp | 32 +++++++++++++++---- tests/cli/cli_contract_test.py | 8 +++-- tests/cpp/tts/test_tokenizer_single_chars.cpp | 24 +++++++------- 11 files changed, 76 insertions(+), 48 deletions(-) diff --git a/app/synthesize.cpp b/app/synthesize.cpp index 4f6becc..be0840c 100644 --- a/app/synthesize.cpp +++ b/app/synthesize.cpp @@ -215,7 +215,8 @@ command_synthesize(int argc, char** argv) { // --tts.sampling-backend cuda. if (parsed.runtime.sampling_backend != nemo_speech::tts::MagpieBackendPreference::Cuda) { - parsed.runtime.sampling_backend = nemo_speech::tts::MagpieBackendPreference::Cpu; + parsed.runtime.sampling_backend = + nemo_speech::tts::MagpieBackendPreference::Cpu; } parsed.runtime.codec_cpu = false; } else { diff --git a/conversion/tts.py b/conversion/tts.py index ce0d3b1..cf73c83 100644 --- a/conversion/tts.py +++ b/conversion/tts.py @@ -109,7 +109,11 @@ def add_metadata( f"rows={sd['final_proj.weight'].shape[0]} expected={expected_logits}" ) n_lt_heads = len( - [k for k in sd if k.startswith("local_transformer_out_projections.") and k.endswith(".weight")] + [ + k + for k in sd + if k.startswith("local_transformer_out_projections.") and k.endswith(".weight") + ] ) if n_lt_heads != n_stacked_codebooks: raise ValueError( diff --git a/server/http/http_server.cpp b/server/http/http_server.cpp index 0dafecc..175644e 100644 --- a/server/http/http_server.cpp +++ b/server/http/http_server.cpp @@ -356,7 +356,8 @@ class TtsPreemptionCoordinator { class TtsPreemptionLease { public: - explicit TtsPreemptionLease(TtsPreemptionCoordinator* coordinator) : coordinator_(coordinator) {} + explicit TtsPreemptionLease(TtsPreemptionCoordinator* coordinator) + : coordinator_(coordinator) {} ~TtsPreemptionLease() { if (coordinator_) coordinator_->release(); @@ -716,8 +717,7 @@ struct Server::Impl { }); } catch (const std::exception&) { - if (this->config.preempt_tts && - this->tts_preemption.superseded(generation)) { + if (this->config.preempt_tts && this->tts_preemption.superseded(generation)) { fail(response, 409, "TTS synthesis was canceled by a newer request"); return; } diff --git a/src/tts/magpietts/decoder.cpp b/src/tts/magpietts/decoder.cpp index 2f7a38f..7663669 100644 --- a/src/tts/magpietts/decoder.cpp +++ b/src/tts/magpietts/decoder.cpp @@ -514,8 +514,7 @@ stack_audio_codes( const std::vector>& audio_codes, const magpietts_hparams& h, std::vector>& stacked) { if ((int)audio_codes.size() != h.audio_codebooks || audio_codes.empty() || - audio_codes[0].empty() || - (int)audio_codes[0].size() % h.frame_stacking_factor != 0) { + audio_codes[0].empty() || (int)audio_codes[0].size() % h.frame_stacking_factor != 0) { return false; } const int raw_len = (int)audio_codes[0].size(); @@ -660,8 +659,8 @@ decoder_eval_impl( } if (cuda_sample) { const bool sampled = MagpieCodebookSampler::runCuda( - model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, h.stacked_audio_codebooks(), - 0); + model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, + h.stacked_audio_codebooks(), 0); ggml_gallocr_free(allocr); ggml_free(ctx); return sampled; @@ -1022,8 +1021,8 @@ decoder_eval_cached_impl( } if (cuda_sample) { const bool sampled = MagpieCodebookSampler::runCuda( - model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, h.stacked_audio_codebooks(), - 0); + model.backend, h, cuda_sample, logits, nullptr, logits_off_floats, + h.stacked_audio_codebooks(), 0); ggml_gallocr_free(allocr); ggml_free(ctx); if (sampled) { diff --git a/src/tts/magpietts/lt.cpp b/src/tts/magpietts/lt.cpp index b2e6bc0..212868e 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -253,8 +253,10 @@ magpietts_model_init_local_transformer_copy( std::vector> copies; if (!magpietts_local_add_tensor_vector( dst.ctx, src.audio_embeddings, dst.audio_embeddings, copies, fp32) || - (src.lt_in_w && !magpietts_local_add_tensor(dst.ctx, src.lt_in_w, &dst.lt_in_w, copies, fp32)) || - (src.lt_in_b && !magpietts_local_add_tensor(dst.ctx, src.lt_in_b, &dst.lt_in_b, copies, fp32)) || + (src.lt_in_w && + !magpietts_local_add_tensor(dst.ctx, src.lt_in_w, &dst.lt_in_w, copies, fp32)) || + (src.lt_in_b && + !magpietts_local_add_tensor(dst.ctx, src.lt_in_b, &dst.lt_in_b, copies, fp32)) || !magpietts_local_add_tensor_vector(dst.ctx, src.lt_out_w, dst.lt_out_w, copies, fp32) || !magpietts_local_add_tensor_vector(dst.ctx, src.lt_out_b, dst.lt_out_b, copies, fp32) || !magpietts_local_copy_transformer_layout(dst.ctx, src.local, dst.local, copies, fp32)) { diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index 5bf55e3..e03cafd 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -721,8 +721,8 @@ magpietts_model_load_impl( h.mask_token_id = gguf_i32(model.gguf, "magpietts.mask_token_id", h.mask_token_id); h.frame_stacking_factor = gguf_i32(model.gguf, "magpietts.frame_stacking_factor", h.frame_stacking_factor); - const int32_t stored_stacked_codebooks = gguf_i32( - model.gguf, "magpietts.stacked_audio_codebooks", h.stacked_audio_codebooks()); + const int32_t stored_stacked_codebooks = + gguf_i32(model.gguf, "magpietts.stacked_audio_codebooks", h.stacked_audio_codebooks()); h.n_embd = gguf_i32(model.gguf, "magpietts.embedding_dim", h.n_embd); h.n_ffn = gguf_i32(model.gguf, "magpietts.ffn_dim", h.n_ffn); h.n_ctx = gguf_i32(model.gguf, "magpietts.context_length", h.n_ctx); @@ -802,9 +802,7 @@ magpietts_model_load_impl( } if (h.frame_stacking_factor < 1) { - fprintf( - stderr, "invalid frame_stacking_factor=%d\n", - h.frame_stacking_factor); + fprintf(stderr, "invalid frame_stacking_factor=%d\n", h.frame_stacking_factor); return false; } if (h.audio_codebooks < 1 || stored_stacked_codebooks != h.stacked_audio_codebooks()) { @@ -901,7 +899,8 @@ magpietts_model_load_impl( return false; } if (model.lt_in_w == nullptr && h.n_embd != h.lt_hidden) { - fprintf(stderr, "local-transformer input projection is missing for incompatible dimensions\n"); + fprintf( + stderr, "local-transformer input projection is missing for incompatible dimensions\n"); return false; } @@ -942,11 +941,12 @@ magpietts_model_load_impl( fprintf( stderr, - "loaded MagpieTTS GGUF: text_vocab=%d audio_codebooks=%d stacked_slots=%d audio_vocab=%d speakers=%d " + "loaded MagpieTTS GGUF: text_vocab=%d audio_codebooks=%d stacked_slots=%d audio_vocab=%d " + "speakers=%d " "attention_prior=%s epsilon=%.4g lookahead=%d start_step=%d advance_threshold=%d " "decay_threshold=%d estimate_layers=%s apply_layers=%s\n", - h.text_vocab_size, h.audio_codebooks, h.stacked_audio_codebooks(), h.audio_vocab_size, h.baked_speakers, - h.apply_attention_prior ? "on" : "off", h.attention_prior_epsilon, + h.text_vocab_size, h.audio_codebooks, h.stacked_audio_codebooks(), h.audio_vocab_size, + h.baked_speakers, h.apply_attention_prior ? "on" : "off", h.attention_prior_epsilon, h.attention_prior_lookahead_window, h.start_prior_after_n_audio_steps, h.attention_prior_advance_threshold, h.attention_prior_decay_threshold, format_i32_list(h.estimate_alignment_from_layers).c_str(), @@ -2149,7 +2149,8 @@ MagpieCodeGenerator::generate( audio_codes[c].push_back(next_codes[c + lane * h.audio_codebooks]); } } - for (int lane = 0; lane < (eos_lane >= 0 ? eos_lane : h.frame_stacking_factor); ++lane) { + for (int lane = 0; lane < (eos_lane >= 0 ? eos_lane : h.frame_stacking_factor); + ++lane) { generated_frames.push_back(codec_frames[(size_t)lane]); } if (eos_lane >= 0) { diff --git a/src/tts/magpietts/model.h b/src/tts/magpietts/model.h index b5511a9..4edf464 100644 --- a/src/tts/magpietts/model.h +++ b/src/tts/magpietts/model.h @@ -53,9 +53,7 @@ struct magpietts_hparams { // Number of model prediction slots in one decoder position. v2602 has // one slot per codec codebook, while v2607 predicts two consecutive codec // frames at once. - int32_t stacked_audio_codebooks() const { - return audio_codebooks * frame_stacking_factor; - } + int32_t stacked_audio_codebooks() const { return audio_codebooks * frame_stacking_factor; } int32_t n_embd = 768; int32_t n_ffn = 3072; diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index 5a7270a..22d3740 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -222,8 +222,7 @@ class MagpieNativeTokenizer::Impl { std::vector languages = {"en-US", "es-ES", "de-DE", "fr-FR", "it-IT", "vi-VN", "hi-IN"}; if (v2607_) { - languages.insert( - languages.end(), {"ar-AE", "ar-SA", "ar-MSA", "ko-KR", "pt-BR"}); + languages.insert(languages.end(), {"ar-AE", "ar-SA", "ar-MSA", "ko-KR", "pt-BR"}); } #ifdef NEMO_SPEECH_TTS_WITH_ZH languages.emplace_back("zh-CN"); diff --git a/src/tts/tokenizer/tokenizer_impl.cpp b/src/tts/tokenizer/tokenizer_impl.cpp index 0e4084e..2989f0e 100644 --- a/src/tts/tokenizer/tokenizer_impl.cpp +++ b/src/tts/tokenizer/tokenizer_impl.cpp @@ -946,8 +946,15 @@ ipa_config_for_language(const std::string& lang) { true}; } if (lang == "pt-br") { - return {"portuguese_Brazilian_phoneme", 1017, "pt_br_prondict", "", "pt-BR", "upper", - "#", true, true}; + return {"portuguese_Brazilian_phoneme", + 1017, + "pt_br_prondict", + "", + "pt-BR", + "upper", + "#", + true, + true}; } throw std::runtime_error("no native IPA tokenizer for language " + lang); } @@ -992,7 +999,15 @@ exact_ipa_tokens(const std::string& tokenizer_name) { }; } if (tokenizer_name == "portuguese_Brazilian_phoneme") { - return {"!", "\"", "#A", "#B", "#C", "#D", "#E", "#F", "#G", "#H", "#I", "#J", "#K", "#L", "#M", "#N", "#O", "#P", "#Q", "#R", "#S", "#T", "#U", "#V", "#W", "#X", "#Y", "#Z", "#À", "#Á", "#Â", "#Ã", "#Ç", "#É", "#Ê", "#Í", "#Ó", "#Ô", "#Õ", "#Ú", "#Ü", "'", "(", ")", ",", "-", ".", "/", ":", ";", "?", "[", "]", "a", "b", "d", "e", "f", "h", "i", "j", "k", "l", "m", "n", "o", "p", "r", "s", "t", "u", "v", "w", "x", "y", "z", "{", "}", "ð", "õ", "ĩ", "ŋ", "ũ", "ɐ", "ɑ", "ɒ", "ɔ", "ə", "ɛ", "ɜ", "ɡ", "ɪ", "ɲ", "ɹ", "ɾ", "ʁ", "ʃ", "ʊ", "ʌ", "ʎ", "ʒ", "ʲ", "ˈ", "ˌ", "ː", "̃", "θ", "ẽ", " ", "", ""}; + return {"!", "\"", "#A", "#B", "#C", "#D", "#E", "#F", "#G", "#H", "#I", "#J", "#K", + "#L", "#M", "#N", "#O", "#P", "#Q", "#R", "#S", "#T", "#U", "#V", "#W", "#X", + "#Y", "#Z", "#À", "#Á", "#Â", "#Ã", "#Ç", "#É", "#Ê", "#Í", "#Ó", "#Ô", "#Õ", + "#Ú", "#Ü", "'", "(", ")", ",", "-", ".", "/", ":", ";", "?", "[", + "]", "a", "b", "d", "e", "f", "h", "i", "j", "k", "l", "m", "n", + "o", "p", "r", "s", "t", "u", "v", "w", "x", "y", "z", "{", "}", + "ð", "õ", "ĩ", "ŋ", "ũ", "ɐ", "ɑ", "ɒ", "ɔ", "ə", "ɛ", "ɜ", "ɡ", + "ɪ", "ɲ", "ɹ", "ɾ", "ʁ", "ʃ", "ʊ", "ʌ", "ʎ", "ʒ", "ʲ", "ˈ", "ˌ", + "ː", "̃", "θ", "ẽ", " ", "", ""}; } return {}; } @@ -1441,7 +1456,10 @@ run_hindi_native(const params& p) { static std::vector arabic_tokens() { - std::vector tokens = {" ", "ء", "آ", "أ", "إ", "ؤ", "ئ", "ا", "ب", "ة", "ت", "ث", "ج", "ح", "خ", "د", "ذ", "ر", "ز", "س", "ش", "ص", "ض", "ط", "ظ", "ع", "غ", "ف", "ق", "ك", "ل", "م", "ن", "ه", "و", "ى", "ي", "ً", "ٌ", "ٍ", "َ", "ُ", "ِ", "ّ", "ٰ", "ْ"}; + std::vector tokens = {" ", "ء", "آ", "أ", "إ", "ؤ", "ئ", "ا", "ب", "ة", "ت", "ث", + "ج", "ح", "خ", "د", "ذ", "ر", "ز", "س", "ش", "ص", "ض", "ط", + "ظ", "ع", "غ", "ف", "ق", "ك", "ل", "م", "ن", "ه", "و", "ى", + "ي", "ً", "ٌ", "ٍ", "َ", "ُ", "ِ", "ّ", "ٰ", "ْ"}; // v2607 pins NeMo ArabicCharsTokenizer charset_version=1, whose mixed-case // Arabic character set deliberately contains this second copy. const std::vector arabic_chars(tokens.begin() + 1, tokens.end()); @@ -1470,7 +1488,8 @@ run_arabic_native(const params& p, int offset, const std::string& tokenizer_name std::vector symbols; for (const auto& symbol : split_utf8(sentence)) { if (symbol == " ") { - if (!symbols.empty() && symbols.back() != symbol) symbols.push_back(symbol); + if (!symbols.empty() && symbols.back() != symbol) + symbols.push_back(symbol); } else if (ids.count(symbol) != 0) { symbols.push_back(symbol); } @@ -1662,7 +1681,8 @@ supports_native(const params& p) { return true; } if ((p.language == "pt-br" || p.language == "ko" || p.language == "ar-ae" || - p.language == "ar-sa" || p.language == "ar-msa") && p.v2607) { + p.language == "ar-sa" || p.language == "ar-msa") && + p.v2607) { return p.language != "pt-br" || fs::is_directory(p.model); } #ifdef NEMO_SPEECH_TTS_WITH_JA diff --git a/tests/cli/cli_contract_test.py b/tests/cli/cli_contract_test.py index 46ddf95..6956473 100644 --- a/tests/cli/cli_contract_test.py +++ b/tests/cli/cli_contract_test.py @@ -143,8 +143,12 @@ def stall_response() -> None: if len(missing_paths) >= 2: error = expect_json_error( run( - binary, "--json", "serve", "--tts.preempt", *missing_arguments, - "--no-warmup" + binary, + "--json", + "serve", + "--tts.preempt", + *missing_arguments, + "--no-warmup", ), 3, "missing_model", diff --git a/tests/cpp/tts/test_tokenizer_single_chars.cpp b/tests/cpp/tts/test_tokenizer_single_chars.cpp index 072503b..f8e029f 100644 --- a/tests/cpp/tts/test_tokenizer_single_chars.cpp +++ b/tests/cpp/tts/test_tokenizer_single_chars.cpp @@ -138,18 +138,18 @@ main(int argc, char** argv) { if (v2607) { bool v2607_ok = true; v2607_ok &= check_tokens(tokenizer, "A", "en-US", {90, 94, 94, 94, 88, 3358}); - v2607_ok &= check_tokens(tokenizer, "مرحبا", "ar-AE", {1329, 1405, 1391, 1387, 1382, - 1381, 1329, 3358}); - v2607_ok &= check_tokens(tokenizer, "مرحبا", "ar-SA", {1493, 1569, 1555, 1551, 1546, - 1545, 1493, 3358}); - v2607_ok &= check_tokens(tokenizer, "مرحبا", "ar-MSA", {1657, 1733, 1719, 1715, 1710, - 1709, 1657, 3358}); - v2607_ok &= check_tokens(tokenizer, "안녕하세요.", "ko-KR", {3212, 3125, 3112, 3211, 3109, - 3125, 3213, 3125, 3128, 3212, - 3108, 3160, 3212, 3130, 3124, - 3022, 2974, 3358}); - v2607_ok &= check_tokens(tokenizer, "Olá!", "pt-BR", {1125, 1082, 1079, 1119, 1070, - 1017, 1125, 3358}); + v2607_ok &= check_tokens( + tokenizer, "مرحبا", "ar-AE", {1329, 1405, 1391, 1387, 1382, 1381, 1329, 3358}); + v2607_ok &= check_tokens( + tokenizer, "مرحبا", "ar-SA", {1493, 1569, 1555, 1551, 1546, 1545, 1493, 3358}); + v2607_ok &= check_tokens( + tokenizer, "مرحبا", "ar-MSA", {1657, 1733, 1719, 1715, 1710, 1709, 1657, 3358}); + v2607_ok &= check_tokens( + tokenizer, "안녕하세요.", "ko-KR", + {3212, 3125, 3112, 3211, 3109, 3125, 3213, 3125, 3128, 3212, 3108, 3160, 3212, 3130, + 3124, 3022, 2974, 3358}); + v2607_ok &= check_tokens( + tokenizer, "Olá!", "pt-BR", {1125, 1082, 1079, 1119, 1070, 1017, 1125, 3358}); return v2607_ok ? 0 : 1; } bool ok = true; From 830db7926a85dbca2b92d37b8de55c8cd1c9cce3 Mon Sep 17 00:00:00 2001 From: Prabhsimran Singh Date: Fri, 21 Aug 2026 21:32:51 +0530 Subject: [PATCH 05/19] refactor(tts): migrate MagpieTTS decoder to the common GGML runtime - add persistent decoder graphs and device-resident KV caches - batch CFG execution and keep sampling on device - compose local codebook inference into a reusable CUDA graph --- ggml-patches/0007-magpietts-nanocodec.patch | 139 ++-- .../0014-cuda-relpos-extensions.patch | 264 ++++++-- ggml-patches/0017-cuda-stream-interop.patch | 58 ++ ggml-patches/README.md | 16 +- scripts/apply-ggml-patches.sh | 8 +- src/runtime/ggml/backend.cpp | 85 ++- src/runtime/ggml/runtime.h | 13 + src/runtime/ggml/session.cpp | 8 + src/runtime/ggml/tensor_container.cpp | 23 + src/tts/magpietts/CMakeLists.txt | 1 + src/tts/magpietts/decoder.cpp | 593 ++++++++++++++++-- src/tts/magpietts/decoder.h | 12 +- src/tts/magpietts/graph.h | 5 + src/tts/magpietts/lt.cpp | 344 ++++++++-- src/tts/magpietts/magpietts.cpp | 2 + src/tts/magpietts/magpietts_cuda_sampling.cu | 503 ++++++++++++--- src/tts/magpietts/magpietts_cuda_sampling.h | 37 ++ src/tts/magpietts/model.cpp | 15 +- tests/cpp/tts/test_magpietts_file.cpp | 121 +++- 19 files changed, 1895 insertions(+), 352 deletions(-) create mode 100644 ggml-patches/0017-cuda-stream-interop.patch diff --git a/ggml-patches/0007-magpietts-nanocodec.patch b/ggml-patches/0007-magpietts-nanocodec.patch index 13dc370..57d4c2b 100644 --- a/ggml-patches/0007-magpietts-nanocodec.patch +++ b/ggml-patches/0007-magpietts-nanocodec.patch @@ -15,7 +15,7 @@ index b54d4a6b..07a94052 100644 @@ -36,6 +37,10 @@ if (CUDAToolkit_FOUND) list(APPEND CMAKE_CUDA_ARCHITECTURES 89-real) endif() - + + if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL "13.0") + list(APPEND CMAKE_CUDA_ARCHITECTURES 110a-real) + endif() @@ -26,7 +26,7 @@ index b54d4a6b..07a94052 100644 @@ -71,16 +76,16 @@ if (CUDAToolkit_FOUND) FetchContent_MakeAvailable(CCCL) endif() - + - # Replace any plain 12X CUDA architectures with their "architecture-specific" equivalents 12Xa. - # 12X is forwards-compatible, 12Xa is not. - # Notably the Blackwell FP4 tensor core instructions are not forwards compatible and therefore need 12Xa. @@ -49,7 +49,7 @@ index b54d4a6b..07a94052 100644 @@ -90,8 +95,8 @@ if (CUDAToolkit_FOUND) set(${ARCHS} ${FIXED_ARCHS}) endforeach() - + - # If we try to compile a "native" build it will use the 12X architectures and fail. - # So we should instead use the native architectures as determined by CMake after replacing 12X with 12Xa. + # If we try to compile a "native" build it may use plain Blackwell architectures and fail. @@ -58,7 +58,7 @@ index b54d4a6b..07a94052 100644 if (CMAKE_CUDA_ARCHITECTURES STREQUAL "native" AND CMAKE_CUDA_ARCHITECTURES_NATIVE MATCHES "^[0-9]+(a|f)?(-real|-virtual)?(;[0-9]+(a|f)?(-real|-virtual)?|;)*$") set(CMAKE_CUDA_ARCHITECTURES ${CMAKE_CUDA_ARCHITECTURES_NATIVE}) diff --git a/src/ggml-cuda/common.cuh b/src/ggml-cuda/common.cuh -index 10817505..e69e710b 100644 +index 9bd40faa..56fe1bb7 100644 --- a/src/ggml-cuda/common.cuh +++ b/src/ggml-cuda/common.cuh @@ -25,6 +25,7 @@ @@ -84,7 +84,7 @@ index 10817505..e69e710b 100644 @@ -315,6 +319,10 @@ static bool amd_wmma_available(const int cc) { return (GGML_CUDA_CC_IS_RDNA4(cc) || GGML_CUDA_CC_IS_RDNA3(cc)); } - + +static bool thor_mma_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && cc == GGML_CUDA_CC_THOR; +} @@ -93,9 +93,9 @@ index 10817505..e69e710b 100644 return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_VOLTA; } @@ -1379,14 +1387,34 @@ struct ggml_backend_cuda_context { - + int64_t last_graph_eviction_sweep = 0; - + + static int64_t cuda_graph_env_ms_to_us(const char * name, int64_t default_ms) { + const char * value = getenv(name); + if (value == nullptr || value[0] == '\0') { @@ -138,7 +138,7 @@ index 8418ba66..20a84516 100644 @@ -1,34 +1,38 @@ #include "conv-transpose-1d.cuh" +#include "convert.cuh" - + -static __global__ void conv_transpose_1d_kernel( +#include + @@ -154,25 +154,25 @@ index 8418ba66..20a84516 100644 if (global_index >= output_size) { return; } - + - int out_index = global_index / dst_ne0; + const int out_t = global_index % dst_ne0; + const int out_c = (global_index / dst_ne0) % dst_ne1; + const int out_b = global_index / (dst_ne0 * dst_ne1); - + float accumulator = 0; - + - for (int c = 0; c < src0_ne2; c++) { - int idx = global_index % dst_ne0; + const int in_end = min(src1_ne0 - 1, out_t / s0); + const int in_start = max(0, (out_t - src0_ne0 + s0) / s0); - + - int kernel_offset = (src0_ne0 * src0_ne1 * c) + (out_index * src0_ne0); - int input_offset = src1_ne0 * c; + for (int c = 0; c < src0_ne2; c++) { + const int kernel_offset = src0_ne0 * (out_c + src0_ne1 * c); + const int input_offset = src1_ne0 * (c + src1_ne1 * out_b); - + - for (int i = 0; i < src1_ne0; i++) { - if (!(idx >= i*s0 && idx < i*s0 + src0_ne0)) { - continue; @@ -180,18 +180,18 @@ index 8418ba66..20a84516 100644 - int weight_idx = idx - i*s0; + for (int i = in_start; i <= in_end; i++) { + const int weight_idx = out_t - i*s0; - + - float kernel_weight = src0[kernel_offset + weight_idx]; - float input_value = src1[input_offset+i]; + const float kernel_weight = ggml_cuda_cast(src0[kernel_offset + weight_idx]); + const float input_value = src1[input_offset+i]; - + accumulator += kernel_weight * input_value; } @@ -37,26 +41,96 @@ static __global__ void conv_transpose_1d_kernel( GGML_UNUSED_VARS(p0, d0, src0_ne3, src1_ne3, dst_ne3, src1_ne1, dst_ne1, src1_ne2, dst_ne2); } - + -static void conv_transpose_1d_f32_f32_cuda( +template +static __global__ void conv_transpose_1d_grouped2_kernel( @@ -243,7 +243,7 @@ index 8418ba66..20a84516 100644 + const T * src0, const float * src1, float * dst, + const bool use_grouped2, cudaStream_t stream) { - + const int num_blocks = (output_size + CUDA_CONV_TRANPOSE_1D_BLOCK_SIZE - 1) / CUDA_CONV_TRANPOSE_1D_BLOCK_SIZE; - conv_transpose_1d_kernel<<>>( - s0,p0,d0,output_size, @@ -286,27 +286,27 @@ index 8418ba66..20a84516 100644 + + return shape_matches && is_nanocodec_up_weight; } - + void ggml_cuda_op_conv_transpose_1d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; - const float * src0_d = (const float *)src0->data; + const void * src0_d = src0->data; - + const ggml_tensor * src1 = dst->src[1]; const float * src1_d = (const float *)src1->data; @@ -64,7 +138,8 @@ void ggml_cuda_op_conv_transpose_1d(ggml_backend_cuda_context & ctx, ggml_tensor float * dst_d = (float *)dst->data; cudaStream_t stream = ctx.stream(); - + - GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); - + GGML_ASSERT(ggml_is_contiguous(src0)); @@ -77,10 +152,19 @@ void ggml_cuda_op_conv_transpose_1d(ggml_backend_cuda_context & ctx, ggml_tensor const int d0 = 1;//opts[4]; - + const int64_t output_size = ggml_nelements(dst); - - conv_transpose_1d_f32_f32_cuda(s0, p0, d0, output_size, @@ -331,13 +331,24 @@ index 8418ba66..20a84516 100644 + } } diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index e25be359..ebdd1451 100644 +index a531ce07..4b4a8488 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -3284,6 +3284,23 @@ static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { +@@ -2485,8 +2485,8 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { + return false; + } + +- //we only support fusion for ncols_dst = 1 +- if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { ++ // MMVF supports a two-column epilogue so paired CFG lanes can share each weight-row load. ++ if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] > 2) { + return false; + } + +@@ -3310,6 +3310,23 @@ static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { return cgraph->nodes[0]; } - + +static const char * ggml_cuda_graph_tensor_name(const ggml_tensor * tensor) { + return tensor && tensor->name[0] ? tensor->name : "(unnamed)"; +} @@ -357,19 +368,19 @@ index e25be359..ebdd1451 100644 + static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { bool res = false; - -@@ -3292,7 +3309,6 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx - + +@@ -3318,7 +3335,6 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx + if (cgraph->uid != 0 && cgraph->uid == graph->uid) { - GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); return false; } -@@ -3820,6 +3837,88 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, +@@ -3891,6 +3907,88 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } - + +static bool ggml_cuda_node_can_be_elided(const struct ggml_cgraph * cgraph, int node_idx, int32_t expected_uses) { + const ggml_tensor * node = cgraph->nodes[node_idx]; + return (node->flags & GGML_TENSOR_FLAG_COMPUTE) != 0 && @@ -454,11 +465,11 @@ index e25be359..ebdd1451 100644 + // try and fuse nodes and return the number of nodes to skip static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { - -@@ -3830,6 +3929,10 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph - + +@@ -3901,6 +3999,10 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph + ggml_tensor * node = cgraph->nodes[i]; - + + if (int n_fused = ggml_cuda_try_fuse_half_snake(cuda_ctx, cgraph, i)) { + return n_fused; + } @@ -466,7 +477,7 @@ index e25be359..ebdd1451 100644 //topk-moe if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || cgraph->nodes[i]->op == GGML_OP_ARGSORT) { -@@ -4475,7 +4578,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, +@@ -4623,7 +4725,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, // Warmup: need at least 2 calls with no property change on the 2nd call if (!properties_changed) { graph->warmup_complete = true; @@ -475,7 +486,7 @@ index e25be359..ebdd1451 100644 use_cuda_graph = true; cuda_graph_update_required = true; } -@@ -4485,7 +4588,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, +@@ -4633,7 +4735,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, if (properties_changed) { // Properties changed - reset warmup, execute directly until stable again graph->warmup_complete = false; @@ -484,7 +495,7 @@ index e25be359..ebdd1451 100644 } else { use_cuda_graph = true; cuda_graph_update_required = graph->instance == nullptr; -@@ -5286,7 +5389,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g +@@ -5434,7 +5536,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g { ggml_type src0_type = op->src[0]->type; ggml_type src1_type = op->src[1]->type; @@ -493,8 +504,8 @@ index e25be359..ebdd1451 100644 return true; } return false; -@@ -5555,6 +5658,12 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t - +@@ -5705,6 +5807,12 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t + { const auto & info = ggml_cuda_info(); + for (int id = 0; id < info.device_count; ++id) { @@ -513,7 +524,7 @@ index aad4c34a..628d8718 100644 @@ -159,6 +159,11 @@ bool ggml_cuda_should_use_mmf(enum ggml_type type, int cc, int warp_size, const return false; } - + + // Thor's tcgen05 paths are provided by CUDA libraries today; avoid ggml's warp-level matrix kernels there. + if (thor_mma_available(cc) && (type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16)) { + return false; @@ -523,13 +534,43 @@ index aad4c34a..628d8718 100644 if (src0_ne[1] <= 1024 && src1_ncols > 512) { return false; diff --git a/src/ggml-cuda/mmvf.cu b/src/ggml-cuda/mmvf.cu -index d9147202..97900cd8 100644 +index d9147202..0c974164 100644 --- a/src/ggml-cuda/mmvf.cu +++ b/src/ggml-cuda/mmvf.cu -@@ -793,9 +793,15 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 +@@ -383,7 +383,10 @@ static void mul_mat_vec_f_switch_fusion( + const dim3 & block_dims, const dim3 & block_nums, const int nbytes_shared, const int ids_stride, const cudaStream_t stream) { + + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; +- if constexpr (ncols_dst == 1) { ++ // The same epilogue addressing works for two adjacent columns. This is especially useful ++ // for classifier-free guidance: both lanes share the weight-row load and still fold the ++ // residual/bias writeback into the projection kernel. ++ if constexpr (ncols_dst <= 2) { + if (has_fusion) { + mul_mat_vec_f<<>> + (x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, +@@ -393,7 +396,7 @@ static void mul_mat_vec_f_switch_fusion( + } + } + +- GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); ++ GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst<=2"); + + mul_mat_vec_f<<>> + (x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, +@@ -650,7 +653,7 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor + + if (fusion) { + GGML_ASSERT( !ids || dst->ne[2] == 1); +- GGML_ASSERT( ids || dst->ne[1] == 1); ++ GGML_ASSERT( ids || dst->ne[1] <= 2); + if (fusion->x_bias) { + GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); + GGML_ASSERT(fusion->x_bias->ne[0] == dst->ne[0]); +@@ -793,9 +796,15 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 } } - + + // Thor's tcgen05 paths are provided by CUDA libraries today; avoid ggml's warp-level vector kernels there. + const bool prefer_cublas_tcgen05 = thor_mma_available(cc); + @@ -542,7 +583,7 @@ index d9147202..97900cd8 100644 if (ampere_mma_available(cc)) { return ne11 <= 3; } -@@ -812,6 +818,9 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 +@@ -812,9 +821,12 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 return ne11 <= 8; case GGML_TYPE_F16: if (GGML_CUDA_CC_IS_NVIDIA(cc)) { @@ -551,8 +592,12 @@ index d9147202..97900cd8 100644 + } const bool src0_small = (src0_ne[1] <= 512 || src0_ne[2]*src0_ne[3] == 1); if (ampere_mma_available(cc)) { - return src0_small && ne11 == 1; -@@ -838,6 +847,9 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 +- return src0_small && ne11 == 1; ++ return src0_small && ne11 <= 2; + } + if (cc >= GGML_CUDA_CC_ADA_LOVELACE) { + return src0_small && ne11 <= 4; +@@ -838,6 +850,9 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 return ne11 <= 8; case GGML_TYPE_BF16: if (GGML_CUDA_CC_IS_NVIDIA(cc)) { @@ -569,7 +614,7 @@ index 384638c1..6f4b219c 100644 @@ -1,6 +1,8 @@ #include "snake.cuh" #include "convert.cuh" - + +#include + // Fused Snake activation: y = x + sin^2(a * x) * inv_b diff --git a/ggml-patches/0014-cuda-relpos-extensions.patch b/ggml-patches/0014-cuda-relpos-extensions.patch index cbf2af9..bac3f72 100644 --- a/ggml-patches/0014-cuda-relpos-extensions.patch +++ b/ggml-patches/0014-cuda-relpos-extensions.patch @@ -1,5 +1,5 @@ diff --git a/include/ggml.h b/include/ggml.h -index fb823571..5cfa65e0 100644 +index fb823571..4671c80d 100644 --- a/include/ggml.h +++ b/include/ggml.h @@ -2436,8 +2436,9 @@ extern "C" { @@ -14,15 +14,16 @@ index fb823571..5cfa65e0 100644 // Q/K/V/P may be non-contiguous views as long as each d_k row is // contiguous — e.g. Q sliced from a fused-QKV projection and K/V read // head-split from a feat-major [n_feat, kv] window (the CUDA op derives -@@ -2460,6 +2461,27 @@ extern "C" { +@@ -2460,6 +2461,46 @@ extern "C" { float scale, bool merge_heads); + // Streaming CUDA variant. K/V contain only the current chunk; cached K/V + // are read directly from a persistent [n_feat*cache_len, slots, 2] F32 -+ // arena using one I32 slot id and circular-cache head per batch item. The -+ // op appends the current chunk at the head and leaves advancing that head -+ // to the caller after the graph run succeeds. ++ // arena using one I32 slot id and circular-cache head per batch item. ++ // ring_heads may be [batch] or [batch,2]; the optional second column is ++ // the number of valid past entries, allowing a fixed graph to skip unused ++ // cache rows. The caller advances the head after a successful graph run. + GGML_API struct ggml_tensor * ggml_fused_relpos_attn_cached( + struct ggml_context * ctx, + struct ggml_tensor * q, @@ -38,12 +39,30 @@ index fb823571..5cfa65e0 100644 + int64_t cache_len, + float scale, + bool merge_heads); ++ ++ // Absolute-position streaming attention. This shares the same persistent ++ // [n_feat*cache_len, slots, 2] F32 K/V arena and fused cache-update kernel ++ // as ggml_fused_relpos_attn_cached, but omits relative-position operands ++ // and biases. Q/K/V are [d_k, chunk_len, n_head, batch]; the current K/V ++ // chunk is attended and appended to the circular cache by the same op. ++ GGML_API struct ggml_tensor * ggml_fused_attn_cached( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * mask, ++ struct ggml_tensor * kv_cache, ++ struct ggml_tensor * slot_ids, ++ struct ggml_tensor * ring_heads, ++ int64_t cache_len, ++ float scale, ++ bool merge_heads); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/src/ggml-cuda/fused-relpos-attn.cu b/src/ggml-cuda/fused-relpos-attn.cu -index f3c4836a..d5640439 100644 +index f3c4836a..e0c0ec27 100644 --- a/src/ggml-cuda/fused-relpos-attn.cu +++ b/src/ggml-cuda/fused-relpos-attn.cu @@ -3,6 +3,8 @@ @@ -853,7 +872,7 @@ index f3c4836a..d5640439 100644 static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { // The target shape fixes dynamic shared memory, so occupancy is invariant // for a given compiled kernel and device. Avoid repeating the CUDA runtime -@@ -310,27 +993,31 @@ static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { +@@ -310,27 +993,32 @@ static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { int blocks = cached[device].load(std::memory_order_relaxed); if (blocks == 0) { CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( @@ -876,7 +895,8 @@ index f3c4836a..d5640439 100644 + const float * __restrict__ Kcache, const float * __restrict__ Vcache, + const int32_t * __restrict__ slot_ids, + const int32_t * __restrict__ ring_heads, -+ int q, int kv, int n_head, int cache_len, float scale, ++ const int32_t * __restrict__ active_lengths, ++ int q, int kv, int n_head, int cache_len, float scale, bool relative, // element strides: x_sq = between queries/keys, x_sh = between heads, // x_sb = between batch items long q_sq, long q_sh, long q_sb, @@ -889,31 +909,40 @@ index f3c4836a..d5640439 100644 extern __shared__ float sh[]; const int dk = blockDim.x; float * Qu = sh; // [dk] -@@ -346,8 +1033,16 @@ static __global__ void fused_relpos_attn_kernel( +@@ -346,11 +1034,20 @@ static __global__ void fused_relpos_attn_kernel( const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; +- const T * Ph = Ppos + (size_t) h * p_sh; +- const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; + const int slot = Cached ? slot_ids[b] : 0; + const int ring_head = Cached ? ring_heads[b] : 0; ++ const int key_begin = active_lengths ? cache_len - active_lengths[b] : 0; + const float * Kch = Cached + ? Kcache + (size_t) slot * cache_ss + (size_t) h * dk + : nullptr; + const float * Vch = Cached + ? Vcache + (size_t) slot * cache_ss + (size_t) h * dk + : nullptr; - const T * Ph = Ppos + (size_t) h * p_sh; -- const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ const T * Ph = relative ? Ppos + (size_t) h * p_sh : nullptr; + const float * Mb = mask ? mask + (size_t) b * m_sb + (size_t) i * m_sq : nullptr; - Qu[d] = Qhi[d] + bu[h * dk + d]; - Qv[d] = Qhi[d] + bv[h * dk + d]; -@@ -364,32 +1059,17 @@ static __global__ void fused_relpos_attn_kernel( +- Qu[d] = Qhi[d] + bu[h * dk + d]; +- Qv[d] = Qhi[d] + bv[h * dk + d]; ++ Qu[d] = Qhi[d] + (relative ? bu[h * dk + d] : 0.0f); ++ Qv[d] = Qhi[d] + (relative ? bv[h * dk + d] : 0.0f); + __syncthreads(); + + // scores. Two layouts: +@@ -363,32 +1060,19 @@ static __global__ void fused_relpos_attn_kernel( + // * otherwise: legacy thread-per-key scalar loop. if (dk == 128) { const int warp = d >> 5, lane = d & 31, nw = dk >> 5; - for (int j = warp; j < kv; j += nw) { +- for (int j = warp; j < kv; j += nw) { - const T * Kj = Kh + (size_t) j * k_sj + lane * 4; ++ for (int j = key_begin + warp; j < kv; j += nw) { const int row = (q - 1) + j - i; // rel-shift index - const T * Pr = Ph + (size_t) row * p_sr + lane * 4; +- const T * Pr = Ph + (size_t) row * p_sr + lane * 4; - float k4[4], p4[4]; - if (sizeof(T) == 2) { - const int2 kr = *(const int2 *) Kj; @@ -936,48 +965,96 @@ index f3c4836a..d5640439 100644 -#pragma unroll - for (int e = 0; e < 4; e++) { - s += k4[e] * Qu[lane * 4 + e] + p4[e] * Qv[lane * 4 + e]; -- } + const float4 k4 = + relpos_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, lane * 4); -+ const float4 p4 = relpos_load4(Pr); + float s = -+ k4.x * Qu[lane * 4 + 0] + p4.x * Qv[lane * 4 + 0] + -+ k4.y * Qu[lane * 4 + 1] + p4.y * Qv[lane * 4 + 1] + -+ k4.z * Qu[lane * 4 + 2] + p4.z * Qv[lane * 4 + 2] + -+ k4.w * Qu[lane * 4 + 3] + p4.w * Qv[lane * 4 + 3]; ++ k4.x * Qu[lane * 4 + 0] + k4.y * Qu[lane * 4 + 1] + ++ k4.z * Qu[lane * 4 + 2] + k4.w * Qu[lane * 4 + 3]; ++ if (relative) { ++ const T * Pr = Ph + (size_t) row * p_sr + lane * 4; ++ const float4 p4 = relpos_load4(Pr); ++ s += p4.x * Qv[lane * 4 + 0] + p4.y * Qv[lane * 4 + 1] + ++ p4.z * Qv[lane * 4 + 2] + p4.w * Qv[lane * 4 + 3]; + } #pragma unroll for (int off = 16; off > 0; off >>= 1) { - s += __shfl_xor_sync(0xffffffff, s, off); -@@ -400,12 +1080,13 @@ static __global__ void fused_relpos_attn_kernel( +@@ -398,15 +1082,39 @@ static __global__ void fused_relpos_attn_kernel( + sc[j] = s * scale + (Mb ? Mb[j] : 0.0f); + } } ++ } else if (dk == 64) { ++ // Preserve one thread per key (important for Magpie's long cache) while reducing each ++ // dot product from 64 scalar loads/instructions to 16 aligned float4 operations. ++ for (int j = key_begin + d; j < kv; j += dk) { ++ const int row = (q - 1) + j - i; ++ float ac = 0.0f; ++ float bd = 0.0f; ++#pragma unroll ++ for (int dd = 0; dd < 64; dd += 4) { ++ const float4 k4 = relpos_load_kv4( ++ Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, dd); ++ const float4 qu4 = relpos_load4(Qu + dd); ++ ac += k4.x * qu4.x + k4.y * qu4.y + k4.z * qu4.z + k4.w * qu4.w; ++ if (relative) { ++ const float4 p4 = relpos_load4(Ph + (size_t) row * p_sr + dd); ++ const float4 qv4 = relpos_load4(Qv + dd); ++ bd += p4.x * qv4.x + p4.y * qv4.y + p4.z * qv4.z + p4.w * qv4.w; ++ } ++ } ++ sc[j] = (ac + bd) * scale + (Mb ? Mb[j] : 0.0f); ++ } } else { - for (int j = d; j < kv; j += dk) { +- for (int j = d; j < kv; j += dk) { - const T * Kj = Kh + (size_t) j * k_sj; ++ for (int j = key_begin + d; j < kv; j += dk) { const int row = (q - 1) + j - i; // rel-shift index - const T * Pr = Ph + (size_t) row * p_sr; +- const T * Pr = Ph + (size_t) row * p_sr; float ac = 0.0f, bd = 0.0f; for (int dd = 0; dd < dk; dd++) { - ac += (float) Kj[dd] * Qu[dd]; +- bd += (float) Pr[dd] * Qv[dd]; + ac += relpos_load_kv( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, dd) * + Qu[dd]; - bd += (float) Pr[dd] * Qv[dd]; ++ if (relative) { ++ const T * Pr = Ph + (size_t) row * p_sr; ++ bd += (float) Pr[dd] * Qv[dd]; ++ } } sc[j] = (ac + bd) * scale + (Mb ? Mb[j] : 0.0f); -@@ -443,7 +1124,10 @@ static __global__ void fused_relpos_attn_kernel( + } +@@ -415,7 +1123,7 @@ static __global__ void fused_relpos_attn_kernel( + + // block max over sc[0..kv) + float lm = -INFINITY; +- for (int j = d; j < kv; j += dk) lm = fmaxf(lm, sc[j]); ++ for (int j = key_begin + d; j < kv; j += dk) lm = fmaxf(lm, sc[j]); + red[d] = lm; + __syncthreads(); + for (int s = dk / 2; s > 0; s >>= 1) { +@@ -427,7 +1135,7 @@ static __global__ void fused_relpos_attn_kernel( + + // exp + block sum + float ls = 0.0f; +- for (int j = d; j < kv; j += dk) { ++ for (int j = key_begin + d; j < kv; j += dk) { + const float e = __expf(sc[j] - m); + sc[j] = e; + ls += e; +@@ -443,7 +1151,10 @@ static __global__ void fused_relpos_attn_kernel( // ctx[d] = inv * sum_j softmax(sc[j]) * V[j,d] (thread d owns output dim d) float c = 0.0f; - for (int j = 0; j < kv; j++) c += sc[j] * (float) Vh[(size_t) j * v_sj + d]; -+ for (int j = 0; j < kv; j++) { ++ for (int j = key_begin; j < kv; j++) { + c += sc[j] * relpos_load_kv( + Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d); + } ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = c * inv; } -@@ -455,6 +1139,10 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor +@@ -455,38 +1166,76 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor const ggml_tensor * bias_u = dst->src[4]; const ggml_tensor * bias_v = dst->src[5]; const ggml_tensor * mask = dst->src[6]; // may be null @@ -985,10 +1062,19 @@ index f3c4836a..d5640439 100644 + const ggml_tensor * slot_ids = dst->src[8]; + const ggml_tensor * ring_heads = dst->src[9]; + const bool cached = kv_cache != nullptr; ++ const bool relative = p != nullptr; GGML_ASSERT(q->type == GGML_TYPE_F32); // K/V/P may be F32 or (all together) F16 — see the kernel comment. -@@ -467,26 +1155,54 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor +- GGML_ASSERT(k->type == v->type && k->type == p->type); ++ GGML_ASSERT(k->type == v->type); ++ GGML_ASSERT(!relative || k->type == p->type); + GGML_ASSERT(k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16); +- GGML_ASSERT(bias_u->type == GGML_TYPE_F32 && bias_v->type == GGML_TYPE_F32); ++ GGML_ASSERT(!relative || (bias_u->type == GGML_TYPE_F32 && bias_v->type == GGML_TYPE_F32)); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + const int d_k = q->ne[0]; const int q_len = q->ne[1]; const int n_head = q->ne[2]; const int batch = q->ne[3]; @@ -998,13 +1084,20 @@ index f3c4836a..d5640439 100644 + const int kv_len = cache_len + chunk_len; GGML_ASSERT((d_k & (d_k - 1)) == 0 && "fused_relpos_attn: d_k must be a power of two"); - GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k && p->ne[0] == d_k); +- GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k && p->ne[0] == d_k); - GGML_ASSERT(k->ne[1] == kv_len && v->ne[1] == kv_len); +- GGML_ASSERT(k->ne[2] == n_head && v->ne[2] == n_head && p->ne[2] == n_head); ++ GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k); + GGML_ASSERT(k->ne[1] == chunk_len && v->ne[1] == chunk_len); - GGML_ASSERT(k->ne[2] == n_head && v->ne[2] == n_head && p->ne[2] == n_head); ++ GGML_ASSERT(k->ne[2] == n_head && v->ne[2] == n_head); GGML_ASSERT(k->ne[3] == batch && v->ne[3] == batch); - GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); - GGML_ASSERT(bias_u->ne[1] == n_head && bias_v->ne[1] == n_head); +- GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); +- GGML_ASSERT(bias_u->ne[1] == n_head && bias_v->ne[1] == n_head); ++ if (relative) { ++ GGML_ASSERT(p->ne[0] == d_k && p->ne[2] == n_head); ++ GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); ++ GGML_ASSERT(bias_u->ne[1] == n_head && bias_v->ne[1] == n_head); ++ } + if (cached) { + GGML_ASSERT(slot_ids != nullptr); + GGML_ASSERT(ring_heads != nullptr); @@ -1015,6 +1108,7 @@ index f3c4836a..d5640439 100644 + GGML_ASSERT(slot_ids->ne[0] == batch); + GGML_ASSERT(ring_heads->type == GGML_TYPE_I32 && ggml_is_contiguous(ring_heads)); + GGML_ASSERT(ring_heads->ne[0] == batch); ++ GGML_ASSERT(ring_heads->ne[1] == 1 || ring_heads->ne[1] == 2); + } else { + GGML_ASSERT(slot_ids == nullptr); + GGML_ASSERT(ring_heads == nullptr); @@ -1050,17 +1144,20 @@ index f3c4836a..d5640439 100644 float scale; memcpy(&scale, dst->op_params, sizeof(scale)); -@@ -501,6 +1217,9 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor +@@ -501,7 +1250,11 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor k_sb = (long)(k->nb[3] / ke); const long v_sj = (long)(v->nb[1] / ke), v_sh = (long)(v->nb[2] / ke), v_sb = (long)(v->nb[3] / ke); +- const long p_sr = (long)(p->nb[1] / ke), p_sh = (long)(p->nb[2] / ke); + const long cache_sj = (long) d_k * n_head; + const long cache_ss = cached ? (long) (kv_cache->nb[1] / sizeof(float)) : 0; + const long cache_sp = cached ? (long) (kv_cache->nb[2] / sizeof(float)) : 0; - const long p_sr = (long)(p->nb[1] / ke), p_sh = (long)(p->nb[2] / ke); ++ const long p_sr = relative ? (long)(p->nb[1] / ke) : 0; ++ const long p_sh = relative ? (long)(p->nb[2] / ke) : 0; const size_t oe = ggml_type_size(dst->type); const long o_si = (long)(dst->nb[1] / oe), o_sh = (long)(dst->nb[2] / oe), -@@ -513,6 +1232,31 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor + o_sb = (long)(dst->nb[3] / oe); +@@ -513,6 +1266,35 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor GGML_ASSERT(shmem <= max_shmem && "fused_relpos_attn: kv window too large for shared memory"); cudaStream_t stream = ctx.stream(); @@ -1069,6 +1166,10 @@ index f3c4836a..d5640439 100644 + const int32_t * active_slots = cached ? (const int32_t *) slot_ids->data : nullptr; + const int32_t * active_ring_heads = + cached ? (const int32_t *) ring_heads->data : nullptr; ++ const int32_t * active_lengths = ++ cached && ring_heads->ne[1] == 2 ++ ? active_ring_heads + ring_heads->ne[0] ++ : nullptr; + + auto update_cache = [&]() { + if (!cached) { @@ -1092,24 +1193,28 @@ index f3c4836a..d5640439 100644 // The cache-aware Nemotron streaming geometry is Q=2, KV=72, H=8, // d_k=128. On SM100 the one-query warp kernel is fastest while its grid -@@ -523,6 +1267,114 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - const bool use_sm100_q2 = +@@ -520,9 +1302,117 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor + // fuse both query rows: halving the block count and reusing K/V then wins. + // cudaOccupancyMaxActiveBlocksPerMultiprocessor uses the compiled kernel's + // actual register count, avoiding a hard-coded batch-size threshold. +- const bool use_sm100_q2 = ++ const bool use_sm100_q2 = relative && device_info.cc == RELPOS_ATTN_CC_SM100 && device_info.warp_size == 32 && - d_k == RELPOS_ATTN_DK_128 && q_len == 2 && kv_len == 72; + d_k == RELPOS_ATTN_DK_128 && q_len == 2 && kv_len == 72 && m_sq == 0; -+ const bool use_register_q2 = ++ const bool use_register_q2 = relative && + GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.cc >= GGML_CUDA_CC_AMPERE && + device_info.warp_size == 32 && d_k == RELPOS_ATTN_DK_128 && q_len == 2 && + kv_len == 72 && cached && cache_len == 70 && chunk_len == 2 && n_head == 8 && + relpos_attn_register_resident_enabled(); -+ const bool use_register_q4 = ++ const bool use_register_q4 = relative && + GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.cc >= GGML_CUDA_CC_AMPERE && + device_info.warp_size == 32 && d_k == RELPOS_ATTN_DK_128 && q_len == 4 && + kv_len == 60 && cached && cache_len == 56 && chunk_len == 4 && n_head == 8 && + relpos_attn_register_resident_enabled(); + const bool common_q = + q_len == 1 || q_len == 2 || q_len == 4 || q_len == 7 || q_len == 14; -+ const bool use_register_common = ++ const bool use_register_common = relative && + GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.cc >= GGML_CUDA_CC_AMPERE && + device_info.warp_size == 32 && d_k == RELPOS_ATTN_DK_128 && common_q && + (cache_len == 56 || cache_len == 70) && cached && chunk_len == q_len && n_head == 8 && @@ -1208,7 +1313,7 @@ index f3c4836a..d5640439 100644 if (use_sm100_q2) { const size_t warp_shmem = ((size_t) 2 * RELPOS_ATTN_DK_128 + kv_len + RELPOS_ATTN_WARPS_128) * sizeof(float); -@@ -531,70 +1383,86 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor +@@ -531,70 +1421,88 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor (size_t) 2 * RELPOS_ATTN_WARPS_128) * sizeof(float); GGML_ASSERT(q2_shmem <= max_shmem); @@ -1308,13 +1413,15 @@ index f3c4836a..d5640439 100644 + constexpr bool Cached = decltype(cache_tag)::value; + fused_relpos_attn_kernel<<>>( + (const float *) q->data, (const T *) k->data, (const T *) v->data, -+ (const T *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, ++ relative ? (const T *) p->data : nullptr, ++ relative ? (const float *) bias_u->data : nullptr, ++ relative ? (const float *) bias_v->data : nullptr, mask ? (const float *) mask->data : nullptr, (float *) dst->data, - q_len, kv_len, n_head, scale, - q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, p_sr, p_sh, o_si, o_sh, o_sb, - m_sb); -+ cache_k, cache_v, active_slots, active_ring_heads, q_len, kv_len, n_head, cache_len, -+ scale, ++ cache_k, cache_v, active_slots, active_ring_heads, active_lengths, q_len, kv_len, ++ n_head, cache_len, scale, relative, + q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, cache_sj, cache_ss, + p_sr, p_sh, o_si, o_sh, o_sb, m_sq, m_sb); + }; @@ -1341,7 +1448,7 @@ index f3c4836a..d5640439 100644 + update_cache(); } diff --git a/src/ggml.c b/src/ggml.c -index 80b5802f..ae1b8db3 100644 +index 80b5802f..09b36241 100644 --- a/src/ggml.c +++ b/src/ggml.c @@ -5377,7 +5377,7 @@ struct ggml_tensor * ggml_flash_attn_ext( @@ -1364,27 +1471,48 @@ index 80b5802f..ae1b8db3 100644 float scale, bool merge_heads) { // Q/K/V/P may be arbitrary-strided views; the CUDA op derives addressing -@@ -5398,10 +5402,11 @@ struct ggml_tensor * ggml_fused_relpos_attn( - GGML_ASSERT(ggml_is_contiguous(bias_u)); - GGML_ASSERT(ggml_is_contiguous(bias_v)); - +@@ -5394,30 +5398,62 @@ struct ggml_tensor * ggml_fused_relpos_attn( + GGML_ASSERT(q->nb[0] == ggml_type_size(q->type)); + GGML_ASSERT(k->nb[0] == ggml_type_size(k->type)); + GGML_ASSERT(v->nb[0] == ggml_type_size(v->type)); +- GGML_ASSERT(p->nb[0] == ggml_type_size(p->type)); +- GGML_ASSERT(ggml_is_contiguous(bias_u)); +- GGML_ASSERT(ggml_is_contiguous(bias_v)); +- - const int64_t d_k = q->ne[0]; - const int64_t kv_len = k->ne[1]; - const int64_t q_len = q->ne[1]; - const int64_t n_head = q->ne[2]; +- +- GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k && p->ne[0] == d_k); +- GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); ++ const bool relative = p != NULL; ++ GGML_ASSERT(relative == (bias_u != NULL)); ++ GGML_ASSERT(relative == (bias_v != NULL)); ++ if (relative) { ++ GGML_ASSERT(p->nb[0] == ggml_type_size(p->type)); ++ GGML_ASSERT(ggml_is_contiguous(bias_u)); ++ GGML_ASSERT(ggml_is_contiguous(bias_v)); ++ } ++ + const int64_t d_k = q->ne[0]; + const int64_t chunk_len = k->ne[1]; + const int64_t kv_len = cache_len + chunk_len; + const int64_t q_len = q->ne[1]; + const int64_t n_head = q->ne[2]; - - GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k && p->ne[0] == d_k); - GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); -@@ -5410,14 +5415,36 @@ struct ggml_tensor * ggml_fused_relpos_attn( ++ ++ GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k); + // The kernel reads rel-pos rows (q_len-1)+j-i for j in [0,kv_len), i in + // [0,q_len) — i.e. rows [0, kv_len+q_len-1) — and takes its head stride // from p->nb, so a longer table (e.g. precomputed for the full chunk // length and reused by shorter tail chunks) is safe. - GGML_ASSERT(p->ne[1] >= kv_len + q_len - 1); // rel-pos length +- GGML_ASSERT(p->ne[1] >= kv_len + q_len - 1); // rel-pos length - GGML_ASSERT(v->ne[1] == kv_len); ++ if (relative) { ++ GGML_ASSERT(p->ne[0] == d_k); ++ GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); ++ GGML_ASSERT(p->ne[1] >= kv_len + q_len - 1); // rel-pos length ++ } + GGML_ASSERT(v->ne[1] == chunk_len); + GGML_ASSERT(cache_len >= 0); + GGML_ASSERT((kv_cache == NULL) == (slot_ids == NULL)); @@ -1398,6 +1526,7 @@ index 80b5802f..ae1b8db3 100644 + GGML_ASSERT(slot_ids->ne[0] == q->ne[3]); + GGML_ASSERT(ring_heads->type == GGML_TYPE_I32 && ggml_is_contiguous(ring_heads)); + GGML_ASSERT(ring_heads->ne[0] == q->ne[3]); ++ GGML_ASSERT(ring_heads->ne[1] == 1 || ring_heads->ne[1] == 2); + } else { + GGML_ASSERT(cache_len == 0); + } @@ -1422,7 +1551,7 @@ index 80b5802f..ae1b8db3 100644 } // Output mirrors q logically: [d_k, q_len, n_head, batch]. With -@@ -5434,6 +5461,7 @@ struct ggml_tensor * ggml_fused_relpos_attn( +@@ -5434,6 +5470,7 @@ struct ggml_tensor * ggml_fused_relpos_attn( } ggml_set_op_params(result, &scale, sizeof(scale)); @@ -1430,7 +1559,7 @@ index 80b5802f..ae1b8db3 100644 result->op = GGML_OP_FUSED_RELPOS_ATTN; result->src[0] = q; -@@ -5443,10 +5471,48 @@ struct ggml_tensor * ggml_fused_relpos_attn( +@@ -5443,10 +5480,65 @@ struct ggml_tensor * ggml_fused_relpos_attn( result->src[4] = bias_u; result->src[5] = bias_v; result->src[6] = mask; @@ -1475,6 +1604,23 @@ index 80b5802f..ae1b8db3 100644 + ctx, q, k, v, p, bias_u, bias_v, mask, kv_cache, slot_ids, ring_heads, cache_len, scale, + merge_heads); +} ++ ++struct ggml_tensor * ggml_fused_attn_cached( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * mask, ++ struct ggml_tensor * kv_cache, ++ struct ggml_tensor * slot_ids, ++ struct ggml_tensor * ring_heads, ++ int64_t cache_len, ++ float scale, ++ bool merge_heads) { ++ return ggml_fused_relpos_attn_impl( ++ ctx, q, k, v, NULL, NULL, NULL, mask, kv_cache, slot_ids, ring_heads, cache_len, scale, ++ merge_heads); ++} + void ggml_flash_attn_ext_set_prec( struct ggml_tensor * a, diff --git a/ggml-patches/0017-cuda-stream-interop.patch b/ggml-patches/0017-cuda-stream-interop.patch new file mode 100644 index 0000000..c8d7ea3 --- /dev/null +++ b/ggml-patches/0017-cuda-stream-interop.patch @@ -0,0 +1,58 @@ +diff --git a/include/ggml-cuda.h b/include/ggml-cuda.h +--- a/include/ggml-cuda.h ++++ b/include/ggml-cuda.h +@@ -23,6 +23,17 @@ GGML_BACKEND_API ggml_backend_t ggml_backend_cuda_init(int device); + + GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend); + ++// Returns the backend's current CUDA/HIP/MUSA stream for in-process device interop. ++// The returned stream is owned by the backend and must not be destroyed by the caller. ++GGML_BACKEND_API void * ggml_backend_cuda_get_stream(ggml_backend_t backend); ++ ++// Returns the backend-owned native CUDA/HIP graph template for a stable GGML graph after its ++// normal warm-up/capture has completed. The opaque handle is borrowed and is intended for runtime ++// graph composition (for example, adding GGML inference as a child node of a larger pipeline ++// graph). It remains owned by the backend and must not be destroyed by the caller. ++GGML_BACKEND_API void * ggml_backend_cuda_get_graph_template( ++ ggml_backend_t backend, const struct ggml_cgraph * cgraph); ++ + // device buffer + GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device); + +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -5514,6 +5514,33 @@ bool ggml_backend_is_cuda(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cuda_guid()); + } + ++void * ggml_backend_cuda_get_stream(ggml_backend_t backend) { ++ if (!ggml_backend_is_cuda(backend)) { ++ return nullptr; ++ } ++ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; ++ return (void *) cuda_ctx->stream(); ++} ++ ++void * ggml_backend_cuda_get_graph_template( ++ ggml_backend_t backend, const struct ggml_cgraph * cgraph) { ++ if (!ggml_backend_is_cuda(backend) || cgraph == nullptr) { ++ return nullptr; ++ } ++#ifdef USE_CUDA_GRAPH ++ ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; ++ const ggml_cuda_graph_key graph_key = ggml_cuda_graph_get_key(cgraph); ++ ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); ++ if (!graph->warmup_complete || graph->graph == nullptr || graph->instance == nullptr) { ++ return nullptr; ++ } ++ return (void *) graph->graph; ++#else ++ GGML_UNUSED(cgraph); ++ return nullptr; ++#endif ++} ++ + int ggml_backend_cuda_get_device_count() { + return ggml_cuda_info().device_count; + } diff --git a/ggml-patches/README.md b/ggml-patches/README.md index 5dc6cb7..8b1d985 100644 --- a/ggml-patches/README.md +++ b/ggml-patches/README.md @@ -140,6 +140,8 @@ stock comparison therefore requires both a pristine ggml checkout and - **0007-magpietts-nanocodec.patch** - adds the CUDA operations used by MagpieTTS and NanoCodec, including grouped transposed convolution and Snake; + extends F16 MMVF and its bias/residual epilogue to two-column inputs for + paired CFG execution; bounds the keyed CUDA graph cache with configurable sweep and idle-eviction intervals; and adds SM110/Jetson Thor architecture handling. @@ -186,12 +188,16 @@ stock comparison therefore requires both a pristine ggml checkout and the performance crossover depends on the GPU and physical batch size. - **0014-cuda-relpos-extensions.patch** - extends fused relative-position - attention for the cache-aware and offline FastConformer paths. The + attention for cache-aware and offline FastConformer paths, and generalizes + the cached op to absolute attention without dummy relative-position + operands. A `[batch,2]` cache metadata tensor carries both circular heads and + active lengths so fixed-shape graphs skip unused cache rows. The `d_k=64` + score path uses aligned `float4` loads. The cache-aware path reads persistent K/V rows directly by state slot and circular head, then overwrites only the rows replaced by the current chunk. Register-resident NVIDIA SM80+ kernels cover the common R=0, 1, 3, 6, and 13 - streaming shapes for both Nemotron cache geometries, with exact-shape kernels - retained where they are faster. The offline mask contract accepts both + streaming shapes for both Nemotron cache geometries, with exact-shape + fallbacks. The offline mask contract accepts both per-batch key-padding masks and `[key, query]` L/R masks. Set `GGML_CUDA_RELPOS_REGISTER_RESIDENT=0` before process start to disable the register-resident specializations without changing the direct @@ -209,6 +215,10 @@ stock comparison therefore requires both a pristine ggml checkout and reshape interleaves those axes for batches larger than one; batch one keeps its original zero-copy path. +- **0017-cuda-stream-interop.patch** - exposes borrowed access to the CUDA + backend's active stream and stable graph templates for composition with + external CUDA work. Both handles remain backend-owned. + ## Regenerating after editing ggml Several patches touch the same ggml files, so regenerating a patch from the diff --git a/scripts/apply-ggml-patches.sh b/scripts/apply-ggml-patches.sh index fb4935f..b4dc00c 100755 --- a/scripts/apply-ggml-patches.sh +++ b/scripts/apply-ggml-patches.sh @@ -3,12 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 # Apply the in-tree ggml patches (ggml-patches/*.patch, in filename order) onto # the vendored ggml submodule. This keeps the submodule pinned to clean -# upstream; our project-specific ggml changes (fused rel-pos attention, NVFP4 -# quantization, norm fusion, dw-conv F16, skinny-q8 GEMM, FastConformer BF16 -# fusions, large-batch CUDA fixes, and MagpieTTS/NanoCodec CUDA ops) are applied -# at build setup time. Later patches also add the NVIDIA SM80+ cached-F16 route -# and a portable direct circular K/V fused-attention path with an SM80+ exact- -# shape register specialization. +# upstream. Project-specific kernels and runtime extensions are applied during +# build setup. # # Patches that still reverse-apply cleanly are skipped. Apply the complete # series to a clean submodule for deterministic setup; later patches may refine diff --git a/src/runtime/ggml/backend.cpp b/src/runtime/ggml/backend.cpp index 886648a..4fca3b7 100644 --- a/src/runtime/ggml/backend.cpp +++ b/src/runtime/ggml/backend.cpp @@ -17,6 +17,12 @@ BackendManager::BackendManager(Params params) { init_backends(); } +BackendManager::BackendManager(Params params, ggml_backend_t borrowed_gpu_backend) { + this->params = params; + borrowed_gpu_backend_ = borrowed_gpu_backend; + init_backends(); +} + void BackendManager::init_backends() { @@ -44,41 +50,57 @@ BackendManager::init_backends() { ggml_backend_dev_t dev = nullptr; if (params.use_gpu) { - int idx = 0; - for (int i = 0; i < dev_count; i++) { - ggml_backend_dev_t dev_cur = ggml_backend_dev_get(i); - GGMLF_LOG_INFO("Device %d: %s\n", i, ggml_backend_dev_name(dev_cur)); - const auto dev_type = ggml_backend_dev_type(dev_cur); - if (dev_type == GGML_BACKEND_DEVICE_TYPE_GPU || - dev_type == GGML_BACKEND_DEVICE_TYPE_IGPU) { - // Register buffer types only for the selected device. - if (idx == params.gpu_device_idx) { - dev = dev_cur; - auto* buft = ggml_backend_dev_buffer_type(dev); - if (buft) { - buft_list.emplace_back(dev, buft); + if (borrowed_gpu_backend_ != nullptr) { + dev = ggml_backend_get_device(borrowed_gpu_backend_); + const auto dev_type = dev ? ggml_backend_dev_type(dev) : GGML_BACKEND_DEVICE_TYPE_CPU; + if (dev == nullptr || + (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && + dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU)) { + throw std::runtime_error("borrowed GPU backend is not a GPU device"); + } + auto* buft = ggml_backend_dev_buffer_type(dev); + if (buft) { + buft_list.emplace_back(dev, buft); + } + gpu_backend = borrowed_gpu_backend_; + GGMLF_LOG_INFO("Using borrowed GPU backend: %s\n", ggml_backend_name(gpu_backend)); + } else { + int idx = 0; + for (int i = 0; i < dev_count; i++) { + ggml_backend_dev_t dev_cur = ggml_backend_dev_get(i); + GGMLF_LOG_INFO("Device %d: %s\n", i, ggml_backend_dev_name(dev_cur)); + const auto dev_type = ggml_backend_dev_type(dev_cur); + if (dev_type == GGML_BACKEND_DEVICE_TYPE_GPU || + dev_type == GGML_BACKEND_DEVICE_TYPE_IGPU) { + // Register buffer types only for the selected device. + if (idx == params.gpu_device_idx) { + dev = dev_cur; + auto* buft = ggml_backend_dev_buffer_type(dev); + if (buft) { + buft_list.emplace_back(dev, buft); + } } - } - if (++idx > params.gpu_device_idx) { - break; + if (++idx > params.gpu_device_idx) { + break; + } } } + if (dev == nullptr) { + throw std::runtime_error( + "use_gpu=true but no matching GPU device found (gpu_device_idx=" + + std::to_string(params.gpu_device_idx) + ")"); + } + GGMLF_LOG_INFO("Using GPU backend: %s\n", ggml_backend_dev_name(dev)); + ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr); + if (backend == nullptr) { + throw std::runtime_error( + std::string("use_gpu=true but ggml_backend_dev_init failed for ") + + ggml_backend_dev_name(dev)); + } + gpu_backend = backend; + backends.emplace_back(backend); } - if (dev == nullptr) { - throw std::runtime_error( - "use_gpu=true but no matching GPU device found (gpu_device_idx=" + - std::to_string(params.gpu_device_idx) + ")"); - } - GGMLF_LOG_INFO("Using GPU backend: %s\n", ggml_backend_dev_name(dev)); - ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr); - if (backend == nullptr) { - throw std::runtime_error( - std::string("use_gpu=true but ggml_backend_dev_init failed for ") + - ggml_backend_dev_name(dev)); - } - gpu_backend = backend; - backends.emplace_back(backend); } for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { @@ -124,7 +146,8 @@ BackendManager::init_backends() { std::vector BackendManager::get_backends() { std::vector handles; - handles.reserve(backends.size()); + handles.reserve(backends.size() + (borrowed_gpu_backend_ ? 1 : 0)); + if (borrowed_gpu_backend_) handles.push_back(borrowed_gpu_backend_); for (const auto& b : backends) handles.push_back(b.get()); return handles; } diff --git a/src/runtime/ggml/runtime.h b/src/runtime/ggml/runtime.h index 7b9cd69..31c052c 100644 --- a/src/runtime/ggml/runtime.h +++ b/src/runtime/ggml/runtime.h @@ -171,6 +171,10 @@ struct Params { class BackendManager { public: explicit BackendManager(Params params); + // Reuse an already initialized GPU backend owned by another pipeline object. The manager + // creates/owns only its auxiliary backends (for example CPU fallback); the borrowed handle + // must outlive this manager and every Session created from it. + BackendManager(Params params, ggml_backend_t borrowed_gpu_backend); // Out-of-line so ggml_backend_ptr's deleter is instantiated in the implementation TU. ~BackendManager(); @@ -194,6 +198,7 @@ class BackendManager { private: Params params; + ggml_backend_t borrowed_gpu_backend_ = nullptr; std::vector backends; // Non-owning alias into `backends`. ggml_backend_t gpu_backend = nullptr; @@ -234,6 +239,10 @@ class TensorContainer { bool has_tensor_by_name(const std::string& name); void cache_tensor(std::string name, ggml_bf_tensor tensor); + // Register storage owned outside this TensorContainer. Imported tensors participate in graph + // construction like ordinary model tensors but are never allocated or freed by the runtime. + ggml_bf_tensor import_tensor(std::string name, ggml_tensor* tensor); + // Declares tensors on the primary device's main buffer type. ggml_bf_tensor create_tensor_1d(std::string name, ggml_type data_type, int64_t ne0); ggml_bf_tensor create_tensor_2d( @@ -285,6 +294,10 @@ class Session { // Allows exact dtype copies and F32-to-F16 conversion. void load_weight(const std::string& gguf_key); + // Import a model tensor whose storage is owned by the embedding pipeline. Call only from + // Module::define_tensors(); the tensor and its backend buffer must outlive this Session. + ggml_bf_tensor import_model_tensor(const std::string& name, ggml_tensor* tensor); + // Invoked before upload so model code can annotate destination storage. // Install before setup(), which loads the model weights. using WeightLoadHook = diff --git a/src/runtime/ggml/session.cpp b/src/runtime/ggml/session.cpp index 76e165c..0e231a9 100644 --- a/src/runtime/ggml/session.cpp +++ b/src/runtime/ggml/session.cpp @@ -303,6 +303,14 @@ Session::load_weight(const std::string& gguf_key) { } } +ggml_bf_tensor +Session::import_model_tensor(const std::string& name, ggml_tensor* tensor) { + if (!model_tensor_container) { + throw std::logic_error("import_model_tensor must be called during Session::setup"); + } + return model_tensor_container->import_tensor(name, tensor); +} + int Session::setup() { std::lock_guard compute_lock(backend_manager_->compute_mutex()); diff --git a/src/runtime/ggml/tensor_container.cpp b/src/runtime/ggml/tensor_container.cpp index 87d6b15..d6bb5ea 100644 --- a/src/runtime/ggml/tensor_container.cpp +++ b/src/runtime/ggml/tensor_container.cpp @@ -136,6 +136,29 @@ TensorContainer::cache_tensor(std::string name, ggml_bf_tensor tensor) { } } +ggml_bf_tensor +TensorContainer::import_tensor(std::string name, ggml_tensor* tensor) { + if (!tensor || !tensor->buffer || !tensor->data) { + throw std::invalid_argument("cannot import an unallocated tensor: " + name); + } + ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(tensor->buffer); + bool supported = false; + for (const auto& entry : buft_list) { + if (entry.second == buft) { + supported = true; + break; + } + } + if (!supported) { + throw std::invalid_argument("imported tensor uses an unregistered buffer type: " + name); + } + ggml_bf_tensor result(tensor, buft); + if (!tensor_lookup.insert(std::make_pair(name, result)).second) { + throw std::runtime_error("duplicate tensor name imported: " + name); + } + return result; +} + ggml_bf_tensor TensorContainer::create_tensor_1d(std::string name, ggml_type data_type, int64_t ne0) { ggml_tensor* meta = ggml_new_tensor_1d(get_temp_ctx(), data_type, ne0); diff --git a/src/tts/magpietts/CMakeLists.txt b/src/tts/magpietts/CMakeLists.txt index 01f1d57..3d99ff4 100644 --- a/src/tts/magpietts/CMakeLists.txt +++ b/src/tts/magpietts/CMakeLists.txt @@ -36,6 +36,7 @@ target_include_directories(nemo_speech_tts PUBLIC target_link_libraries(nemo_speech_tts PUBLIC ggml nemo_speech_common + nemo_speech_runtime_ggml nemo_speech_tts_nanocodec_obj nemo_speech_tts_preproc nemo_speech_tts_tokenizer_obj diff --git a/src/tts/magpietts/decoder.cpp b/src/tts/magpietts/decoder.cpp index 7663669..56ec3c5 100644 --- a/src/tts/magpietts/decoder.cpp +++ b/src/tts/magpietts/decoder.cpp @@ -6,9 +6,13 @@ #include #include #include +#include #include +#include +#include #include +#include "../../runtime/ggml/runtime.h" #include "graph.h" #include "nvtx_utils.h" @@ -399,6 +403,385 @@ read_alignment_outputs( return true; } +namespace { + +constexpr int kMagpieCfgLanes = 2; + +bool +runtime_layer_selected(const std::vector& layers, int layer) { + return layers.empty() || + std::find(layers.begin(), layers.end(), static_cast(layer)) != layers.end(); +} + +std::string +runtime_kv_name(int layer) { + return "magpietts.decoder.runtime.kv." + std::to_string(layer); +} + +// Single-token decoder graph with external cross-K/V and persistent self-K/V storage. +class PersistentDecoderModule final : public ggml_runtime::Module { + public: + PersistentDecoderModule( + const magpietts_model& model, const DecoderCrossKvCache& cross_kv, int text_len, + int cache_len) + : model_(model), cross_kv_(cross_kv), text_len_(text_len), cache_len_(cache_len) { + for (int layer = 0; layer < static_cast(model_.decoder.layers.size()); ++layer) { + if (model_.decoder.layers[layer].has_cross && + runtime_layer_selected(model_.decoder.estimate_alignment_from_layers, layer)) { + ++alignment_count_; + } + } + } + + void define_tensors(ggml_runtime::Session* session) override { + std::unordered_set seen; + int imported = 0; + auto import = [&](ggml_tensor* tensor) { + if (!tensor || !seen.insert(tensor).second) { + return; + } + session->import_model_tensor( + "magpietts.decoder.external." + std::to_string(imported++), tensor); + }; + + for (ggml_tensor* embedding : model_.audio_embeddings) import(embedding); + import(model_.decoder.pos_emb); + import(model_.decoder.norm_out); + for (const magpietts_layer& layer : model_.decoder.layers) { + import(layer.norm_self); + import(layer.self_qkv); + import(layer.self_o); + import(layer.norm_xattn_query); + import(layer.cross_q); + import(layer.cross_o); + import(layer.norm_ff); + for (ggml_tensor* tensor : layer.ff_proj) import(tensor); + for (ggml_tensor* tensor : layer.ff_out) import(tensor); + } + import(cross_kv_.memory_k); + import(cross_kv_.memory_v); + + session->model_tensor_container->create_tensor_1d( + "magpietts.decoder.runtime.slot_ids", GGML_TYPE_I32, kMagpieCfgLanes); + for (int layer = 0; layer < model_.hparams.n_dec_layer; ++layer) { + session->model_tensor_container->create_tensor_3d( + runtime_kv_name(layer), GGML_TYPE_F32, + static_cast(model_.hparams.n_embd) * cache_len_, kMagpieCfgLanes, 2); + } + } + + ggml_runtime::TensorBag build_graph( + ggml_runtime::Session* session, ggml_runtime::TensorBag inputs, + ggml_runtime::TensorContainer* tc) override { + if (inputs.tensor_count() != 4) { + throw std::runtime_error("Magpie persistent decoder expects four inputs"); + } + const auto tokens = inputs.get_tensor(0); + const auto position = inputs.get_tensor(1); + const auto cache_meta = inputs.get_tensor(2); + const auto prior = inputs.get_tensor(3); + const auto bf_ctx = tc->get_ctx_of_buffer_type(tokens.buft); + ggml_context* ctx = bf_ctx.ctx; + const magpietts_hparams& h = model_.hparams; + const magpietts_transformer& tr = model_.decoder; + + ggml_tensor* audio = nullptr; + for (int codebook = 0; codebook < h.stacked_audio_codebooks(); ++codebook) { + ggml_tensor* token = ggml_view_1d( + ctx, tokens.tensor, 1, static_cast(codebook) * sizeof(int32_t)); + ggml_tensor* embedding = ggml_get_rows(ctx, model_.audio_embeddings[codebook], token); + audio = audio ? ggml_add(ctx, audio, embedding) : embedding; + } + audio = ggml_scale(ctx, audio, 1.0f / static_cast(h.stacked_audio_codebooks())); + audio = ggml_add(ctx, audio, ggml_get_rows(ctx, tr.pos_emb, position.tensor)); + // Store CFG lanes as projection columns for two-column MMVF. + ggml_tensor* x = ggml_concat(ctx, audio, audio, 1); // [E,CFG=2] + + const int64_t d_head = tr.n_embd / tr.n_head; + auto slots = session->model_tensor_container->get_tensor_by_name( + "magpietts.decoder.runtime.slot_ids"); + std::vector alignment_outputs; + + for (int layer_index = 0; layer_index < static_cast(tr.layers.size()); + ++layer_index) { + const magpietts_layer& layer = tr.layers[layer_index]; + ggml_tensor* residual = x; + ggml_tensor* cur = layer_norm(ctx, x, layer.norm_self); + ggml_tensor* qkv = linear(ctx, layer.self_qkv, cur); + const size_t element = ggml_element_size(qkv); + auto split_heads = [&](size_t offset) { + // Expose Q/K/V as [d_head,1,n_head,B] views without staging copies. + return ggml_view_4d( + ctx, qkv, d_head, 1, tr.n_head, kMagpieCfgLanes, qkv->nb[1], + static_cast(d_head) * element, qkv->nb[1], offset); + }; + ggml_tensor* q = split_heads(0); + ggml_tensor* k = split_heads(static_cast(tr.n_embd) * element); + ggml_tensor* v = split_heads(static_cast(2 * tr.n_embd) * element); + auto kv = session->model_tensor_container->get_tensor_by_name( + runtime_kv_name(layer_index)); + ggml_tensor* heads = ggml_fused_attn_cached( + ctx, q, k, v, nullptr, kv.tensor, slots.tensor, cache_meta.tensor, cache_len_, + 1.0f / std::sqrt(static_cast(d_head)), true); + ggml_tensor* merged = ggml_reshape_2d( + ctx, ggml_permute(ctx, heads, 0, 2, 1, 3), tr.n_embd, + kMagpieCfgLanes); + x = ggml_add(ctx, residual, linear(ctx, layer.self_o, merged)); + + // Text cross-attention applies only to the conditional lane. + if (tr.has_cross && layer.has_cross) { + ggml_tensor* cond = ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], 0); + ggml_tensor* uncond = + ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], x->nb[1]); + ggml_tensor* cross_in = layer_norm(ctx, cond, layer.norm_xattn_query); + ggml_tensor* last_attn = nullptr; + const bool apply_prior = tr.apply_attention_prior && + runtime_layer_selected(tr.apply_prior_to_layers, layer_index); + const bool collect = runtime_layer_selected( + tr.estimate_alignment_from_layers, layer_index); + ggml_tensor* cross = cross_attention_cached( + ctx, tr, layer, cross_kv_, layer_index, cross_in, + apply_prior ? prior.tensor : nullptr, collect ? &last_attn : nullptr, true); + cond = ggml_add(ctx, cond, cross); + x = ggml_concat(ctx, cond, uncond, 1); + if (last_attn) { + alignment_outputs.push_back(last_attn); + } + } + + residual = x; + cur = layer_norm(ctx, x, layer.norm_ff); + cur = causal_conv1d(ctx, cur, layer.ff_proj); + cur = ggml_gelu(ctx, cur); + cur = causal_conv1d(ctx, cur, layer.ff_out); + x = ggml_add(ctx, residual, cur); + } + if (tr.norm_out) x = layer_norm(ctx, x, tr.norm_out); + + ggml_tensor* cond = ggml_cont_2d( + ctx, ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], 0), tr.n_embd, 1); + ggml_tensor* uncond = ggml_cont_2d( + ctx, ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], x->nb[1]), tr.n_embd, 1); + ggml_set_name(cond, "magpietts_decoder_runtime_hidden_cond"); + ggml_set_name(uncond, "magpietts_decoder_runtime_hidden_uncond"); + ggml_runtime::TensorBag outputs; + outputs.add_tensor({cond, bf_ctx.buft}); + outputs.add_tensor({uncond, bf_ctx.buft}); + if (alignment_outputs.size() != alignment_count_) { + throw std::runtime_error("Magpie persistent decoder alignment topology changed"); + } + if (!alignment_outputs.empty()) { + ggml_tensor* sum = nullptr; + for (ggml_tensor* alignment : alignment_outputs) { + ggml_tensor* per_layer = alignment; + if (tr.n_cross_head > 1) { + per_layer = ggml_reshape_1d( + ctx, ggml_sum_rows(ctx, ggml_cont(ctx, ggml_transpose(ctx, alignment))), + text_len_); + } else { + per_layer = ggml_reshape_1d(ctx, alignment, text_len_); + } + sum = sum ? ggml_add(ctx, sum, per_layer) : per_layer; + } + ggml_tensor* mean = ggml_scale( + ctx, sum, + 1.0f / + static_cast(alignment_outputs.size() * tr.n_cross_head)); + ggml_set_name(mean, "magpietts_decoder_runtime_alignment_mean"); + outputs.add_tensor({mean, bf_ctx.buft}); + } + return outputs; + } + + void set_data(ggml_runtime::Session* session) override { + const int32_t slots[kMagpieCfgLanes] = {0, 1}; + auto slot_ids = session->model_tensor_container->get_tensor_by_name( + "magpietts.decoder.runtime.slot_ids"); + ggml_backend_tensor_set(slot_ids.tensor, slots, 0, sizeof(slots)); + for (int layer = 0; layer < model_.hparams.n_dec_layer; ++layer) { + auto kv = session->model_tensor_container->get_tensor_by_name(runtime_kv_name(layer)); + ggml_backend_tensor_memset(kv.tensor, 0, 0, ggml_nbytes(kv.tensor)); + } + } + + size_t alignment_count() const { return alignment_count_; } + + private: + const magpietts_model& model_; + const DecoderCrossKvCache& cross_kv_; + int text_len_ = 0; + int cache_len_ = 0; + size_t alignment_count_ = 0; +}; + +} // namespace + +class MagpieDecoder::PersistentDecoderRuntime { + public: + PersistentDecoderRuntime( + const magpietts_model& model, const DecoderCrossKvCache& cross_kv, int text_len) + : model_(model), cross_kv_(&cross_kv), text_len_(text_len), + cache_len_( + model.hparams.baked_context_length + + (model.hparams.max_decoder_steps + model.hparams.frame_stacking_factor - 1) / + model.hparams.frame_stacking_factor - + 1), + backend_manager_(ggml_runtime::Params{true, 0, nullptr}, model.backend), + module_(model, cross_kv, text_len, cache_len_), + session_(backend_manager_, &module_, nullptr) { + if (cache_len_ <= 0 || cache_len_ >= model.hparams.n_ctx) { + throw std::runtime_error("invalid persistent Magpie decoder cache length"); + } + session_.set_run_cache_capacity(1); + session_.setup(); + } + + bool matches(const DecoderCrossKvCache* cross_kv, int text_len) const { + return cross_kv == cross_kv_ && text_len == text_len_; + } + + bool sequence_matches(int n_tokens) const { return n_tokens == n_tokens_; } + + void seed(const DecoderKvCache& cond, const DecoderKvCache& uncond) { + if (cond.n_tokens <= 0 || cond.n_tokens != uncond.n_tokens || cond.n_tokens > cache_len_) { + throw std::runtime_error("cannot seed persistent decoder from incompatible KV caches"); + } + ggml_context* ctx = new_graph_context(); + const size_t element = sizeof(float); + const size_t source_layer_bytes = + static_cast(cond.n_ctx) * cond.n_embd * element; + const size_t copy_elements = static_cast(cond.n_tokens) * cond.n_embd; + const size_t destination_token = static_cast(cache_len_ - cond.n_tokens); + for (int layer = 0; layer < cond.n_layers; ++layer) { + auto arena = session_.model_tensor_container->get_tensor_by_name(runtime_kv_name(layer)); + for (int plane = 0; plane < 2; ++plane) { + ggml_tensor* dst_base = arena.tensor; + ggml_tensor* cond_src_base = plane == 0 ? cond.memory_k : cond.memory_v; + ggml_tensor* uncond_src_base = plane == 0 ? uncond.memory_k : uncond.memory_v; + const ggml_tensor* sources[kMagpieCfgLanes] = {cond_src_base, uncond_src_base}; + for (int lane = 0; lane < kMagpieCfgLanes; ++lane) { + ggml_tensor* src = ggml_view_1d( + ctx, const_cast(sources[lane]), copy_elements, + static_cast(layer) * source_layer_bytes); + const size_t dst_offset = + static_cast(plane) * dst_base->nb[2] + + static_cast(lane) * dst_base->nb[1] + + destination_token * cond.n_embd * element; + ggml_tensor* dst = ggml_view_1d(ctx, dst_base, copy_elements, dst_offset); + ggml_backend_tensor_copy_async(model_.backend, model_.backend, src, dst); + } + } + } + ggml_backend_synchronize(model_.backend); + ggml_free(ctx); + n_tokens_ = cond.n_tokens; + valid_tokens_ = cond.n_tokens; + ring_head_ = 0; + } + + bool eval( + const std::vector>& audio_codes, DecoderKvCache& cond_kv, + DecoderKvCache& uncond_kv, decoder_result& cond_result, decoder_result& uncond_result, + magpietts_backend_tensor* cond_hidden_out, + magpietts_backend_tensor* uncond_hidden_out, + const magpietts_decoder_attention* attention) { + const ggml_nvtx::range nvtx_range("magpietts_persistent_decoder_eval"); + const magpietts_hparams& h = model_.hparams; + if (!cond_hidden_out || !uncond_hidden_out || !cond_hidden_out->tensor || + !uncond_hidden_out->tensor || static_cast(audio_codes.size()) != h.audio_codebooks || + audio_codes.empty()) { + return false; + } + const size_t raw_len = audio_codes[0].size(); + if (raw_len == 0 || raw_len % h.frame_stacking_factor != 0) return false; + for (const auto& codes : audio_codes) { + if (codes.size() != raw_len) return false; + } + const int total_len = h.baked_context_length + + static_cast(raw_len / h.frame_stacking_factor); + if (total_len != n_tokens_ + 1 || n_tokens_ >= cache_len_) return false; + + std::vector tokens(static_cast(h.stacked_audio_codebooks())); + const size_t frame_start = raw_len - static_cast(h.frame_stacking_factor); + for (int lane = 0; lane < h.frame_stacking_factor; ++lane) { + for (int codebook = 0; codebook < h.audio_codebooks; ++codebook) { + tokens[static_cast(codebook + lane * h.audio_codebooks)] = + audio_codes[static_cast(codebook)][frame_start + lane]; + } + } + const int32_t position = n_tokens_; + // Column-major [CFG lane, {ring head, valid length}]. The fused attention kernel reads + // only the active suffix while the graph and arena shapes remain constant. + const int32_t cache_meta[kMagpieCfgLanes * 2] = { + ring_head_, ring_head_, valid_tokens_, valid_tokens_}; + std::vector log_prior(static_cast(text_len_), 0.0f); + if (attention && attention->prior) { + if (static_cast(attention->prior->size()) != text_len_) return false; + for (int i = 0; i < text_len_; ++i) { + log_prior[static_cast(i)] = + std::log(std::max((*attention->prior)[static_cast(i)], 1.0e-20f)); + } + } + + std::vector inputs = { + {"magpietts.decoder.runtime.tokens", GGML_TYPE_I32, tokens.data(), + {h.stacked_audio_codebooks()}}, + {"magpietts.decoder.runtime.position", GGML_TYPE_I32, &position, {1}}, + {"magpietts.decoder.runtime.cache_meta", GGML_TYPE_I32, cache_meta, + {kMagpieCfgLanes, 2}}, + {"magpietts.decoder.runtime.prior", GGML_TYPE_F32, log_prior.data(), {text_len_}}}; + + ggml_runtime::DeviceTensor cond_device; + ggml_runtime::DeviceTensor uncond_device; + std::vector alignment(static_cast(text_len_)); + const bool has_alignment = module_.alignment_count() > 0; + std::vector outputs(has_alignment ? 3 : 2); + outputs[0].index = 0; + outputs[0].device_tensor = &cond_device; + outputs[1].index = 1; + outputs[1].device_tensor = &uncond_device; + if (has_alignment) { + outputs[2].index = 2; + outputs[2].host_buffer = alignment.data(); + outputs[2].nbytes = alignment.size() * sizeof(float); + } + session_.run(inputs, outputs); + ggml_backend_tensor_copy_async( + model_.backend, model_.backend, cond_device.tensor, cond_hidden_out->tensor); + ggml_backend_tensor_copy_async( + model_.backend, model_.backend, uncond_device.tensor, uncond_hidden_out->tensor); + + if (attention && attention->alignment_scores) { + *attention->alignment_scores = alignment; + } + + ++n_tokens_; + valid_tokens_ = std::min(cache_len_, valid_tokens_ + 1); + ring_head_ = (ring_head_ + 1) % cache_len_; + cond_kv.n_tokens = n_tokens_; + uncond_kv.n_tokens = n_tokens_; + cond_result.hidden_last.clear(); + uncond_result.hidden_last.clear(); + return true; + } + + private: + const magpietts_model& model_; + const DecoderCrossKvCache* cross_kv_ = nullptr; + int text_len_ = 0; + int cache_len_ = 0; + int n_tokens_ = 0; + int valid_tokens_ = 0; + int ring_head_ = 0; + ggml_runtime::BackendManager backend_manager_; + PersistentDecoderModule module_; + ggml_runtime::Session session_; +}; + +MagpieDecoder::MagpieDecoder(const magpietts_model& model) : model_(model) {} + +MagpieDecoder::~MagpieDecoder() = default; + bool MagpieDecoder::eval( const std::vector& text_cond, int text_len, @@ -447,6 +830,48 @@ MagpieDecoder::evalCachedPair( const magpietts_backend_tensor* text_cond_device, magpietts_backend_tensor* cond_hidden_out, magpietts_backend_tensor* uncond_hidden_out, DecoderCrossKvCache* cond_cross_kv, const magpietts_decoder_attention* attention) const { + const bool persistent_candidate = + cuda_sample == nullptr && cond_hidden_out != nullptr && uncond_hidden_out != nullptr && + cond_cross_kv != nullptr && cond_cross_kv->validFor(model_, text_len) && + model_.hparams.dec_kernel == 1 && magpietts_backend_is_cuda(model_.backend) && + cond_kv.n_tokens > 0 && cond_kv.n_tokens == uncond_kv.n_tokens; + if (persistent_candidate) { + try { + if (persistent_runtime_ && !persistent_runtime_->matches(cond_cross_kv, text_len)) { + // Cross-cache address or shape changes require a new graph. + persistent_runtime_.reset(); + cond_kv.clear(); + uncond_kv.clear(); + } + if (cond_kv.n_tokens > 0 && !persistent_runtime_) { + persistent_runtime_ = + std::make_unique(model_, *cond_cross_kv, text_len); + persistent_runtime_->seed(cond_kv, uncond_kv); + fprintf( + stderr, + "MagpieTTS decoder runtime: fixed-shape CUDA graph, CFG batch=2, " + "device K/V arena enabled\n"); + } + if (persistent_runtime_ && + !persistent_runtime_->sequence_matches(cond_kv.n_tokens)) { + // Reuse the graph and reseed the cache suffix. + persistent_runtime_->seed(cond_kv, uncond_kv); + } + if (persistent_runtime_ && persistent_runtime_->eval( + audio_codes, cond_kv, uncond_kv, cond_result, + uncond_result, cond_hidden_out, uncond_hidden_out, + attention)) { + return true; + } + persistent_runtime_.reset(); + cond_kv.clear(); + uncond_kv.clear(); + } + catch (const std::exception& e) { + fprintf(stderr, "MagpieTTS persistent decoder failed: %s\n", e.what()); + return false; + } + } return decoder_eval_cached_pair_impl( model_, text_cond, text_len, audio_codes, speaker, threads, cond_kv, uncond_kv, cond_result, uncond_result, cuda_sample, text_cond_device, cond_hidden_out, uncond_hidden_out, @@ -617,10 +1042,14 @@ decoder_eval_impl( ggml_set_name(dec_out, "magpietts_decoder_out"); ggml_set_output(dec_out); - ggml_tensor* logits = linear(ctx, model.final_proj_w, dec_out, model.final_proj_b); - logits = ggml_cont(ctx, ggml_cast(ctx, logits, GGML_TYPE_F32)); - ggml_set_name(logits, "magpietts_decoder_logits"); - ggml_set_output(logits); + const bool compute_logits = cuda_sample || result.logits_required; + ggml_tensor* logits = nullptr; + if (compute_logits) { + logits = linear(ctx, model.final_proj_w, dec_out, model.final_proj_b); + logits = ggml_cont(ctx, ggml_cast(ctx, logits, GGML_TYPE_F32)); + ggml_set_name(logits, "magpietts_decoder_logits"); + ggml_set_output(logits); + } ggml_gallocr_t allocr = nullptr; const size_t hidden_off = (size_t)h.n_embd * (total_len - 1) * sizeof(float); @@ -630,7 +1059,9 @@ decoder_eval_impl( ggml_set_name(hidden_last, "magpietts_decoder_hidden_last"); ggml_set_output(hidden_last); } - ggml_build_forward_expand(gf, logits); + if (logits) { + ggml_build_forward_expand(gf, logits); + } ggml_build_forward_expand(gf, dec_out); if (hidden_last) { ggml_build_forward_expand(gf, hidden_last); @@ -671,12 +1102,16 @@ decoder_eval_impl( return true; } - result.logits_last.resize(logits_last_size); + if (result.logits_required) { + result.logits_last.resize(logits_last_size); + } result.hidden_last.resize(h.n_embd); const size_t logits_off = logits_off_floats * sizeof(float); - magpietts_backend_tensor_get_staged( - model, output_staging, logits, result.logits_last.data(), logits_off, - result.logits_last.size() * sizeof(float)); + if (result.logits_required) { + magpietts_backend_tensor_get_staged( + model, output_staging, logits, result.logits_last.data(), logits_off, + result.logits_last.size() * sizeof(float)); + } magpietts_backend_tensor_get_staged( model, output_staging, dec_out, result.hidden_last.data(), hidden_off, result.hidden_last.size() * sizeof(float)); @@ -761,10 +1196,15 @@ decoder_eval_pair_impl( ggml_set_name(dec_out_cond, "magpietts_decoder_out_cond"); ggml_set_output(dec_out_cond); - ggml_tensor* logits_cond = linear(ctx, model.final_proj_w, dec_out_cond, model.final_proj_b); - logits_cond = ggml_cont(ctx, ggml_cast(ctx, logits_cond, GGML_TYPE_F32)); - ggml_set_name(logits_cond, "magpietts_decoder_logits_cond"); - ggml_set_output(logits_cond); + const bool compute_logits = + cuda_sample || cond_result.logits_required || uncond_result.logits_required; + ggml_tensor* logits_cond = nullptr; + if (compute_logits) { + logits_cond = linear(ctx, model.final_proj_w, dec_out_cond, model.final_proj_b); + logits_cond = ggml_cont(ctx, ggml_cast(ctx, logits_cond, GGML_TYPE_F32)); + ggml_set_name(logits_cond, "magpietts_decoder_logits_cond"); + ggml_set_output(logits_cond); + } ggml_tensor* dec_out_uncond = transformer_forward(ctx, model.decoder, dec_in_uncond, pos, nullptr); @@ -772,15 +1212,21 @@ decoder_eval_pair_impl( ggml_set_name(dec_out_uncond, "magpietts_decoder_out_uncond"); ggml_set_output(dec_out_uncond); - ggml_tensor* logits_uncond = - linear(ctx, model.final_proj_w, dec_out_uncond, model.final_proj_b); - logits_uncond = ggml_cont(ctx, ggml_cast(ctx, logits_uncond, GGML_TYPE_F32)); - ggml_set_name(logits_uncond, "magpietts_decoder_logits_uncond"); - ggml_set_output(logits_uncond); + ggml_tensor* logits_uncond = nullptr; + if (compute_logits) { + logits_uncond = linear(ctx, model.final_proj_w, dec_out_uncond, model.final_proj_b); + logits_uncond = ggml_cont(ctx, ggml_cast(ctx, logits_uncond, GGML_TYPE_F32)); + ggml_set_name(logits_uncond, "magpietts_decoder_logits_uncond"); + ggml_set_output(logits_uncond); + } - ggml_build_forward_expand(gf, logits_cond); + if (logits_cond) { + ggml_build_forward_expand(gf, logits_cond); + } ggml_build_forward_expand(gf, dec_out_cond); - ggml_build_forward_expand(gf, logits_uncond); + if (logits_uncond) { + ggml_build_forward_expand(gf, logits_uncond); + } ggml_build_forward_expand(gf, dec_out_uncond); for (ggml_tensor* t : alignment_outputs) { ggml_build_forward_expand(gf, t); @@ -838,20 +1284,28 @@ decoder_eval_pair_impl( return true; } - cond_result.logits_last.resize(logits_last_size); + if (cond_result.logits_required) { + cond_result.logits_last.resize(logits_last_size); + } cond_result.hidden_last.resize(h.n_embd); - uncond_result.logits_last.resize(logits_last_size); + if (uncond_result.logits_required) { + uncond_result.logits_last.resize(logits_last_size); + } uncond_result.hidden_last.resize(h.n_embd); const size_t logits_off = logits_off_floats * sizeof(float); - magpietts_backend_tensor_get_staged( - model, output_staging, logits_cond, cond_result.logits_last.data(), logits_off, - cond_result.logits_last.size() * sizeof(float)); + if (cond_result.logits_required) { + magpietts_backend_tensor_get_staged( + model, output_staging, logits_cond, cond_result.logits_last.data(), logits_off, + cond_result.logits_last.size() * sizeof(float)); + } magpietts_backend_tensor_get_staged( model, output_staging, dec_out_cond, cond_result.hidden_last.data(), hidden_off, cond_result.hidden_last.size() * sizeof(float)); - magpietts_backend_tensor_get_staged( - model, output_staging, logits_uncond, uncond_result.logits_last.data(), logits_off, - uncond_result.logits_last.size() * sizeof(float)); + if (uncond_result.logits_required) { + magpietts_backend_tensor_get_staged( + model, output_staging, logits_uncond, uncond_result.logits_last.data(), logits_off, + uncond_result.logits_last.size() * sizeof(float)); + } magpietts_backend_tensor_get_staged( model, output_staging, dec_out_uncond, uncond_result.hidden_last.data(), hidden_off, uncond_result.hidden_last.size() * sizeof(float)); @@ -981,11 +1435,15 @@ decoder_eval_cached_impl( ggml_set_name(dec_out, "magpietts_decoder_out_cached"); ggml_set_output(dec_out); - ggml_tensor* logits = linear(ctx, model.final_proj_w, dec_out, model.final_proj_b); - logits = ggml_cont(ctx, ggml_cast(ctx, logits, GGML_TYPE_F32)); - ggml_set_name(logits, "magpietts_decoder_logits_cached"); - ggml_set_output(logits); - ggml_build_forward_expand(gf, logits); + const bool compute_logits = cuda_sample || result.logits_required; + ggml_tensor* logits = nullptr; + if (compute_logits) { + logits = linear(ctx, model.final_proj_w, dec_out, model.final_proj_b); + logits = ggml_cont(ctx, ggml_cast(ctx, logits, GGML_TYPE_F32)); + ggml_set_name(logits, "magpietts_decoder_logits_cached"); + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + } ggml_build_forward_expand(gf, dec_out); for (ggml_tensor* t : alignment_outputs) { ggml_build_forward_expand(gf, t); @@ -1037,12 +1495,16 @@ decoder_eval_cached_impl( return true; } - result.logits_last.resize(logits_last_size); + if (result.logits_required) { + result.logits_last.resize(logits_last_size); + } result.hidden_last.resize(h.n_embd); const size_t logits_off = logits_off_floats * sizeof(float); - magpietts_backend_tensor_get_staged( - model, output_staging, logits, result.logits_last.data(), logits_off, - result.logits_last.size() * sizeof(float)); + if (result.logits_required) { + magpietts_backend_tensor_get_staged( + model, output_staging, logits, result.logits_last.data(), logits_off, + result.logits_last.size() * sizeof(float)); + } magpietts_backend_tensor_get_staged( model, output_staging, dec_out, result.hidden_last.data(), hidden_off, result.hidden_last.size() * sizeof(float)); @@ -1196,10 +1658,15 @@ decoder_eval_cached_pair_impl( ggml_set_name(dec_out_cond, "magpietts_decoder_out_cond_cached"); ggml_set_output(dec_out_cond); - ggml_tensor* logits_cond = linear(ctx, model.final_proj_w, dec_out_cond, model.final_proj_b); - logits_cond = ggml_cont(ctx, ggml_cast(ctx, logits_cond, GGML_TYPE_F32)); - ggml_set_name(logits_cond, "magpietts_decoder_logits_cond_cached"); - ggml_set_output(logits_cond); + const bool compute_logits = + cuda_sample || cond_result.logits_required || uncond_result.logits_required; + ggml_tensor* logits_cond = nullptr; + if (compute_logits) { + logits_cond = linear(ctx, model.final_proj_w, dec_out_cond, model.final_proj_b); + logits_cond = ggml_cont(ctx, ggml_cast(ctx, logits_cond, GGML_TYPE_F32)); + ggml_set_name(logits_cond, "magpietts_decoder_logits_cond_cached"); + ggml_set_output(logits_cond); + } ggml_tensor* dec_out_uncond = transformer_forward_cached( ctx, gf, model.decoder, dec_in_uncond, pos, nullptr, uncond_kv, nullptr, n_past); @@ -1207,15 +1674,21 @@ decoder_eval_cached_pair_impl( ggml_set_name(dec_out_uncond, "magpietts_decoder_out_uncond_cached"); ggml_set_output(dec_out_uncond); - ggml_tensor* logits_uncond = - linear(ctx, model.final_proj_w, dec_out_uncond, model.final_proj_b); - logits_uncond = ggml_cont(ctx, ggml_cast(ctx, logits_uncond, GGML_TYPE_F32)); - ggml_set_name(logits_uncond, "magpietts_decoder_logits_uncond_cached"); - ggml_set_output(logits_uncond); + ggml_tensor* logits_uncond = nullptr; + if (compute_logits) { + logits_uncond = linear(ctx, model.final_proj_w, dec_out_uncond, model.final_proj_b); + logits_uncond = ggml_cont(ctx, ggml_cast(ctx, logits_uncond, GGML_TYPE_F32)); + ggml_set_name(logits_uncond, "magpietts_decoder_logits_uncond_cached"); + ggml_set_output(logits_uncond); + } - ggml_build_forward_expand(gf, logits_cond); + if (logits_cond) { + ggml_build_forward_expand(gf, logits_cond); + } ggml_build_forward_expand(gf, dec_out_cond); - ggml_build_forward_expand(gf, logits_uncond); + if (logits_uncond) { + ggml_build_forward_expand(gf, logits_uncond); + } ggml_build_forward_expand(gf, dec_out_uncond); for (ggml_tensor* t : alignment_outputs) { ggml_build_forward_expand(gf, t); @@ -1279,20 +1752,28 @@ decoder_eval_cached_pair_impl( return true; } - cond_result.logits_last.resize(logits_last_size); + if (cond_result.logits_required) { + cond_result.logits_last.resize(logits_last_size); + } cond_result.hidden_last.resize(h.n_embd); - uncond_result.logits_last.resize(logits_last_size); + if (uncond_result.logits_required) { + uncond_result.logits_last.resize(logits_last_size); + } uncond_result.hidden_last.resize(h.n_embd); const size_t logits_off = logits_off_floats * sizeof(float); - magpietts_backend_tensor_get_staged( - model, output_staging, logits_cond, cond_result.logits_last.data(), logits_off, - cond_result.logits_last.size() * sizeof(float)); + if (cond_result.logits_required) { + magpietts_backend_tensor_get_staged( + model, output_staging, logits_cond, cond_result.logits_last.data(), logits_off, + cond_result.logits_last.size() * sizeof(float)); + } magpietts_backend_tensor_get_staged( model, output_staging, dec_out_cond, cond_result.hidden_last.data(), hidden_off, cond_result.hidden_last.size() * sizeof(float)); - magpietts_backend_tensor_get_staged( - model, output_staging, logits_uncond, uncond_result.logits_last.data(), logits_off, - uncond_result.logits_last.size() * sizeof(float)); + if (uncond_result.logits_required) { + magpietts_backend_tensor_get_staged( + model, output_staging, logits_uncond, uncond_result.logits_last.data(), logits_off, + uncond_result.logits_last.size() * sizeof(float)); + } magpietts_backend_tensor_get_staged( model, output_staging, dec_out_uncond, uncond_result.hidden_last.data(), hidden_off, uncond_result.hidden_last.size() * sizeof(float)); diff --git a/src/tts/magpietts/decoder.h b/src/tts/magpietts/decoder.h index 996bc2e..a298053 100644 --- a/src/tts/magpietts/decoder.h +++ b/src/tts/magpietts/decoder.h @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 #pragma once +#include + #include "model.h" namespace nemo_speech::tts { @@ -61,6 +63,7 @@ class DecoderCrossKvCache { }; struct decoder_result { + bool logits_required = true; std::vector logits_last; std::vector hidden_last; std::vector cross_attn_last; @@ -85,7 +88,11 @@ struct magpietts_cuda_sample_request { class MagpieDecoder { public: - explicit MagpieDecoder(const magpietts_model& model) : model_(model) {} + explicit MagpieDecoder(const magpietts_model& model); + ~MagpieDecoder(); + + MagpieDecoder(const MagpieDecoder&) = delete; + MagpieDecoder& operator=(const MagpieDecoder&) = delete; bool eval( const std::vector& text_cond, int text_len, @@ -123,8 +130,11 @@ class MagpieDecoder { const magpietts_decoder_attention* attention = nullptr) const; private: + class PersistentDecoderRuntime; + const magpietts_model& model_; mutable MagpiePinnedHostScratch output_staging_; + mutable std::unique_ptr persistent_runtime_; }; class MagpieCodebookSampler { diff --git a/src/tts/magpietts/graph.h b/src/tts/magpietts/graph.h index 097a9f5..ebf8da7 100644 --- a/src/tts/magpietts/graph.h +++ b/src/tts/magpietts/graph.h @@ -23,6 +23,11 @@ ggml_tensor* cross_attention( ggml_context* ctx, const magpietts_transformer& tr, const magpietts_layer& layer, ggml_tensor* x, ggml_tensor* memory, ggml_tensor* attn_prior = nullptr, ggml_tensor** last_attn = nullptr); +ggml_tensor* cross_attention_cached( + ggml_context* ctx, const magpietts_transformer& tr, const magpietts_layer& layer, + const DecoderCrossKvCache& cross_kv, int layer_index, ggml_tensor* x, + ggml_tensor* attn_prior = nullptr, ggml_tensor** last_attn = nullptr, + bool prior_is_log = false); ggml_tensor* transformer_forward( ggml_context* ctx, const magpietts_transformer& tr, ggml_tensor* x, ggml_tensor* pos, ggml_tensor* cond, ggml_tensor* attn_prior = nullptr, diff --git a/src/tts/magpietts/lt.cpp b/src/tts/magpietts/lt.cpp index 212868e..f700347 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -13,6 +13,10 @@ #include "graph.h" #include "nvtx_utils.h" +#if defined(MAGPIETTS_CUDA_SAMPLING) +#include "ggml-cuda.h" +#endif + namespace nemo_speech::tts { class LocalTransformerGraph { @@ -454,6 +458,107 @@ local_transformer_forward_cached_fixed_pos( return tr.norm_out ? layer_norm(ctx, x, tr.norm_out) : x; } +static ggml_tensor* +local_self_attention_cached_pair( + ggml_context* ctx, ggml_cgraph* gf, const magpietts_transformer& tr, + const magpietts_layer& layer, DecoderKvCache& cond_cache, DecoderKvCache& uncond_cache, + int layer_index, int n_past, ggml_tensor* x) { + constexpr int kCfgLanes = 2; + const int64_t n_embd = tr.n_embd; + const int64_t n_head = tr.n_head; + const int64_t d_head = n_embd / n_head; + const int64_t n_tok = x->ne[1]; + const int64_t n_total = n_past + n_tok; + + ggml_tensor* qkv = linear(ctx, layer.self_qkv, x); + const size_t element = ggml_element_size(qkv); + auto qkv_slice = [&](size_t offset) { + return ggml_view_3d( + ctx, qkv, n_embd, n_tok, kCfgLanes, qkv->nb[1], qkv->nb[2], offset); + }; + ggml_tensor* qcur = qkv_slice(0); + ggml_tensor* kcur = qkv_slice(static_cast(n_embd) * element); + ggml_tensor* vcur = qkv_slice(static_cast(2 * n_embd) * element); + + const size_t cache_element = ggml_element_size(cond_cache.memory_k); + const size_t layer_offset = + static_cast(layer_index) * cond_cache.n_ctx * n_embd * cache_element; + const size_t write_offset = + layer_offset + static_cast(n_past) * n_embd * cache_element; + DecoderKvCache* caches[kCfgLanes] = {&cond_cache, &uncond_cache}; + for (int lane = 0; lane < kCfgLanes; ++lane) { + ggml_tensor* k_lane = ggml_view_2d( + ctx, kcur, n_embd, n_tok, kcur->nb[1], static_cast(lane) * kcur->nb[2]); + ggml_tensor* v_lane = ggml_view_2d( + ctx, vcur, n_embd, n_tok, vcur->nb[1], static_cast(lane) * vcur->nb[2]); + ggml_tensor* k_dst = + ggml_view_1d(ctx, caches[lane]->memory_k, n_tok * n_embd, write_offset); + ggml_tensor* v_dst = + ggml_view_1d(ctx, caches[lane]->memory_v, n_tok * n_embd, write_offset); + ggml_tensor* k_copy = ggml_cpy(ctx, k_lane, k_dst); + ggml_set_name(k_copy, "magpietts_local_pair_kv_copy_k"); + ggml_build_forward_expand(gf, k_copy); + ggml_tensor* v_copy = ggml_cpy(ctx, v_lane, v_dst); + ggml_set_name(v_copy, "magpietts_local_pair_kv_copy_v"); + ggml_build_forward_expand(gf, v_copy); + } + + ggml_tensor* q = ggml_permute( + ctx, ggml_cont_4d(ctx, qcur, d_head, n_head, n_tok, kCfgLanes), 0, 2, 1, 3); + ggml_tensor* k_cond = ggml_view_2d( + ctx, cond_cache.memory_k, n_embd, n_total, n_embd * cache_element, layer_offset); + ggml_tensor* k_uncond = ggml_view_2d( + ctx, uncond_cache.memory_k, n_embd, n_total, n_embd * cache_element, layer_offset); + ggml_tensor* k_pair = ggml_concat(ctx, k_cond, k_uncond, 2); + ggml_tensor* k = ggml_permute( + ctx, ggml_reshape_4d(ctx, k_pair, d_head, n_head, n_total, kCfgLanes), 0, 2, 1, 3); + ggml_tensor* scores = ggml_scale( + ctx, ggml_mul_mat(ctx, k, q), 1.0f / std::sqrt(static_cast(d_head))); + // Each local graph evaluates exactly one new position and exposes only the populated cache + // prefix, so all keys are causal-valid and no diagonal mask is needed. + ggml_tensor* probs = ggml_soft_max(ctx, scores); + + ggml_tensor* v_cond = ggml_view_2d( + ctx, cond_cache.memory_v, n_embd, n_total, n_embd * cache_element, layer_offset); + ggml_tensor* v_uncond = ggml_view_2d( + ctx, uncond_cache.memory_v, n_embd, n_total, n_embd * cache_element, layer_offset); + ggml_tensor* v_pair = ggml_concat(ctx, v_cond, v_uncond, 2); + ggml_tensor* v_trans = ggml_cont_4d( + ctx, + ggml_permute( + ctx, ggml_reshape_4d(ctx, v_pair, d_head, n_head, n_total, kCfgLanes), 1, 2, 0, + 3), + n_total, d_head, n_head, kCfgLanes); + ggml_tensor* weighted = ggml_mul_mat(ctx, v_trans, probs); + ggml_tensor* merged = ggml_permute(ctx, weighted, 0, 2, 1, 3); + ggml_tensor* out = ggml_cont_3d(ctx, merged, n_embd, n_tok, kCfgLanes); + return linear(ctx, layer.self_o, out); +} + +static ggml_tensor* +local_transformer_forward_cached_pair_fixed_pos( + ggml_context* ctx, ggml_cgraph* gf, const magpietts_transformer& tr, ggml_tensor* x, + ggml_tensor* pos_emb, DecoderKvCache& cond_cache, DecoderKvCache& uncond_cache, int n_past) { + pos_emb = ggml_cont(ctx, ggml_cast(ctx, pos_emb, GGML_TYPE_F32)); + x = ggml_add(ctx, x, pos_emb); + for (int il = 0; il < static_cast(tr.layers.size()); ++il) { + const magpietts_layer& layer = tr.layers[static_cast(il)]; + ggml_tensor* residual = x; + ggml_tensor* cur = layer_norm(ctx, x, layer.norm_self); + cur = local_self_attention_cached_pair( + ctx, gf, tr, layer, cond_cache, uncond_cache, il, n_past, cur); + x = ggml_add(ctx, residual, cur); + + residual = x; + cur = layer_norm(ctx, x, layer.norm_ff); + cur = causal_conv1d(ctx, cur, layer.ff_proj); + cur = ggml_gelu(ctx, cur); + cur = causal_conv1d(ctx, cur, layer.ff_out); + x = ggml_add(ctx, residual, cur); + } + return tr.norm_out ? layer_norm(ctx, x, tr.norm_out) : x; +} + static bool local_transformer_graph_init( const magpietts_model& model, bool pair, int codebook_idx, local_transformer_graph_bank& bank, @@ -503,24 +608,20 @@ local_transformer_graph_init( pair ? "magpietts_build_local_transformer_pair_graph" : "magpietts_build_local_transformer_graph"); - ggml_tensor* cur_cond = nullptr; - ggml_tensor* cur_uncond = nullptr; + ggml_tensor* input_cond = nullptr; + ggml_tensor* input_uncond = nullptr; if (codebook_idx == 0) { graph.dec_cond = ggml_new_tensor_2d(graph.ctx, GGML_TYPE_F32, h.n_embd, 1); ggml_set_name( graph.dec_cond, pair ? "magpietts_local_transformer_dec_cond" : "magpietts_local_transformer_dec_last"); ggml_set_input(graph.dec_cond); - cur_cond = model.lt_in_w - ? linear(graph.ctx, model.lt_in_w, graph.dec_cond, model.lt_in_b) - : graph.dec_cond; + input_cond = graph.dec_cond; if (pair) { graph.dec_uncond = ggml_new_tensor_2d(graph.ctx, GGML_TYPE_F32, h.n_embd, 1); ggml_set_name(graph.dec_uncond, "magpietts_local_transformer_dec_uncond"); ggml_set_input(graph.dec_uncond); - cur_uncond = model.lt_in_w - ? linear(graph.ctx, model.lt_in_w, graph.dec_uncond, model.lt_in_b) - : graph.dec_uncond; + input_uncond = graph.dec_uncond; } } else { const std::string name = "magpietts_local_transformer_prev_code"; @@ -529,10 +630,8 @@ local_transformer_graph_init( ggml_set_input(graph.prev_token); ggml_tensor* emb = ggml_get_rows( graph.ctx, model.audio_embeddings[codebook_idx - 1], graph.prev_token); - cur_cond = model.lt_in_w ? linear(graph.ctx, model.lt_in_w, emb, model.lt_in_b) : emb; - if (pair) { - cur_uncond = cur_cond; - } + input_cond = emb; + if (pair) input_uncond = emb; } graph.pos_emb = ggml_view_2d( @@ -540,30 +639,45 @@ local_transformer_graph_init( (size_t)codebook_idx * model.local.pos_emb->nb[1]); ggml_set_name(graph.pos_emb, "magpietts_local_transformer_pos_emb"); - DecoderKvCache& cond_cache = pair ? bank.cond_cache : bank.single_cache; - ggml_tensor* out_cond = local_transformer_forward_cached_fixed_pos( - graph.ctx, graph.gf, model.local, cur_cond, graph.pos_emb, cond_cache, codebook_idx); - graph.logits_cond = - linear(graph.ctx, model.lt_out_w[codebook_idx], out_cond, model.lt_out_b[codebook_idx]); - graph.logits_cond = - ggml_cont(graph.ctx, ggml_cast(graph.ctx, graph.logits_cond, GGML_TYPE_F32)); - ggml_set_name( - graph.logits_cond, pair ? "magpietts_local_transformer_logits_cond" - : "magpietts_local_transformer_logits"); - ggml_set_output(graph.logits_cond); - ggml_build_forward_expand(graph.gf, graph.logits_cond); - if (pair) { - ggml_tensor* out_uncond = local_transformer_forward_cached_fixed_pos( - graph.ctx, graph.gf, model.local, cur_uncond, graph.pos_emb, bank.uncond_cache, - codebook_idx); - graph.logits_uncond = linear( - graph.ctx, model.lt_out_w[codebook_idx], out_uncond, model.lt_out_b[codebook_idx]); - graph.logits_uncond = - ggml_cont(graph.ctx, ggml_cast(graph.ctx, graph.logits_uncond, GGML_TYPE_F32)); + ggml_tensor* pair_input = ggml_concat(graph.ctx, input_cond, input_uncond, 2); + ggml_tensor* cur_pair = model.lt_in_w + ? linear( + graph.ctx, model.lt_in_w, pair_input, model.lt_in_b) + : pair_input; + ggml_tensor* out_pair = local_transformer_forward_cached_pair_fixed_pos( + graph.ctx, graph.gf, model.local, cur_pair, graph.pos_emb, bank.cond_cache, + bank.uncond_cache, codebook_idx); + ggml_tensor* logits_pair = linear( + graph.ctx, model.lt_out_w[codebook_idx], out_pair, model.lt_out_b[codebook_idx]); + logits_pair = + ggml_cont(graph.ctx, ggml_cast(graph.ctx, logits_pair, GGML_TYPE_F32)); + graph.logits_cond = ggml_view_2d( + graph.ctx, logits_pair, h.audio_vocab_size, 1, logits_pair->nb[1], 0); + graph.logits_uncond = ggml_view_2d( + graph.ctx, logits_pair, h.audio_vocab_size, 1, logits_pair->nb[1], + logits_pair->nb[2]); + ggml_set_name(graph.logits_cond, "magpietts_local_transformer_logits_cond"); ggml_set_name(graph.logits_uncond, "magpietts_local_transformer_logits_uncond"); + ggml_set_output(graph.logits_cond); ggml_set_output(graph.logits_uncond); + ggml_build_forward_expand(graph.gf, graph.logits_cond); ggml_build_forward_expand(graph.gf, graph.logits_uncond); + } else { + ggml_tensor* cur_cond = model.lt_in_w + ? linear( + graph.ctx, model.lt_in_w, input_cond, model.lt_in_b) + : input_cond; + ggml_tensor* out_cond = local_transformer_forward_cached_fixed_pos( + graph.ctx, graph.gf, model.local, cur_cond, graph.pos_emb, bank.single_cache, + codebook_idx); + graph.logits_cond = linear( + graph.ctx, model.lt_out_w[codebook_idx], out_cond, model.lt_out_b[codebook_idx]); + graph.logits_cond = + ggml_cont(graph.ctx, ggml_cast(graph.ctx, graph.logits_cond, GGML_TYPE_F32)); + ggml_set_name(graph.logits_cond, "magpietts_local_transformer_logits"); + ggml_set_output(graph.logits_cond); + ggml_build_forward_expand(graph.gf, graph.logits_cond); } tag_graph_first_node(graph.gf); } @@ -710,13 +824,41 @@ local_transformer_graph_eval_cuda( fprintf(stderr, "CUDA local-transformer eval requires a CUDA sampler\n"); return false; } + const bool building_sequence = + magpietts_cuda_sampler_sequence_build_active(cuda_sample.sampler); { const ggml_nvtx::range nvtx_inputs("magpietts_local_transformer_graph_set_device_inputs"); if (graph.codebook_idx == 0) { - ggml_backend_tensor_copy(cond_hidden, graph.dec_cond); + if (building_sequence) { + char error[256] = {}; + if (!magpietts_cuda_sampler_sequence_add_device_copy( + cuda_sample.sampler, cond_hidden->data, graph.dec_cond->data, + ggml_nbytes(graph.dec_cond), error, sizeof(error))) { + fprintf( + stderr, "CUDA local-transformer conditional input node failed: %s\n", + error[0] ? error : "unknown error"); + return false; + } + } else { + ggml_backend_tensor_copy_async( + model.backend, model.backend, cond_hidden, graph.dec_cond); + } if (graph.pair) { - ggml_backend_tensor_copy(uncond_hidden, graph.dec_uncond); + if (building_sequence) { + char error[256] = {}; + if (!magpietts_cuda_sampler_sequence_add_device_copy( + cuda_sample.sampler, uncond_hidden->data, graph.dec_uncond->data, + ggml_nbytes(graph.dec_uncond), error, sizeof(error))) { + fprintf( + stderr, "CUDA local-transformer unconditional input node failed: %s\n", + error[0] ? error : "unknown error"); + return false; + } + } else { + ggml_backend_tensor_copy_async( + model.backend, model.backend, uncond_hidden, graph.dec_uncond); + } } } else { if (!graph.prev_token->data) { @@ -743,7 +885,21 @@ local_transformer_graph_eval_cuda( ggml_status status = GGML_STATUS_FAILED; { const ggml_nvtx::range nvtx_compute("magpietts_local_transformer_graph_compute_cuda"); - status = ggml_backend_graph_compute(model.backend, graph.gf); + if (building_sequence) { + void* graph_template = + ggml_backend_cuda_get_graph_template(model.backend, graph.gf); + char error[256] = {}; + if (!graph_template || !magpietts_cuda_sampler_sequence_add_ggml_graph( + cuda_sample.sampler, graph_template, error, sizeof(error))) { + fprintf( + stderr, "CUDA local-transformer child graph is unavailable: %s\n", + error[0] ? error : "GGML graph has not completed warm-up"); + return false; + } + status = GGML_STATUS_SUCCESS; + } else { + status = ggml_backend_graph_compute_async(model.backend, graph.gf); + } } if (status != GGML_STATUS_SUCCESS) { fprintf( @@ -760,18 +916,15 @@ local_transformer_graph_eval_cuda( return false; } - ggml_backend_synchronize(model.backend); const size_t off = 0; const float* logits_cond = (const float*)graph.logits_cond->data + off; const float* logits_uncond = graph.pair ? (const float*)graph.logits_uncond->data + off : nullptr; char error[256] = {}; - const bool ok = magpietts_cuda_sample_codebooks_device( + const bool ok = magpietts_cuda_sample_codebooks_device_configured( cuda_sample.sampler, logits_cond, logits_uncond, 1, model.hparams.audio_vocab_size, - model.hparams.audio_codebook_size, model.hparams.audio_eos_id, cuda_sample.use_cfg, - cuda_sample.cfg_scale, cuda_sample.temperature, cuda_sample.top_k, - cuda_sample.forbid_audio_eos, cuda_sample.seed, cuda_sample.frame_index, codebook_idx, - codebook_idx, error, sizeof(error)); + model.hparams.audio_codebook_size, model.hparams.audio_eos_id, codebook_idx, codebook_idx, + error, sizeof(error)); if (!ok) { fprintf( stderr, "CUDA local-transformer sampling failed: %s\n", @@ -893,31 +1046,102 @@ sample_local_codebooks_cuda_impl( if (!local_graphs.beginFrame(model, use_cfg)) { return false; } + char stream_error[256] = {}; + if (!magpietts_cuda_sampler_bind_stream( + cuda_sampler, ggml_backend_cuda_get_stream(model.backend), stream_error, + sizeof(stream_error))) { + fprintf( + stderr, "CUDA local-transformer stream binding failed: %s\n", + stream_error[0] ? stream_error : "unknown error"); + return false; + } codes.clear(); argmax_codes.clear(); - for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { - magpietts_cuda_sample_request cuda_sample; -#if defined(MAGPIETTS_CUDA_SAMPLING) - cuda_sample.sampler = cuda_sampler; -#else - (void)cuda_sampler; -#endif - cuda_sample.use_cfg = use_cfg; - cuda_sample.cfg_scale = cfg_scale; - cuda_sample.temperature = temperature; - cuda_sample.top_k = top_k; - cuda_sample.forbid_audio_eos = forbid_audio_eos; - cuda_sample.seed = seed; - cuda_sample.frame_index = frame_index; - const bool ok = local_transformer_graph_bank_eval_cuda( - model, local_graphs, use_cfg, cond_hidden, uncond_hidden, c, c, threads, cuda_sample); - if (!ok) { - return false; + char error[256] = {}; + if (!magpietts_cuda_sampler_configure( + cuda_sampler, use_cfg, cfg_scale, temperature, top_k, forbid_audio_eos, seed, + frame_index, error, sizeof(error))) { + fprintf( + stderr, "CUDA local-transformer sampler config failed: %s\n", + error[0] ? error : "unknown error"); + return false; + } + + magpietts_cuda_sample_request cuda_sample; + cuda_sample.sampler = cuda_sampler; + cuda_sample.use_cfg = use_cfg; + cuda_sample.cfg_scale = cfg_scale; + cuda_sample.temperature = temperature; + cuda_sample.top_k = top_k; + cuda_sample.forbid_audio_eos = forbid_audio_eos; + cuda_sample.seed = seed; + cuda_sample.frame_index = frame_index; + auto run_chain = [&]() { + for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { + if (!local_transformer_graph_bank_eval_cuda( + model, local_graphs, use_cfg, cond_hidden, uncond_hidden, c, c, threads, + cuda_sample)) { + return false; + } } + return true; + }; + + bool chain_ok = false; + if (magpietts_cuda_sampler_sequence_is_ready(cuda_sampler)) { + chain_ok = magpietts_cuda_sampler_sequence_launch( + cuda_sampler, error, sizeof(error)); + } else if ( + magpietts_cuda_sampler_sequence_is_warm(cuda_sampler) && + !magpietts_cuda_sampler_sequence_is_disabled(cuda_sampler)) { + // Compose the codebook graphs, sampling, and token feedback. + bool build_ok = magpietts_cuda_sampler_sequence_begin_build( + cuda_sampler, error, sizeof(error)); + if (build_ok) { + build_ok = magpietts_cuda_sampler_upload_config( + cuda_sampler, error, sizeof(error)) && + run_chain(); + } + if (build_ok) { + build_ok = magpietts_cuda_sampler_sequence_finish_build_and_launch( + cuda_sampler, error, sizeof(error)); + } else { + magpietts_cuda_sampler_sequence_abort_build(cuda_sampler); + } + if (build_ok) { + fprintf( + stderr, + "MagpieTTS local codebook chain: composed one reusable CUDA graph (%d " + "codebooks)\n", + h.stacked_audio_codebooks()); + chain_ok = true; + } else { + fprintf( + stderr, "warning: CUDA local graph composition unavailable; using async chain: %s\n", + error[0] ? error : "unknown error"); + magpietts_cuda_sampler_sequence_disable(cuda_sampler); + chain_ok = magpietts_cuda_sampler_upload_config( + cuda_sampler, error, sizeof(error)) && + run_chain(); + } + } else { + chain_ok = magpietts_cuda_sampler_upload_config( + cuda_sampler, error, sizeof(error)) && + run_chain(); + if (chain_ok && !magpietts_cuda_sampler_sequence_is_disabled(cuda_sampler)) { + // Initialize the per-codebook graphs before composing them. + magpietts_cuda_sampler_sequence_mark_warm(cuda_sampler); + } + } + if (!chain_ok) { + fprintf( + stderr, "CUDA local-transformer chain failed: %s\n", + error[0] ? error : "unknown error"); + return false; } codes.assign((size_t)h.stacked_audio_codebooks(), 0); argmax_codes.assign((size_t)h.stacked_audio_codebooks(), 0); - char error[256] = {}; + error[0] = '\0'; if (!magpietts_cuda_copy_sampled_codebooks( cuda_sampler, h.stacked_audio_codebooks(), codes.data(), argmax_codes.data(), error, sizeof(error))) { diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index 33073f3..483f491 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1318,6 +1318,8 @@ stream_magpie_to_audio( decoder_result cond; decoder_result uncond; + cond.logits_required = !params.use_local_transformer; + uncond.logits_required = !params.use_local_transformer; const bool forbid_eos = step * h.frame_stacking_factor < h.min_generated_frames; std::vector next_codes; std::vector argmax_codes; diff --git a/src/tts/magpietts/magpietts_cuda_sampling.cu b/src/tts/magpietts/magpietts_cuda_sampling.cu index ad0ecf3..3172eb0 100644 --- a/src/tts/magpietts/magpietts_cuda_sampling.cu +++ b/src/tts/magpietts/magpietts_cuda_sampling.cu @@ -2,6 +2,8 @@ // SPDX-License-Identifier: Apache-2.0 #include +#include + #include #include #include @@ -10,14 +12,39 @@ static constexpr int MAGPIETTS_CUDA_MAX_VOCAB = 4096; static constexpr int MAGPIETTS_CUDA_BLOCK_SIZE = 256; +static constexpr int MAGPIETTS_CUDA_SMALL_VOCAB = 2048; +static constexpr int MAGPIETTS_CUDA_SMALL_ITEMS_PER_THREAD = + MAGPIETTS_CUDA_SMALL_VOCAB / MAGPIETTS_CUDA_BLOCK_SIZE; +static constexpr int MAGPIETTS_CUDA_MAX_ITEMS_PER_THREAD = + MAGPIETTS_CUDA_MAX_VOCAB / MAGPIETTS_CUDA_BLOCK_SIZE; + +struct alignas(16) magpietts_cuda_sampling_config { + float cfg_scale = 1.0f; + float temperature = 0.0f; + int top_k = 1; + int frame_index = 0; + uint64_t seed = 0; + int use_cfg = 0; + int forbid_audio_eos = 0; +}; struct magpietts_cuda_sampler { int codebooks = 0; cudaStream_t stream = nullptr; + cudaStream_t owned_stream = nullptr; + bool uses_external_stream = false; int32_t* d_codes = nullptr; int32_t* d_argmax = nullptr; int32_t* d_top_ids = nullptr; float* d_top_vals = nullptr; + magpietts_cuda_sampling_config* h_config = nullptr; + magpietts_cuda_sampling_config* d_config = nullptr; + cudaGraph_t sequence_graph = nullptr; + cudaGraphExec_t sequence_exec = nullptr; + cudaGraphNode_t sequence_tail = nullptr; + bool sequence_warm = false; + bool sequence_build_active = false; + bool sequence_disabled = false; }; bool @@ -98,28 +125,26 @@ sampled_logit( return logit; } -static __device__ __forceinline__ bool -better_logit(int lhs_id, float lhs, int rhs_id, float rhs) { - return lhs > rhs || (lhs == rhs && lhs_id < rhs_id); -} - +template __global__ void magpietts_sample_codebooks_kernel( const float* logits_cond, const float* logits_uncond, int codebooks, int vocab_size, - int audio_codebook_size, int audio_eos_id, bool use_cfg, float cfg_scale, float temperature, - int top_k, bool forbid_audio_eos, uint64_t seed, int frame_index, int codebook_offset, - int output_offset, int32_t* top_ids_scratch, float* top_vals_scratch, int32_t* codes_out, - int32_t* argmax_out) { + int audio_codebook_size, int audio_eos_id, const magpietts_cuda_sampling_config* config_ptr, + int codebook_offset, int output_offset, int32_t* top_ids_scratch, float* top_vals_scratch, + int32_t* codes_out, int32_t* argmax_out) { const int c = blockIdx.x; if (c >= codebooks) { return; } - __shared__ float s_vals[MAGPIETTS_CUDA_BLOCK_SIZE]; - __shared__ int s_ids[MAGPIETTS_CUDA_BLOCK_SIZE]; + const magpietts_cuda_sampling_config config = *config_ptr; + + using block_sort = cub::BlockRadixSort< + float, MAGPIETTS_CUDA_BLOCK_SIZE, items_per_thread, int>; + __shared__ typename block_sort::TempStorage sort_storage; __shared__ double s_sums[MAGPIETTS_CUDA_BLOCK_SIZE]; - int k = top_k < vocab_size ? top_k : vocab_size; + int k = config.top_k < vocab_size ? config.top_k : vocab_size; if (k < 1) { k = 1; } @@ -128,62 +153,47 @@ magpietts_sample_codebooks_kernel( int32_t* top_ids = top_ids_scratch + (size_t)c * vocab_size; float* top_vals = top_vals_scratch + (size_t)c * vocab_size; - for (int rank = 0; rank < k; ++rank) { - float local_val = -INFINITY; - int local_id = vocab_size; - - for (int id = threadIdx.x; id < vocab_size; id += blockDim.x) { - bool selected = false; - for (int prev = 0; prev < rank; ++prev) { - selected = selected || id == top_ids[prev]; - } - if (selected) { - continue; - } - const float logit = sampled_logit( - logits_cond, logits_uncond, off, id, audio_codebook_size, audio_eos_id, use_cfg, - cfg_scale, forbid_audio_eos); - if (better_logit(id, logit, local_id, local_val)) { - local_val = logit; - local_id = id; - } - } - - s_vals[threadIdx.x] = local_val; - s_ids[threadIdx.x] = local_id; - __syncthreads(); - - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (threadIdx.x < stride && - better_logit( - s_ids[threadIdx.x + stride], s_vals[threadIdx.x + stride], s_ids[threadIdx.x], - s_vals[threadIdx.x])) { - s_vals[threadIdx.x] = s_vals[threadIdx.x + stride]; - s_ids[threadIdx.x] = s_ids[threadIdx.x + stride]; - } - __syncthreads(); - } - - if (threadIdx.x == 0) { - top_vals[rank] = s_vals[0]; - top_ids[rank] = s_ids[0] == vocab_size ? 0 : s_ids[0]; + // Sort once per codebook with work independent of top-k. + float thread_vals[items_per_thread]; + int thread_ids[items_per_thread]; +#pragma unroll + for (int item = 0; item < items_per_thread; ++item) { + const int id = (int)threadIdx.x * items_per_thread + item; + thread_vals[item] = id < vocab_size + ? sampled_logit( + logits_cond, logits_uncond, off, id, audio_codebook_size, + audio_eos_id, config.use_cfg != 0, config.cfg_scale, + config.forbid_audio_eos != 0) + : -INFINITY; + thread_ids[item] = id; + } + block_sort(sort_storage).SortDescendingBlockedToStriped(thread_vals, thread_ids); + __syncthreads(); + +#pragma unroll + for (int item = 0; item < items_per_thread; ++item) { + const int rank = item * MAGPIETTS_CUDA_BLOCK_SIZE + (int)threadIdx.x; + if (rank < k) { + top_vals[rank] = thread_vals[item]; + top_ids[rank] = thread_ids[item]; } - __syncthreads(); } + __syncthreads(); if (threadIdx.x != 0) { s_sums[threadIdx.x] = 0.0; } int sampled = top_ids[0]; - if (temperature <= 0.0f) { + if (config.temperature <= 0.0f) { sampled = top_ids[0]; } else { const float max_logit = top_vals[0]; double local_sum = 0.0; for (int i = threadIdx.x; i < k; i += blockDim.x) { if (isfinite(top_vals[i])) { - local_sum += exp((double)(top_vals[i] - max_logit) / (double)temperature); + local_sum += + exp((double)(top_vals[i] - max_logit) / (double)config.temperature); } } s_sums[threadIdx.x] = local_sum; @@ -195,11 +205,13 @@ magpietts_sample_codebooks_kernel( __syncthreads(); } if (threadIdx.x == 0 && s_sums[0] > 0.0) { - const double target = uniform01(seed, frame_index, codebook_offset + c) * s_sums[0]; + const double target = + uniform01(config.seed, config.frame_index, codebook_offset + c) * s_sums[0]; double acc = 0.0; for (int i = 0; i < k; ++i) { if (isfinite(top_vals[i])) { - acc += exp((double)(top_vals[i] - max_logit) / (double)temperature); + acc += + exp((double)(top_vals[i] - max_logit) / (double)config.temperature); } if (target <= acc) { sampled = top_ids[i]; @@ -222,21 +234,22 @@ magpietts_cuda_sampler_create(int codebooks) { } magpietts_cuda_sampler* sampler = new magpietts_cuda_sampler; sampler->codebooks = codebooks; - cudaError_t err = cudaStreamCreateWithFlags(&sampler->stream, cudaStreamNonBlocking); + cudaError_t err = cudaStreamCreateWithFlags(&sampler->owned_stream, cudaStreamNonBlocking); if (err != cudaSuccess) { delete sampler; return nullptr; } + sampler->stream = sampler->owned_stream; err = cudaMalloc(&sampler->d_codes, (size_t)codebooks * sizeof(int32_t)); if (err != cudaSuccess) { - cudaStreamDestroy(sampler->stream); + cudaStreamDestroy(sampler->owned_stream); delete sampler; return nullptr; } err = cudaMalloc(&sampler->d_argmax, (size_t)codebooks * sizeof(int32_t)); if (err != cudaSuccess) { cudaFree(sampler->d_codes); - cudaStreamDestroy(sampler->stream); + cudaStreamDestroy(sampler->owned_stream); delete sampler; return nullptr; } @@ -245,7 +258,7 @@ magpietts_cuda_sampler_create(int codebooks) { if (err != cudaSuccess) { cudaFree(sampler->d_argmax); cudaFree(sampler->d_codes); - cudaStreamDestroy(sampler->stream); + cudaStreamDestroy(sampler->owned_stream); delete sampler; return nullptr; } @@ -255,7 +268,31 @@ magpietts_cuda_sampler_create(int codebooks) { cudaFree(sampler->d_top_ids); cudaFree(sampler->d_argmax); cudaFree(sampler->d_codes); - cudaStreamDestroy(sampler->stream); + cudaStreamDestroy(sampler->owned_stream); + delete sampler; + return nullptr; + } + err = cudaHostAlloc( + reinterpret_cast(&sampler->h_config), sizeof(*sampler->h_config), + cudaHostAllocDefault); + if (err != cudaSuccess) { + cudaFree(sampler->d_top_vals); + cudaFree(sampler->d_top_ids); + cudaFree(sampler->d_argmax); + cudaFree(sampler->d_codes); + cudaStreamDestroy(sampler->owned_stream); + delete sampler; + return nullptr; + } + *sampler->h_config = magpietts_cuda_sampling_config{}; + err = cudaMalloc(&sampler->d_config, sizeof(*sampler->d_config)); + if (err != cudaSuccess) { + cudaFreeHost(sampler->h_config); + cudaFree(sampler->d_top_vals); + cudaFree(sampler->d_top_ids); + cudaFree(sampler->d_argmax); + cudaFree(sampler->d_codes); + cudaStreamDestroy(sampler->owned_stream); delete sampler; return nullptr; } @@ -267,14 +304,255 @@ magpietts_cuda_sampler_free(magpietts_cuda_sampler* sampler) { if (!sampler) { return; } + if (sampler->sequence_exec) { + cudaGraphExecDestroy(sampler->sequence_exec); + } + if (sampler->sequence_graph) { + cudaGraphDestroy(sampler->sequence_graph); + } + cudaFree(sampler->d_config); + cudaFreeHost(sampler->h_config); cudaFree(sampler->d_codes); cudaFree(sampler->d_argmax); cudaFree(sampler->d_top_ids); cudaFree(sampler->d_top_vals); - cudaStreamDestroy(sampler->stream); + cudaStreamDestroy(sampler->owned_stream); delete sampler; } +bool +magpietts_cuda_sampler_bind_stream( + magpietts_cuda_sampler* sampler, void* stream, char* error, size_t error_size) { + if (!sampler || !stream) { + set_error(error, error_size, "invalid CUDA sampler stream binding"); + return false; + } + sampler->stream = (cudaStream_t)stream; + sampler->uses_external_stream = true; + if (error && error_size > 0) { + error[0] = '\0'; + } + return true; +} + +bool +magpietts_cuda_sampler_configure( + magpietts_cuda_sampler* sampler, bool use_cfg, float cfg_scale, float temperature, int top_k, + bool forbid_audio_eos, uint64_t seed, int frame_index, char* error, size_t error_size) { + if (!sampler || !sampler->h_config || !sampler->d_config) { + set_error(error, error_size, "invalid CUDA sampler configuration"); + return false; + } + sampler->h_config->cfg_scale = cfg_scale; + sampler->h_config->temperature = temperature; + sampler->h_config->top_k = top_k; + sampler->h_config->frame_index = frame_index; + sampler->h_config->seed = seed; + sampler->h_config->use_cfg = use_cfg ? 1 : 0; + sampler->h_config->forbid_audio_eos = forbid_audio_eos ? 1 : 0; + if (error && error_size > 0) error[0] = '\0'; + return true; +} + +bool +magpietts_cuda_sampler_upload_config( + magpietts_cuda_sampler* sampler, char* error, size_t error_size) { + if (!sampler || !sampler->h_config || !sampler->d_config) { + set_error(error, error_size, "invalid CUDA sampler config upload"); + return false; + } + cudaError_t err = cudaSuccess; + if (sampler->sequence_build_active) { + cudaGraphNode_t node = nullptr; + const cudaGraphNode_t* deps = sampler->sequence_tail ? &sampler->sequence_tail : nullptr; + const size_t dependency_count = sampler->sequence_tail ? 1 : 0; + err = cudaGraphAddMemcpyNode1D( + &node, sampler->sequence_graph, deps, dependency_count, sampler->d_config, + sampler->h_config, sizeof(*sampler->d_config), cudaMemcpyHostToDevice); + if (err == cudaSuccess) sampler->sequence_tail = node; + } else { + err = cudaMemcpyAsync( + sampler->d_config, sampler->h_config, sizeof(*sampler->d_config), + cudaMemcpyHostToDevice, sampler->stream); + } + if (err != cudaSuccess) { + set_error(error, error_size, "failed to upload CUDA sampler config", err); + return false; + } + if (error && error_size > 0) error[0] = '\0'; + return true; +} + +bool +magpietts_cuda_sampler_sequence_is_warm(const magpietts_cuda_sampler* sampler) { + return sampler && sampler->sequence_warm; +} + +bool +magpietts_cuda_sampler_sequence_is_ready(const magpietts_cuda_sampler* sampler) { + return sampler && sampler->sequence_exec; +} + +bool +magpietts_cuda_sampler_sequence_is_disabled(const magpietts_cuda_sampler* sampler) { + return !sampler || sampler->sequence_disabled; +} + +bool +magpietts_cuda_sampler_sequence_build_active(const magpietts_cuda_sampler* sampler) { + return sampler && sampler->sequence_build_active; +} + +void +magpietts_cuda_sampler_sequence_mark_warm(magpietts_cuda_sampler* sampler) { + if (sampler) sampler->sequence_warm = true; +} + +bool +magpietts_cuda_sampler_sequence_begin_build( + magpietts_cuda_sampler* sampler, char* error, size_t error_size) { + if (!sampler || !sampler->stream || sampler->sequence_exec || sampler->sequence_disabled || + sampler->sequence_build_active) { + set_error(error, error_size, "invalid CUDA local sequence graph-build state"); + return false; + } + cudaGraph_t graph = nullptr; + const cudaError_t err = cudaGraphCreate(&graph, 0); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to create CUDA local sequence graph", err); + return false; + } + sampler->sequence_graph = graph; + sampler->sequence_tail = nullptr; + sampler->sequence_build_active = true; + if (error && error_size > 0) error[0] = '\0'; + return true; +} + +bool +magpietts_cuda_sampler_sequence_finish_build_and_launch( + magpietts_cuda_sampler* sampler, char* error, size_t error_size) { + if (!sampler || !sampler->sequence_build_active) { + set_error(error, error_size, "CUDA local sequence graph build is not active"); + return false; + } + cudaGraph_t graph = sampler->sequence_graph; + cudaError_t err = cudaSuccess; + sampler->sequence_build_active = false; + sampler->sequence_tail = nullptr; + if (!graph) { + set_error(error, error_size, "CUDA local sequence graph is empty"); + return false; + } + cudaGraphExec_t exec = nullptr; + err = cudaGraphInstantiate(&exec, graph, nullptr, nullptr, 0); + if (err != cudaSuccess) { + cudaGraphDestroy(graph); + sampler->sequence_graph = nullptr; + set_error(error, error_size, "failed to instantiate CUDA local sequence graph", err); + return false; + } + err = cudaGraphLaunch(exec, sampler->stream); + if (err != cudaSuccess) { + cudaGraphExecDestroy(exec); + cudaGraphDestroy(graph); + sampler->sequence_graph = nullptr; + set_error(error, error_size, "failed to launch composed CUDA local sequence", err); + return false; + } + sampler->sequence_exec = exec; + if (error && error_size > 0) error[0] = '\0'; + return true; +} + +void +magpietts_cuda_sampler_sequence_abort_build(magpietts_cuda_sampler* sampler) { + if (!sampler || !sampler->sequence_build_active) return; + if (sampler->sequence_graph) cudaGraphDestroy(sampler->sequence_graph); + sampler->sequence_graph = nullptr; + sampler->sequence_tail = nullptr; + sampler->sequence_build_active = false; +} + +void +magpietts_cuda_sampler_sequence_disable(magpietts_cuda_sampler* sampler) { + if (!sampler) return; + magpietts_cuda_sampler_sequence_abort_build(sampler); + if (sampler->sequence_exec) { + cudaGraphExecDestroy(sampler->sequence_exec); + sampler->sequence_exec = nullptr; + } + if (sampler->sequence_graph) { + cudaGraphDestroy(sampler->sequence_graph); + sampler->sequence_graph = nullptr; + } + sampler->sequence_tail = nullptr; + sampler->sequence_disabled = true; +} + +bool +magpietts_cuda_sampler_sequence_launch( + magpietts_cuda_sampler* sampler, char* error, size_t error_size) { + if (!sampler || !sampler->sequence_exec) { + set_error(error, error_size, "CUDA local sequence graph is not ready"); + return false; + } + const cudaError_t err = cudaGraphLaunch(sampler->sequence_exec, sampler->stream); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to launch CUDA local sequence graph", err); + return false; + } + if (error && error_size > 0) error[0] = '\0'; + return true; +} + +bool +magpietts_cuda_sampler_sequence_add_ggml_graph( + magpietts_cuda_sampler* sampler, void* graph_template, char* error, size_t error_size) { + if (!sampler || !sampler->sequence_build_active || !sampler->sequence_graph || + !graph_template) { + set_error(error, error_size, "invalid GGML child graph for CUDA local sequence"); + return false; + } + cudaGraphNode_t node = nullptr; + const cudaGraphNode_t* deps = sampler->sequence_tail ? &sampler->sequence_tail : nullptr; + const size_t dependency_count = sampler->sequence_tail ? 1 : 0; + const cudaError_t err = cudaGraphAddChildGraphNode( + &node, sampler->sequence_graph, deps, dependency_count, + reinterpret_cast(graph_template)); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to add GGML child graph to CUDA local sequence", err); + return false; + } + sampler->sequence_tail = node; + if (error && error_size > 0) error[0] = '\0'; + return true; +} + +bool +magpietts_cuda_sampler_sequence_add_device_copy( + magpietts_cuda_sampler* sampler, const void* src_device, void* dst_device, size_t bytes, + char* error, size_t error_size) { + if (!sampler || !sampler->sequence_build_active || !sampler->sequence_graph || !src_device || + !dst_device || bytes == 0) { + set_error(error, error_size, "invalid device copy for CUDA local sequence"); + return false; + } + cudaGraphNode_t node = nullptr; + const cudaGraphNode_t* deps = sampler->sequence_tail ? &sampler->sequence_tail : nullptr; + const size_t dependency_count = sampler->sequence_tail ? 1 : 0; + const cudaError_t err = cudaGraphAddMemcpyNode1D( + &node, sampler->sequence_graph, deps, dependency_count, dst_device, src_device, bytes, + cudaMemcpyDeviceToDevice); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to add device copy to CUDA local sequence", err); + return false; + } + sampler->sequence_tail = node; + if (error && error_size > 0) error[0] = '\0'; + return true; +} + bool magpietts_cuda_sample_codebooks( magpietts_cuda_sampler* sampler, const float* logits_cond, const float* logits_uncond, @@ -302,6 +580,22 @@ magpietts_cuda_sample_codebooks_device( int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, bool use_cfg, float cfg_scale, float temperature, int top_k, bool forbid_audio_eos, uint64_t seed, int frame_index, int codebook_offset, int output_offset, char* error, size_t error_size) { + if (!magpietts_cuda_sampler_configure( + sampler, use_cfg, cfg_scale, temperature, top_k, forbid_audio_eos, seed, frame_index, + error, error_size) || + !magpietts_cuda_sampler_upload_config(sampler, error, error_size)) { + return false; + } + return magpietts_cuda_sample_codebooks_device_configured( + sampler, logits_cond, logits_uncond, codebooks, vocab_size, audio_codebook_size, + audio_eos_id, codebook_offset, output_offset, error, error_size); +} + +bool +magpietts_cuda_sample_codebooks_device_configured( + magpietts_cuda_sampler* sampler, const float* logits_cond, const float* logits_uncond, + int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, + int codebook_offset, int output_offset, char* error, size_t error_size) { if (!sampler || !logits_cond) { set_error(error, error_size, "invalid CUDA sampler device arguments"); return false; @@ -319,11 +613,59 @@ magpietts_cuda_sample_codebooks_device( return false; } - magpietts_sample_codebooks_kernel<<stream>>>( - logits_cond, logits_uncond, codebooks, vocab_size, audio_codebook_size, audio_eos_id, - use_cfg, cfg_scale, temperature, top_k, forbid_audio_eos, seed, frame_index, - codebook_offset, output_offset, sampler->d_top_ids, sampler->d_top_vals, sampler->d_codes, - sampler->d_argmax); + if (sampler->sequence_build_active) { + magpietts_cuda_sampling_config* config = sampler->d_config; + int32_t* top_ids = sampler->d_top_ids; + float* top_vals = sampler->d_top_vals; + int32_t* codes = sampler->d_codes; + int32_t* argmax = sampler->d_argmax; + void* kernel_args[] = { + &logits_cond, &logits_uncond, &codebooks, &vocab_size, + &audio_codebook_size, &audio_eos_id, &config, &codebook_offset, + &output_offset, &top_ids, &top_vals, &codes, + &argmax}; + cudaKernelNodeParams params{}; + params.func = vocab_size <= MAGPIETTS_CUDA_SMALL_VOCAB + ? reinterpret_cast( + magpietts_sample_codebooks_kernel< + MAGPIETTS_CUDA_SMALL_ITEMS_PER_THREAD>) + : reinterpret_cast( + magpietts_sample_codebooks_kernel< + MAGPIETTS_CUDA_MAX_ITEMS_PER_THREAD>); + params.gridDim = dim3((unsigned int)codebooks, 1, 1); + params.blockDim = dim3(MAGPIETTS_CUDA_BLOCK_SIZE, 1, 1); + params.sharedMemBytes = 0; + params.kernelParams = kernel_args; + params.extra = nullptr; + cudaGraphNode_t node = nullptr; + const cudaGraphNode_t* deps = sampler->sequence_tail ? &sampler->sequence_tail : nullptr; + const size_t dependency_count = sampler->sequence_tail ? 1 : 0; + const cudaError_t err = cudaGraphAddKernelNode( + &node, sampler->sequence_graph, deps, dependency_count, ¶ms); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to add sampler to CUDA local sequence", err); + return false; + } + sampler->sequence_tail = node; + if (error && error_size > 0) error[0] = '\0'; + return true; + } + + if (vocab_size <= MAGPIETTS_CUDA_SMALL_VOCAB) { + magpietts_sample_codebooks_kernel + <<stream>>>( + logits_cond, logits_uncond, codebooks, vocab_size, audio_codebook_size, + audio_eos_id, sampler->d_config, codebook_offset, output_offset, + sampler->d_top_ids, sampler->d_top_vals, sampler->d_codes, + sampler->d_argmax); + } else { + magpietts_sample_codebooks_kernel + <<stream>>>( + logits_cond, logits_uncond, codebooks, vocab_size, audio_codebook_size, + audio_eos_id, sampler->d_config, codebook_offset, output_offset, + sampler->d_top_ids, sampler->d_top_vals, sampler->d_codes, + sampler->d_argmax); + } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { set_error(error, error_size, "failed to launch CUDA sampler", err); @@ -349,10 +691,19 @@ magpietts_cuda_copy_sampled_code_to_device( return false; } - cudaError_t err = cudaStreamSynchronize(sampler->stream); - if (err != cudaSuccess) { - set_error(error, error_size, "failed to synchronize CUDA sampled code", err); - return false; + if (sampler->sequence_build_active) { + return magpietts_cuda_sampler_sequence_add_device_copy( + sampler, sampler->d_codes + codebook, dst_device, sizeof(int32_t), error, + error_size); + } + + cudaError_t err = cudaSuccess; + if (!sampler->uses_external_stream) { + err = cudaStreamSynchronize(sampler->stream); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to synchronize CUDA sampled code", err); + return false; + } } err = cudaMemcpyAsync( dst_device, sampler->d_codes + codebook, sizeof(int32_t), cudaMemcpyDeviceToDevice, @@ -361,10 +712,12 @@ magpietts_cuda_copy_sampled_code_to_device( set_error(error, error_size, "failed to copy CUDA sampled code to device", err); return false; } - err = cudaStreamSynchronize(sampler->stream); - if (err != cudaSuccess) { - set_error(error, error_size, "failed to synchronize CUDA sampled-code copy", err); - return false; + if (!sampler->uses_external_stream) { + err = cudaStreamSynchronize(sampler->stream); + if (err != cudaSuccess) { + set_error(error, error_size, "failed to synchronize CUDA sampled-code copy", err); + return false; + } } if (error && error_size > 0) { error[0] = '\0'; diff --git a/src/tts/magpietts/magpietts_cuda_sampling.h b/src/tts/magpietts/magpietts_cuda_sampling.h index 0925b6b..1660276 100644 --- a/src/tts/magpietts/magpietts_cuda_sampling.h +++ b/src/tts/magpietts/magpietts_cuda_sampling.h @@ -11,6 +11,37 @@ magpietts_cuda_sampler* magpietts_cuda_sampler_create(int codebooks); void magpietts_cuda_sampler_free(magpietts_cuda_sampler* sampler); bool magpietts_cuda_device_is_uma(void); +// Bind to a caller-owned stream. +bool magpietts_cuda_sampler_bind_stream( + magpietts_cuda_sampler* sampler, void* stream, char* error, size_t error_size); + +// Configure per-frame values stored in a stable device buffer. +bool magpietts_cuda_sampler_configure( + magpietts_cuda_sampler* sampler, bool use_cfg, float cfg_scale, float temperature, int top_k, + bool forbid_audio_eos, uint64_t seed, int frame_index, char* error, size_t error_size); +bool magpietts_cuda_sampler_upload_config( + magpietts_cuda_sampler* sampler, char* error, size_t error_size); + +// Compose and launch the local-transformer sequence as a CUDA graph. +bool magpietts_cuda_sampler_sequence_is_warm(const magpietts_cuda_sampler* sampler); +bool magpietts_cuda_sampler_sequence_is_ready(const magpietts_cuda_sampler* sampler); +bool magpietts_cuda_sampler_sequence_is_disabled(const magpietts_cuda_sampler* sampler); +bool magpietts_cuda_sampler_sequence_build_active(const magpietts_cuda_sampler* sampler); +void magpietts_cuda_sampler_sequence_mark_warm(magpietts_cuda_sampler* sampler); +bool magpietts_cuda_sampler_sequence_begin_build( + magpietts_cuda_sampler* sampler, char* error, size_t error_size); +bool magpietts_cuda_sampler_sequence_finish_build_and_launch( + magpietts_cuda_sampler* sampler, char* error, size_t error_size); +void magpietts_cuda_sampler_sequence_abort_build(magpietts_cuda_sampler* sampler); +void magpietts_cuda_sampler_sequence_disable(magpietts_cuda_sampler* sampler); +bool magpietts_cuda_sampler_sequence_launch( + magpietts_cuda_sampler* sampler, char* error, size_t error_size); +bool magpietts_cuda_sampler_sequence_add_ggml_graph( + magpietts_cuda_sampler* sampler, void* graph_template, char* error, size_t error_size); +bool magpietts_cuda_sampler_sequence_add_device_copy( + magpietts_cuda_sampler* sampler, const void* src_device, void* dst_device, size_t bytes, + char* error, size_t error_size); + bool magpietts_cuda_sample_codebooks( magpietts_cuda_sampler* sampler, const float* logits_cond, const float* logits_uncond, int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, bool use_cfg, @@ -24,6 +55,12 @@ bool magpietts_cuda_sample_codebooks_device( float cfg_scale, float temperature, int top_k, bool forbid_audio_eos, uint64_t seed, int frame_index, int codebook_offset, int output_offset, char* error, size_t error_size); +// Launch using the most recently uploaded configuration. +bool magpietts_cuda_sample_codebooks_device_configured( + magpietts_cuda_sampler* sampler, const float* logits_cond, const float* logits_uncond, + int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, + int codebook_offset, int output_offset, char* error, size_t error_size); + bool magpietts_cuda_copy_sampled_code_to_device( magpietts_cuda_sampler* sampler, int codebook, void* dst_device, char* error, size_t error_size); diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index e03cafd..712060a 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -1146,11 +1146,11 @@ cross_attention( return linear(ctx, layer.cross_o, out); } -static ggml_tensor* +ggml_tensor* cross_attention_cached( ggml_context* ctx, const magpietts_transformer& tr, const magpietts_layer& layer, const DecoderCrossKvCache& cross_kv, int layer_index, ggml_tensor* x, ggml_tensor* attn_prior, - ggml_tensor** last_attn) { + ggml_tensor** last_attn, bool prior_is_log) { const int64_t d_head = tr.n_cross_dhead; const int64_t n_head = tr.n_cross_head; const int64_t cross_dim = d_head * n_head; @@ -1170,8 +1170,13 @@ cross_attention_cached( 0, 2, 1, 3); ggml_tensor* kq = ggml_mul_mat(ctx, kh, qh); kq = ggml_scale(ctx, kq, 1.0f / std::sqrt((float)d_head)); - ggml_tensor* kq_soft = ggml_soft_max(ctx, kq); - if (attn_prior) { + ggml_tensor* kq_soft = nullptr; + if (attn_prior && prior_is_log) { + kq_soft = ggml_soft_max(ctx, ggml_add(ctx, kq, ggml_repeat(ctx, attn_prior, kq))); + } else { + kq_soft = ggml_soft_max(ctx, kq); + } + if (attn_prior && !prior_is_log) { ggml_tensor* prior = ggml_repeat(ctx, attn_prior, kq_soft); kq_soft = ggml_mul(ctx, kq_soft, prior); ggml_tensor* normalizer = ggml_repeat(ctx, ggml_sum_rows(ctx, kq_soft), kq_soft); @@ -2002,6 +2007,8 @@ MagpieCodeGenerator::generate( decoder_result cond; decoder_result uncond; + cond.logits_required = !params.use_local_transformer; + uncond.logits_required = !params.use_local_transformer; if (use_cuda_sampling) { if (params.use_local_transformer) { const bool decode_ok = diff --git a/tests/cpp/tts/test_magpietts_file.cpp b/tests/cpp/tts/test_magpietts_file.cpp index f78d78a..8f458b3 100644 --- a/tests/cpp/tts/test_magpietts_file.cpp +++ b/tests/cpp/tts/test_magpietts_file.cpp @@ -28,13 +28,17 @@ usage(const char* argv0) { std::fprintf( stderr, "usage: %s --tts.magpie-model magpie.gguf --tts.codec-model codec.gguf " - "--tts.tokens-file tokens.txt --tts.wav-out out.wav [options]\n" + "(--tts.tokens-file tokens.txt | --tts.text TEXT --tts.tokenizer-model-dir DIR) " + "--tts.wav-out out.wav [options]\n" "\n" "required:\n" " --tts.magpie-model PATH\n" " --tts.codec-model PATH\n" " --tts.tokens LIST Comma/space-separated token IDs\n" " --tts.tokens-file PATH File containing token IDs\n" + " --tts.text TEXT Natural text to tokenize and synthesize\n" + " --tts.tokenizer-model-dir DIR\n" + " Extracted Magpie tokenizer assets for --tts.text\n" " --tts.wav-out PATH Output WAV written after synthesis completes\n" "\n" "options:\n" @@ -43,6 +47,7 @@ usage(const char* argv0) { " --threads N CPU threads (default 4)\n" " --tts.codec-threads N Codec CPU threads (default --threads)\n" " --tts.chunk-frames N Codec frames per internal callback chunk (default 3)\n" + " --tts.language-code LANG Tokenizer language (default en-US)\n" " --tts.lt-backend auto|cpu|cuda\n" " Local-transformer backend\n" " --tts.lt-fp32 Run the local transformer entirely in FP32\n" @@ -56,6 +61,8 @@ usage(const char* argv0) { " --tts.no-local-transformer Sample directly from decoder final projection\n" " --tts.no-kv-cache Recompute decoder prefix each frame\n" " --tts.no-stateful-codec Disable fast layer-state codec\n" + " --warmup-runs N Untimed synthesis runs before measurement (default 0)\n" + " --runs N Number of measured synthesis runs (default 1)\n" " --benchmark Print timing summary\n" " --verbose Print detailed runtime logs\n", argv0); @@ -165,6 +172,8 @@ write_wav(const std::string& path, int sample_rate, const std::vector& int main(int argc, char** argv) { tts::magpie_stream_params params; + int warmup_runs = 0; + int measured_runs = 1; for (int i = 1; i < argc; ++i) { const std::string arg = argv[i]; @@ -184,6 +193,10 @@ main(int argc, char** argv) { params.tokens = tts::parse_token_list(need_value(arg.c_str())); } else if (arg == "--tts.tokens-file") { params.tokens_file = need_value(arg.c_str()); + } else if (arg == "--tts.text") { + params.text = need_value(arg.c_str()); + } else if (arg == "--tts.tokenizer-model-dir") { + params.tokenizer_model_dir = need_value(arg.c_str()); } else if (arg == "--tts.wav-out" || arg == "--tts.output") { params.wav_out = need_value(arg.c_str()); } else if (arg == "--tts.speaker") { @@ -196,6 +209,8 @@ main(int argc, char** argv) { params.codec_threads = parse_int(need_value(arg.c_str()), "--tts.codec-threads"); } else if (arg == "--tts.chunk-frames") { params.chunk_frames = parse_int(need_value(arg.c_str()), "--tts.chunk-frames"); + } else if (arg == "--tts.language-code") { + params.language_code = need_value(arg.c_str()); } else if (arg == "--tts.lt-backend") { if (!tts::parse_backend_preference(need_value(arg.c_str()), params.lt_backend)) { std::fprintf(stderr, "--tts.lt-backend must be auto, cpu, or cuda\n"); @@ -226,6 +241,10 @@ main(int argc, char** argv) { params.use_kv_cache = false; } else if (arg == "--tts.no-stateful-codec") { params.use_stateful_codec = false; + } else if (arg == "--warmup-runs") { + warmup_runs = parse_int(need_value(arg.c_str()), "--warmup-runs"); + } else if (arg == "--runs") { + measured_runs = parse_int(need_value(arg.c_str()), "--runs"); } else if (arg == "--benchmark") { params.benchmark = true; } else if (arg == "--verbose") { @@ -243,13 +262,51 @@ main(int argc, char** argv) { if (!params.tokens_file.empty()) { params.tokens = tts::parse_token_list(tts::read_file(params.tokens_file)); } + if (!params.text.empty() && !params.tokens.empty()) { + std::fprintf(stderr, "provide either --tts.text or token IDs, not both\n"); + return 1; + } + if (!params.text.empty() && params.tokenizer_model_dir.empty()) { + std::fprintf(stderr, "--tts.text requires --tts.tokenizer-model-dir\n"); + return 1; + } + + std::vector> token_chunks; + if (!params.text.empty()) { + try { + tts::MagpieNativeTokenizer tokenizer(params.tokenizer_model_dir); + const tts::MagpieTokenizationResult tokenized = + tokenizer.tokenize(params.text, params.language_code); + params.tokens = tokenized.tokens; + token_chunks.reserve(tokenized.chunks.size()); + for (const auto& chunk : tokenized.chunks) { + token_chunks.push_back(chunk.tokens); + } + std::fprintf( + stderr, "benchmark_text=\"%s\" language=%s tokenizer=%s tokens=%zu chunks=%zu\n", + params.text.c_str(), tokenized.language.c_str(), tokenized.tokenizer_name.c_str(), + params.tokens.size(), token_chunks.size()); + } + catch (const std::exception& e) { + std::fprintf(stderr, "failed to tokenize benchmark text: %s\n", e.what()); + return 1; + } + } else if (!params.tokens.empty()) { + token_chunks.push_back(params.tokens); + std::fprintf(stderr, "benchmark_input=synthetic_token_ids tokens=%zu\n", params.tokens.size()); + } params.warmup_tokens = params.tokens; + params.warmup_token_chunks = token_chunks; if (params.magpie_model.empty() || params.codec_model.empty() || params.tokens.empty() || params.wav_out.empty()) { usage(argv[0]); return 1; } + if (warmup_runs < 0 || measured_runs <= 0) { + std::fprintf(stderr, "--warmup-runs must be non-negative and --runs must be positive\n"); + return 1; + } tts::MagpieStreamingRuntime runtime; if (!runtime.load( @@ -257,17 +314,51 @@ main(int argc, char** argv) { return 1; } + for (int run = 0; run < warmup_runs; ++run) { + tts::stream_run_metrics warmup_metrics; + const bool ok = runtime.synthesize( + params, token_chunks, [](const std::vector&) { return true; }, + warmup_metrics, "warmup", false); + if (!ok) { + return 1; + } + } + std::vector pcm; tts::stream_run_metrics metrics; - const bool ok = runtime.synthesize( - params, params.tokens, - [&](const std::vector& bytes) { - pcm.insert(pcm.end(), bytes.begin(), bytes.end()); - return true; - }, - metrics); - if (!ok) { - return 1; + double e2e_rtfx_sum = 0.0; + double ttfa_ms_sum = 0.0; + double decoder_ttft_ms_sum = 0.0; + double decoder_itl_ms_sum = 0.0; + double encoder_ms_sum = 0.0; + for (int run = 0; run < measured_runs; ++run) { + pcm.clear(); + tts::stream_run_metrics run_metrics; + const bool ok = runtime.synthesize( + params, token_chunks, + [&](const std::vector& bytes) { + pcm.insert(pcm.end(), bytes.begin(), bytes.end()); + return true; + }, + run_metrics, "riva_tts", false); + if (!ok) { + return 1; + } + metrics = run_metrics; + e2e_rtfx_sum += run_metrics.e2e_rtfx; + ttfa_ms_sum += run_metrics.e2e.first_event_ms; + decoder_ttft_ms_sum += run_metrics.decoder.first_event_ms; + decoder_itl_ms_sum += run_metrics.decoder.inter_event_avg_ms(); + encoder_ms_sum += run_metrics.encoder_ms; + if (params.benchmark && measured_runs > 1) { + std::fprintf( + stderr, + "benchmark_run=%d/%d frames=%d e2e_rtfx=%.2f ttfa_ms=%.2f " + "encoder_ms=%.2f decoder_ttft_ms=%.2f decoder_itl_avg_ms=%.2f\n", + run + 1, measured_runs, run_metrics.generated_frames, run_metrics.e2e_rtfx, + run_metrics.e2e.first_event_ms, run_metrics.encoder_ms, + run_metrics.decoder.first_event_ms, run_metrics.decoder.inter_event_avg_ms()); + } } if (!write_wav(params.wav_out, runtime.sampleRate(), pcm)) { return 1; @@ -287,6 +378,16 @@ main(int argc, char** argv) { metrics.generated_frames, metrics.chunks, (unsigned long long)metrics.samples_written, metrics.e2e_rtfx, metrics.e2e.first_event_ms, metrics.decoder.first_event_ms, metrics.codec.first_event_ms); + if (measured_runs > 1) { + const double divisor = (double)measured_runs; + std::fprintf( + stderr, + "benchmark_mean runs=%d e2e_rtfx=%.2f ttfa_ms=%.2f encoder_ms=%.2f " + "decoder_ttft_ms=%.2f decoder_itl_avg_ms=%.2f\n", + measured_runs, e2e_rtfx_sum / divisor, ttfa_ms_sum / divisor, + encoder_ms_sum / divisor, decoder_ttft_ms_sum / divisor, + decoder_itl_ms_sum / divisor); + } } return 0; } From 675b8bd4f66277e91d8140ba99364710e2ea4e96 Mon Sep 17 00:00:00 2001 From: Prabhsimran Singh Date: Sat, 22 Aug 2026 14:55:23 +0530 Subject: [PATCH 06/19] perf(tts): optimize magpie decoding with cached attention - batch conditional and unconditional CFG lanes in one forward pass - keep LT K/V caches resident on the device across codebooks - generalize fused GGML attention for cached and relative-position modes - optimize projection and LayerNorm scheduling for frame-stacked decoding --- CMakeLists.txt | 9 +- ggml-patches/0007-magpietts-nanocodec.patch | 22 +- ...014-cuda-fused-attention-extensions.patch} | 1991 +++++++++++++---- .../0015-cuda-ctc-batch-fusions.patch | 31 +- ggml-patches/README.md | 234 +- src/tts/magpietts/lt.cpp | 269 ++- tests/cpp/tts/CMakeLists.txt | 5 + .../tts/test_magpietts_cached_attention.cpp | 196 ++ 8 files changed, 2033 insertions(+), 724 deletions(-) rename ggml-patches/{0014-cuda-relpos-extensions.patch => 0014-cuda-fused-attention-extensions.patch} (51%) create mode 100644 tests/cpp/tts/test_magpietts_cached_attention.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 580e6dc..8295001 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -147,12 +147,9 @@ option(NEMO_SPEECH_CUBLAS_SHIM "Build the in-tree drop-in cuBLAS shim (libcu # norm/BF16 fusions, CUDA graph and launch fixes) are transparent ggml-cuda internals. option(NEMO_SPEECH_GGML_PATCHED "Linked ggml has the project ASR patches applied (fused rel-pos op, F16 dw-conv)" ON) -# Fused rel-pos attention CUDA op (GGML_OP_FUSED_RELPOS_ATTN). Replaces the -# unfused 3-GEMM + rel-shift + softmax attention sequence with one kernel in the -# non-cache-aware encoder path. Defaults ON whenever GGML_CUDA + a patched ggml -# are present (it links a patch-only op symbol); forced OFF otherwise, where the -# encoder runs the unfused stock-ggml path. Kept as an explicit option so it can -# be toggled OFF for debugging / WER bisection (the kernel stores k/v/p F16). +# Relative-position mode of the fused CUDA attention op. Replaces the unfused +# encoder attention sequence with one kernel. Defaults ON with CUDA and the +# patched ggml; the encoder otherwise uses stock ggml ops. if(GGML_CUDA AND NEMO_SPEECH_GGML_PATCHED) option(NEMO_SPEECH_FUSED_RELPOS_ATTN "Use the fused rel-pos attention CUDA op in the encoder" ON) else() diff --git a/ggml-patches/0007-magpietts-nanocodec.patch b/ggml-patches/0007-magpietts-nanocodec.patch index 57d4c2b..23f8c8b 100644 --- a/ggml-patches/0007-magpietts-nanocodec.patch +++ b/ggml-patches/0007-magpietts-nanocodec.patch @@ -534,7 +534,7 @@ index aad4c34a..628d8718 100644 if (src0_ne[1] <= 1024 && src1_ncols > 512) { return false; diff --git a/src/ggml-cuda/mmvf.cu b/src/ggml-cuda/mmvf.cu -index d9147202..0c974164 100644 +index d9147202..a347447c 100644 --- a/src/ggml-cuda/mmvf.cu +++ b/src/ggml-cuda/mmvf.cu @@ -383,7 +383,10 @@ static void mul_mat_vec_f_switch_fusion( @@ -558,7 +558,19 @@ index d9147202..0c974164 100644 mul_mat_vec_f<<>> (x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, -@@ -650,7 +653,7 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor +@@ -436,6 +439,11 @@ void launch_mul_mat_vec_f_cuda( + block_size_best = block_size; + } + } ++ if constexpr (std::is_same_v) { ++ if (warp_size == 32 && ncols >= 768 && ncols_dst <= 2) { ++ block_size_best = 96; ++ } ++ } + + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + +@@ -650,7 +658,7 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor if (fusion) { GGML_ASSERT( !ids || dst->ne[2] == 1); @@ -567,7 +579,7 @@ index d9147202..0c974164 100644 if (fusion->x_bias) { GGML_ASSERT(fusion->x_bias->type == GGML_TYPE_F32); GGML_ASSERT(fusion->x_bias->ne[0] == dst->ne[0]); -@@ -793,9 +796,15 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 +@@ -793,9 +801,15 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 } } @@ -583,7 +595,7 @@ index d9147202..0c974164 100644 if (ampere_mma_available(cc)) { return ne11 <= 3; } -@@ -812,9 +821,12 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 +@@ -812,9 +826,12 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 return ne11 <= 8; case GGML_TYPE_F16: if (GGML_CUDA_CC_IS_NVIDIA(cc)) { @@ -597,7 +609,7 @@ index d9147202..0c974164 100644 } if (cc >= GGML_CUDA_CC_ADA_LOVELACE) { return src0_small && ne11 <= 4; -@@ -838,6 +850,9 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 +@@ -838,6 +855,9 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 return ne11 <= 8; case GGML_TYPE_BF16: if (GGML_CUDA_CC_IS_NVIDIA(cc)) { diff --git a/ggml-patches/0014-cuda-relpos-extensions.patch b/ggml-patches/0014-cuda-fused-attention-extensions.patch similarity index 51% rename from ggml-patches/0014-cuda-relpos-extensions.patch rename to ggml-patches/0014-cuda-fused-attention-extensions.patch index bac3f72..4d9de3f 100644 --- a/ggml-patches/0014-cuda-relpos-extensions.patch +++ b/ggml-patches/0014-cuda-fused-attention-extensions.patch @@ -1,8 +1,34 @@ diff --git a/include/ggml.h b/include/ggml.h -index fb823571..4671c80d 100644 +index fb823571..d403756e 100644 --- a/include/ggml.h +++ b/include/ggml.h -@@ -2436,8 +2436,9 @@ extern "C" { +@@ -583,7 +583,7 @@ extern "C" { + + GGML_OP_GLU, + +- GGML_OP_FUSED_RELPOS_ATTN, ++ GGML_OP_FUSED_ATTN, + + GGML_OP_COUNT, + }; +@@ -2423,12 +2423,11 @@ extern "C" { + struct ggml_tensor * a, + struct ggml_tensor * sinks); + +- // Fused FastConformer relative-position multi-head attention. +- // Replaces the content (K*Qu) + position (P*Qv with rel-shift) + softmax + +- // context (attn*V) op sequence with one kernel. CUDA-only (CPU/other +- // backends report unsupported and the unfused graph runs instead). ++ // Fused CUDA multi-head attention with optional relative-position terms ++ // and an optional persistent circular K/V cache. CPU and other backends ++ // report this op as unsupported. + // +- // Logical shapes, head dim ne[0]=d_k fastest: ++ // Logical shapes, with head dimension d_k as ne[0]: + // q [d_k, q_len, n_head, batch] pre-bias query (Qu/Qv added inside) + // k [d_k, kv_len, n_head, batch] + // v [d_k, kv_len, n_head, batch] +@@ -2436,8 +2435,9 @@ extern "C" { // pos_len >= kv_len + q_len - 1 // bias_u [d_k, n_head] pos_bias_u (content term) // bias_v [d_k, n_head] pos_bias_v (position term) @@ -14,14 +40,14 @@ index fb823571..4671c80d 100644 // Q/K/V/P may be non-contiguous views as long as each d_k row is // contiguous — e.g. Q sliced from a fused-QKV projection and K/V read // head-split from a feat-major [n_feat, kv] window (the CUDA op derives -@@ -2460,6 +2461,46 @@ extern "C" { +@@ -2460,6 +2460,44 @@ extern "C" { float scale, bool merge_heads); + // Streaming CUDA variant. K/V contain only the current chunk; cached K/V + // are read directly from a persistent [n_feat*cache_len, slots, 2] F32 + // arena using one I32 slot id and circular-cache head per batch item. -+ // ring_heads may be [batch] or [batch,2]; the optional second column is ++ // cache_state may be [batch] or [batch,2]; the optional second column is + // the number of valid past entries, allowing a fixed graph to skip unused + // cache rows. The caller advances the head after a successful graph run. + GGML_API struct ggml_tensor * ggml_fused_relpos_attn_cached( @@ -35,16 +61,14 @@ index fb823571..4671c80d 100644 + struct ggml_tensor * mask, + struct ggml_tensor * kv_cache, + struct ggml_tensor * slot_ids, -+ struct ggml_tensor * ring_heads, ++ struct ggml_tensor * cache_state, + int64_t cache_len, + float scale, + bool merge_heads); + -+ // Absolute-position streaming attention. This shares the same persistent -+ // [n_feat*cache_len, slots, 2] F32 K/V arena and fused cache-update kernel -+ // as ggml_fused_relpos_attn_cached, but omits relative-position operands -+ // and biases. Q/K/V are [d_k, chunk_len, n_head, batch]; the current K/V -+ // chunk is attended and appended to the circular cache by the same op. ++ // Standard cached attention. Q/K/V are [d_k, chunk_len, n_head, batch]. ++ // The current K/V chunk is attended and appended to the shared persistent ++ // [n_feat*cache_len, slots, 2] F32 circular cache. + GGML_API struct ggml_tensor * ggml_fused_attn_cached( + struct ggml_context * ctx, + struct ggml_tensor * q, @@ -53,7 +77,7 @@ index fb823571..4671c80d 100644 + struct ggml_tensor * mask, + struct ggml_tensor * kv_cache, + struct ggml_tensor * slot_ids, -+ struct ggml_tensor * ring_heads, ++ struct ggml_tensor * cache_state, + int64_t cache_len, + float scale, + bool merge_heads); @@ -61,37 +85,90 @@ index fb823571..4671c80d 100644 // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, -diff --git a/src/ggml-cuda/fused-relpos-attn.cu b/src/ggml-cuda/fused-relpos-attn.cu -index f3c4836a..e0c0ec27 100644 ---- a/src/ggml-cuda/fused-relpos-attn.cu -+++ b/src/ggml-cuda/fused-relpos-attn.cu -@@ -3,6 +3,8 @@ - #include "fused-relpos-attn.cuh" +diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c +index cce0e5a8..43de6346 100644 +--- a/src/ggml-cpu/ggml-cpu.c ++++ b/src/ggml-cpu/ggml-cpu.c +@@ -1988,9 +1988,9 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm + { + ggml_compute_forward_flash_attn_ext(params, tensor); + } break; +- case GGML_OP_FUSED_RELPOS_ATTN: ++ case GGML_OP_FUSED_ATTN: + { +- GGML_ABORT("FUSED_RELPOS_ATTN has no CPU path (CUDA-only)"); ++ GGML_ABORT("FUSED_ATTN has no CPU path (CUDA-only)"); + } break; + case GGML_OP_FLASH_ATTN_BACK: + { +diff --git a/src/ggml-cpu/ggml-cpu.cpp b/src/ggml-cpu/ggml-cpu.cpp +index 7429cc45..0c570639 100644 +--- a/src/ggml-cpu/ggml-cpu.cpp ++++ b/src/ggml-cpu/ggml-cpu.cpp +@@ -439,7 +439,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st + } - #include -+#include + switch (op->op) { +- case GGML_OP_FUSED_RELPOS_ATTN: ++ case GGML_OP_FUSED_ATTN: + return false; // CUDA-only; CPU falls back to the unfused graph + case GGML_OP_CPY: + case GGML_OP_SET_ROWS: +diff --git a/src/ggml-cuda/fused-attention.cu b/src/ggml-cuda/fused-attention.cu +new file mode 100644 +index 00000000..0a1c1766 +--- /dev/null ++++ b/src/ggml-cuda/fused-attention.cu +@@ -0,0 +1,1650 @@ ++// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. ++// SPDX-License-Identifier: Apache-2.0 ++#include "fused-attention.cuh" ++ ++#include +#include - - // Fused FastConformer relative-position multi-head attention. - // -@@ -17,7 +19,8 @@ - // K/V window — as long as d_k stays innermost-contiguous (asserted by the op - // constructor; the vectorized loads rely on it). P is [d_k, pos_len, n_head] - // with pos_len = kv + q - 1 rows addressable; bu/bv are [d_k, n_head]; --// mask is [kv] (shared) or [kv, batch] (per-stream) additive (0 / -inf), or NULL. ++ ++// Fused multi-head attention with optional relative-position terms and a ++// persistent circular K/V cache. ++// ++// The generic path uses one block per (head, query, batch). Scores, optional ++// relative-position terms, softmax, and the attn*V context are computed in one ++// kernel, so intermediate score matrices are not written to global memory. ++// ++// Operand addressing is fully stride-driven (strides read from each tensor's ++// nb[] by the host wrapper, in elements): Q/K/V may be non-contiguous views — ++// e.g. the Q slice of a fused-QKV projection, or a feat-major [n_feat, kv] ++// K/V window — as long as d_k stays innermost-contiguous (asserted by the op ++// constructor; the vectorized loads rely on it). P is [d_k, pos_len, n_head] ++// with pos_len = kv + q - 1 rows addressable; bu/bv are [d_k, n_head]; +// mask is [kv] (shared), [kv, q] (offline per-query), or [kv, batch] +// (streaming per-stream) additive (0 / -inf), or NULL. - // Output ctx is [d_k, q, n_head, batch] logical; with the merge_heads op flag - // its memory layout is head-merged ([d_k+h*d_k] innermost, i.e. a plain - // (n_feat, q, batch) matrix), so the output projection consumes it without a -@@ -41,6 +44,17 @@ static __device__ __forceinline__ float relpos_warp_sum(float value) { - return value; - } - -+// Unlike relpos_warp_sum(), return the complete sum to every lane. The ++// Output ctx is [d_k, q, n_head, batch] logical; with the merge_heads op flag ++// its memory layout is head-merged ([d_k+h*d_k] innermost, i.e. a plain ++// (n_feat, q, batch) matrix), so the output projection consumes it without a ++// permute copy. ++// ++// Requires d_k to be a power of two (the softmax reduction halves blockDim.x). ++ ++// K/V/P are templated: F16 operands halve the dominant re-read traffic when ++// the caller stages them; F32 operands skip the staging casts entirely. All ++// math stays in F32 either way. ++ ++static constexpr int RELPOS_ATTN_DK_128 = 128; ++static constexpr int RELPOS_ATTN_WARPS_128 = RELPOS_ATTN_DK_128 / 32; ++static constexpr int RELPOS_ATTN_CC_SM100 = 1000; ++ ++static __device__ __forceinline__ float attention_warp_sum(float value) { ++#pragma unroll ++ for (int offset = 16; offset > 0; offset >>= 1) { ++ value += __shfl_down_sync(0xffffffff, value, offset); ++ } ++ return value; ++} ++ ++// Unlike attention_warp_sum(), return the complete sum to every lane. The +// register-resident Q=2 kernel needs each lane to retain the softmax weight +// for the four V features it accumulates. -+static __device__ __forceinline__ float relpos_warp_all_sum(float value) { ++static __device__ __forceinline__ float attention_warp_all_sum(float value) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value += __shfl_xor_sync(0xffffffff, value, offset); @@ -99,15 +176,28 @@ index f3c4836a..e0c0ec27 100644 + return value; +} + - static __device__ __forceinline__ float4 relpos_load4(const float * ptr) { - return *reinterpret_cast(ptr); - } -@@ -53,21 +67,98 @@ static __device__ __forceinline__ float4 relpos_load4(const half * ptr) { - __low2float(values[1]), __high2float(values[1])); - } - ++static __device__ __forceinline__ float4 attention_load4(const float * ptr) { ++ return *reinterpret_cast(ptr); ++} ++ ++static __device__ __forceinline__ float4 attention_load4(const half * ptr) { ++ const int2 packed = *reinterpret_cast(ptr); ++ const half2 * values = reinterpret_cast(&packed); ++ return make_float4( ++ __low2float(values[0]), __high2float(values[0]), ++ __low2float(values[1]), __high2float(values[1])); ++} ++ ++static __device__ __forceinline__ float2 attention_load2(const float * ptr) { ++ return *reinterpret_cast(ptr); ++} ++ ++static __device__ __forceinline__ float2 attention_load2(const half * ptr) { ++ return __half22float2(*reinterpret_cast(ptr)); ++} ++ +template -+static __device__ __forceinline__ float4 relpos_load_kv4( ++static __device__ __forceinline__ float4 attention_load_kv4( + const T * chunk_head, const float * cache_head, int j, int cache_len, + int ring_head, long chunk_sj, long cache_sj, int d4) { + if constexpr (Cached) { @@ -116,15 +206,15 @@ index f3c4836a..e0c0ec27 100644 + if (physical_j >= cache_len) { + physical_j -= cache_len; + } -+ return relpos_load4(cache_head + (size_t) physical_j * cache_sj + d4); ++ return attention_load4(cache_head + (size_t) physical_j * cache_sj + d4); + } -+ return relpos_load4(chunk_head + (size_t) (j - cache_len) * chunk_sj + d4); ++ return attention_load4(chunk_head + (size_t) (j - cache_len) * chunk_sj + d4); + } -+ return relpos_load4(chunk_head + (size_t) j * chunk_sj + d4); ++ return attention_load4(chunk_head + (size_t) j * chunk_sj + d4); +} + +template -+static __device__ __forceinline__ float relpos_load_kv( ++static __device__ __forceinline__ float attention_load_kv( + const T * chunk_head, const float * cache_head, int j, int cache_len, + int ring_head, long chunk_sj, long cache_sj, int d) { + if constexpr (Cached) { @@ -140,12 +230,29 @@ index f3c4836a..e0c0ec27 100644 + return (float) chunk_head[(size_t) j * chunk_sj + d]; +} + ++template ++static __device__ __forceinline__ float2 attention_load_kv2( ++ const T * chunk_head, const float * cache_head, int j, int cache_len, ++ int ring_head, long chunk_sj, long cache_sj, int d2) { ++ if constexpr (Cached) { ++ if (j < cache_len) { ++ int physical_j = ring_head + j; ++ if (physical_j >= cache_len) { ++ physical_j -= cache_len; ++ } ++ return attention_load2(cache_head + (size_t) physical_j * cache_sj + d2); ++ } ++ return attention_load2(chunk_head + (size_t) (j - cache_len) * chunk_sj + d2); ++ } ++ return attention_load2(chunk_head + (size_t) j * chunk_sj + d2); ++} ++ +// Each thread owns one feature and appends only the current chunk to its +// circular cache. ring_heads[b] is the oldest physical row before this step, +// so those rows are exactly the ones the new chunk replaces. K and V use +// separate planes of the same persistent arena. +template -+static __global__ void fused_relpos_attn_update_cache_kernel( ++static __global__ void fused_attention_update_cache_kernel( + float * __restrict__ arena, const T * __restrict__ K, + const T * __restrict__ V, const int32_t * __restrict__ slot_ids, + const int32_t * __restrict__ ring_heads, @@ -179,34 +286,154 @@ index f3c4836a..e0c0ec27 100644 + } +} + - // SM100 streaming specialization for d_k=128 and q=2. Keeping one block per - // query preserves the two rows' parallelism while the complete grid fits in a - // resident wave. Compared with the generic shared-memory reduction tree, the - // score-producing warps retain their local maxima and max/sum use only four - // warp partials. This removes fourteen block barriers and 124 scratch floats. --template +template - static __global__ void fused_relpos_attn_warp_128_kernel( - const float * __restrict__ Q, const T * __restrict__ K, - const T * __restrict__ V, const T * __restrict__ Ppos, - const float * __restrict__ bu, const float * __restrict__ bv, - const float * __restrict__ mask, float * __restrict__ ctx, -- int kv, float scale, ++static __global__ void fused_cached_attention_q1_d64_kernel( ++ const float * __restrict__ Q, const T * __restrict__ K, ++ const T * __restrict__ V, const float * __restrict__ mask, ++ float * __restrict__ ctx, const float * __restrict__ Kcache, ++ const float * __restrict__ Vcache, const int32_t * __restrict__ slot_ids, ++ const int32_t * __restrict__ ring_heads, ++ const int32_t * __restrict__ active_lengths, ++ int kv, int cache_len, float scale, ++ long q_sh, long q_sb, long k_sj, long k_sh, long k_sb, ++ long v_sj, long v_sh, long v_sb, long cache_sj, long cache_ss, ++ long o_sh, long o_sb, long m_sb) { ++ constexpr int d_k = 64; ++ constexpr int warps = 8; ++ constexpr int value_parts = 4; ++ extern __shared__ float sh[]; ++ float * query = sh; ++ float * scores = query + d_k; ++ float * reduce = scores + kv; ++ float * value_partials = reduce + warps; ++ ++ const int h = blockIdx.x; ++ const int b = blockIdx.z; ++ const int tid = threadIdx.x; ++ const int warp = tid >> 5; ++ const int lane = tid & 31; ++ const float * Qh = Q + (size_t) b * q_sb + (size_t) h * q_sh; ++ const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; ++ const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; ++ const int slot = Cached ? slot_ids[b] : 0; ++ const int ring_head = Cached ? ring_heads[b] : 0; ++ const int key_begin = active_lengths ? cache_len - active_lengths[b] : 0; ++ const float * Kch = Cached ++ ? Kcache + (size_t) slot * cache_ss + (size_t) h * d_k ++ : nullptr; ++ const float * Vch = Cached ++ ? Vcache + (size_t) slot * cache_ss + (size_t) h * d_k ++ : nullptr; ++ const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ ++ if (tid < d_k) { ++ query[tid] = Qh[tid]; ++ } ++ __syncthreads(); ++ ++ const int d2 = lane * 2; ++ const float2 q2 = attention_load2(query + d2); ++ float local_max = -INFINITY; ++ for (int j = key_begin + warp; j < kv; j += warps) { ++ const float2 k2 = attention_load_kv2( ++ Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, d2); ++ float score = k2.x * q2.x + k2.y * q2.y; ++ score = attention_warp_sum(score); ++ if (lane == 0) { ++ score = score * scale + (Mb ? Mb[j] : 0.0f); ++ scores[j] = score; ++ local_max = fmaxf(local_max, score); ++ } ++ } ++ if (lane == 0) { ++ reduce[warp] = local_max; ++ } ++ __syncthreads(); ++ if (tid == 0) { ++ float maximum = reduce[0]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ maximum = fmaxf(maximum, reduce[w]); ++ } ++ reduce[0] = maximum; ++ } ++ __syncthreads(); ++ const float maximum = reduce[0]; ++ ++ float local_sum = 0.0f; ++ for (int j = key_begin + tid; j < kv; j += blockDim.x) { ++ const float weight = __expf(scores[j] - maximum); ++ scores[j] = weight; ++ local_sum += weight; ++ } ++ local_sum = attention_warp_sum(local_sum); ++ if (lane == 0) { ++ reduce[warp] = local_sum; ++ } ++ __syncthreads(); ++ if (tid == 0) { ++ float sum = reduce[0]; ++#pragma unroll ++ for (int w = 1; w < warps; ++w) { ++ sum += reduce[w]; ++ } ++ reduce[0] = sum; ++ } ++ __syncthreads(); ++ ++ const int d = tid & (d_k - 1); ++ const int part = tid / d_k; ++ float context = 0.0f; ++ for (int j = key_begin + part; j < kv; j += value_parts) { ++ context += scores[j] * attention_load_kv( ++ Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d); ++ } ++ value_partials[part * d_k + d] = context; ++ __syncthreads(); ++ if (tid < d_k) { ++ context = value_partials[tid]; ++#pragma unroll ++ for (int part_index = 1; part_index < value_parts; ++part_index) { ++ context += value_partials[part_index * d_k + tid]; ++ } ++ ctx[(size_t) b * o_sb + (size_t) h * o_sh + tid] = context / reduce[0]; ++ } ++} ++ ++// One-block-per-query specialization for d_k=128 and q=2. Score-producing ++// warps retain reduction state to reduce synchronization and scratch use. ++template ++static __global__ void fused_relpos_attn_warp_128_kernel( ++ const float * __restrict__ Q, const T * __restrict__ K, ++ const T * __restrict__ V, const T * __restrict__ Ppos, ++ const float * __restrict__ bu, const float * __restrict__ bv, ++ const float * __restrict__ mask, float * __restrict__ ctx, + const float * __restrict__ Kcache, const float * __restrict__ Vcache, + const int32_t * __restrict__ slot_ids, + const int32_t * __restrict__ ring_heads, + int kv, int cache_len, float scale, - long q_sq, long q_sh, long q_sb, - long k_sj, long k_sh, long k_sb, - long v_sj, long v_sh, long v_sb, ++ long q_sq, long q_sh, long q_sb, ++ long k_sj, long k_sh, long k_sb, ++ long v_sj, long v_sh, long v_sb, + long cache_sj, long cache_ss, - long p_sr, long p_sh, - long o_si, long o_sh, long o_sb, long m_sb) { - extern __shared__ float sh[]; -@@ -86,6 +177,14 @@ static __global__ void fused_relpos_attn_warp_128_kernel( - const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; - const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; - const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; ++ long p_sr, long p_sh, ++ long o_si, long o_sh, long o_sb, long m_sb) { ++ extern __shared__ float sh[]; ++ float * Qu = sh; ++ float * Qv = Qu + RELPOS_ATTN_DK_128; ++ float * sc = Qv + RELPOS_ATTN_DK_128; ++ float * red = sc + kv; ++ ++ const int h = blockIdx.x; ++ const int i = blockIdx.y; ++ const int b = blockIdx.z; ++ const int d = threadIdx.x; ++ const int warp = d >> 5; ++ const int lane = d & 31; ++ ++ const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; ++ const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; ++ const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; + const int slot = Cached ? slot_ids[b] : 0; + const int ring_head = Cached ? ring_heads[b] : 0; + const float * Kch = Cached @@ -215,57 +442,122 @@ index f3c4836a..e0c0ec27 100644 + const float * Vch = Cached + ? Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128 + : nullptr; - const T * Ph = Ppos + (size_t) h * p_sh; - const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; - -@@ -98,7 +197,8 @@ static __global__ void fused_relpos_attn_warp_128_kernel( - const float4 qv4 = *reinterpret_cast(Qv + d4); - float produced_max = -INFINITY; - for (int j = warp; j < kv; j += RELPOS_ATTN_WARPS_128) { -- const float4 k4 = relpos_load4(Kh + (size_t) j * k_sj + d4); -+ const float4 k4 = relpos_load_kv4( ++ const T * Ph = Ppos + (size_t) h * p_sh; ++ const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ ++ Qu[d] = Qhi[d] + bu[h * RELPOS_ATTN_DK_128 + d]; ++ Qv[d] = Qhi[d] + bv[h * RELPOS_ATTN_DK_128 + d]; ++ __syncthreads(); ++ ++ const int d4 = lane * 4; ++ const float4 qu4 = *reinterpret_cast(Qu + d4); ++ const float4 qv4 = *reinterpret_cast(Qv + d4); ++ float produced_max = -INFINITY; ++ for (int j = warp; j < kv; j += RELPOS_ATTN_WARPS_128) { ++ const float4 k4 = attention_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, d4); - const int row = 1 + j - i; - const float4 p4 = relpos_load4(Ph + (size_t) row * p_sr + d4); - float score = -@@ -154,7 +254,9 @@ static __global__ void fused_relpos_attn_warp_128_kernel( - - float context = 0.0f; - for (int j = 0; j < kv; ++j) { -- context += sc[j] * (float) Vh[(size_t) j * v_sj + d]; ++ const int row = 1 + j - i; ++ const float4 p4 = attention_load4(Ph + (size_t) row * p_sr + d4); ++ float score = ++ k4.x * qu4.x + p4.x * qv4.x + ++ k4.y * qu4.y + p4.y * qv4.y + ++ k4.z * qu4.z + p4.z * qv4.z + ++ k4.w * qu4.w + p4.w * qv4.w; ++ score = attention_warp_sum(score); ++ if (lane == 0) { ++ sc[j] = score * scale + (Mb ? Mb[j] : 0.0f); ++ produced_max = fmaxf(produced_max, sc[j]); ++ } ++ } ++ if (lane == 0) { ++ red[warp] = produced_max; ++ } ++ __syncthreads(); ++ ++ if (d == 0) { ++ float maximum = red[0]; ++#pragma unroll ++ for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { ++ maximum = fmaxf(maximum, red[w]); ++ } ++ red[0] = maximum; ++ } ++ __syncthreads(); ++ const float maximum = red[0]; ++ ++ // red[] is reused below. Ensure every warp has captured red[0] first; ++ // otherwise warp 0 can overwrite the maximum while another warp loads it. ++ __syncthreads(); ++ float local_sum = 0.0f; ++ for (int j = d; j < kv; j += RELPOS_ATTN_DK_128) { ++ const float weight = __expf(sc[j] - maximum); ++ sc[j] = weight; ++ local_sum += weight; ++ } ++ local_sum = attention_warp_sum(local_sum); ++ if (lane == 0) { ++ red[warp] = local_sum; ++ } ++ __syncthreads(); ++ if (d == 0) { ++ float sum = red[0]; ++#pragma unroll ++ for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { ++ sum += red[w]; ++ } ++ red[0] = sum; ++ } ++ __syncthreads(); ++ ++ float context = 0.0f; ++ for (int j = 0; j < kv; ++j) { + context += -+ sc[j] * relpos_load_kv( ++ sc[j] * attention_load_kv( + Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d); - } - ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = context / red[0]; - } -@@ -163,16 +265,20 @@ static __global__ void fused_relpos_attn_warp_128_kernel( - // block computes both query rows. K and V are then loaded once and reused; - // only the two relative-position rows differ. The reduction order matches the - // one-query specialization so changing batch size does not change numerics. --template ++ } ++ ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = context / red[0]; ++} ++ ++// Once the one-block-per-query grid spills into another occupancy wave, one ++// block computes both query rows. K and V are then loaded once and reused; ++// only the two relative-position rows differ. The reduction order matches the ++// one-query specialization so changing batch size does not change numerics. +template - static __global__ void fused_relpos_attn_q2_warp_128_kernel( - const float * __restrict__ Q, const T * __restrict__ K, - const T * __restrict__ V, const T * __restrict__ Ppos, - const float * __restrict__ bu, const float * __restrict__ bv, - const float * __restrict__ mask, float * __restrict__ ctx, -- int kv, float scale, ++static __global__ void fused_relpos_attn_q2_warp_128_kernel( ++ const float * __restrict__ Q, const T * __restrict__ K, ++ const T * __restrict__ V, const T * __restrict__ Ppos, ++ const float * __restrict__ bu, const float * __restrict__ bv, ++ const float * __restrict__ mask, float * __restrict__ ctx, + const float * __restrict__ Kcache, const float * __restrict__ Vcache, + const int32_t * __restrict__ slot_ids, + const int32_t * __restrict__ ring_heads, + int kv, int cache_len, float scale, - long q_sq, long q_sh, long q_sb, - long k_sj, long k_sh, long k_sb, - long v_sj, long v_sh, long v_sb, ++ long q_sq, long q_sh, long q_sb, ++ long k_sj, long k_sh, long k_sb, ++ long v_sj, long v_sh, long v_sb, + long cache_sj, long cache_ss, - long p_sr, long p_sh, - long o_si, long o_sh, long o_sb, long m_sb) { - extern __shared__ float sh[]; -@@ -195,6 +301,14 @@ static __global__ void fused_relpos_attn_q2_warp_128_kernel( - const float * Qh1 = Qh0 + q_sq; - const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; - const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; ++ long p_sr, long p_sh, ++ long o_si, long o_sh, long o_sb, long m_sb) { ++ extern __shared__ float sh[]; ++ float * Qu0 = sh; ++ float * Qv0 = Qu0 + RELPOS_ATTN_DK_128; ++ float * Qu1 = Qv0 + RELPOS_ATTN_DK_128; ++ float * Qv1 = Qu1 + RELPOS_ATTN_DK_128; ++ float * sc0 = Qv1 + RELPOS_ATTN_DK_128; ++ float * sc1 = sc0 + kv; ++ float * red0 = sc1 + kv; ++ float * red1 = red0 + RELPOS_ATTN_WARPS_128; ++ ++ const int h = blockIdx.x; ++ const int b = blockIdx.z; ++ const int d = threadIdx.x; ++ const int warp = d >> 5; ++ const int lane = d & 31; ++ ++ const float * Qh0 = Q + (size_t) b * q_sb + (size_t) h * q_sh; ++ const float * Qh1 = Qh0 + q_sq; ++ const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; ++ const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; + const int slot = Cached ? slot_ids[b] : 0; + const int ring_head = Cached ? ring_heads[b] : 0; + const float * Kch = Cached @@ -274,45 +566,118 @@ index f3c4836a..e0c0ec27 100644 + const float * Vch = Cached + ? Vcache + (size_t) slot * cache_ss + (size_t) h * RELPOS_ATTN_DK_128 + : nullptr; - const T * Ph = Ppos + (size_t) h * p_sh; - const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; - -@@ -214,7 +328,8 @@ static __global__ void fused_relpos_attn_q2_warp_128_kernel( - float produced_max0 = -INFINITY; - float produced_max1 = -INFINITY; - for (int j = warp; j < kv; j += RELPOS_ATTN_WARPS_128) { -- const float4 k4 = relpos_load4(Kh + (size_t) j * k_sj + d4); -+ const float4 k4 = relpos_load_kv4( ++ const T * Ph = Ppos + (size_t) h * p_sh; ++ const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ ++ const float bias_u = bu[h * RELPOS_ATTN_DK_128 + d]; ++ const float bias_v = bv[h * RELPOS_ATTN_DK_128 + d]; ++ Qu0[d] = Qh0[d] + bias_u; ++ Qv0[d] = Qh0[d] + bias_v; ++ Qu1[d] = Qh1[d] + bias_u; ++ Qv1[d] = Qh1[d] + bias_v; ++ __syncthreads(); ++ ++ const int d4 = lane * 4; ++ const float4 qu04 = *reinterpret_cast(Qu0 + d4); ++ const float4 qv04 = *reinterpret_cast(Qv0 + d4); ++ const float4 qu14 = *reinterpret_cast(Qu1 + d4); ++ const float4 qv14 = *reinterpret_cast(Qv1 + d4); ++ float produced_max0 = -INFINITY; ++ float produced_max1 = -INFINITY; ++ for (int j = warp; j < kv; j += RELPOS_ATTN_WARPS_128) { ++ const float4 k4 = attention_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, d4); - const float4 p04 = relpos_load4(Ph + (size_t) (j + 1) * p_sr + d4); - const float4 p14 = relpos_load4(Ph + (size_t) j * p_sr + d4); - float score0 = -@@ -292,7 +407,9 @@ static __global__ void fused_relpos_attn_q2_warp_128_kernel( - float context0 = 0.0f; - float context1 = 0.0f; - for (int j = 0; j < kv; ++j) { -- const float value = (float) Vh[(size_t) j * v_sj + d]; ++ const float4 p04 = attention_load4(Ph + (size_t) (j + 1) * p_sr + d4); ++ const float4 p14 = attention_load4(Ph + (size_t) j * p_sr + d4); ++ float score0 = ++ k4.x * qu04.x + p04.x * qv04.x + ++ k4.y * qu04.y + p04.y * qv04.y + ++ k4.z * qu04.z + p04.z * qv04.z + ++ k4.w * qu04.w + p04.w * qv04.w; ++ float score1 = ++ k4.x * qu14.x + p14.x * qv14.x + ++ k4.y * qu14.y + p14.y * qv14.y + ++ k4.z * qu14.z + p14.z * qv14.z + ++ k4.w * qu14.w + p14.w * qv14.w; ++ score0 = attention_warp_sum(score0); ++ score1 = attention_warp_sum(score1); ++ if (lane == 0) { ++ const float additive_mask = Mb ? Mb[j] : 0.0f; ++ sc0[j] = score0 * scale + additive_mask; ++ sc1[j] = score1 * scale + additive_mask; ++ produced_max0 = fmaxf(produced_max0, sc0[j]); ++ produced_max1 = fmaxf(produced_max1, sc1[j]); ++ } ++ } ++ if (lane == 0) { ++ red0[warp] = produced_max0; ++ red1[warp] = produced_max1; ++ } ++ __syncthreads(); ++ ++ if (d == 0) { ++ float maximum0 = red0[0]; ++ float maximum1 = red1[0]; ++#pragma unroll ++ for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { ++ maximum0 = fmaxf(maximum0, red0[w]); ++ maximum1 = fmaxf(maximum1, red1[w]); ++ } ++ red0[0] = maximum0; ++ red1[0] = maximum1; ++ } ++ __syncthreads(); ++ const float maximum0 = red0[0]; ++ const float maximum1 = red1[0]; ++ __syncthreads(); ++ ++ float local_sum0 = 0.0f; ++ float local_sum1 = 0.0f; ++ for (int j = d; j < kv; j += RELPOS_ATTN_DK_128) { ++ const float weight0 = __expf(sc0[j] - maximum0); ++ const float weight1 = __expf(sc1[j] - maximum1); ++ sc0[j] = weight0; ++ sc1[j] = weight1; ++ local_sum0 += weight0; ++ local_sum1 += weight1; ++ } ++ local_sum0 = attention_warp_sum(local_sum0); ++ local_sum1 = attention_warp_sum(local_sum1); ++ if (lane == 0) { ++ red0[warp] = local_sum0; ++ red1[warp] = local_sum1; ++ } ++ __syncthreads(); ++ if (d == 0) { ++ float sum0 = red0[0]; ++ float sum1 = red1[0]; ++#pragma unroll ++ for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { ++ sum0 += red0[w]; ++ sum1 += red1[w]; ++ } ++ red0[0] = sum0; ++ red1[0] = sum1; ++ } ++ __syncthreads(); ++ ++ float context0 = 0.0f; ++ float context1 = 0.0f; ++ for (int j = 0; j < kv; ++j) { + const float value = -+ relpos_load_kv( ++ attention_load_kv( + Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d); - context0 += sc0[j] * value; - context1 += sc1[j] * value; - } -@@ -301,7 +418,573 @@ static __global__ void fused_relpos_attn_q2_warp_128_kernel( - ctx[out + o_si] = context1 / red1[0]; - } - -+// Portable NVIDIA SM80+ exact-shape Q=2/KV=72 specialization. One -+// 256-thread CTA owns a (batch, head) pair and each of its eight warps owns -+// nine consecutive KV columns through -+// both scoring and value accumulation. The two sets of nine softmax weights -+// remain in registers instead of making a score/weight round trip through -+// shared memory. Only per-warp denominators and partial context vectors cross -+// warps. __launch_bounds__ requests four resident blocks per multiprocessor. -+// -+// Max subtraction keeps exponent evaluation stable for the model's score range; -+// the per-warp maxima cross through one small shared-memory exchange. - template ++ context0 += sc0[j] * value; ++ context1 += sc1[j] * value; ++ } ++ const size_t out = (size_t) b * o_sb + (size_t) h * o_sh + d; ++ ctx[out] = context0 / red0[0]; ++ ctx[out + o_si] = context1 / red1[0]; ++} ++ ++// Register-resident Q=2, KV=72 specialization. Each warp owns consecutive ++// cache positions and reuses K/V across both query rows. ++template +__global__ __launch_bounds__(256, 4) void fused_relpos_attn_q2_register_kernel( + const float * __restrict__ Q, const T * __restrict__ K, + const T * __restrict__ V, const T * __restrict__ Ppos, @@ -357,10 +722,10 @@ index f3c4836a..e0c0ec27 100644 + const T * Ph = Ppos + (size_t) h * p_sh; + const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; + -+ const float4 q04 = relpos_load4(Qh0 + d4); -+ const float4 q14 = relpos_load4(Qh1 + d4); -+ const float4 bu4 = relpos_load4(bu + (size_t) h * RELPOS_ATTN_DK_128 + d4); -+ const float4 bv4 = relpos_load4(bv + (size_t) h * RELPOS_ATTN_DK_128 + d4); ++ const float4 q04 = attention_load4(Qh0 + d4); ++ const float4 q14 = attention_load4(Qh1 + d4); ++ const float4 bu4 = attention_load4(bu + (size_t) h * RELPOS_ATTN_DK_128 + d4); ++ const float4 bv4 = attention_load4(bv + (size_t) h * RELPOS_ATTN_DK_128 + d4); + const float4 qc0 = make_float4( + q04.x + bu4.x, q04.y + bu4.y, q04.z + bu4.z, q04.w + bu4.w); + const float4 qc1 = make_float4( @@ -374,20 +739,20 @@ index f3c4836a..e0c0ec27 100644 + float scores1[keys_per_warp]; + float local_max0 = -INFINITY; + float local_max1 = -INFINITY; -+ float4 p14 = relpos_load4(Ph + (size_t) j0 * p_sr + d4); ++ float4 p14 = attention_load4(Ph + (size_t) j0 * p_sr + d4); +#pragma unroll + for (int c = 0; c < keys_per_warp; ++c) { + const int j = j0 + c; -+ const float4 p04 = relpos_load4(Ph + (size_t) (j + 1) * p_sr + d4); ++ const float4 p04 = attention_load4(Ph + (size_t) (j + 1) * p_sr + d4); + float4 k4; + if (c < 7 || j < cache_len) { + int physical_j = ring_head + j; + if (physical_j >= cache_len) { + physical_j -= cache_len; + } -+ k4 = relpos_load4(Kch + (size_t) physical_j * cache_sj + d4); ++ k4 = attention_load4(Kch + (size_t) physical_j * cache_sj + d4); + } else { -+ k4 = relpos_load4(Kh + (size_t) (j - cache_len) * k_sj + d4); ++ k4 = attention_load4(Kh + (size_t) (j - cache_len) * k_sj + d4); + } + float score0 = + k4.x * qc0.x + p04.x * qp0.x + @@ -399,8 +764,8 @@ index f3c4836a..e0c0ec27 100644 + k4.y * qc1.y + p14.y * qp1.y + + k4.z * qc1.z + p14.z * qp1.z + + k4.w * qc1.w + p14.w * qp1.w; -+ score0 = relpos_warp_all_sum(score0); -+ score1 = relpos_warp_all_sum(score1); ++ score0 = attention_warp_all_sum(score0); ++ score1 = attention_warp_all_sum(score1); + const float additive_mask = Mb ? Mb[j] : 0.0f; + score0 = score0 * scale + additive_mask; + score1 = score1 * scale + additive_mask; @@ -452,9 +817,9 @@ index f3c4836a..e0c0ec27 100644 + if (physical_j >= cache_len) { + physical_j -= cache_len; + } -+ v4 = relpos_load4(Vch + (size_t) physical_j * cache_sj + d4); ++ v4 = attention_load4(Vch + (size_t) physical_j * cache_sj + d4); + } else { -+ v4 = relpos_load4(Vh + (size_t) (j - cache_len) * v_sj + d4); ++ v4 = attention_load4(Vh + (size_t) (j - cache_len) * v_sj + d4); + } + acc0.x = fmaf(scores0[c], v4.x, acc0.x); + acc0.y = fmaf(scores0[c], v4.y, acc0.y); @@ -490,9 +855,8 @@ index f3c4836a..e0c0ec27 100644 + context / denominator; +} + -+// Nemotron 3.5 R=3 specialization: Q=4, cache=56, KV=60. Six warps own -+// ten consecutive keys each and reuse every K/V load across all four query -+// rows. Scores and softmax weights remain register-resident. ++// Register-resident Q=4, cache=56, KV=60 specialization. Each warp reuses K/V ++// across all query rows. +template +__global__ __launch_bounds__(192, 3) void fused_relpos_attn_q4_register_kernel( + const float * __restrict__ Q, const T * __restrict__ K, @@ -541,12 +905,12 @@ index f3c4836a..e0c0ec27 100644 + const T * Ph = Ppos + (size_t) h * p_sh; + const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; + -+ const float4 q04 = relpos_load4(Qh0 + d4); -+ const float4 q14 = relpos_load4(Qh1 + d4); -+ const float4 q24 = relpos_load4(Qh2 + d4); -+ const float4 q34 = relpos_load4(Qh3 + d4); -+ const float4 bu4 = relpos_load4(bu + (size_t) h * RELPOS_ATTN_DK_128 + d4); -+ const float4 bv4 = relpos_load4(bv + (size_t) h * RELPOS_ATTN_DK_128 + d4); ++ const float4 q04 = attention_load4(Qh0 + d4); ++ const float4 q14 = attention_load4(Qh1 + d4); ++ const float4 q24 = attention_load4(Qh2 + d4); ++ const float4 q34 = attention_load4(Qh3 + d4); ++ const float4 bu4 = attention_load4(bu + (size_t) h * RELPOS_ATTN_DK_128 + d4); ++ const float4 bv4 = attention_load4(bv + (size_t) h * RELPOS_ATTN_DK_128 + d4); +#define RELPOS_ADD_BIAS(q, bias) \ + make_float4((q).x + (bias).x, (q).y + (bias).y, (q).z + (bias).z, (q).w + (bias).w) + const float4 qc0 = RELPOS_ADD_BIAS(q04, bu4); @@ -570,19 +934,19 @@ index f3c4836a..e0c0ec27 100644 +#pragma unroll + for (int c = 0; c < keys_per_warp; ++c) { + const int j = j0 + c; -+ const float4 k4 = relpos_load_kv4( ++ const float4 k4 = attention_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, d4); -+ const float4 p04 = relpos_load4(Ph + (size_t) (j + 3) * p_sr + d4); -+ const float4 p14 = relpos_load4(Ph + (size_t) (j + 2) * p_sr + d4); -+ const float4 p24 = relpos_load4(Ph + (size_t) (j + 1) * p_sr + d4); -+ const float4 p34 = relpos_load4(Ph + (size_t) j * p_sr + d4); ++ const float4 p04 = attention_load4(Ph + (size_t) (j + 3) * p_sr + d4); ++ const float4 p14 = attention_load4(Ph + (size_t) (j + 2) * p_sr + d4); ++ const float4 p24 = attention_load4(Ph + (size_t) (j + 1) * p_sr + d4); ++ const float4 p34 = attention_load4(Ph + (size_t) j * p_sr + d4); +#define RELPOS_SCORE(k, qc, p, qp) \ + ((k).x * (qc).x + (p).x * (qp).x + (k).y * (qc).y + (p).y * (qp).y + \ + (k).z * (qc).z + (p).z * (qp).z + (k).w * (qc).w + (p).w * (qp).w) -+ float score0 = relpos_warp_all_sum(RELPOS_SCORE(k4, qc0, p04, qp0)); -+ float score1 = relpos_warp_all_sum(RELPOS_SCORE(k4, qc1, p14, qp1)); -+ float score2 = relpos_warp_all_sum(RELPOS_SCORE(k4, qc2, p24, qp2)); -+ float score3 = relpos_warp_all_sum(RELPOS_SCORE(k4, qc3, p34, qp3)); ++ float score0 = attention_warp_all_sum(RELPOS_SCORE(k4, qc0, p04, qp0)); ++ float score1 = attention_warp_all_sum(RELPOS_SCORE(k4, qc1, p14, qp1)); ++ float score2 = attention_warp_all_sum(RELPOS_SCORE(k4, qc2, p24, qp2)); ++ float score3 = attention_warp_all_sum(RELPOS_SCORE(k4, qc3, p34, qp3)); +#undef RELPOS_SCORE + const float additive_mask = Mb ? Mb[j] : 0.0f; + score0 = score0 * scale + additive_mask; @@ -648,7 +1012,7 @@ index f3c4836a..e0c0ec27 100644 +#pragma unroll + for (int c = 0; c < keys_per_warp; ++c) { + const int j = j0 + c; -+ const float4 v4 = relpos_load_kv4( ++ const float4 v4 = attention_load_kv4( + Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d4); +#define RELPOS_ACCUM(acc, weight, value) \ + do { \ @@ -688,9 +1052,8 @@ index f3c4836a..e0c0ec27 100644 + } +} + -+// Register-resident paired-query path for the cache-aware geometries exposed -+// by the Nemotron streaming models. Pairing halves the CTA count for larger -+// right-context presets while reusing every K/V load across two query rows. ++// Paired-query path for common cache-aware geometries. Pairing reduces the ++// block count and reuses K/V across query rows. +template +__global__ __launch_bounds__(256, 4) void fused_relpos_attn_common_q2_register_kernel( + const float * __restrict__ Q, const T * __restrict__ K, @@ -739,10 +1102,10 @@ index f3c4836a..e0c0ec27 100644 + const T * Ph = Ppos + (size_t) h * p_sh; + const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; + -+ const float4 q04 = relpos_load4(Qh0 + d4); -+ const float4 q14 = relpos_load4(Qh1 + d4); -+ const float4 bu4 = relpos_load4(bu + (size_t) h * RELPOS_ATTN_DK_128 + d4); -+ const float4 bv4 = relpos_load4(bv + (size_t) h * RELPOS_ATTN_DK_128 + d4); ++ const float4 q04 = attention_load4(Qh0 + d4); ++ const float4 q14 = attention_load4(Qh1 + d4); ++ const float4 bu4 = attention_load4(bu + (size_t) h * RELPOS_ATTN_DK_128 + d4); ++ const float4 bv4 = attention_load4(bv + (size_t) h * RELPOS_ATTN_DK_128 + d4); +#define RELPOS_ADD_BIAS(q, bias) \ + make_float4((q).x + (bias).x, (q).y + (bias).y, (q).z + (bias).z, (q).w + (bias).w) + const float4 qc0 = RELPOS_ADD_BIAS(q04, bu4); @@ -761,17 +1124,17 @@ index f3c4836a..e0c0ec27 100644 + float score0 = -INFINITY; + float score1 = -INFINITY; + if (j < kv_len) { -+ const float4 k4 = relpos_load_kv4( ++ const float4 k4 = attention_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, d4); + const int pos0 = j + q_len - 1 - query0; -+ const float4 p04 = relpos_load4(Ph + (size_t) pos0 * p_sr + d4); ++ const float4 p04 = attention_load4(Ph + (size_t) pos0 * p_sr + d4); +#define RELPOS_SCORE(k, qc, p, qp) \ + ((k).x * (qc).x + (p).x * (qp).x + (k).y * (qc).y + (p).y * (qp).y + \ + (k).z * (qc).z + (p).z * (qp).z + (k).w * (qc).w + (p).w * (qp).w) -+ score0 = relpos_warp_all_sum(RELPOS_SCORE(k4, qc0, p04, qp0)); ++ score0 = attention_warp_all_sum(RELPOS_SCORE(k4, qc0, p04, qp0)); + if (has_query1) { -+ const float4 p14 = relpos_load4(Ph + (size_t) (pos0 - 1) * p_sr + d4); -+ score1 = relpos_warp_all_sum(RELPOS_SCORE(k4, qc1, p14, qp1)); ++ const float4 p14 = attention_load4(Ph + (size_t) (pos0 - 1) * p_sr + d4); ++ score1 = attention_warp_all_sum(RELPOS_SCORE(k4, qc1, p14, qp1)); + } +#undef RELPOS_SCORE + const float additive_mask = Mb ? Mb[j] : 0.0f; @@ -826,7 +1189,7 @@ index f3c4836a..e0c0ec27 100644 + const int j = j0 + c; + if (j >= kv_len) + continue; -+ const float4 v4 = relpos_load_kv4( ++ const float4 v4 = attention_load_kv4( + Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d4); +#define RELPOS_ACCUM(acc, weight, value) \ + do { \ @@ -860,61 +1223,55 @@ index f3c4836a..e0c0ec27 100644 + } +} + -+static bool relpos_attn_register_resident_enabled() { -+ static const bool enabled = []() { -+ const char * value = std::getenv("GGML_CUDA_RELPOS_REGISTER_RESIDENT"); -+ return value == nullptr || std::atoi(value) != 0; -+ }(); -+ return enabled; -+} -+ +template - static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { - // The target shape fixes dynamic shared memory, so occupancy is invariant - // for a given compiled kernel and device. Avoid repeating the CUDA runtime -@@ -310,27 +993,32 @@ static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { - int blocks = cached[device].load(std::memory_order_relaxed); - if (blocks == 0) { - CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( -- &blocks, fused_relpos_attn_warp_128_kernel, RELPOS_ATTN_DK_128, shmem)); ++static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { ++ // Dynamic shared memory is fixed for a given kernel shape, so cache the ++ // occupancy query per device. ++ static std::atomic cached[GGML_CUDA_MAX_DEVICES] = {}; ++ int blocks = cached[device].load(std::memory_order_relaxed); ++ if (blocks == 0) { ++ CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &blocks, fused_relpos_attn_warp_128_kernel, RELPOS_ATTN_DK_128, shmem)); - GGML_ASSERT(blocks > 0); - cached[device].store(blocks, std::memory_order_relaxed); - } - return blocks; - } - --template ++ GGML_ASSERT(blocks > 0); ++ cached[device].store(blocks, std::memory_order_relaxed); ++ } ++ return blocks; ++} ++ +template - static __global__ void fused_relpos_attn_kernel( - const float * __restrict__ Q, const T * __restrict__ K, - const T * __restrict__ V, const T * __restrict__ Ppos, - const float * __restrict__ bu, const float * __restrict__ bv, - const float * __restrict__ mask, float * __restrict__ ctx, -- int q, int kv, int n_head, float scale, ++static __global__ void fused_attention_kernel( ++ const float * __restrict__ Q, const T * __restrict__ K, ++ const T * __restrict__ V, const T * __restrict__ Ppos, ++ const float * __restrict__ bu, const float * __restrict__ bv, ++ const float * __restrict__ mask, float * __restrict__ ctx, + const float * __restrict__ Kcache, const float * __restrict__ Vcache, + const int32_t * __restrict__ slot_ids, + const int32_t * __restrict__ ring_heads, + const int32_t * __restrict__ active_lengths, + int q, int kv, int n_head, int cache_len, float scale, bool relative, - // element strides: x_sq = between queries/keys, x_sh = between heads, - // x_sb = between batch items - long q_sq, long q_sh, long q_sb, - long k_sj, long k_sh, long k_sb, - long v_sj, long v_sh, long v_sb, ++ // element strides: x_sq = between queries/keys, x_sh = between heads, ++ // x_sb = between batch items ++ long q_sq, long q_sh, long q_sb, ++ long k_sj, long k_sh, long k_sb, ++ long v_sj, long v_sh, long v_sb, + long cache_sj, long cache_ss, - long p_sr, long p_sh, -- long o_si, long o_sh, long o_sb, long m_sb) { ++ long p_sr, long p_sh, + long o_si, long o_sh, long o_sb, long m_sq, long m_sb) { - extern __shared__ float sh[]; - const int dk = blockDim.x; - float * Qu = sh; // [dk] -@@ -346,11 +1034,20 @@ static __global__ void fused_relpos_attn_kernel( - const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; - const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; - const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; -- const T * Ph = Ppos + (size_t) h * p_sh; -- const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; ++ extern __shared__ float sh[]; ++ const int dk = blockDim.x; ++ float * Qu = sh; // [dk] ++ float * Qv = sh + dk; // [dk] ++ float * sc = sh + 2 * dk; // [kv] ++ float * red = sh + 2 * dk + kv; // [dk] reduction scratch ++ ++ const int h = blockIdx.x; // head ++ const int i = blockIdx.y; // query ++ const int b = blockIdx.z; // batch ++ const int d = threadIdx.x; // head dim 0..dk-1 ++ ++ const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; ++ const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; ++ const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; + const int slot = Cached ? slot_ids[b] : 0; + const int ring_head = Cached ? ring_heads[b] : 0; + const int key_begin = active_lengths ? cache_len - active_lengths[b] : 0; @@ -926,63 +1283,43 @@ index f3c4836a..e0c0ec27 100644 + : nullptr; + const T * Ph = relative ? Ppos + (size_t) h * p_sh : nullptr; + const float * Mb = mask ? mask + (size_t) b * m_sb + (size_t) i * m_sq : nullptr; - -- Qu[d] = Qhi[d] + bu[h * dk + d]; -- Qv[d] = Qhi[d] + bv[h * dk + d]; ++ + Qu[d] = Qhi[d] + (relative ? bu[h * dk + d] : 0.0f); + Qv[d] = Qhi[d] + (relative ? bv[h * dk + d] : 0.0f); - __syncthreads(); - - // scores. Two layouts: -@@ -363,32 +1060,19 @@ static __global__ void fused_relpos_attn_kernel( - // * otherwise: legacy thread-per-key scalar loop. - if (dk == 128) { - const int warp = d >> 5, lane = d & 31, nw = dk >> 5; -- for (int j = warp; j < kv; j += nw) { -- const T * Kj = Kh + (size_t) j * k_sj + lane * 4; ++ __syncthreads(); ++ ++ // scores. Two layouts: ++ // * dk == 128 (the FastConformer case): WARP-COOPERATIVE — each warp owns ++ // a key j and the 32 lanes split the 128 dims 4-a-piece with one ++ // vectorized row load + shuffle reduction. The original ++ // thread-per-key loop left dk-kv threads idle (kv ~= 50 < 128) and ++ // issued dk scalar loads per row, which made the kernel ++ // load-issue-bound (F16 operands alone changed nothing). ++ // * otherwise: legacy thread-per-key scalar loop. ++ if (dk == 128) { ++ const int warp = d >> 5, lane = d & 31, nw = dk >> 5; + for (int j = key_begin + warp; j < kv; j += nw) { - const int row = (q - 1) + j - i; // rel-shift index -- const T * Pr = Ph + (size_t) row * p_sr + lane * 4; -- float k4[4], p4[4]; -- if (sizeof(T) == 2) { -- const int2 kr = *(const int2 *) Kj; -- const int2 pr = *(const int2 *) Pr; -- const half2 * kh = (const half2 *) &kr; -- const half2 * ph = (const half2 *) ≺ -- k4[0] = __low2float(kh[0]); k4[1] = __high2float(kh[0]); -- k4[2] = __low2float(kh[1]); k4[3] = __high2float(kh[1]); -- p4[0] = __low2float(ph[0]); p4[1] = __high2float(ph[0]); -- p4[2] = __low2float(ph[1]); p4[3] = __high2float(ph[1]); -- } else { -- const float4 kr = *(const float4 *) Kj; -- const float4 pr = *(const float4 *) Pr; -- k4[0] = ((const float *) &kr)[0]; k4[1] = ((const float *) &kr)[1]; -- k4[2] = ((const float *) &kr)[2]; k4[3] = ((const float *) &kr)[3]; -- p4[0] = ((const float *) &pr)[0]; p4[1] = ((const float *) &pr)[1]; -- p4[2] = ((const float *) &pr)[2]; p4[3] = ((const float *) &pr)[3]; -- } -- float s = 0.0f; --#pragma unroll -- for (int e = 0; e < 4; e++) { -- s += k4[e] * Qu[lane * 4 + e] + p4[e] * Qv[lane * 4 + e]; ++ const int row = (q - 1) + j - i; // rel-shift index + const float4 k4 = -+ relpos_load_kv4( ++ attention_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, lane * 4); + float s = + k4.x * Qu[lane * 4 + 0] + k4.y * Qu[lane * 4 + 1] + + k4.z * Qu[lane * 4 + 2] + k4.w * Qu[lane * 4 + 3]; + if (relative) { + const T * Pr = Ph + (size_t) row * p_sr + lane * 4; -+ const float4 p4 = relpos_load4(Pr); ++ const float4 p4 = attention_load4(Pr); + s += p4.x * Qv[lane * 4 + 0] + p4.y * Qv[lane * 4 + 1] + + p4.z * Qv[lane * 4 + 2] + p4.w * Qv[lane * 4 + 3]; - } - #pragma unroll - for (int off = 16; off > 0; off >>= 1) { -@@ -398,15 +1082,39 @@ static __global__ void fused_relpos_attn_kernel( - sc[j] = s * scale + (Mb ? Mb[j] : 0.0f); - } - } ++ } ++#pragma unroll ++ for (int off = 16; off > 0; off >>= 1) { ++ s += __shfl_xor_sync(0xffffffff, s, off); ++ } ++ if (lane == 0) { ++ sc[j] = s * scale + (Mb ? Mb[j] : 0.0f); ++ } ++ } + } else if (dk == 64) { + // Preserve one thread per key (important for Magpie's long cache) while reducing each + // dot product from 64 scalar loads/instructions to 16 aligned float4 operations. @@ -992,107 +1329,108 @@ index f3c4836a..e0c0ec27 100644 + float bd = 0.0f; +#pragma unroll + for (int dd = 0; dd < 64; dd += 4) { -+ const float4 k4 = relpos_load_kv4( ++ const float4 k4 = attention_load_kv4( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, dd); -+ const float4 qu4 = relpos_load4(Qu + dd); ++ const float4 qu4 = attention_load4(Qu + dd); + ac += k4.x * qu4.x + k4.y * qu4.y + k4.z * qu4.z + k4.w * qu4.w; + if (relative) { -+ const float4 p4 = relpos_load4(Ph + (size_t) row * p_sr + dd); -+ const float4 qv4 = relpos_load4(Qv + dd); ++ const float4 p4 = attention_load4(Ph + (size_t) row * p_sr + dd); ++ const float4 qv4 = attention_load4(Qv + dd); + bd += p4.x * qv4.x + p4.y * qv4.y + p4.z * qv4.z + p4.w * qv4.w; + } + } + sc[j] = (ac + bd) * scale + (Mb ? Mb[j] : 0.0f); + } - } else { -- for (int j = d; j < kv; j += dk) { -- const T * Kj = Kh + (size_t) j * k_sj; ++ } else { + for (int j = key_begin + d; j < kv; j += dk) { - const int row = (q - 1) + j - i; // rel-shift index -- const T * Pr = Ph + (size_t) row * p_sr; - float ac = 0.0f, bd = 0.0f; - for (int dd = 0; dd < dk; dd++) { -- ac += (float) Kj[dd] * Qu[dd]; -- bd += (float) Pr[dd] * Qv[dd]; -+ ac += relpos_load_kv( ++ const int row = (q - 1) + j - i; // rel-shift index ++ float ac = 0.0f, bd = 0.0f; ++ for (int dd = 0; dd < dk; dd++) { ++ ac += attention_load_kv( + Kh, Kch, j, cache_len, ring_head, k_sj, cache_sj, dd) * + Qu[dd]; + if (relative) { + const T * Pr = Ph + (size_t) row * p_sr; + bd += (float) Pr[dd] * Qv[dd]; + } - } - sc[j] = (ac + bd) * scale + (Mb ? Mb[j] : 0.0f); - } -@@ -415,7 +1123,7 @@ static __global__ void fused_relpos_attn_kernel( - - // block max over sc[0..kv) - float lm = -INFINITY; -- for (int j = d; j < kv; j += dk) lm = fmaxf(lm, sc[j]); ++ } ++ sc[j] = (ac + bd) * scale + (Mb ? Mb[j] : 0.0f); ++ } ++ } ++ __syncthreads(); ++ ++ // block max over sc[0..kv) ++ float lm = -INFINITY; + for (int j = key_begin + d; j < kv; j += dk) lm = fmaxf(lm, sc[j]); - red[d] = lm; - __syncthreads(); - for (int s = dk / 2; s > 0; s >>= 1) { -@@ -427,7 +1135,7 @@ static __global__ void fused_relpos_attn_kernel( - - // exp + block sum - float ls = 0.0f; -- for (int j = d; j < kv; j += dk) { ++ red[d] = lm; ++ __syncthreads(); ++ for (int s = dk / 2; s > 0; s >>= 1) { ++ if (d < s) red[d] = fmaxf(red[d], red[d + s]); ++ __syncthreads(); ++ } ++ const float m = red[0]; ++ __syncthreads(); ++ ++ // exp + block sum ++ float ls = 0.0f; + for (int j = key_begin + d; j < kv; j += dk) { - const float e = __expf(sc[j] - m); - sc[j] = e; - ls += e; -@@ -443,7 +1151,10 @@ static __global__ void fused_relpos_attn_kernel( - - // ctx[d] = inv * sum_j softmax(sc[j]) * V[j,d] (thread d owns output dim d) - float c = 0.0f; -- for (int j = 0; j < kv; j++) c += sc[j] * (float) Vh[(size_t) j * v_sj + d]; ++ const float e = __expf(sc[j] - m); ++ sc[j] = e; ++ ls += e; ++ } ++ red[d] = ls; ++ __syncthreads(); ++ for (int s = dk / 2; s > 0; s >>= 1) { ++ if (d < s) red[d] += red[d + s]; ++ __syncthreads(); ++ } ++ const float inv = 1.0f / red[0]; ++ __syncthreads(); ++ ++ // ctx[d] = inv * sum_j softmax(sc[j]) * V[j,d] (thread d owns output dim d) ++ float c = 0.0f; + for (int j = key_begin; j < kv; j++) { -+ c += sc[j] * relpos_load_kv( ++ c += sc[j] * attention_load_kv( + Vh, Vch, j, cache_len, ring_head, v_sj, cache_sj, d); + } - ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = c * inv; - } - -@@ -455,38 +1166,76 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - const ggml_tensor * bias_u = dst->src[4]; - const ggml_tensor * bias_v = dst->src[5]; - const ggml_tensor * mask = dst->src[6]; // may be null ++ ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = c * inv; ++} ++ ++void ggml_cuda_op_fused_attention(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ++ const ggml_tensor * q = dst->src[0]; ++ const ggml_tensor * k = dst->src[1]; ++ const ggml_tensor * v = dst->src[2]; ++ const ggml_tensor * p = dst->src[3]; ++ const ggml_tensor * bias_u = dst->src[4]; ++ const ggml_tensor * bias_v = dst->src[5]; ++ const ggml_tensor * mask = dst->src[6]; // may be null + const ggml_tensor * kv_cache = dst->src[7]; + const ggml_tensor * slot_ids = dst->src[8]; -+ const ggml_tensor * ring_heads = dst->src[9]; ++ const ggml_tensor * cache_state = dst->src[9]; + const bool cached = kv_cache != nullptr; + const bool relative = p != nullptr; - - GGML_ASSERT(q->type == GGML_TYPE_F32); - // K/V/P may be F32 or (all together) F16 — see the kernel comment. -- GGML_ASSERT(k->type == v->type && k->type == p->type); ++ ++ GGML_ASSERT(q->type == GGML_TYPE_F32); ++ // K/V/P may be F32 or (all together) F16 — see the kernel comment. + GGML_ASSERT(k->type == v->type); + GGML_ASSERT(!relative || k->type == p->type); - GGML_ASSERT(k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16); -- GGML_ASSERT(bias_u->type == GGML_TYPE_F32 && bias_v->type == GGML_TYPE_F32); ++ GGML_ASSERT(k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16); + GGML_ASSERT(!relative || (bias_u->type == GGML_TYPE_F32 && bias_v->type == GGML_TYPE_F32)); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - const int d_k = q->ne[0]; - const int q_len = q->ne[1]; - const int n_head = q->ne[2]; - const int batch = q->ne[3]; -- const int kv_len = k->ne[1]; ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ ++ const int d_k = q->ne[0]; ++ const int q_len = q->ne[1]; ++ const int n_head = q->ne[2]; ++ const int batch = q->ne[3]; + const int chunk_len = k->ne[1]; + const int cache_len = cached ? ggml_get_op_params_i32(dst, 1) : 0; + const int kv_len = cache_len + chunk_len; - - GGML_ASSERT((d_k & (d_k - 1)) == 0 && "fused_relpos_attn: d_k must be a power of two"); -- GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k && p->ne[0] == d_k); -- GGML_ASSERT(k->ne[1] == kv_len && v->ne[1] == kv_len); -- GGML_ASSERT(k->ne[2] == n_head && v->ne[2] == n_head && p->ne[2] == n_head); ++ ++ GGML_ASSERT((d_k & (d_k - 1)) == 0 && "fused_attention: d_k must be a power of two"); + GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k); + GGML_ASSERT(k->ne[1] == chunk_len && v->ne[1] == chunk_len); + GGML_ASSERT(k->ne[2] == n_head && v->ne[2] == n_head); - GGML_ASSERT(k->ne[3] == batch && v->ne[3] == batch); -- GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); -- GGML_ASSERT(bias_u->ne[1] == n_head && bias_v->ne[1] == n_head); ++ GGML_ASSERT(k->ne[3] == batch && v->ne[3] == batch); + if (relative) { + GGML_ASSERT(p->ne[0] == d_k && p->ne[2] == n_head); + GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); @@ -1100,27 +1438,23 @@ index f3c4836a..e0c0ec27 100644 + } + if (cached) { + GGML_ASSERT(slot_ids != nullptr); -+ GGML_ASSERT(ring_heads != nullptr); ++ GGML_ASSERT(cache_state != nullptr); + GGML_ASSERT(kv_cache->type == GGML_TYPE_F32 && ggml_is_contiguous(kv_cache)); + GGML_ASSERT(kv_cache->ne[0] == (int64_t) d_k * n_head * cache_len); + GGML_ASSERT(kv_cache->ne[2] == 2 && kv_cache->ne[3] == 1); + GGML_ASSERT(slot_ids->type == GGML_TYPE_I32 && ggml_is_contiguous(slot_ids)); + GGML_ASSERT(slot_ids->ne[0] == batch); -+ GGML_ASSERT(ring_heads->type == GGML_TYPE_I32 && ggml_is_contiguous(ring_heads)); -+ GGML_ASSERT(ring_heads->ne[0] == batch); -+ GGML_ASSERT(ring_heads->ne[1] == 1 || ring_heads->ne[1] == 2); ++ GGML_ASSERT(cache_state->type == GGML_TYPE_I32 && ggml_is_contiguous(cache_state)); ++ GGML_ASSERT(cache_state->ne[0] == batch); ++ GGML_ASSERT(cache_state->ne[1] == 1 || cache_state->ne[1] == 2); + } else { + GGML_ASSERT(slot_ids == nullptr); -+ GGML_ASSERT(ring_heads == nullptr); ++ GGML_ASSERT(cache_state == nullptr); + } - if (mask != nullptr) { - GGML_ASSERT(mask->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(mask)); - GGML_ASSERT(mask->ne[0] == kv_len); -- // Shared across the batch (ne[1]==1) or one key-mask column per -- // stream (ne[1]==batch, the cache-aware layout). -- GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == batch); -- GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); ++ if (mask != nullptr) { ++ GGML_ASSERT(mask->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_is_contiguous(mask)); ++ GGML_ASSERT(mask->ne[0] == kv_len); + if (cached) { + // Streaming: shared across the batch or one key-mask column per + // stream. @@ -1131,44 +1465,53 @@ index f3c4836a..e0c0ec27 100644 + GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == q_len); + GGML_ASSERT(mask->ne[2] == 1 && (mask->ne[3] == 1 || mask->ne[3] == batch)); + } - } ++ } + const long m_sq = + (mask != nullptr && !cached && mask->ne[1] == q_len && q_len > 1) + ? (long) (mask->nb[1] / sizeof(float)) + : 0; - const long m_sb = -- (mask != nullptr && mask->ne[1] == batch && batch > 1) ? (long) (mask->nb[1] / sizeof(float)) : 0; ++ const long m_sb = + (mask != nullptr && batch > 1 && ((cached && mask->ne[1] == batch) || (!cached && mask->ne[3] == batch))) + ? (long) ((cached ? mask->nb[1] : mask->nb[3]) / sizeof(float)) + : 0; - - float scale; - memcpy(&scale, dst->op_params, sizeof(scale)); -@@ -501,7 +1250,11 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - k_sb = (long)(k->nb[3] / ke); - const long v_sj = (long)(v->nb[1] / ke), v_sh = (long)(v->nb[2] / ke), - v_sb = (long)(v->nb[3] / ke); -- const long p_sr = (long)(p->nb[1] / ke), p_sh = (long)(p->nb[2] / ke); ++ ++ float scale; ++ memcpy(&scale, dst->op_params, sizeof(scale)); ++ ++ // Element strides from tensor byte strides. d_k rows must be contiguous ++ // (constructor invariant), everything else is free-form. ++ const size_t qe = ggml_type_size(q->type); ++ const size_t ke = ggml_type_size(k->type); ++ const long q_sq = (long)(q->nb[1] / qe), q_sh = (long)(q->nb[2] / qe), ++ q_sb = (long)(q->nb[3] / qe); ++ const long k_sj = (long)(k->nb[1] / ke), k_sh = (long)(k->nb[2] / ke), ++ k_sb = (long)(k->nb[3] / ke); ++ const long v_sj = (long)(v->nb[1] / ke), v_sh = (long)(v->nb[2] / ke), ++ v_sb = (long)(v->nb[3] / ke); + const long cache_sj = (long) d_k * n_head; + const long cache_ss = cached ? (long) (kv_cache->nb[1] / sizeof(float)) : 0; + const long cache_sp = cached ? (long) (kv_cache->nb[2] / sizeof(float)) : 0; + const long p_sr = relative ? (long)(p->nb[1] / ke) : 0; + const long p_sh = relative ? (long)(p->nb[2] / ke) : 0; - const size_t oe = ggml_type_size(dst->type); - const long o_si = (long)(dst->nb[1] / oe), o_sh = (long)(dst->nb[2] / oe), - o_sb = (long)(dst->nb[3] / oe); -@@ -513,6 +1266,35 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - GGML_ASSERT(shmem <= max_shmem && "fused_relpos_attn: kv window too large for shared memory"); - - cudaStream_t stream = ctx.stream(); ++ const size_t oe = ggml_type_size(dst->type); ++ const long o_si = (long)(dst->nb[1] / oe), o_sh = (long)(dst->nb[2] / oe), ++ o_sb = (long)(dst->nb[3] / oe); ++ ++ const int device = ggml_cuda_get_device(); ++ const auto & device_info = ggml_cuda_info().devices[device]; ++ const size_t shmem = ((size_t) 3 * d_k + kv_len) * sizeof(float); ++ const size_t max_shmem = device_info.smpb; ++ GGML_ASSERT(shmem <= max_shmem && "fused_attention: kv window too large for shared memory"); ++ ++ cudaStream_t stream = ctx.stream(); + const float * cache_k = cached ? (const float *) kv_cache->data : nullptr; + const float * cache_v = cached ? cache_k + cache_sp : nullptr; + const int32_t * active_slots = cached ? (const int32_t *) slot_ids->data : nullptr; + const int32_t * active_ring_heads = -+ cached ? (const int32_t *) ring_heads->data : nullptr; ++ cached ? (const int32_t *) cache_state->data : nullptr; + const int32_t * active_lengths = -+ cached && ring_heads->ne[1] == 2 -+ ? active_ring_heads + ring_heads->ne[0] ++ cached && cache_state->ne[1] == 2 ++ ? active_ring_heads + cache_state->ne[0] + : nullptr; + + auto update_cache = [&]() { @@ -1177,48 +1520,72 @@ index f3c4836a..e0c0ec27 100644 + } + const dim3 update_grid((d_k * n_head + 255) / 256, batch, 2); + if (k->type == GGML_TYPE_F16) { -+ fused_relpos_attn_update_cache_kernel<<>>( ++ fused_attention_update_cache_kernel<<>>( + (float *) kv_cache->data, (const half *) k->data, (const half *) v->data, + active_slots, active_ring_heads, cache_len, chunk_len, d_k * n_head, d_k, + cache_ss, cache_sp, + k_sj, k_sh, k_sb, v_sj, v_sh, v_sb); + } else { -+ fused_relpos_attn_update_cache_kernel<<>>( ++ fused_attention_update_cache_kernel<<>>( + (float *) kv_cache->data, (const float *) k->data, (const float *) v->data, + active_slots, active_ring_heads, cache_len, chunk_len, d_k * n_head, d_k, + cache_ss, cache_sp, + k_sj, k_sh, k_sb, v_sj, v_sh, v_sb); + } + }; - - // The cache-aware Nemotron streaming geometry is Q=2, KV=72, H=8, - // d_k=128. On SM100 the one-query warp kernel is fastest while its grid -@@ -520,9 +1302,117 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - // fuse both query rows: halving the block count and reusing K/V then wins. - // cudaOccupancyMaxActiveBlocksPerMultiprocessor uses the compiled kernel's - // actual register count, avoiding a hard-coded batch-size threshold. -- const bool use_sm100_q2 = ++ ++ const bool use_cached_q1_warp64 = !relative && cached && ++ GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.warp_size == 32 && ++ d_k == 64 && q_len == 1; ++ if (use_cached_q1_warp64) { ++ constexpr int threads = 256; ++ constexpr int warps = threads / 32; ++ constexpr int value_partials = 4 * 64; ++ const size_t specialized_shmem = ++ ((size_t) 64 + kv_len + warps + value_partials) * sizeof(float); ++ GGML_ASSERT(specialized_shmem <= max_shmem); ++ const dim3 specialized_grid(n_head, 1, batch); ++ if (k->type == GGML_TYPE_F16) { ++ fused_cached_attention_q1_d64_kernel ++ <<>>( ++ (const float *) q->data, (const half *) k->data, (const half *) v->data, ++ mask ? (const float *) mask->data : nullptr, (float *) dst->data, ++ cache_k, cache_v, active_slots, active_ring_heads, active_lengths, ++ kv_len, cache_len, scale, q_sh, q_sb, k_sj, k_sh, k_sb, ++ v_sj, v_sh, v_sb, cache_sj, cache_ss, o_sh, o_sb, m_sb); ++ } else { ++ fused_cached_attention_q1_d64_kernel ++ <<>>( ++ (const float *) q->data, (const float *) k->data, (const float *) v->data, ++ mask ? (const float *) mask->data : nullptr, (float *) dst->data, ++ cache_k, cache_v, active_slots, active_ring_heads, active_lengths, ++ kv_len, cache_len, scale, q_sh, q_sb, k_sj, k_sh, k_sb, ++ v_sj, v_sh, v_sb, cache_sj, cache_ss, o_sh, o_sb, m_sb); ++ } ++ update_cache(); ++ return; ++ } ++ ++ // Select fixed-shape relative-position kernels when their geometry and ++ // device requirements are met. For Q=2, KV=72, occupancy determines ++ // whether one block processes one or two query rows. + const bool use_sm100_q2 = relative && - device_info.cc == RELPOS_ATTN_CC_SM100 && device_info.warp_size == 32 && -- d_k == RELPOS_ATTN_DK_128 && q_len == 2 && kv_len == 72; ++ device_info.cc == RELPOS_ATTN_CC_SM100 && device_info.warp_size == 32 && + d_k == RELPOS_ATTN_DK_128 && q_len == 2 && kv_len == 72 && m_sq == 0; + const bool use_register_q2 = relative && + GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.cc >= GGML_CUDA_CC_AMPERE && + device_info.warp_size == 32 && d_k == RELPOS_ATTN_DK_128 && q_len == 2 && -+ kv_len == 72 && cached && cache_len == 70 && chunk_len == 2 && n_head == 8 && -+ relpos_attn_register_resident_enabled(); ++ kv_len == 72 && cached && cache_len == 70 && chunk_len == 2 && n_head == 8; + const bool use_register_q4 = relative && + GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.cc >= GGML_CUDA_CC_AMPERE && + device_info.warp_size == 32 && d_k == RELPOS_ATTN_DK_128 && q_len == 4 && -+ kv_len == 60 && cached && cache_len == 56 && chunk_len == 4 && n_head == 8 && -+ relpos_attn_register_resident_enabled(); ++ kv_len == 60 && cached && cache_len == 56 && chunk_len == 4 && n_head == 8; + const bool common_q = + q_len == 1 || q_len == 2 || q_len == 4 || q_len == 7 || q_len == 14; + const bool use_register_common = relative && + GGML_CUDA_CC_IS_NVIDIA(device_info.cc) && device_info.cc >= GGML_CUDA_CC_AMPERE && + device_info.warp_size == 32 && d_k == RELPOS_ATTN_DK_128 && common_q && -+ (cache_len == 56 || cache_len == 70) && cached && chunk_len == q_len && n_head == 8 && -+ relpos_attn_register_resident_enabled(); ++ (cache_len == 56 || cache_len == 70) && cached && chunk_len == q_len && n_head == 8; + if (use_register_q4) { + constexpr int register_warps = 6; + constexpr int register_queries = 4; @@ -1310,16 +1677,14 @@ index f3c4836a..e0c0ec27 100644 + update_cache(); + return; + } - if (use_sm100_q2) { - const size_t warp_shmem = - ((size_t) 2 * RELPOS_ATTN_DK_128 + kv_len + RELPOS_ATTN_WARPS_128) * sizeof(float); -@@ -531,70 +1421,88 @@ void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor - (size_t) 2 * RELPOS_ATTN_WARPS_128) * sizeof(float); - GGML_ASSERT(q2_shmem <= max_shmem); - -- const int max_single_blocks_per_sm = k->type == GGML_TYPE_F16 -- ? relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem) -- : relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem); ++ if (use_sm100_q2) { ++ const size_t warp_shmem = ++ ((size_t) 2 * RELPOS_ATTN_DK_128 + kv_len + RELPOS_ATTN_WARPS_128) * sizeof(float); ++ const size_t q2_shmem = ++ ((size_t) 4 * RELPOS_ATTN_DK_128 + (size_t) 2 * kv_len + ++ (size_t) 2 * RELPOS_ATTN_WARPS_128) * sizeof(float); ++ GGML_ASSERT(q2_shmem <= max_shmem); ++ + int max_single_blocks_per_sm; + if (k->type == GGML_TYPE_F16) { + max_single_blocks_per_sm = cached @@ -1330,39 +1695,26 @@ index f3c4836a..e0c0ec27 100644 + ? relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem) + : relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem); + } - const int64_t single_query_blocks = (int64_t) n_head * q_len * batch; - const int64_t single_wave_blocks = (int64_t) device_info.nsm * max_single_blocks_per_sm; - const bool fuse_queries = single_query_blocks > single_wave_blocks; - const dim3 tuned_grid(n_head, fuse_queries ? 1 : q_len, batch); - -- if (k->type == GGML_TYPE_F16) { ++ const int64_t single_query_blocks = (int64_t) n_head * q_len * batch; ++ const int64_t single_wave_blocks = (int64_t) device_info.nsm * max_single_blocks_per_sm; ++ const bool fuse_queries = single_query_blocks > single_wave_blocks; ++ const dim3 tuned_grid(n_head, fuse_queries ? 1 : q_len, batch); ++ + auto launch_tuned = [&](auto scalar, auto cache_tag) { + using T = decltype(scalar); + constexpr bool Cached = decltype(cache_tag)::value; - if (fuse_queries) { -- fused_relpos_attn_q2_warp_128_kernel<<< ++ if (fuse_queries) { + fused_relpos_attn_q2_warp_128_kernel<<< - tuned_grid, RELPOS_ATTN_DK_128, q2_shmem, stream>>>( -- (const float *) q->data, (const half *) k->data, (const half *) v->data, -- (const half *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, -- mask ? (const float *) mask->data : nullptr, (float *) dst->data, -- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, -- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); ++ tuned_grid, RELPOS_ATTN_DK_128, q2_shmem, stream>>>( + (const float *) q->data, (const T *) k->data, (const T *) v->data, + (const T *) p->data, (const float *) bias_u->data, + (const float *) bias_v->data, mask ? (const float *) mask->data : nullptr, + (float *) dst->data, cache_k, cache_v, active_slots, active_ring_heads, kv_len, + cache_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, + cache_sj, cache_ss, p_sr, p_sh, o_si, o_sh, o_sb, m_sb); - } else { -- fused_relpos_attn_warp_128_kernel<<< ++ } else { + fused_relpos_attn_warp_128_kernel<<< - tuned_grid, RELPOS_ATTN_DK_128, warp_shmem, stream>>>( -- (const float *) q->data, (const half *) k->data, (const half *) v->data, -- (const half *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, -- mask ? (const float *) mask->data : nullptr, (float *) dst->data, -- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, -- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); ++ tuned_grid, RELPOS_ATTN_DK_128, warp_shmem, stream>>>( + (const float *) q->data, (const T *) k->data, (const T *) v->data, + (const T *) p->data, (const float *) bias_u->data, + (const float *) bias_v->data, mask ? (const float *) mask->data : nullptr, @@ -1376,50 +1728,28 @@ index f3c4836a..e0c0ec27 100644 + launch_tuned(half{}, std::true_type{}); + } else { + launch_tuned(half{}, std::false_type{}); - } - } else { -- if (fuse_queries) { -- fused_relpos_attn_q2_warp_128_kernel<<< -- tuned_grid, RELPOS_ATTN_DK_128, q2_shmem, stream>>>( -- (const float *) q->data, (const float *) k->data, (const float *) v->data, -- (const float *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, -- mask ? (const float *) mask->data : nullptr, (float *) dst->data, -- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, -- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); ++ } ++ } else { + if (cached) { + launch_tuned(float{}, std::true_type{}); - } else { -- fused_relpos_attn_warp_128_kernel<<< -- tuned_grid, RELPOS_ATTN_DK_128, warp_shmem, stream>>>( -- (const float *) q->data, (const float *) k->data, (const float *) v->data, -- (const float *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, -- mask ? (const float *) mask->data : nullptr, (float *) dst->data, -- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, -- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); ++ } else { + launch_tuned(float{}, std::false_type{}); - } - } ++ } ++ } + update_cache(); - return; - } - - const dim3 grid(n_head, q_len, batch); -- if (k->type == GGML_TYPE_F16) { -- fused_relpos_attn_kernel<<>>( -- (const float *) q->data, (const half *) k->data, (const half *) v->data, -- (const half *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, ++ return; ++ } ++ ++ const dim3 grid(n_head, q_len, batch); + auto launch_generic = [&](auto scalar, auto cache_tag) { + using T = decltype(scalar); + constexpr bool Cached = decltype(cache_tag)::value; -+ fused_relpos_attn_kernel<<>>( ++ fused_attention_kernel<<>>( + (const float *) q->data, (const T *) k->data, (const T *) v->data, + relative ? (const T *) p->data : nullptr, + relative ? (const float *) bias_u->data : nullptr, + relative ? (const float *) bias_v->data : nullptr, - mask ? (const float *) mask->data : nullptr, (float *) dst->data, -- q_len, kv_len, n_head, scale, -- q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, p_sr, p_sh, o_si, o_sh, o_sb, -- m_sb); ++ mask ? (const float *) mask->data : nullptr, (float *) dst->data, + cache_k, cache_v, active_slots, active_ring_heads, active_lengths, q_len, kv_len, + n_head, cache_len, scale, relative, + q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, cache_sj, cache_ss, @@ -1431,7 +1761,625 @@ index f3c4836a..e0c0ec27 100644 + } else { + launch_generic(half{}, std::false_type{}); + } - } else { ++ } else { ++ if (cached) { ++ launch_generic(float{}, std::true_type{}); ++ } else { ++ launch_generic(float{}, std::false_type{}); ++ } ++ } ++ update_cache(); ++} +diff --git a/src/ggml-cuda/fused-relpos-attn.cuh b/src/ggml-cuda/fused-attention.cuh +similarity index 64% +rename from src/ggml-cuda/fused-relpos-attn.cuh +rename to src/ggml-cuda/fused-attention.cuh +index 1543fa18..48e23c67 100644 +--- a/src/ggml-cuda/fused-relpos-attn.cuh ++++ b/src/ggml-cuda/fused-attention.cuh +@@ -2,4 +2,4 @@ + // SPDX-License-Identifier: Apache-2.0 + #include "common.cuh" + +-void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); ++void ggml_cuda_op_fused_attention(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +diff --git a/src/ggml-cuda/fused-relpos-attn.cu b/src/ggml-cuda/fused-relpos-attn.cu +deleted file mode 100644 +index f3c4836a..00000000 +--- a/src/ggml-cuda/fused-relpos-attn.cu ++++ /dev/null +@@ -1,600 +0,0 @@ +-// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +-// SPDX-License-Identifier: Apache-2.0 +-#include "fused-relpos-attn.cuh" +- +-#include +- +-// Fused FastConformer relative-position multi-head attention. +-// +-// One block per (head, query, batch); blockDim.x = d_k threads (one per head +-// dim). Scores, rel-shifted position term, scale+mask, two-pass softmax, and +-// the attn*V context are all computed in-kernel, so the rel-shift matrix and +-// the score matrix are never materialized in global memory. +-// +-// Operand addressing is fully stride-driven (strides read from each tensor's +-// nb[] by the host wrapper, in elements): Q/K/V may be non-contiguous views — +-// e.g. the Q slice of a fused-QKV projection, or a feat-major [n_feat, kv] +-// K/V window — as long as d_k stays innermost-contiguous (asserted by the op +-// constructor; the vectorized loads rely on it). P is [d_k, pos_len, n_head] +-// with pos_len = kv + q - 1 rows addressable; bu/bv are [d_k, n_head]; +-// mask is [kv] (shared) or [kv, batch] (per-stream) additive (0 / -inf), or NULL. +-// Output ctx is [d_k, q, n_head, batch] logical; with the merge_heads op flag +-// its memory layout is head-merged ([d_k+h*d_k] innermost, i.e. a plain +-// (n_feat, q, batch) matrix), so the output projection consumes it without a +-// permute copy. +-// +-// Requires d_k to be a power of two (the softmax reduction halves blockDim.x). +- +-// K/V/P are templated: F16 operands halve the dominant re-read traffic when +-// the caller stages them; F32 operands skip the staging casts entirely. All +-// math stays in F32 either way. +- +-static constexpr int RELPOS_ATTN_DK_128 = 128; +-static constexpr int RELPOS_ATTN_WARPS_128 = RELPOS_ATTN_DK_128 / 32; +-static constexpr int RELPOS_ATTN_CC_SM100 = 1000; +- +-static __device__ __forceinline__ float relpos_warp_sum(float value) { +-#pragma unroll +- for (int offset = 16; offset > 0; offset >>= 1) { +- value += __shfl_down_sync(0xffffffff, value, offset); +- } +- return value; +-} +- +-static __device__ __forceinline__ float4 relpos_load4(const float * ptr) { +- return *reinterpret_cast(ptr); +-} +- +-static __device__ __forceinline__ float4 relpos_load4(const half * ptr) { +- const int2 packed = *reinterpret_cast(ptr); +- const half2 * values = reinterpret_cast(&packed); +- return make_float4( +- __low2float(values[0]), __high2float(values[0]), +- __low2float(values[1]), __high2float(values[1])); +-} +- +-// SM100 streaming specialization for d_k=128 and q=2. Keeping one block per +-// query preserves the two rows' parallelism while the complete grid fits in a +-// resident wave. Compared with the generic shared-memory reduction tree, the +-// score-producing warps retain their local maxima and max/sum use only four +-// warp partials. This removes fourteen block barriers and 124 scratch floats. +-template +-static __global__ void fused_relpos_attn_warp_128_kernel( +- const float * __restrict__ Q, const T * __restrict__ K, +- const T * __restrict__ V, const T * __restrict__ Ppos, +- const float * __restrict__ bu, const float * __restrict__ bv, +- const float * __restrict__ mask, float * __restrict__ ctx, +- int kv, float scale, +- long q_sq, long q_sh, long q_sb, +- long k_sj, long k_sh, long k_sb, +- long v_sj, long v_sh, long v_sb, +- long p_sr, long p_sh, +- long o_si, long o_sh, long o_sb, long m_sb) { +- extern __shared__ float sh[]; +- float * Qu = sh; +- float * Qv = Qu + RELPOS_ATTN_DK_128; +- float * sc = Qv + RELPOS_ATTN_DK_128; +- float * red = sc + kv; +- +- const int h = blockIdx.x; +- const int i = blockIdx.y; +- const int b = blockIdx.z; +- const int d = threadIdx.x; +- const int warp = d >> 5; +- const int lane = d & 31; +- +- const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; +- const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; +- const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; +- const T * Ph = Ppos + (size_t) h * p_sh; +- const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; +- +- Qu[d] = Qhi[d] + bu[h * RELPOS_ATTN_DK_128 + d]; +- Qv[d] = Qhi[d] + bv[h * RELPOS_ATTN_DK_128 + d]; +- __syncthreads(); +- +- const int d4 = lane * 4; +- const float4 qu4 = *reinterpret_cast(Qu + d4); +- const float4 qv4 = *reinterpret_cast(Qv + d4); +- float produced_max = -INFINITY; +- for (int j = warp; j < kv; j += RELPOS_ATTN_WARPS_128) { +- const float4 k4 = relpos_load4(Kh + (size_t) j * k_sj + d4); +- const int row = 1 + j - i; +- const float4 p4 = relpos_load4(Ph + (size_t) row * p_sr + d4); +- float score = +- k4.x * qu4.x + p4.x * qv4.x + +- k4.y * qu4.y + p4.y * qv4.y + +- k4.z * qu4.z + p4.z * qv4.z + +- k4.w * qu4.w + p4.w * qv4.w; +- score = relpos_warp_sum(score); +- if (lane == 0) { +- sc[j] = score * scale + (Mb ? Mb[j] : 0.0f); +- produced_max = fmaxf(produced_max, sc[j]); +- } +- } +- if (lane == 0) { +- red[warp] = produced_max; +- } +- __syncthreads(); +- +- if (d == 0) { +- float maximum = red[0]; +-#pragma unroll +- for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { +- maximum = fmaxf(maximum, red[w]); +- } +- red[0] = maximum; +- } +- __syncthreads(); +- const float maximum = red[0]; +- +- // red[] is reused below. Ensure every warp has captured red[0] first; +- // otherwise warp 0 can overwrite the maximum while another warp loads it. +- __syncthreads(); +- float local_sum = 0.0f; +- for (int j = d; j < kv; j += RELPOS_ATTN_DK_128) { +- const float weight = __expf(sc[j] - maximum); +- sc[j] = weight; +- local_sum += weight; +- } +- local_sum = relpos_warp_sum(local_sum); +- if (lane == 0) { +- red[warp] = local_sum; +- } +- __syncthreads(); +- if (d == 0) { +- float sum = red[0]; +-#pragma unroll +- for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { +- sum += red[w]; +- } +- red[0] = sum; +- } +- __syncthreads(); +- +- float context = 0.0f; +- for (int j = 0; j < kv; ++j) { +- context += sc[j] * (float) Vh[(size_t) j * v_sj + d]; +- } +- ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = context / red[0]; +-} +- +-// Once the one-block-per-query grid spills into another occupancy wave, one +-// block computes both query rows. K and V are then loaded once and reused; +-// only the two relative-position rows differ. The reduction order matches the +-// one-query specialization so changing batch size does not change numerics. +-template +-static __global__ void fused_relpos_attn_q2_warp_128_kernel( +- const float * __restrict__ Q, const T * __restrict__ K, +- const T * __restrict__ V, const T * __restrict__ Ppos, +- const float * __restrict__ bu, const float * __restrict__ bv, +- const float * __restrict__ mask, float * __restrict__ ctx, +- int kv, float scale, +- long q_sq, long q_sh, long q_sb, +- long k_sj, long k_sh, long k_sb, +- long v_sj, long v_sh, long v_sb, +- long p_sr, long p_sh, +- long o_si, long o_sh, long o_sb, long m_sb) { +- extern __shared__ float sh[]; +- float * Qu0 = sh; +- float * Qv0 = Qu0 + RELPOS_ATTN_DK_128; +- float * Qu1 = Qv0 + RELPOS_ATTN_DK_128; +- float * Qv1 = Qu1 + RELPOS_ATTN_DK_128; +- float * sc0 = Qv1 + RELPOS_ATTN_DK_128; +- float * sc1 = sc0 + kv; +- float * red0 = sc1 + kv; +- float * red1 = red0 + RELPOS_ATTN_WARPS_128; +- +- const int h = blockIdx.x; +- const int b = blockIdx.z; +- const int d = threadIdx.x; +- const int warp = d >> 5; +- const int lane = d & 31; +- +- const float * Qh0 = Q + (size_t) b * q_sb + (size_t) h * q_sh; +- const float * Qh1 = Qh0 + q_sq; +- const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; +- const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; +- const T * Ph = Ppos + (size_t) h * p_sh; +- const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; +- +- const float bias_u = bu[h * RELPOS_ATTN_DK_128 + d]; +- const float bias_v = bv[h * RELPOS_ATTN_DK_128 + d]; +- Qu0[d] = Qh0[d] + bias_u; +- Qv0[d] = Qh0[d] + bias_v; +- Qu1[d] = Qh1[d] + bias_u; +- Qv1[d] = Qh1[d] + bias_v; +- __syncthreads(); +- +- const int d4 = lane * 4; +- const float4 qu04 = *reinterpret_cast(Qu0 + d4); +- const float4 qv04 = *reinterpret_cast(Qv0 + d4); +- const float4 qu14 = *reinterpret_cast(Qu1 + d4); +- const float4 qv14 = *reinterpret_cast(Qv1 + d4); +- float produced_max0 = -INFINITY; +- float produced_max1 = -INFINITY; +- for (int j = warp; j < kv; j += RELPOS_ATTN_WARPS_128) { +- const float4 k4 = relpos_load4(Kh + (size_t) j * k_sj + d4); +- const float4 p04 = relpos_load4(Ph + (size_t) (j + 1) * p_sr + d4); +- const float4 p14 = relpos_load4(Ph + (size_t) j * p_sr + d4); +- float score0 = +- k4.x * qu04.x + p04.x * qv04.x + +- k4.y * qu04.y + p04.y * qv04.y + +- k4.z * qu04.z + p04.z * qv04.z + +- k4.w * qu04.w + p04.w * qv04.w; +- float score1 = +- k4.x * qu14.x + p14.x * qv14.x + +- k4.y * qu14.y + p14.y * qv14.y + +- k4.z * qu14.z + p14.z * qv14.z + +- k4.w * qu14.w + p14.w * qv14.w; +- score0 = relpos_warp_sum(score0); +- score1 = relpos_warp_sum(score1); +- if (lane == 0) { +- const float additive_mask = Mb ? Mb[j] : 0.0f; +- sc0[j] = score0 * scale + additive_mask; +- sc1[j] = score1 * scale + additive_mask; +- produced_max0 = fmaxf(produced_max0, sc0[j]); +- produced_max1 = fmaxf(produced_max1, sc1[j]); +- } +- } +- if (lane == 0) { +- red0[warp] = produced_max0; +- red1[warp] = produced_max1; +- } +- __syncthreads(); +- +- if (d == 0) { +- float maximum0 = red0[0]; +- float maximum1 = red1[0]; +-#pragma unroll +- for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { +- maximum0 = fmaxf(maximum0, red0[w]); +- maximum1 = fmaxf(maximum1, red1[w]); +- } +- red0[0] = maximum0; +- red1[0] = maximum1; +- } +- __syncthreads(); +- const float maximum0 = red0[0]; +- const float maximum1 = red1[0]; +- __syncthreads(); +- +- float local_sum0 = 0.0f; +- float local_sum1 = 0.0f; +- for (int j = d; j < kv; j += RELPOS_ATTN_DK_128) { +- const float weight0 = __expf(sc0[j] - maximum0); +- const float weight1 = __expf(sc1[j] - maximum1); +- sc0[j] = weight0; +- sc1[j] = weight1; +- local_sum0 += weight0; +- local_sum1 += weight1; +- } +- local_sum0 = relpos_warp_sum(local_sum0); +- local_sum1 = relpos_warp_sum(local_sum1); +- if (lane == 0) { +- red0[warp] = local_sum0; +- red1[warp] = local_sum1; +- } +- __syncthreads(); +- if (d == 0) { +- float sum0 = red0[0]; +- float sum1 = red1[0]; +-#pragma unroll +- for (int w = 1; w < RELPOS_ATTN_WARPS_128; ++w) { +- sum0 += red0[w]; +- sum1 += red1[w]; +- } +- red0[0] = sum0; +- red1[0] = sum1; +- } +- __syncthreads(); +- +- float context0 = 0.0f; +- float context1 = 0.0f; +- for (int j = 0; j < kv; ++j) { +- const float value = (float) Vh[(size_t) j * v_sj + d]; +- context0 += sc0[j] * value; +- context1 += sc1[j] * value; +- } +- const size_t out = (size_t) b * o_sb + (size_t) h * o_sh + d; +- ctx[out] = context0 / red0[0]; +- ctx[out + o_si] = context1 / red1[0]; +-} +- +-template +-static int relpos_attn_warp_128_max_blocks_per_sm(int device, size_t shmem) { +- // The target shape fixes dynamic shared memory, so occupancy is invariant +- // for a given compiled kernel and device. Avoid repeating the CUDA runtime +- // query in every attention layer and graph execution. +- static std::atomic cached[GGML_CUDA_MAX_DEVICES] = {}; +- int blocks = cached[device].load(std::memory_order_relaxed); +- if (blocks == 0) { +- CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor( +- &blocks, fused_relpos_attn_warp_128_kernel, RELPOS_ATTN_DK_128, shmem)); +- GGML_ASSERT(blocks > 0); +- cached[device].store(blocks, std::memory_order_relaxed); +- } +- return blocks; +-} +- +-template +-static __global__ void fused_relpos_attn_kernel( +- const float * __restrict__ Q, const T * __restrict__ K, +- const T * __restrict__ V, const T * __restrict__ Ppos, +- const float * __restrict__ bu, const float * __restrict__ bv, +- const float * __restrict__ mask, float * __restrict__ ctx, +- int q, int kv, int n_head, float scale, +- // element strides: x_sq = between queries/keys, x_sh = between heads, +- // x_sb = between batch items +- long q_sq, long q_sh, long q_sb, +- long k_sj, long k_sh, long k_sb, +- long v_sj, long v_sh, long v_sb, +- long p_sr, long p_sh, +- long o_si, long o_sh, long o_sb, long m_sb) { +- extern __shared__ float sh[]; +- const int dk = blockDim.x; +- float * Qu = sh; // [dk] +- float * Qv = sh + dk; // [dk] +- float * sc = sh + 2 * dk; // [kv] +- float * red = sh + 2 * dk + kv; // [dk] reduction scratch +- +- const int h = blockIdx.x; // head +- const int i = blockIdx.y; // query +- const int b = blockIdx.z; // batch +- const int d = threadIdx.x; // head dim 0..dk-1 +- +- const float * Qhi = Q + (size_t) b * q_sb + (size_t) h * q_sh + (size_t) i * q_sq; +- const T * Kh = K + (size_t) b * k_sb + (size_t) h * k_sh; +- const T * Vh = V + (size_t) b * v_sb + (size_t) h * v_sh; +- const T * Ph = Ppos + (size_t) h * p_sh; +- const float * Mb = mask ? mask + (size_t) b * m_sb : nullptr; +- +- Qu[d] = Qhi[d] + bu[h * dk + d]; +- Qv[d] = Qhi[d] + bv[h * dk + d]; +- __syncthreads(); +- +- // scores. Two layouts: +- // * dk == 128 (the FastConformer case): WARP-COOPERATIVE — each warp owns +- // a key j and the 32 lanes split the 128 dims 4-a-piece with one +- // vectorized row load + shuffle reduction. The original +- // thread-per-key loop left dk-kv threads idle (kv ~= 50 < 128) and +- // issued dk scalar loads per row, which made the kernel +- // load-issue-bound (F16 operands alone changed nothing). +- // * otherwise: legacy thread-per-key scalar loop. +- if (dk == 128) { +- const int warp = d >> 5, lane = d & 31, nw = dk >> 5; +- for (int j = warp; j < kv; j += nw) { +- const T * Kj = Kh + (size_t) j * k_sj + lane * 4; +- const int row = (q - 1) + j - i; // rel-shift index +- const T * Pr = Ph + (size_t) row * p_sr + lane * 4; +- float k4[4], p4[4]; +- if (sizeof(T) == 2) { +- const int2 kr = *(const int2 *) Kj; +- const int2 pr = *(const int2 *) Pr; +- const half2 * kh = (const half2 *) &kr; +- const half2 * ph = (const half2 *) ≺ +- k4[0] = __low2float(kh[0]); k4[1] = __high2float(kh[0]); +- k4[2] = __low2float(kh[1]); k4[3] = __high2float(kh[1]); +- p4[0] = __low2float(ph[0]); p4[1] = __high2float(ph[0]); +- p4[2] = __low2float(ph[1]); p4[3] = __high2float(ph[1]); +- } else { +- const float4 kr = *(const float4 *) Kj; +- const float4 pr = *(const float4 *) Pr; +- k4[0] = ((const float *) &kr)[0]; k4[1] = ((const float *) &kr)[1]; +- k4[2] = ((const float *) &kr)[2]; k4[3] = ((const float *) &kr)[3]; +- p4[0] = ((const float *) &pr)[0]; p4[1] = ((const float *) &pr)[1]; +- p4[2] = ((const float *) &pr)[2]; p4[3] = ((const float *) &pr)[3]; +- } +- float s = 0.0f; +-#pragma unroll +- for (int e = 0; e < 4; e++) { +- s += k4[e] * Qu[lane * 4 + e] + p4[e] * Qv[lane * 4 + e]; +- } +-#pragma unroll +- for (int off = 16; off > 0; off >>= 1) { +- s += __shfl_xor_sync(0xffffffff, s, off); +- } +- if (lane == 0) { +- sc[j] = s * scale + (Mb ? Mb[j] : 0.0f); +- } +- } +- } else { +- for (int j = d; j < kv; j += dk) { +- const T * Kj = Kh + (size_t) j * k_sj; +- const int row = (q - 1) + j - i; // rel-shift index +- const T * Pr = Ph + (size_t) row * p_sr; +- float ac = 0.0f, bd = 0.0f; +- for (int dd = 0; dd < dk; dd++) { +- ac += (float) Kj[dd] * Qu[dd]; +- bd += (float) Pr[dd] * Qv[dd]; +- } +- sc[j] = (ac + bd) * scale + (Mb ? Mb[j] : 0.0f); +- } +- } +- __syncthreads(); +- +- // block max over sc[0..kv) +- float lm = -INFINITY; +- for (int j = d; j < kv; j += dk) lm = fmaxf(lm, sc[j]); +- red[d] = lm; +- __syncthreads(); +- for (int s = dk / 2; s > 0; s >>= 1) { +- if (d < s) red[d] = fmaxf(red[d], red[d + s]); +- __syncthreads(); +- } +- const float m = red[0]; +- __syncthreads(); +- +- // exp + block sum +- float ls = 0.0f; +- for (int j = d; j < kv; j += dk) { +- const float e = __expf(sc[j] - m); +- sc[j] = e; +- ls += e; +- } +- red[d] = ls; +- __syncthreads(); +- for (int s = dk / 2; s > 0; s >>= 1) { +- if (d < s) red[d] += red[d + s]; +- __syncthreads(); +- } +- const float inv = 1.0f / red[0]; +- __syncthreads(); +- +- // ctx[d] = inv * sum_j softmax(sc[j]) * V[j,d] (thread d owns output dim d) +- float c = 0.0f; +- for (int j = 0; j < kv; j++) c += sc[j] * (float) Vh[(size_t) j * v_sj + d]; +- ctx[(size_t) b * o_sb + (size_t) i * o_si + (size_t) h * o_sh + d] = c * inv; +-} +- +-void ggml_cuda_op_fused_relpos_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +- const ggml_tensor * q = dst->src[0]; +- const ggml_tensor * k = dst->src[1]; +- const ggml_tensor * v = dst->src[2]; +- const ggml_tensor * p = dst->src[3]; +- const ggml_tensor * bias_u = dst->src[4]; +- const ggml_tensor * bias_v = dst->src[5]; +- const ggml_tensor * mask = dst->src[6]; // may be null +- +- GGML_ASSERT(q->type == GGML_TYPE_F32); +- // K/V/P may be F32 or (all together) F16 — see the kernel comment. +- GGML_ASSERT(k->type == v->type && k->type == p->type); +- GGML_ASSERT(k->type == GGML_TYPE_F32 || k->type == GGML_TYPE_F16); +- GGML_ASSERT(bias_u->type == GGML_TYPE_F32 && bias_v->type == GGML_TYPE_F32); +- GGML_ASSERT(dst->type == GGML_TYPE_F32); +- +- const int d_k = q->ne[0]; +- const int q_len = q->ne[1]; +- const int n_head = q->ne[2]; +- const int batch = q->ne[3]; +- const int kv_len = k->ne[1]; +- +- GGML_ASSERT((d_k & (d_k - 1)) == 0 && "fused_relpos_attn: d_k must be a power of two"); +- GGML_ASSERT(k->ne[0] == d_k && v->ne[0] == d_k && p->ne[0] == d_k); +- GGML_ASSERT(k->ne[1] == kv_len && v->ne[1] == kv_len); +- GGML_ASSERT(k->ne[2] == n_head && v->ne[2] == n_head && p->ne[2] == n_head); +- GGML_ASSERT(k->ne[3] == batch && v->ne[3] == batch); +- GGML_ASSERT(bias_u->ne[0] == d_k && bias_v->ne[0] == d_k); +- GGML_ASSERT(bias_u->ne[1] == n_head && bias_v->ne[1] == n_head); +- if (mask != nullptr) { +- GGML_ASSERT(mask->type == GGML_TYPE_F32); +- GGML_ASSERT(ggml_is_contiguous(mask)); +- GGML_ASSERT(mask->ne[0] == kv_len); +- // Shared across the batch (ne[1]==1) or one key-mask column per +- // stream (ne[1]==batch, the cache-aware layout). +- GGML_ASSERT(mask->ne[1] == 1 || mask->ne[1] == batch); +- GGML_ASSERT(mask->ne[2] == 1 && mask->ne[3] == 1); +- } +- const long m_sb = +- (mask != nullptr && mask->ne[1] == batch && batch > 1) ? (long) (mask->nb[1] / sizeof(float)) : 0; +- +- float scale; +- memcpy(&scale, dst->op_params, sizeof(scale)); +- +- // Element strides from tensor byte strides. d_k rows must be contiguous +- // (constructor invariant), everything else is free-form. +- const size_t qe = ggml_type_size(q->type); +- const size_t ke = ggml_type_size(k->type); +- const long q_sq = (long)(q->nb[1] / qe), q_sh = (long)(q->nb[2] / qe), +- q_sb = (long)(q->nb[3] / qe); +- const long k_sj = (long)(k->nb[1] / ke), k_sh = (long)(k->nb[2] / ke), +- k_sb = (long)(k->nb[3] / ke); +- const long v_sj = (long)(v->nb[1] / ke), v_sh = (long)(v->nb[2] / ke), +- v_sb = (long)(v->nb[3] / ke); +- const long p_sr = (long)(p->nb[1] / ke), p_sh = (long)(p->nb[2] / ke); +- const size_t oe = ggml_type_size(dst->type); +- const long o_si = (long)(dst->nb[1] / oe), o_sh = (long)(dst->nb[2] / oe), +- o_sb = (long)(dst->nb[3] / oe); +- +- const int device = ggml_cuda_get_device(); +- const auto & device_info = ggml_cuda_info().devices[device]; +- const size_t shmem = ((size_t) 3 * d_k + kv_len) * sizeof(float); +- const size_t max_shmem = device_info.smpb; +- GGML_ASSERT(shmem <= max_shmem && "fused_relpos_attn: kv window too large for shared memory"); +- +- cudaStream_t stream = ctx.stream(); +- +- // The cache-aware Nemotron streaming geometry is Q=2, KV=72, H=8, +- // d_k=128. On SM100 the one-query warp kernel is fastest while its grid +- // fits in one resident wave. If that grid exceeds its measured occupancy, +- // fuse both query rows: halving the block count and reusing K/V then wins. +- // cudaOccupancyMaxActiveBlocksPerMultiprocessor uses the compiled kernel's +- // actual register count, avoiding a hard-coded batch-size threshold. +- const bool use_sm100_q2 = +- device_info.cc == RELPOS_ATTN_CC_SM100 && device_info.warp_size == 32 && +- d_k == RELPOS_ATTN_DK_128 && q_len == 2 && kv_len == 72; +- if (use_sm100_q2) { +- const size_t warp_shmem = +- ((size_t) 2 * RELPOS_ATTN_DK_128 + kv_len + RELPOS_ATTN_WARPS_128) * sizeof(float); +- const size_t q2_shmem = +- ((size_t) 4 * RELPOS_ATTN_DK_128 + (size_t) 2 * kv_len + +- (size_t) 2 * RELPOS_ATTN_WARPS_128) * sizeof(float); +- GGML_ASSERT(q2_shmem <= max_shmem); +- +- const int max_single_blocks_per_sm = k->type == GGML_TYPE_F16 +- ? relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem) +- : relpos_attn_warp_128_max_blocks_per_sm(device, warp_shmem); +- const int64_t single_query_blocks = (int64_t) n_head * q_len * batch; +- const int64_t single_wave_blocks = (int64_t) device_info.nsm * max_single_blocks_per_sm; +- const bool fuse_queries = single_query_blocks > single_wave_blocks; +- const dim3 tuned_grid(n_head, fuse_queries ? 1 : q_len, batch); +- +- if (k->type == GGML_TYPE_F16) { +- if (fuse_queries) { +- fused_relpos_attn_q2_warp_128_kernel<<< +- tuned_grid, RELPOS_ATTN_DK_128, q2_shmem, stream>>>( +- (const float *) q->data, (const half *) k->data, (const half *) v->data, +- (const half *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, +- mask ? (const float *) mask->data : nullptr, (float *) dst->data, +- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, +- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); +- } else { +- fused_relpos_attn_warp_128_kernel<<< +- tuned_grid, RELPOS_ATTN_DK_128, warp_shmem, stream>>>( +- (const float *) q->data, (const half *) k->data, (const half *) v->data, +- (const half *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, +- mask ? (const float *) mask->data : nullptr, (float *) dst->data, +- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, +- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); +- } +- } else { +- if (fuse_queries) { +- fused_relpos_attn_q2_warp_128_kernel<<< +- tuned_grid, RELPOS_ATTN_DK_128, q2_shmem, stream>>>( +- (const float *) q->data, (const float *) k->data, (const float *) v->data, +- (const float *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, +- mask ? (const float *) mask->data : nullptr, (float *) dst->data, +- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, +- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); +- } else { +- fused_relpos_attn_warp_128_kernel<<< +- tuned_grid, RELPOS_ATTN_DK_128, warp_shmem, stream>>>( +- (const float *) q->data, (const float *) k->data, (const float *) v->data, +- (const float *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, +- mask ? (const float *) mask->data : nullptr, (float *) dst->data, +- kv_len, scale, q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, +- p_sr, p_sh, o_si, o_sh, o_sb, m_sb); +- } +- } +- return; +- } +- +- const dim3 grid(n_head, q_len, batch); +- if (k->type == GGML_TYPE_F16) { +- fused_relpos_attn_kernel<<>>( +- (const float *) q->data, (const half *) k->data, (const half *) v->data, +- (const half *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, +- mask ? (const float *) mask->data : nullptr, (float *) dst->data, +- q_len, kv_len, n_head, scale, +- q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, p_sr, p_sh, o_si, o_sh, o_sb, +- m_sb); +- } else { - fused_relpos_attn_kernel<<>>( - (const float *) q->data, (const float *) k->data, (const float *) v->data, - (const float *) p->data, (const float *) bias_u->data, (const float *) bias_v->data, @@ -1439,24 +2387,72 @@ index f3c4836a..e0c0ec27 100644 - q_len, kv_len, n_head, scale, - q_sq, q_sh, q_sb, k_sj, k_sh, k_sb, v_sj, v_sh, v_sb, p_sr, p_sh, o_si, o_sh, o_sb, - m_sb); -+ if (cached) { -+ launch_generic(float{}, std::true_type{}); -+ } else { -+ launch_generic(float{}, std::false_type{}); -+ } - } -+ update_cache(); - } +- } +-} +diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu +index 1b7e26e6..b9700eae 100644 +--- a/src/ggml-cuda/ggml-cuda.cu ++++ b/src/ggml-cuda/ggml-cuda.cu +@@ -25,7 +25,7 @@ + #include "ggml-cuda/diagmask.cuh" + #include "ggml-cuda/diag.cuh" + #include "ggml-cuda/fattn.cuh" +-#include "ggml-cuda/fused-relpos-attn.cuh" ++#include "ggml-cuda/fused-attention.cuh" + #include "ggml-cuda/getrows.cuh" + #include "ggml-cuda/im2col.cuh" + #include "ggml-cuda/mmf.cuh" +@@ -3214,8 +3214,8 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_FLASH_ATTN_EXT: + ggml_cuda_flash_attn_ext(ctx, dst); + break; +- case GGML_OP_FUSED_RELPOS_ATTN: +- ggml_cuda_op_fused_relpos_attn(ctx, dst); ++ case GGML_OP_FUSED_ATTN: ++ ggml_cuda_op_fused_attention(ctx, dst); + break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_cuda_cross_entropy_loss(ctx, dst); +@@ -5975,7 +5975,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g + #endif // GGML_USE_MUSA + case GGML_OP_FLASH_ATTN_EXT: + return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); +- case GGML_OP_FUSED_RELPOS_ATTN: ++ case GGML_OP_FUSED_ATTN: + return (op->src[0]->ne[0] & (op->src[0]->ne[0] - 1)) == 0; // d_k power of two + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: diff --git a/src/ggml.c b/src/ggml.c -index 80b5802f..09b36241 100644 +index 80b5802f..4f5aeb09 100644 --- a/src/ggml.c +++ b/src/ggml.c -@@ -5377,7 +5377,7 @@ struct ggml_tensor * ggml_flash_attn_ext( +@@ -1079,7 +1079,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + + "GLU", + +- "FUSED_RELPOS_ATTN", ++ "FUSED_ATTN", + }; + + static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +@@ -1191,7 +1191,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + + "glu(x)", + +- "fused_relpos_attn(q,k,v,p)", ++ "fused_attn(q,k,v,p)", + }; + + static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); +@@ -5375,9 +5375,9 @@ struct ggml_tensor * ggml_flash_attn_ext( + return result; + } - // ggml_fused_relpos_attn +-// ggml_fused_relpos_attn ++// ggml_fused_attention -struct ggml_tensor * ggml_fused_relpos_attn( -+static struct ggml_tensor * ggml_fused_relpos_attn_impl( ++static struct ggml_tensor * ggml_fused_attention_impl( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, @@ -1466,7 +2462,7 @@ index 80b5802f..09b36241 100644 struct ggml_tensor * mask, + struct ggml_tensor * kv_cache, + struct ggml_tensor * slot_ids, -+ struct ggml_tensor * ring_heads, ++ struct ggml_tensor * cache_state, + int64_t cache_len, float scale, bool merge_heads) { @@ -1516,7 +2512,7 @@ index 80b5802f..09b36241 100644 + GGML_ASSERT(v->ne[1] == chunk_len); + GGML_ASSERT(cache_len >= 0); + GGML_ASSERT((kv_cache == NULL) == (slot_ids == NULL)); -+ GGML_ASSERT((kv_cache == NULL) == (ring_heads == NULL)); ++ GGML_ASSERT((kv_cache == NULL) == (cache_state == NULL)); + if (kv_cache) { + GGML_ASSERT(cache_len > 0); + GGML_ASSERT(kv_cache->type == GGML_TYPE_F32 && ggml_is_contiguous(kv_cache)); @@ -1524,9 +2520,9 @@ index 80b5802f..09b36241 100644 + GGML_ASSERT(kv_cache->ne[2] == 2 && kv_cache->ne[3] == 1); + GGML_ASSERT(slot_ids->type == GGML_TYPE_I32 && ggml_is_contiguous(slot_ids)); + GGML_ASSERT(slot_ids->ne[0] == q->ne[3]); -+ GGML_ASSERT(ring_heads->type == GGML_TYPE_I32 && ggml_is_contiguous(ring_heads)); -+ GGML_ASSERT(ring_heads->ne[0] == q->ne[3]); -+ GGML_ASSERT(ring_heads->ne[1] == 1 || ring_heads->ne[1] == 2); ++ GGML_ASSERT(cache_state->type == GGML_TYPE_I32 && ggml_is_contiguous(cache_state)); ++ GGML_ASSERT(cache_state->ne[0] == q->ne[3]); ++ GGML_ASSERT(cache_state->ne[1] == 1 || cache_state->ne[1] == 2); + } else { + GGML_ASSERT(cache_len == 0); + } @@ -1551,21 +2547,24 @@ index 80b5802f..09b36241 100644 } // Output mirrors q logically: [d_k, q_len, n_head, batch]. With -@@ -5434,6 +5470,7 @@ struct ggml_tensor * ggml_fused_relpos_attn( +@@ -5434,8 +5470,9 @@ struct ggml_tensor * ggml_fused_relpos_attn( } ggml_set_op_params(result, &scale, sizeof(scale)); + ggml_set_op_params_i32(result, 1, (int32_t) cache_len); - result->op = GGML_OP_FUSED_RELPOS_ATTN; +- result->op = GGML_OP_FUSED_RELPOS_ATTN; ++ result->op = GGML_OP_FUSED_ATTN; result->src[0] = q; + result->src[1] = k; + result->src[2] = v; @@ -5443,10 +5480,65 @@ struct ggml_tensor * ggml_fused_relpos_attn( result->src[4] = bias_u; result->src[5] = bias_v; result->src[6] = mask; + result->src[7] = kv_cache; + result->src[8] = slot_ids; -+ result->src[9] = ring_heads; ++ result->src[9] = cache_state; return result; } @@ -1581,7 +2580,7 @@ index 80b5802f..09b36241 100644 + struct ggml_tensor * mask, + float scale, + bool merge_heads) { -+ return ggml_fused_relpos_attn_impl( ++ return ggml_fused_attention_impl( + ctx, q, k, v, p, bias_u, bias_v, mask, NULL, NULL, NULL, 0, scale, merge_heads); +} + @@ -1596,12 +2595,12 @@ index 80b5802f..09b36241 100644 + struct ggml_tensor * mask, + struct ggml_tensor * kv_cache, + struct ggml_tensor * slot_ids, -+ struct ggml_tensor * ring_heads, ++ struct ggml_tensor * cache_state, + int64_t cache_len, + float scale, + bool merge_heads) { -+ return ggml_fused_relpos_attn_impl( -+ ctx, q, k, v, p, bias_u, bias_v, mask, kv_cache, slot_ids, ring_heads, cache_len, scale, ++ return ggml_fused_attention_impl( ++ ctx, q, k, v, p, bias_u, bias_v, mask, kv_cache, slot_ids, cache_state, cache_len, scale, + merge_heads); +} + @@ -1613,12 +2612,12 @@ index 80b5802f..09b36241 100644 + struct ggml_tensor * mask, + struct ggml_tensor * kv_cache, + struct ggml_tensor * slot_ids, -+ struct ggml_tensor * ring_heads, ++ struct ggml_tensor * cache_state, + int64_t cache_len, + float scale, + bool merge_heads) { -+ return ggml_fused_relpos_attn_impl( -+ ctx, q, k, v, NULL, NULL, NULL, mask, kv_cache, slot_ids, ring_heads, cache_len, scale, ++ return ggml_fused_attention_impl( ++ ctx, q, k, v, NULL, NULL, NULL, mask, kv_cache, slot_ids, cache_state, cache_len, scale, + merge_heads); +} + diff --git a/ggml-patches/0015-cuda-ctc-batch-fusions.patch b/ggml-patches/0015-cuda-ctc-batch-fusions.patch index fb92c3f..1f33287 100644 --- a/ggml-patches/0015-cuda-ctc-batch-fusions.patch +++ b/ggml-patches/0015-cuda-ctc-batch-fusions.patch @@ -284,7 +284,27 @@ index 555c951..e30c025 100644 template static __global__ void group_norm_f32(const float * x, float * dst, const int group_size, const int ne_elements, const float eps) { // blockIdx.x: num_groups idx -@@ -343,10 +399,9 @@ static void norm_mul_add_cuda( +@@ -334,7 +390,18 @@ static void norm_mul_add_cuda( + const int64_t stride_sample, const float eps, cudaStream_t stream) { + const dim3 blocks_num(nrows, nchannels, nsamples); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; +- if (ncols < 1024) { ++ if (ncols == 768 && WARP_SIZE == 32) { ++ constexpr int block_size = 384; ++ if (add) { ++ norm_mul_add_f32 ++ <<>>( ++ x, mul, add, dst, ncols, stride_row, stride_channel, stride_sample, eps); ++ } else { ++ norm_mul_add_f32 ++ <<>>( ++ x, mul, add, dst, ncols, stride_row, stride_channel, stride_sample, eps); ++ } ++ } else if (ncols < 1024) { + const dim3 block_dims(WARP_SIZE, 1, 1); + if (add) { + norm_mul_add_f32<<>>( +@@ -343,10 +410,8 @@ static void norm_mul_add_cuda( norm_mul_add_f32<<>>( x, mul, add, dst, ncols, stride_row, stride_channel, stride_sample, eps); } @@ -293,12 +313,11 @@ index 555c951..e30c025 100644 - // while reducing the affine LayerNorm reduction from 32 warps/block - // to 8. Retain the established 1024-thread path on older devices. + } else if (ncols == 1024 && GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_AMPERE) { -+ // Four elements per thread keeps high occupancy while reducing the -+ // affine LayerNorm reduction from 32 warps/block to 8. ++ // Four elements per thread balances reduction work and occupancy. const dim3 block_dims(256, 1, 1); if (add) { norm_mul_add_f32<256, true, T><<>>( -@@ -519,7 +574,7 @@ void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +@@ -519,7 +584,7 @@ void ggml_cuda_op_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { // or add_tensor), mirroring ggml_cuda_op_rms_norm_fused(_add). Eligibility // (row-vector operands, F32, contiguity) is enforced in ggml_cuda_can_fuse. void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_tensor * mul_tensor, @@ -307,7 +326,7 @@ index 555c951..e30c025 100644 const ggml_tensor * norm_src = (ggml_tensor *) dst->src[0]; float eps = 0.0f; memcpy(&eps, dst->op_params, sizeof(float)); -@@ -563,13 +618,19 @@ void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, +@@ -563,13 +628,19 @@ void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, const int64_t s02 = norm_src->nb[2] / ts0; const int64_t s03 = norm_src->nb[3] / ts0; @@ -334,7 +353,7 @@ index 555c951..e30c025 100644 } else { norm_mul_add_cuda( src0_d, mul_d, add_d, dst_d, -@@ -577,6 +638,50 @@ void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, +@@ -577,6 +648,50 @@ void ggml_cuda_op_norm_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst, } } diff --git a/ggml-patches/README.md b/ggml-patches/README.md index 8b1d985..409a75e 100644 --- a/ggml-patches/README.md +++ b/ggml-patches/README.md @@ -54,170 +54,68 @@ stock comparison therefore requires both a pristine ggml checkout and ## Patches -- **0001-fused-relpos-attn.patch** - adds `GGML_OP_FUSED_RELPOS_ATTN`, a fused - FastConformer relative-position attention op (content + position-with-rel-shift - + scale + mask + softmax + context in one CUDA kernel). Touches `ggml.h`, - `ggml.c`, the CPU backend (unsupported stub + `supports_op` false); adds - `ggml-cuda/fused-relpos-attn.{cu,cuh}`. The kernel takes K/V/P in F32 or F16 - (it is re-read per query, so F16 halves the dominant traffic; math stays F32) - and uses a warp-cooperative, vectorized score loop for d_k == 128 (the - thread-per-key scalar loop is load-issue-bound and left as the generic - fallback). For the SM100 streaming shape (`d_k=128`, `q=2`, `kv=72`), an - occupancy query selects between one block per query and a two-query block - that reuses K/V after the grid exceeds one resident wave. Other CUDA shapes - retain the generic fused kernel. - The op is stride-general: Q/K/V/P may be non-contiguous views - (only d_k rows must be contiguous; the CUDA op derives all addressing from - tensor `nb[]`), `merge_heads` emits a head-merged output layout whose - permute view is a contiguous `(n_feat, q, batch)` matrix, and the rel-pos - table length assert is `>=` so one precomputed table serves shorter tail - chunks — together this lets the cache-aware streaming encoder call the - kernel directly on the fused-QKV output and feat-major K/V windows with - zero staging copies. Wired into the encoder behind - `NEMO_SPEECH_FUSED_RELPOS_ATTN`. - -- **0002-nvfp4-warp-quantizer.patch** - reworks the CUDA NVFP4 MMQ activation - quantizer. Upstream uses one thread per 16-element sub-block (serial) plus a - 5-candidate scale search, which dominates small-batch streaming GEMMs. Adds a - single-pass warp-cooperative kernel (~8x threads, MXFP4-style) gated by - `NEMO_SPEECH_NVFP4_WARP`, a `NEMO_SPEECH_NVFP4_SCALE_SEARCH` width knob (1..5), and a - byte-exact self-check (`NEMO_SPEECH_NVFP4_SELFCHECK`). The native-FP4 MMQ path is - Blackwell-1200-only; these knobs are inert on architectures using the generic - path. - -- **0003-norm-mul-add-fusion.patch** - fused LayerNorm kernel - (`norm_mul_add_f32`): `GGML_OP_NORM` + row-vector gamma `MUL` + optional - row-vector beta `ADD` in one launch (upstream only fuses `RMS_NORM`). - Restricted to the classic affine pattern (ne0-length contiguous vectors). - Eligibility + dispatch live in patch 0006. Graph code must emit non-inplace - mul/add for the fusion to match (inplace ops are views). - -- **0004-conv2d-dw-f16-kernel.patch** - `conv2d-dw.cu` accepts an F16 kernel - (weights) with F32 input/output (templated kernel type). Lets the encoder's - depthwise convs run the direct CUDA kernel instead of im2col + GEMM while - keeping converter-produced F16 conv weights. - -- **0005-skinny-q8-gemm.patch** - adds `ggml-cuda/skinny-q8.{cu,cuh}`: a - Q8_0 x F32 GEMM specialized for skinny activations (9 <= N <= 64, the - streaming-encoder shape where mul_mat_q runs latency-bound). int8 tensor-core - `mma.m16n8k32` with per-q8-block scaling, K128 two-buffer cp.async pipeline, - once-per-tensor weight repack into aligned planes (cached; weight buffers - only), warp-coalesced activation quantizer, deterministic K-split reduction - for small-M shapes, one grid.z launch for all 64-column outer-batch tiles, - and an optional fused row-vector bias epilogue. It is enabled by default; logical - per-sequence-width dispatch prevents outer batch size from selecting different - math (`GGML_SKINNY_Q8_OUTER_BATCH=1` opts into dense outer-batch flattening - - use with `GGML_SKINNY_Q8_INPLACE=0` under a multi-stream scheduler). Accepts - serialized tensor-planar Q8 weights - (`GGML_TENSOR_FLAG_Q8_PLANAR`, see 0006) without a runtime repack. Kill - switch: `GGML_SKINNY_Q8=0`. The repack is in-place by default (reuses the - weight buffer, saving the ~1.07 GB cudaMalloc duplicate on parakeet-xxl), - which is correct and fast for the streaming-ASR encoder runtime. Two - caveats for the llama.cpp NMT decoder, which the NMT pipeline handles by - setting an env var process-wide before any service warms up: (1) the - in-place D2D memcpy is a stream-ordering hazard under llama.cpp's - multi-stream graph-split scheduling (it corrupts the GEMM even though the - repacked bytes are correct), and (2) the kernel is tuned for the encoder's - N=9..64 shape, so for the decoder's N=1 decode it is ~17% slower than stock - mmvq. So: NMT without ASR sets `GGML_SKINNY_Q8=0` (disable skinny entirely; - correct and faster); NMT alongside ASR sets `GGML_SKINNY_Q8_INPLACE=0` - (keep the encoder's skinny, force the safe separate-buffer layout). - -- **0006-cuda-dispatch-wiring.patch** - `ggml-cuda.cu` + `common.cuh` + - `mmvq.cu`/`vecdotq.cuh` + the `ggml.h` flag: the CUDA-backend wiring for - the patches above (fused-relpos op dispatch and `supports_op`, - `GGML_OP_NORM` fusion acceptance + eval-loop dispatch, skinny-q8 mul_mat - dispatch hook and the skinny GEMM+bias fusion), plus the streaming Q8 - weight/epilogue work: the serialized tensor-planar Q8 layout - (`GGML_TENSOR_FLAG_Q8_PLANAR`; planar Q8_0 vec-dot + planar MMVQ dispatch, - `VDR_Q8_0_Q8_1_MMVQ` 2->4; GGUFs produced with - `convert_model.py --outtype q8_0 --q8-layout planar`) and the Q8 - narrow bias/SiLU epilogue fusion for the two-frame streaming chunk - (`GGML_CUDA_Q8_NARROW_EPILOGUE`, default on: MUL_MAT+UNARY and - MUL_MAT+ADD+SILU fusion with broadcast Linear-bias support, plus - flattened-outer-batch MMVQ eligibility). Also carries the documented - sm_110 finding on the Blackwell FP4 gate in `common.cuh`. - -- **0007-magpietts-nanocodec.patch** - adds the CUDA operations used by - MagpieTTS and NanoCodec, including grouped transposed convolution and Snake; - extends F16 MMVF and its bias/residual epilogue to two-column inputs for - paired CFG execution; - bounds the keyed CUDA graph cache with configurable sweep and idle-eviction - intervals; and adds SM110/Jetson Thor architecture handling. - -- **0008-cublas-bf16-projections.patch** - recognizes shared F32/F16/BF16 - weights broadcast over contiguous outer activation dimensions and presents - `[K,T,B,...]` as one `[K,T*B*...]` cuBLAS GEMM. For BF16 projections it also - folds row bias and optional SiLU/rounding into the required output conversion. - Native BF16 epilogues dispatch only on NVIDIA SM80+; older architectures keep - the established conversion and elementwise paths. - -- **0009-fastconformer-cuda-fusions.patch** - adds the CUDA sigmoid GLU used by - the convolution module, fuses Macaron `residual + scale * ff`, and lets fused - affine LayerNorm write the BF16 projection input directly. The graph rewrites - are behind `NEMO_SPEECH_FASTCONFORMER_CUDA_FUSIONS`; BF16 output requires - NVIDIA SM80+, and the 256-thread specialization for 1024-wide LayerNorm rows - is selected only on SM90+. - -- **0010-cuda-pad-large-batch-grid.patch** - flattens CUDA PAD's tensor-slice - launch into `grid.x`. Upstream maps `ne2 * ne3` onto `grid.z`, which exceeds - CUDA's 65,535-block z-dimension limit for Nemotron's 256-channel causal - subsampling tensors at batch sizes of 256 or larger. The flattened launch - preserves the same indexing while allowing the large batches required for - throughput sweeps. +- **0001-fused-relpos-attn.patch** - adds relative-position fused attention for + CUDA, including stride-aware inputs, F16 K/V/P storage, head-merged output, + and the `NEMO_SPEECH_FUSED_RELPOS_ATTN` encoder path. + +- **0002-nvfp4-warp-quantizer.patch** - adds warp-cooperative NVFP4 activation + quantization with configurable scale search and an optional self-check. + +- **0003-norm-mul-add-fusion.patch** - fuses affine LayerNorm with row-vector + scale and optional bias. + +- **0004-conv2d-dw-f16-kernel.patch** - supports F16 weights in the direct + depthwise-convolution CUDA kernel with F32 input and output. + +- **0005-skinny-q8-gemm.patch** - adds a Q8_0 x F32 GEMM for skinny streaming + activations, including planar weights, deterministic K-split reduction, and + an optional bias epilogue. `GGML_SKINNY_Q8` controls dispatch; + `GGML_SKINNY_Q8_INPLACE=0` uses separate repack storage when required by a + multi-stream scheduler. + +- **0006-cuda-dispatch-wiring.patch** - wires fused attention, affine + LayerNorm, skinny-Q8, planar-Q8, and narrow bias/SiLU epilogues into the CUDA + backend. + +- **0007-magpietts-nanocodec.patch** - adds grouped transposed convolution and + Snake for MagpieTTS and NanoCodec, two-column MMVF epilogues for paired CFG, + bounded CUDA graph caching, and CUDA architecture handling. + +- **0008-cublas-bf16-projections.patch** - flattens contiguous outer activation + dimensions into shared-weight cuBLAS GEMMs and folds supported BF16 projection + epilogues into output conversion. + +- **0009-fastconformer-cuda-fusions.patch** - adds sigmoid GLU, Macaron + residual, affine LayerNorm conversion, and BF16 projection fusions for + FastConformer. + +- **0010-cuda-pad-large-batch-grid.patch** - flattens CUDA PAD launches into + `grid.x` so large batch dimensions do not exceed the `grid.z` limit. - **0011-cuda-graph-shape-key.patch** - keys cached CUDA graph executables by - the host graph identity plus a structural signature containing node count and - endpoint tensor descriptors. This prevents allocator reuse from associating - a new batch shape/topology with an incompatible executable while avoiding a - full graph scan on every dispatch. - -- **0012-cuda-streaming-cache-copies.patch** - recognizes inner-contiguous F32 - cache-tail views and materializes all batch planes with one pitched - `cudaMemcpy2DAsync`; adds aligned float4 fast paths for gathering and - scattering large indexed state-arena rows, including multi-plane K/V arenas. - The shape/alignment guards keep all other COPY, GET_ROWS, and SET_ROWS cases - on their existing kernels. - -- **0013-cuda-cached-f16-cublas.patch** - optional cached-F16/cuBLAS path for - skinny Q8 projections on NVIDIA SM80+. It expands immutable Q8 weights once, - converts only the live activation, and retains FP32 accumulation/output. - Runtime selection is controlled by `GGML_SKINNY_Q8_CUBLAS_F16` and its - minimum-N threshold; cuBLAS chooses the implementation for the active GPU. - Keep it opt-in because the F16 cache consumes additional device memory and - the performance crossover depends on the GPU and physical batch size. - -- **0014-cuda-relpos-extensions.patch** - extends fused relative-position - attention for cache-aware and offline FastConformer paths, and generalizes - the cached op to absolute attention without dummy relative-position - operands. A `[batch,2]` cache metadata tensor carries both circular heads and - active lengths so fixed-shape graphs skip unused cache rows. The `d_k=64` - score path uses aligned `float4` loads. The - cache-aware path reads persistent K/V rows directly by state slot and - circular head, then overwrites only the rows replaced by the current chunk. - Register-resident NVIDIA SM80+ kernels cover the common R=0, 1, 3, 6, and 13 - streaming shapes for both Nemotron cache geometries, with exact-shape - fallbacks. The offline mask contract accepts both - per-batch key-padding masks and `[key, query]` L/R masks. Set - `GGML_CUDA_RELPOS_REGISTER_RESIDENT=0` before process start to - disable the register-resident specializations without changing the direct - circular-cache path. - -- **0015-cuda-ctc-batch-fusions.patch** - reduces large-batch FastConformer - overhead by fusing BatchNorm and BatchNorm+transpose+SiLU graph patterns, - extending SiLU and affine LayerNorm output conversion to F16, and folding - bias and residual addition into the cached-F16 cuBLASLt projection. The - eligibility checks preserve the unfused path for unsupported layouts, - precisions, and GPUs. - -- **0016-fix-batched-conv1d-layout.patch** - restores the batch and output-channel - axes after the flattened Conv1D matrix multiplication. The upstream direct - reshape interleaves those axes for batches larger than one; batch one keeps - its original zero-copy path. - -- **0017-cuda-stream-interop.patch** - exposes borrowed access to the CUDA - backend's active stream and stable graph templates for composition with - external CUDA work. Both handles remain backend-owned. + graph identity and structural shape so allocator reuse cannot select an + incompatible executable. + +- **0012-cuda-streaming-cache-copies.patch** - adds pitched cache-tail copies + and vectorized indexed state-arena transfers for streaming K/V state. + +- **0013-cuda-cached-f16-cublas.patch** - adds an optional cached-F16/cuBLAS + path for skinny Q8 projections while retaining F32 accumulation and output. + +- **0014-cuda-fused-attention-extensions.patch** - generalizes fused attention + to standard and relative-position modes, persistent circular K/V caches, + active-length state, streaming FastConformer shapes, and Magpie's cached + single-query shape. + +- **0015-cuda-ctc-batch-fusions.patch** - adds BatchNorm, transpose/SiLU, + affine LayerNorm, and cached-F16 projection epilogues used by batched + FastConformer and Magpie. + +- **0016-fix-batched-conv1d-layout.patch** - restores batch and output-channel + axes after flattened Conv1D matrix multiplication. + +- **0017-cuda-stream-interop.patch** - exposes borrowed access to the active + CUDA stream and stable graph templates for external graph composition. ## Regenerating after editing ggml @@ -260,14 +158,14 @@ same way. Patch 0013 is intentionally layered on top of 0005. To regenerate it without folding the generic skinny-Q8 implementation into the cached-F16 patch, use a temporary ggml worktree, apply and stage patches 0001 through 0012 as the -baseline, then copy in only the cached-F16 changes and the SM100 CMake target -correction and run `git diff` against that staged baseline for `CMakeLists.txt` -and `skinny-q8.cu`. +baseline, then copy in only the cached-F16 changes and the CUDA architecture +target correction, and run `git diff` against that staged baseline for +`CMakeLists.txt` and `skinny-q8.cu`. Patch 0014 is intentionally layered on top of 0001 and 0013. Apply and stage -patches 0001 through 0013 in a temporary worktree, copy the edited -`include/ggml.h`, `src/ggml.c`, and `src/ggml-cuda/fused-relpos-attn.cu` into -that worktree, then generate 0014 with `git diff` against the staged baseline. +patches 0001 through 0013 in a temporary worktree, copy the edited attention +API, dispatch, and CUDA implementation into that worktree, then generate 0014 +with `git diff` against the staged baseline. Generate patches with `git diff` only (GNU `diff`/editors can strip the leading space on blank context lines, which `git apply` rejects as corrupt). diff --git a/src/tts/magpietts/lt.cpp b/src/tts/magpietts/lt.cpp index f700347..dca2032 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -19,6 +19,31 @@ namespace nemo_speech::tts { +class LocalTransformerCudaAttentionCache { + public: + LocalTransformerCudaAttentionCache() = default; + ~LocalTransformerCudaAttentionCache(); + + LocalTransformerCudaAttentionCache(const LocalTransformerCudaAttentionCache&) = delete; + LocalTransformerCudaAttentionCache& operator=(const LocalTransformerCudaAttentionCache&) = + delete; + LocalTransformerCudaAttentionCache(LocalTransformerCudaAttentionCache&& other) noexcept; + LocalTransformerCudaAttentionCache& operator=( + LocalTransformerCudaAttentionCache&& other) noexcept; + + bool init(const magpietts_model& model, int lane_count); + void reset(); + + ggml_context* ctx = nullptr; + ggml_backend_buffer_t buffer = nullptr; + std::vector layers; + std::vector cache_states; + ggml_tensor* slot_ids = nullptr; + int n_ctx = 0; + int n_embd = 0; + int lanes = 0; +}; + class LocalTransformerGraph { public: LocalTransformerGraph() = default; @@ -39,6 +64,7 @@ class LocalTransformerGraph { ggml_tensor* dec_uncond = nullptr; ggml_tensor* pos_emb = nullptr; ggml_tensor* prev_token = nullptr; + ggml_tensor* cache_state = nullptr; ggml_tensor* logits_cond = nullptr; ggml_tensor* logits_uncond = nullptr; @@ -68,6 +94,7 @@ class LocalTransformerGraphBank { DecoderKvCache single_cache; DecoderKvCache cond_cache; DecoderKvCache uncond_cache; + LocalTransformerCudaAttentionCache pair_cuda_attention_cache; }; using local_transformer_graph = LocalTransformerGraph; @@ -315,6 +342,122 @@ magpietts_model_init_local_transformer_fp32( return magpietts_model_init_local_transformer_copy(src, dst, use_cuda, true); } +LocalTransformerCudaAttentionCache::~LocalTransformerCudaAttentionCache() { + reset(); +} + +LocalTransformerCudaAttentionCache::LocalTransformerCudaAttentionCache( + LocalTransformerCudaAttentionCache&& other) noexcept { + *this = std::move(other); +} + +LocalTransformerCudaAttentionCache& +LocalTransformerCudaAttentionCache::operator=(LocalTransformerCudaAttentionCache&& other) noexcept { + if (this != &other) { + reset(); + ctx = other.ctx; + buffer = other.buffer; + layers = std::move(other.layers); + cache_states = std::move(other.cache_states); + slot_ids = other.slot_ids; + n_ctx = other.n_ctx; + n_embd = other.n_embd; + lanes = other.lanes; + other.ctx = nullptr; + other.buffer = nullptr; + other.slot_ids = nullptr; + other.n_ctx = 0; + other.n_embd = 0; + other.lanes = 0; + } + return *this; +} + +void +LocalTransformerCudaAttentionCache::reset() { + if (buffer) { + ggml_backend_buffer_free(buffer); + buffer = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + layers.clear(); + cache_states.clear(); + slot_ids = nullptr; + n_ctx = 0; + n_embd = 0; + lanes = 0; +} + +bool +LocalTransformerCudaAttentionCache::init(const magpietts_model& model, int lane_count) { + const magpietts_hparams& h = model.hparams; + if (ctx) { + if (n_ctx == h.lt_ctx && n_embd == h.lt_hidden && lanes == lane_count && + static_cast(layers.size()) == h.lt_layers && + static_cast(cache_states.size()) == h.stacked_audio_codebooks()) { + return true; + } + reset(); + } + + ggml_init_params params = { + /*.mem_size =*/ggml_tensor_overhead() * + static_cast(h.lt_layers + h.stacked_audio_codebooks() + 1), + /*.mem_buffer =*/nullptr, + /*.no_alloc =*/true, + }; + ctx = ggml_init(params); + if (!ctx) { + fprintf(stderr, "failed to allocate local CUDA attention cache context\n"); + return false; + } + + layers.reserve(static_cast(h.lt_layers)); + for (int layer = 0; layer < h.lt_layers; ++layer) { + ggml_tensor* arena = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, static_cast(h.lt_hidden) * h.lt_ctx, lane_count, 2); + const std::string name = "magpietts_local_cuda_kv_" + std::to_string(layer); + ggml_set_name(arena, name.c_str()); + layers.push_back(arena); + } + slot_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, lane_count); + ggml_set_name(slot_ids, "magpietts_local_cuda_slot_ids"); + cache_states.reserve(static_cast(h.stacked_audio_codebooks())); + for (int codebook = 0; codebook < h.stacked_audio_codebooks(); ++codebook) { + ggml_tensor* state_tensor = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, lane_count, 2); + const std::string name = "magpietts_local_cuda_cache_state_" + std::to_string(codebook); + ggml_set_name(state_tensor, name.c_str()); + cache_states.push_back(state_tensor); + } + + buffer = ggml_backend_alloc_ctx_tensors(ctx, model.backend); + if (!buffer) { + fprintf(stderr, "failed to allocate local CUDA attention cache buffer\n"); + reset(); + return false; + } + std::vector slots(static_cast(lane_count)); + for (int lane = 0; lane < lane_count; ++lane) slots[static_cast(lane)] = lane; + ggml_backend_tensor_set(slot_ids, slots.data(), 0, slots.size() * sizeof(int32_t)); + for (ggml_tensor* arena : layers) { + ggml_backend_tensor_memset(arena, 0, 0, ggml_nbytes(arena)); + } + std::vector state(static_cast(lane_count) * 2); + for (int codebook = 0; codebook < h.stacked_audio_codebooks(); ++codebook) { + std::fill(state.begin(), state.end(), codebook); + ggml_backend_tensor_set( + cache_states[static_cast(codebook)], state.data(), 0, + state.size() * sizeof(int32_t)); + } + n_ctx = h.lt_ctx; + n_embd = h.lt_hidden; + lanes = lane_count; + return true; +} + LocalTransformerGraph::~LocalTransformerGraph() { reset(); } @@ -334,6 +477,7 @@ LocalTransformerGraph::operator=(LocalTransformerGraph&& other) noexcept { dec_uncond = other.dec_uncond; pos_emb = other.pos_emb; prev_token = other.prev_token; + cache_state = other.cache_state; logits_cond = other.logits_cond; logits_uncond = other.logits_uncond; codebook_idx = other.codebook_idx; @@ -349,6 +493,7 @@ LocalTransformerGraph::operator=(LocalTransformerGraph&& other) noexcept { other.dec_uncond = nullptr; other.pos_emb = nullptr; other.prev_token = nullptr; + other.cache_state = nullptr; other.logits_cond = nullptr; other.logits_uncond = nullptr; other.codebook_idx = -1; @@ -373,6 +518,7 @@ LocalTransformerGraph::reset() { dec_uncond = nullptr; pos_emb = nullptr; prev_token = nullptr; + cache_state = nullptr; logits_cond = nullptr; logits_uncond = nullptr; codebook_idx = -1; @@ -399,6 +545,7 @@ LocalTransformerGraphBank::operator=(LocalTransformerGraphBank&& other) noexcept single_cache = std::move(other.single_cache); cond_cache = std::move(other.cond_cache); uncond_cache = std::move(other.uncond_cache); + pair_cuda_attention_cache = std::move(other.pair_cuda_attention_cache); } return *this; } @@ -410,12 +557,16 @@ LocalTransformerGraphBank::reset() { single_cache.reset(); cond_cache.reset(); uncond_cache.reset(); + pair_cuda_attention_cache.reset(); } bool LocalTransformerGraphBank::beginFrame(const magpietts_model& model, bool pair) { const auto& h = model.hparams; if (pair) { + if (magpietts_backend_is_cuda(model.backend)) { + return pair_cuda_attention_cache.init(model, 2); + } if (!cond_cache.init( model.backend, h.lt_layers, h.lt_ctx, h.lt_hidden, "local conditional") || !uncond_cache.init( @@ -433,6 +584,33 @@ LocalTransformerGraphBank::beginFrame(const magpietts_model& model, bool pair) { return true; } +static ggml_tensor* +local_self_attention_cuda_cached_pair( + ggml_context* ctx, const magpietts_transformer& tr, const magpietts_layer& layer, + LocalTransformerCudaAttentionCache& cache, int layer_index, ggml_tensor* cache_state, + ggml_tensor* x) { + constexpr int kCfgLanes = 2; + const int64_t n_embd = tr.n_embd; + const int64_t d_head = n_embd / tr.n_head; + ggml_tensor* qkv = + ggml_reshape_3d(ctx, linear(ctx, layer.self_qkv, x), 3 * n_embd, 1, kCfgLanes); + const size_t element = ggml_element_size(qkv); + auto split_heads = [&](size_t offset) { + return ggml_view_4d( + ctx, qkv, d_head, 1, tr.n_head, kCfgLanes, qkv->nb[1], + static_cast(d_head) * element, qkv->nb[2], offset); + }; + ggml_tensor* q = split_heads(0); + ggml_tensor* k = split_heads(static_cast(n_embd) * element); + ggml_tensor* v = split_heads(static_cast(2 * n_embd) * element); + ggml_tensor* heads = ggml_fused_attn_cached( + ctx, q, k, v, nullptr, cache.layers[static_cast(layer_index)], cache.slot_ids, + cache_state, cache.n_ctx, 1.0f / std::sqrt(static_cast(d_head)), true); + ggml_tensor* merged = + ggml_reshape_2d(ctx, ggml_permute(ctx, heads, 0, 2, 1, 3), n_embd, kCfgLanes); + return linear(ctx, layer.self_o, merged); +} + static ggml_tensor* local_transformer_forward_cached_fixed_pos( ggml_context* ctx, ggml_cgraph* gf, const magpietts_transformer& tr, ggml_tensor* x, @@ -473,8 +651,7 @@ local_self_attention_cached_pair( ggml_tensor* qkv = linear(ctx, layer.self_qkv, x); const size_t element = ggml_element_size(qkv); auto qkv_slice = [&](size_t offset) { - return ggml_view_3d( - ctx, qkv, n_embd, n_tok, kCfgLanes, qkv->nb[1], qkv->nb[2], offset); + return ggml_view_3d(ctx, qkv, n_embd, n_tok, kCfgLanes, qkv->nb[1], qkv->nb[2], offset); }; ggml_tensor* qcur = qkv_slice(0); ggml_tensor* kcur = qkv_slice(static_cast(n_embd) * element); @@ -483,8 +660,7 @@ local_self_attention_cached_pair( const size_t cache_element = ggml_element_size(cond_cache.memory_k); const size_t layer_offset = static_cast(layer_index) * cond_cache.n_ctx * n_embd * cache_element; - const size_t write_offset = - layer_offset + static_cast(n_past) * n_embd * cache_element; + const size_t write_offset = layer_offset + static_cast(n_past) * n_embd * cache_element; DecoderKvCache* caches[kCfgLanes] = {&cond_cache, &uncond_cache}; for (int lane = 0; lane < kCfgLanes; ++lane) { ggml_tensor* k_lane = ggml_view_2d( @@ -503,8 +679,8 @@ local_self_attention_cached_pair( ggml_build_forward_expand(gf, v_copy); } - ggml_tensor* q = ggml_permute( - ctx, ggml_cont_4d(ctx, qcur, d_head, n_head, n_tok, kCfgLanes), 0, 2, 1, 3); + ggml_tensor* q = + ggml_permute(ctx, ggml_cont_4d(ctx, qcur, d_head, n_head, n_tok, kCfgLanes), 0, 2, 1, 3); ggml_tensor* k_cond = ggml_view_2d( ctx, cond_cache.memory_k, n_embd, n_total, n_embd * cache_element, layer_offset); ggml_tensor* k_uncond = ggml_view_2d( @@ -512,8 +688,8 @@ local_self_attention_cached_pair( ggml_tensor* k_pair = ggml_concat(ctx, k_cond, k_uncond, 2); ggml_tensor* k = ggml_permute( ctx, ggml_reshape_4d(ctx, k_pair, d_head, n_head, n_total, kCfgLanes), 0, 2, 1, 3); - ggml_tensor* scores = ggml_scale( - ctx, ggml_mul_mat(ctx, k, q), 1.0f / std::sqrt(static_cast(d_head))); + ggml_tensor* scores = + ggml_scale(ctx, ggml_mul_mat(ctx, k, q), 1.0f / std::sqrt(static_cast(d_head))); // Each local graph evaluates exactly one new position and exposes only the populated cache // prefix, so all keys are causal-valid and no diagonal mask is needed. ggml_tensor* probs = ggml_soft_max(ctx, scores); @@ -526,8 +702,7 @@ local_self_attention_cached_pair( ggml_tensor* v_trans = ggml_cont_4d( ctx, ggml_permute( - ctx, ggml_reshape_4d(ctx, v_pair, d_head, n_head, n_total, kCfgLanes), 1, 2, 0, - 3), + ctx, ggml_reshape_4d(ctx, v_pair, d_head, n_head, n_total, kCfgLanes), 1, 2, 0, 3), n_total, d_head, n_head, kCfgLanes); ggml_tensor* weighted = ggml_mul_mat(ctx, v_trans, probs); ggml_tensor* merged = ggml_permute(ctx, weighted, 0, 2, 1, 3); @@ -538,15 +713,20 @@ local_self_attention_cached_pair( static ggml_tensor* local_transformer_forward_cached_pair_fixed_pos( ggml_context* ctx, ggml_cgraph* gf, const magpietts_transformer& tr, ggml_tensor* x, - ggml_tensor* pos_emb, DecoderKvCache& cond_cache, DecoderKvCache& uncond_cache, int n_past) { + ggml_tensor* pos_emb, DecoderKvCache& cond_cache, DecoderKvCache& uncond_cache, + LocalTransformerCudaAttentionCache* cuda_attention_cache, ggml_tensor* cache_state, + int n_past) { pos_emb = ggml_cont(ctx, ggml_cast(ctx, pos_emb, GGML_TYPE_F32)); x = ggml_add(ctx, x, pos_emb); for (int il = 0; il < static_cast(tr.layers.size()); ++il) { const magpietts_layer& layer = tr.layers[static_cast(il)]; ggml_tensor* residual = x; ggml_tensor* cur = layer_norm(ctx, x, layer.norm_self); - cur = local_self_attention_cached_pair( - ctx, gf, tr, layer, cond_cache, uncond_cache, il, n_past, cur); + cur = cuda_attention_cache + ? local_self_attention_cuda_cached_pair( + ctx, tr, layer, *cuda_attention_cache, il, cache_state, cur) + : local_self_attention_cached_pair( + ctx, gf, tr, layer, cond_cache, uncond_cache, il, n_past, cur); x = ggml_add(ctx, residual, cur); residual = x; @@ -592,6 +772,12 @@ local_transformer_graph_init( graph.codebook_idx = codebook_idx; graph.seq_len = 1; graph.pair = pair; + const bool cuda_cached_attention = pair && magpietts_backend_is_cuda(model.backend); + + if (cuda_cached_attention) { + graph.cache_state = + bank.pair_cuda_attention_cache.cache_states[static_cast(codebook_idx)]; + } if (!model.local.pos_emb || model.local.pos_emb->ne[0] != model.local.n_embd || model.local.pos_emb->ne[1] <= codebook_idx) { @@ -631,7 +817,8 @@ local_transformer_graph_init( ggml_tensor* emb = ggml_get_rows( graph.ctx, model.audio_embeddings[codebook_idx - 1], graph.prev_token); input_cond = emb; - if (pair) input_uncond = emb; + if (pair) + input_uncond = emb; } graph.pos_emb = ggml_view_2d( @@ -640,23 +827,24 @@ local_transformer_graph_init( ggml_set_name(graph.pos_emb, "magpietts_local_transformer_pos_emb"); if (pair) { - ggml_tensor* pair_input = ggml_concat(graph.ctx, input_cond, input_uncond, 2); - ggml_tensor* cur_pair = model.lt_in_w - ? linear( - graph.ctx, model.lt_in_w, pair_input, model.lt_in_b) - : pair_input; + const int pair_dim = cuda_cached_attention ? 1 : 2; + ggml_tensor* pair_input = ggml_concat(graph.ctx, input_cond, input_uncond, pair_dim); + ggml_tensor* cur_pair = + model.lt_in_w ? linear(graph.ctx, model.lt_in_w, pair_input, model.lt_in_b) + : pair_input; ggml_tensor* out_pair = local_transformer_forward_cached_pair_fixed_pos( graph.ctx, graph.gf, model.local, cur_pair, graph.pos_emb, bank.cond_cache, - bank.uncond_cache, codebook_idx); + bank.uncond_cache, + cuda_cached_attention ? &bank.pair_cuda_attention_cache : nullptr, + graph.cache_state, codebook_idx); ggml_tensor* logits_pair = linear( graph.ctx, model.lt_out_w[codebook_idx], out_pair, model.lt_out_b[codebook_idx]); - logits_pair = - ggml_cont(graph.ctx, ggml_cast(graph.ctx, logits_pair, GGML_TYPE_F32)); - graph.logits_cond = ggml_view_2d( - graph.ctx, logits_pair, h.audio_vocab_size, 1, logits_pair->nb[1], 0); + logits_pair = ggml_cont(graph.ctx, ggml_cast(graph.ctx, logits_pair, GGML_TYPE_F32)); + graph.logits_cond = + ggml_view_2d(graph.ctx, logits_pair, h.audio_vocab_size, 1, logits_pair->nb[1], 0); graph.logits_uncond = ggml_view_2d( graph.ctx, logits_pair, h.audio_vocab_size, 1, logits_pair->nb[1], - logits_pair->nb[2]); + cuda_cached_attention ? logits_pair->nb[1] : logits_pair->nb[2]); ggml_set_name(graph.logits_cond, "magpietts_local_transformer_logits_cond"); ggml_set_name(graph.logits_uncond, "magpietts_local_transformer_logits_uncond"); ggml_set_output(graph.logits_cond); @@ -664,10 +852,9 @@ local_transformer_graph_init( ggml_build_forward_expand(graph.gf, graph.logits_cond); ggml_build_forward_expand(graph.gf, graph.logits_uncond); } else { - ggml_tensor* cur_cond = model.lt_in_w - ? linear( - graph.ctx, model.lt_in_w, input_cond, model.lt_in_b) - : input_cond; + ggml_tensor* cur_cond = + model.lt_in_w ? linear(graph.ctx, model.lt_in_w, input_cond, model.lt_in_b) + : input_cond; ggml_tensor* out_cond = local_transformer_forward_cached_fixed_pos( graph.ctx, graph.gf, model.local, cur_cond, graph.pos_emb, bank.single_cache, codebook_idx); @@ -886,8 +1073,7 @@ local_transformer_graph_eval_cuda( { const ggml_nvtx::range nvtx_compute("magpietts_local_transformer_graph_compute_cuda"); if (building_sequence) { - void* graph_template = - ggml_backend_cuda_get_graph_template(model.backend, graph.gf); + void* graph_template = ggml_backend_cuda_get_graph_template(model.backend, graph.gf); char error[256] = {}; if (!graph_template || !magpietts_cuda_sampler_sequence_add_ggml_graph( cuda_sample.sampler, graph_template, error, sizeof(error))) { @@ -1089,17 +1275,15 @@ sample_local_codebooks_cuda_impl( bool chain_ok = false; if (magpietts_cuda_sampler_sequence_is_ready(cuda_sampler)) { - chain_ok = magpietts_cuda_sampler_sequence_launch( - cuda_sampler, error, sizeof(error)); + chain_ok = magpietts_cuda_sampler_sequence_launch(cuda_sampler, error, sizeof(error)); } else if ( magpietts_cuda_sampler_sequence_is_warm(cuda_sampler) && !magpietts_cuda_sampler_sequence_is_disabled(cuda_sampler)) { // Compose the codebook graphs, sampling, and token feedback. - bool build_ok = magpietts_cuda_sampler_sequence_begin_build( - cuda_sampler, error, sizeof(error)); + bool build_ok = + magpietts_cuda_sampler_sequence_begin_build(cuda_sampler, error, sizeof(error)); if (build_ok) { - build_ok = magpietts_cuda_sampler_upload_config( - cuda_sampler, error, sizeof(error)) && + build_ok = magpietts_cuda_sampler_upload_config(cuda_sampler, error, sizeof(error)) && run_chain(); } if (build_ok) { @@ -1117,17 +1301,16 @@ sample_local_codebooks_cuda_impl( chain_ok = true; } else { fprintf( - stderr, "warning: CUDA local graph composition unavailable; using async chain: %s\n", + stderr, + "warning: CUDA local graph composition unavailable; using async chain: %s\n", error[0] ? error : "unknown error"); magpietts_cuda_sampler_sequence_disable(cuda_sampler); - chain_ok = magpietts_cuda_sampler_upload_config( - cuda_sampler, error, sizeof(error)) && + chain_ok = magpietts_cuda_sampler_upload_config(cuda_sampler, error, sizeof(error)) && run_chain(); } } else { - chain_ok = magpietts_cuda_sampler_upload_config( - cuda_sampler, error, sizeof(error)) && - run_chain(); + chain_ok = + magpietts_cuda_sampler_upload_config(cuda_sampler, error, sizeof(error)) && run_chain(); if (chain_ok && !magpietts_cuda_sampler_sequence_is_disabled(cuda_sampler)) { // Initialize the per-codebook graphs before composing them. magpietts_cuda_sampler_sequence_mark_warm(cuda_sampler); diff --git a/tests/cpp/tts/CMakeLists.txt b/tests/cpp/tts/CMakeLists.txt index c06c038..d78656b 100644 --- a/tests/cpp/tts/CMakeLists.txt +++ b/tests/cpp/tts/CMakeLists.txt @@ -11,6 +11,11 @@ target_link_libraries(test_magpietts_attention_prior PRIVATE nemo_speech_tts_mag add_executable(test_magpietts_frame_stacking test_magpietts_frame_stacking.cpp) target_link_libraries(test_magpietts_frame_stacking PRIVATE nemo_speech_tts_magpietts) +if(GGML_CUDA AND NEMO_SPEECH_GGML_PATCHED) + add_executable(test_magpietts_cached_attention test_magpietts_cached_attention.cpp) + target_link_libraries(test_magpietts_cached_attention PRIVATE nemo_speech_tts_magpietts) +endif() + add_executable(test_magpietts_asr test_magpietts_asr.cpp) target_link_libraries(test_magpietts_asr PRIVATE tests_cpp_wer diff --git a/tests/cpp/tts/test_magpietts_cached_attention.cpp b/tests/cpp/tts/test_magpietts_cached_attention.cpp new file mode 100644 index 0000000..8d81513 --- /dev/null +++ b/tests/cpp/tts/test_magpietts_cached_attention.cpp @@ -0,0 +1,196 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include +#include +#include + +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +namespace { + +constexpr int kHeadDim = 64; +constexpr int kHeads = 12; +constexpr int kBatch = 2; +constexpr int kCacheLength = 18; +constexpr int kFeatures = kHeadDim * kHeads; + +size_t +qkv_index(int d, int head, int batch) { + return static_cast(d + kHeadDim * (head + kHeads * batch)); +} + +size_t +cache_index(int d, int head, int position, int slot, int plane) { + return static_cast( + d + kHeadDim * head + kFeatures * position + kFeatures * kCacheLength * slot + + kFeatures * kCacheLength * kBatch * plane); +} + +std::vector +reference_attention( + const std::vector& q, const std::vector& k, const std::vector& v, + const std::vector& cache, const int32_t* ring_heads, const int32_t* active_lengths) { + std::vector output(q.size()); + const float scale = 1.0f / std::sqrt(static_cast(kHeadDim)); + for (int batch = 0; batch < kBatch; ++batch) { + const int key_begin = kCacheLength - active_lengths[batch]; + for (int head = 0; head < kHeads; ++head) { + std::vector scores; + scores.reserve(static_cast(active_lengths[batch] + 1)); + float maximum = -INFINITY; + for (int logical = key_begin; logical <= kCacheLength; ++logical) { + float score = 0.0f; + for (int d = 0; d < kHeadDim; ++d) { + const float key = + logical == kCacheLength + ? k[qkv_index(d, head, batch)] + : cache[cache_index( + d, head, (ring_heads[batch] + logical) % kCacheLength, batch, 0)]; + score += q[qkv_index(d, head, batch)] * key; + } + score *= scale; + scores.push_back(score); + maximum = std::max(maximum, score); + } + + float denominator = 0.0f; + for (float& score : scores) { + score = std::exp(score - maximum); + denominator += score; + } + for (int d = 0; d < kHeadDim; ++d) { + float context = 0.0f; + for (int logical = key_begin; logical <= kCacheLength; ++logical) { + const size_t score_index = static_cast(logical - key_begin); + const float value = + logical == kCacheLength + ? v[qkv_index(d, head, batch)] + : cache[cache_index( + d, head, (ring_heads[batch] + logical) % kCacheLength, batch, 1)]; + context += scores[score_index] * value; + } + output[qkv_index(d, head, batch)] = context / denominator; + } + } + } + return output; +} + +} // namespace + +int +main() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "FAIL: CUDA backend unavailable\n"); + return 1; + } + + const size_t tensor_count = 7; + ggml_init_params params = { + /*.mem_size =*/ggml_tensor_overhead() * tensor_count + ggml_graph_overhead(), + /*.mem_buffer =*/nullptr, + /*.no_alloc =*/true, + }; + ggml_context* ctx = ggml_init(params); + if (!ctx) { + std::fprintf(stderr, "FAIL: ggml context allocation\n"); + ggml_backend_free(backend); + return 1; + } + + ggml_tensor* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kHeadDim, 1, kHeads, kBatch); + ggml_tensor* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kHeadDim, 1, kHeads, kBatch); + ggml_tensor* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kHeadDim, 1, kHeads, kBatch); + ggml_tensor* cache = + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kFeatures * kCacheLength, kBatch, 2); + ggml_tensor* slot_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, kBatch); + ggml_tensor* cache_state = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kBatch, 2); + ggml_tensor* output = ggml_fused_attn_cached( + ctx, q, k, v, nullptr, cache, slot_ids, cache_state, kCacheLength, + 1.0f / std::sqrt(static_cast(kHeadDim)), true); + ggml_set_output(output); + + ggml_cgraph* graph = ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); + ggml_build_forward_expand(graph, output); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + std::fprintf(stderr, "FAIL: CUDA tensor allocation\n"); + ggml_free(ctx); + ggml_backend_free(backend); + return 1; + } + + std::vector q_data(ggml_nelements(q)); + std::vector k_data(ggml_nelements(k)); + std::vector v_data(ggml_nelements(v)); + std::vector cache_data(ggml_nelements(cache)); + for (size_t i = 0; i < q_data.size(); ++i) { + q_data[i] = 0.25f * std::sin(0.013f * static_cast(i + 1)); + k_data[i] = 0.30f * std::cos(0.017f * static_cast(i + 3)); + v_data[i] = 0.35f * std::sin(0.019f * static_cast(i + 5)); + } + for (size_t i = 0; i < cache_data.size(); ++i) { + cache_data[i] = 0.40f * std::cos(0.007f * static_cast(i + 7)); + } + const int32_t slots[kBatch] = {0, 1}; + const int32_t state[kBatch * 2] = {0, 11, 0, 7}; + const std::vector reference = + reference_attention(q_data, k_data, v_data, cache_data, state, state + kBatch); + + ggml_backend_tensor_set(q, q_data.data(), 0, q_data.size() * sizeof(float)); + ggml_backend_tensor_set(k, k_data.data(), 0, k_data.size() * sizeof(float)); + ggml_backend_tensor_set(v, v_data.data(), 0, v_data.size() * sizeof(float)); + ggml_backend_tensor_set(cache, cache_data.data(), 0, cache_data.size() * sizeof(float)); + ggml_backend_tensor_set(slot_ids, slots, 0, sizeof(slots)); + ggml_backend_tensor_set(cache_state, state, 0, sizeof(state)); + + const enum ggml_status status = ggml_backend_graph_compute(backend, graph); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "FAIL: CUDA graph compute (%d)\n", static_cast(status)); + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + return 1; + } + + std::vector actual(reference.size()); + std::vector updated_cache(cache_data.size()); + ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); + ggml_backend_tensor_get(cache, updated_cache.data(), 0, updated_cache.size() * sizeof(float)); + + float max_error = 0.0f; + for (size_t i = 0; i < actual.size(); ++i) { + max_error = std::max(max_error, std::fabs(actual[i] - reference[i])); + } + bool cache_ok = true; + for (int batch = 0; batch < kBatch; ++batch) { + for (int head = 0; head < kHeads; ++head) { + for (int d = 0; d < kHeadDim; ++d) { + const size_t source = qkv_index(d, head, batch); + cache_ok &= + updated_cache[cache_index(d, head, state[batch], batch, 0)] == k_data[source]; + cache_ok &= + updated_cache[cache_index(d, head, state[batch], batch, 1)] == v_data[source]; + } + } + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + ggml_backend_free(backend); + + if (max_error > 2.0e-5f) { + std::fprintf(stderr, "FAIL: cached attention max error %.8g\n", max_error); + return 1; + } + if (!cache_ok) { + std::fprintf(stderr, "FAIL: circular cache update mismatch\n"); + return 1; + } + std::printf("cached attention max error %.8g\n", max_error); + return 0; +} From fe5850e4e368cfc663baae93b5e01815e9202750 Mon Sep 17 00:00:00 2001 From: Prabhsmran Singh Date: Mon, 24 Aug 2026 15:55:10 +0530 Subject: [PATCH 07/19] perf(tts): enable optimized magpie streaming default --- app/serve.cpp | 7 ------- app/synthesize.cpp | 7 ------- config/server.example.yaml | 4 ++-- config/tts.example.yaml | 4 ++-- docs/tts/configuration.md | 4 ++-- src/tts/magpietts/config.cpp | 4 ++-- src/tts/magpietts/magpietts.cpp | 11 ++++++++--- src/tts/magpietts/magpietts.h | 2 +- src/tts/magpietts/model.cpp | 16 ++++++++++++++-- src/tts/magpietts/runtime.h | 2 +- tests/cpp/tts/test_magpietts_asr.cpp | 4 ++-- .../cpp/tts/test_magpietts_cached_attention.cpp | 14 ++++++++++++++ tests/cpp/tts/test_magpietts_file.cpp | 9 +++++---- 13 files changed, 53 insertions(+), 35 deletions(-) diff --git a/app/serve.cpp b/app/serve.cpp index 0cefccf..5635471 100644 --- a/app/serve.cpp +++ b/app/serve.cpp @@ -501,13 +501,6 @@ run_server(int argc, char** argv) { tts_config.runtime.codec_cpu = gpu < 0; } else { tts_config.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cuda; - // CUDA decoding does not imply CUDA local-transformer sampling. Keep sampling - // on CPU unless the TTS-specific setting explicitly requests CUDA. - if (tts_config.runtime.sampling_backend != - nemo_speech::tts::MagpieBackendPreference::Cuda) { - tts_config.runtime.sampling_backend = - nemo_speech::tts::MagpieBackendPreference::Cpu; - } } } tts_config.runtime.magpie_model = magpie_path; diff --git a/app/synthesize.cpp b/app/synthesize.cpp index be0840c..4ca4898 100644 --- a/app/synthesize.cpp +++ b/app/synthesize.cpp @@ -211,13 +211,6 @@ command_synthesize(int argc, char** argv) { parsed.runtime.codec_cpu = true; } else if (cuda_device) { parsed.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cuda; - // Preserve CPU sampling by default; CUDA sampling is opt-in via - // --tts.sampling-backend cuda. - if (parsed.runtime.sampling_backend != - nemo_speech::tts::MagpieBackendPreference::Cuda) { - parsed.runtime.sampling_backend = - nemo_speech::tts::MagpieBackendPreference::Cpu; - } parsed.runtime.codec_cpu = false; } else { parsed.runtime.lt_backend = nemo_speech::tts::MagpieBackendPreference::Cpu; diff --git a/config/server.example.yaml b/config/server.example.yaml index 3e6dbac..187baad 100644 --- a/config/server.example.yaml +++ b/config/server.example.yaml @@ -65,7 +65,7 @@ tts: steps: -1 # -1 = model default top-k: -1 # -1 = model default - chunk-frames: 3 + chunk-frames: 4 codec-queue-depth: 4 codec-history-frames: -1 codec-future-frames: 1 @@ -79,7 +79,7 @@ tts: codec-cpu: false lt-backend: auto # auto | cpu | cuda - sampling-backend: auto # auto | cpu | cuda + sampling-backend: auto # CUDA with a CUDA Magpie/LT path; otherwise CPU uma-mode: auto # auto | off | on longform: auto # auto | off | on diff --git a/config/tts.example.yaml b/config/tts.example.yaml index cb10bb2..9e1b694 100644 --- a/config/tts.example.yaml +++ b/config/tts.example.yaml @@ -26,7 +26,7 @@ tts: # temperature: 0.6 # omitted = model default # cfg-scale: 2.5 # omitted = model default - chunk-frames: 3 + chunk-frames: 4 codec-queue-depth: 4 codec-history-frames: -1 codec-future-frames: 1 @@ -40,7 +40,7 @@ tts: codec-cpu: false lt-backend: auto # auto | cpu | cuda - sampling-backend: auto # auto | cpu | cuda + sampling-backend: auto # CUDA with a CUDA Magpie/LT path; otherwise CPU uma-mode: auto # auto | off | on longform: auto # auto | off | on diff --git a/docs/tts/configuration.md b/docs/tts/configuration.md index 97ccb67..b7ec8f0 100644 --- a/docs/tts/configuration.md +++ b/docs/tts/configuration.md @@ -151,7 +151,7 @@ All keys nest under `tts.`. Defaults shown; CLI alias listed where one exists. | key | CLI alias | default | meaning | |---|---|---|---| -| `tts.chunk-frames` | - | `3` | codec frames per streamed audio chunk | +| `tts.chunk-frames` | - | `4` | codec frames per streamed audio chunk | | `tts.codec-queue-depth` | - | `4` | codec worker queue depth | | `tts.codec-history-frames` | - | `-1` | rolling codec history frames | | `tts.codec-future-frames` | - | `1` | rolling codec future frames | @@ -168,7 +168,7 @@ All keys nest under `tts.`. Defaults shown; CLI alias listed where one exists. | `tts.codec-threads` | - | `0` | codec CPU threads; `0` = use `threads` | | `tts.lt-backend` | - | `auto` | local-transformer backend: `auto`/`cpu`/`cuda` | | `tts.lt-fp32` | `--tts.local-transformer-fp32` | `false` | run the local transformer in FP32 | -| `tts.sampling-backend` | - | `auto` | sampling backend: `auto`/`cpu`/`cuda` | +| `tts.sampling-backend` | - | `auto` | sampling backend; `auto` uses CUDA when the Magpie/local-transformer path is CUDA, otherwise CPU | | `tts.uma-mode` | - | `auto` | CUDA managed memory: `auto`/`off`/`on` | | `tts.longform` | - | `auto` | sentence-chunk longform mode: `auto`/`off`/`on` | diff --git a/src/tts/magpietts/config.cpp b/src/tts/magpietts/config.cpp index ef84e32..eb2472d 100644 --- a/src/tts/magpietts/config.cpp +++ b/src/tts/magpietts/config.cpp @@ -191,7 +191,7 @@ register_runtime_config(common::ParameterParser& p, MagpieRuntimeConfig& c) { [&c](const std::string& value) { c.sampling_backend = runtime_backend_from_string("tts.sampling-backend", value); }, - "Sampling backend: auto, cpu, or cuda"); + "Sampling backend: auto (CUDA when compatible with the Magpie/LT path), cpu, or cuda"); p.Register( "uma-mode", [&c](const std::string& value) { @@ -284,7 +284,7 @@ register_stream_params(common::ParameterParser& p, magpie_stream_params& c) { [&c](const std::string& value) { c.sampling_backend = stream_backend_from_string("tts.sampling-backend", value); }, - "Sampling backend: auto, cpu, or cuda"); + "Sampling backend: auto (CUDA when compatible with the Magpie/LT path), cpu, or cuda"); p.Register( "uma-mode", [&c](const std::string& value) { diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index 483f491..d1b0176 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -911,11 +911,12 @@ struct codec_stream_worker { if (end <= read_idx) { return out; } + const int chunk_end = std::min(end, read_idx + chunk_size); out.history_frames = 0; out.chunk_index = chunks_done; - out.final_read = is_last_token_in && last_token_id <= end; - out.frames.assign(audio_codes.begin() + read_idx, audio_codes.begin() + end); - read_idx = end; + out.final_read = is_last_token_in && last_token_id <= chunk_end; + out.frames.assign(audio_codes.begin() + read_idx, audio_codes.begin() + chunk_end); + read_idx = chunk_end; has_room.notify_one(); return out; } @@ -1090,6 +1091,10 @@ stream_magpie_to_audio( if (!magpietts_resolve_sampling_backend(magpie, params.sampling_backend, use_cuda_sampling)) { return false; } + if (params.sampling_backend == MAGPIETTS_BACKEND_AUTO && params.use_local_transformer && + !use_cuda_lt) { + use_cuda_sampling = false; + } if (params.use_local_transformer && !use_cuda_lt && use_cuda_sampling) { fprintf( stderr, diff --git a/src/tts/magpietts/magpietts.h b/src/tts/magpietts/magpietts.h index 41534ae..bbdd825 100644 --- a/src/tts/magpietts/magpietts.h +++ b/src/tts/magpietts/magpietts.h @@ -50,7 +50,7 @@ struct magpie_stream_params { int seed = -1; int steps = -1; int top_k = -1; - int chunk_frames = 3; + int chunk_frames = 4; int codec_queue_depth = 4; int codec_history_frames = -1; int codec_future_frames = 1; diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index 712060a..4ea6432 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -1836,12 +1836,17 @@ bool magpietts_resolve_sampling_backend( const magpietts_model& model, magpietts_backend_preference requested, bool& use_cuda_sampling) { use_cuda_sampling = false; - if (requested == MAGPIETTS_BACKEND_AUTO || requested == MAGPIETTS_BACKEND_CPU) { + if (requested == MAGPIETTS_BACKEND_CPU) { return true; } #if defined(MAGPIETTS_CUDA_SAMPLING) - if (!magpietts_backend_is_cuda(model.backend)) { + const bool cuda_backend = magpietts_backend_is_cuda(model.backend); + if (requested == MAGPIETTS_BACKEND_AUTO) { + use_cuda_sampling = cuda_backend; + return true; + } + if (!cuda_backend) { fprintf( stderr, "--sampling-backend cuda requires a CUDA ggml backend; current backend is %s\n", ggml_backend_name(model.backend)); @@ -1851,6 +1856,9 @@ magpietts_resolve_sampling_backend( return true; #else (void)model; + if (requested == MAGPIETTS_BACKEND_AUTO) { + return true; + } fprintf( stderr, "--sampling-backend cuda requires building MagpieTTS with GGML_CUDA=ON and CUDAToolkit\n"); @@ -1876,6 +1884,10 @@ MagpieCodeGenerator::generate( if (!magpietts_resolve_sampling_backend(model, params.sampling_backend, use_cuda_sampling)) { return false; } + if (params.sampling_backend == MAGPIETTS_BACKEND_AUTO && params.use_local_transformer && + !use_cuda_lt) { + use_cuda_sampling = false; + } if (params.use_local_transformer && !use_cuda_lt && use_cuda_sampling) { fprintf( stderr, diff --git a/src/tts/magpietts/runtime.h b/src/tts/magpietts/runtime.h index 38e53b0..2b71287 100644 --- a/src/tts/magpietts/runtime.h +++ b/src/tts/magpietts/runtime.h @@ -37,7 +37,7 @@ struct MagpieRuntimeConfig { int seed = -1; int steps = -1; int top_k = -1; - int chunk_frames = 3; + int chunk_frames = 4; int codec_queue_depth = 4; int codec_history_frames = -1; int codec_future_frames = 1; diff --git a/tests/cpp/tts/test_magpietts_asr.cpp b/tests/cpp/tts/test_magpietts_asr.cpp index ac7f0a9..90668d5 100644 --- a/tests/cpp/tts/test_magpietts_asr.cpp +++ b/tests/cpp/tts/test_magpietts_asr.cpp @@ -46,7 +46,7 @@ struct Args { int top_k = -1; int threads = 4; int codec_threads = 0; - int chunk_frames = 3; + int chunk_frames = 4; int gpu = 0; int chunk_ms = 160; int right_ctx = 1; @@ -98,7 +98,7 @@ usage(const char* argv0) { " --tts.sampling-backend auto|cpu|cuda\n" " --threads N CPU threads (default 4)\n" " --tts.codec-threads N Codec CPU threads (default --threads)\n" - " --tts.chunk-frames N Codec frames per internal audio chunk (default 3)\n" + " --tts.chunk-frames N Codec frames per internal audio chunk (default 4)\n" " --tts.codec-cpu Force NanoCodec decoder onto CPU backend\n" " --gpu N ASR GPU index (default 0; -1 for CPU)\n" " --chunk-ms N ASR streaming chunk size (default 160)\n" diff --git a/tests/cpp/tts/test_magpietts_cached_attention.cpp b/tests/cpp/tts/test_magpietts_cached_attention.cpp index 8d81513..704c1b6 100644 --- a/tests/cpp/tts/test_magpietts_cached_attention.cpp +++ b/tests/cpp/tts/test_magpietts_cached_attention.cpp @@ -8,6 +8,7 @@ #include "ggml-backend.h" #include "ggml-cuda.h" #include "ggml.h" +#include "tts/magpietts/model.h" namespace { @@ -89,6 +90,19 @@ main() { return 1; } + nemo_speech::tts::magpietts_model backend_model; + backend_model.backend = backend; + bool use_cuda_sampling = false; + if (!nemo_speech::tts::magpietts_resolve_sampling_backend( + backend_model, nemo_speech::tts::MAGPIETTS_BACKEND_AUTO, use_cuda_sampling) || + !use_cuda_sampling) { + std::fprintf(stderr, "FAIL: automatic sampling did not select CUDA\n"); + backend_model.backend = nullptr; + ggml_backend_free(backend); + return 1; + } + backend_model.backend = nullptr; + const size_t tensor_count = 7; ggml_init_params params = { /*.mem_size =*/ggml_tensor_overhead() * tensor_count + ggml_graph_overhead(), diff --git a/tests/cpp/tts/test_magpietts_file.cpp b/tests/cpp/tts/test_magpietts_file.cpp index 8f458b3..81ead34 100644 --- a/tests/cpp/tts/test_magpietts_file.cpp +++ b/tests/cpp/tts/test_magpietts_file.cpp @@ -46,7 +46,7 @@ usage(const char* argv0) { " --tts.steps N Max decoder frames\n" " --threads N CPU threads (default 4)\n" " --tts.codec-threads N Codec CPU threads (default --threads)\n" - " --tts.chunk-frames N Codec frames per internal callback chunk (default 3)\n" + " --tts.chunk-frames N Codec frames per internal callback chunk (default 4)\n" " --tts.language-code LANG Tokenizer language (default en-US)\n" " --tts.lt-backend auto|cpu|cuda\n" " Local-transformer backend\n" @@ -293,7 +293,8 @@ main(int argc, char** argv) { } } else if (!params.tokens.empty()) { token_chunks.push_back(params.tokens); - std::fprintf(stderr, "benchmark_input=synthetic_token_ids tokens=%zu\n", params.tokens.size()); + std::fprintf( + stderr, "benchmark_input=synthetic_token_ids tokens=%zu\n", params.tokens.size()); } params.warmup_tokens = params.tokens; params.warmup_token_chunks = token_chunks; @@ -317,8 +318,8 @@ main(int argc, char** argv) { for (int run = 0; run < warmup_runs; ++run) { tts::stream_run_metrics warmup_metrics; const bool ok = runtime.synthesize( - params, token_chunks, [](const std::vector&) { return true; }, - warmup_metrics, "warmup", false); + params, token_chunks, [](const std::vector&) { return true; }, warmup_metrics, + "warmup", false); if (!ok) { return 1; } From 44f46b298fcb2c0b74a56a214b10bc4e13f7c357 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Wed, 26 Aug 2026 05:52:20 +0000 Subject: [PATCH 08/19] Tokenizer fixes Signed-off-by: Anand Joseph --- conversion/tts.py | 4 + conversion/tts_tokenizer_profiles.py | 178 +++++ docs/tts/models.md | 9 + src/tts/magpietts/magpietts.cpp | 11 + src/tts/magpietts/magpietts.h | 2 + src/tts/magpietts/model.cpp | 17 + src/tts/magpietts/model.h | 19 + src/tts/magpietts/runtime.cpp | 16 + src/tts/magpietts/runtime.h | 2 + src/tts/synthesizer.cpp | 9 + src/tts/tokenizer/mandarin_tokenizer.cpp | 19 +- src/tts/tokenizer/mandarin_tokenizer.h | 4 +- src/tts/tokenizer/tokenizer.cpp | 628 ++++++++++++++++-- src/tts/tokenizer/tokenizer.h | 2 + src/tts/tokenizer/tokenizer_impl.cpp | 365 ++++------ .../conversion/tts_tokenizer_profiles_test.py | 80 +++ tests/cpp/tts/CMakeLists.txt | 20 + tests/cpp/tts/test_grpc_tts_config.cpp | 3 +- .../cpp/tts/test_magpietts_frame_stacking.cpp | 21 + tests/cpp/tts/test_tokenizer_mandarin.cpp | 6 + tests/cpp/tts/test_tokenizer_single_chars.cpp | 72 +- 21 files changed, 1177 insertions(+), 310 deletions(-) create mode 100644 conversion/tts_tokenizer_profiles.py create mode 100644 tests/conversion/tts_tokenizer_profiles_test.py diff --git a/conversion/tts.py b/conversion/tts.py index cf73c83..adc7f75 100644 --- a/conversion/tts.py +++ b/conversion/tts.py @@ -24,6 +24,7 @@ import torch from .source import extract_archive, find_checkpoint_files, load_state_dict, read_checkpoint_config +from .tts_tokenizer_profiles import tokenizer_profile SPECIAL_AUDIO_TOKENS = 8 SPEAKER_NAMES = ["John", "Sofia", "Aria", "Jason", "Leo"] @@ -93,6 +94,7 @@ def add_metadata( decoder = cfg["decoder"] lt_hidden = int(cfg.get("local_transformer_hidden_dim", 256)) frame_stacking = int(cfg.get("frame_stacking_factor", 1)) + profile = tokenizer_profile(cfg, text_vocab, frame_stacking) n_stacked_codebooks = int( len([k for k in sd if k.startswith("audio_embeddings.") and k.endswith(".weight")]) ) @@ -128,6 +130,7 @@ def add_metadata( "model_type": cfg.get("model_type"), "nemo_target": cfg.get("target"), "nemo_version": cfg.get("nemo_version"), + "tokenizer_profile": profile, "codec_model": cfg.get("codecmodel_path"), "text_vocab_size": text_vocab, "audio_codebooks": n_codebooks, @@ -151,6 +154,7 @@ def add_metadata( ) writer.add_string("magpietts.nemo_target", str(cfg.get("target", ""))) writer.add_string("magpietts.nemo_version", str(cfg.get("nemo_version", ""))) + writer.add_string("magpietts.tokenizer_profile", profile) writer.add_string("magpietts.model_type", str(cfg.get("model_type", ""))) writer.add_string("magpietts.codec_model", str(cfg.get("codecmodel_path", ""))) writer.add_string("magpietts.config_json", json.dumps(cfg, ensure_ascii=False, sort_keys=True)) diff --git a/conversion/tts_tokenizer_profiles.py b/conversion/tts_tokenizer_profiles.py new file mode 100644 index 0000000..dd9b1d8 --- /dev/null +++ b/conversion/tts_tokenizer_profiles.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Known MagpieTTS tokenizer layouts shared by conversion checks.""" + +from __future__ import annotations + +from typing import Any + +TOKENIZER_ORDERS = { + "v2602": ( + "english_phoneme", + "spanish_phoneme", + "german_phoneme", + "mandarin_phoneme", + "japanese_phoneme", + "french_chartokenizer", + "hindi_chartokenizer", + "italian_phoneme", + "vietnamese_phoneme", + "text_ce_tokenizer", + ), + "v2607": ( + "english_phoneme", + "text_ce_tokenizer", + "spanish_phoneme", + "german_phoneme", + "mandarin_phoneme", + "japanese_phoneme", + "portuguese_Brazilian_phoneme", + "hindi_phoneme", + "arabic_AE_chartokenizer", + "arabic_SA_chartokenizer", + "arabic_MSA_chartokenizer", + "french_chartokenizer", + "italian_chartokenizer", + "vietnamese_chartokenizer", + "korean_chartokenizer", + ), +} + +TOKENIZER_PROFILE_DIMENSIONS = { + "v2602": (2362, 1), + "v2607": (3359, 2), +} + +TOKENIZER_PROFILE_NEMO_VERSIONS = { + "v2602": "2.6.0rc0", + "v2607": "2.8.0rc0", +} + +IPA_TARGET = ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer" +) +BYT5_TARGET = "AutoTokenizer" +TOKENIZER_TARGETS = { + "v2602": { + "english_phoneme": IPA_TARGET, + "spanish_phoneme": IPA_TARGET, + "german_phoneme": IPA_TARGET, + "mandarin_phoneme": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ChinesePhonemesTokenizer" + ), + "japanese_phoneme": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "JapanesePhonemeTokenizer" + ), + "french_chartokenizer": BYT5_TARGET, + "hindi_chartokenizer": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "HindiCharsTokenizer" + ), + "italian_phoneme": BYT5_TARGET, + "vietnamese_phoneme": BYT5_TARGET, + "text_ce_tokenizer": BYT5_TARGET, + }, + "v2607": { + "english_phoneme": IPA_TARGET, + "text_ce_tokenizer": BYT5_TARGET, + "spanish_phoneme": IPA_TARGET, + "german_phoneme": IPA_TARGET, + "mandarin_phoneme": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ChinesePhonemesTokenizer" + ), + "japanese_phoneme": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "JapanesePhonemeTokenizer" + ), + "portuguese_Brazilian_phoneme": IPA_TARGET, + "hindi_phoneme": IPA_TARGET, + "arabic_AE_chartokenizer": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ArabicCharsTokenizer" + ), + "arabic_SA_chartokenizer": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ArabicCharsTokenizer" + ), + "arabic_MSA_chartokenizer": ( + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ArabicCharsTokenizer" + ), + "french_chartokenizer": BYT5_TARGET, + "italian_chartokenizer": BYT5_TARGET, + "vietnamese_chartokenizer": BYT5_TARGET, + "korean_chartokenizer": BYT5_TARGET, + }, +} + +V2607_LANGUAGE_MAPPING = { + "en": ["english_phoneme"], + "de": ["german_phoneme"], + "es": ["spanish_phoneme"], + "fr": ["french_chartokenizer"], + "it": ["italian_chartokenizer"], + "vi": ["vietnamese_chartokenizer"], + "zh": ["mandarin_phoneme"], + "hi": ["hindi_phoneme"], + "ja": ["japanese_phoneme"], + "pt-BR": ["portuguese_Brazilian_phoneme"], + "ko": ["korean_chartokenizer"], + "ar-AE": ["arabic_AE_chartokenizer"], + "ar-SA": ["arabic_SA_chartokenizer"], + "ar-MSA": ["arabic_MSA_chartokenizer"], +} + + +def tokenizer_profile(cfg: dict[str, Any], text_vocab: int, frame_stacking: int) -> str: + tokenizers = cfg.get("text_tokenizers") + if not isinstance(tokenizers, dict): + raise ValueError("Magpie config has no text_tokenizers mapping") + order = tuple(tokenizers) + matches = [name for name, expected in TOKENIZER_ORDERS.items() if order == expected] + if len(matches) != 1: + raise ValueError(f"unsupported Magpie tokenizer layout: {list(order)}") + profile = matches[0] + if str(cfg.get("nemo_version", "")) != TOKENIZER_PROFILE_NEMO_VERSIONS[profile]: + raise ValueError( + f"unsupported {profile} nemo_version: {cfg.get('nemo_version')!r}" + ) + for name, target in TOKENIZER_TARGETS[profile].items(): + if tokenizers[name].get("_target_") != target: + raise ValueError( + f"unsupported {profile} tokenizer target for {name}: " + f"{tokenizers[name].get('_target_')!r}" + ) + expected_japanese_case = "upper" if profile == "v2602" else "lower" + japanese_case = ( + tokenizers["japanese_phoneme"].get("g2p", {}).get("ascii_letter_case") + ) + if japanese_case != expected_japanese_case: + raise ValueError( + f"unsupported {profile} Japanese ascii_letter_case: {japanese_case!r}" + ) + if profile == "v2607": + hindi = tokenizers["hindi_phoneme"] + if ( + hindi.get("locale") != "hi-IN" + or hindi.get("g2p", {}).get("locale") != "hi-IN" + ): + raise ValueError("unsupported v2607 Hindi IPA locale") + if ( + tokenizers["portuguese_Brazilian_phoneme"].get("locale_specific_punct") + is not False + ): + raise ValueError("unsupported v2607 Portuguese punctuation mode") + if cfg.get("language_to_tokenizer_mapping") != V2607_LANGUAGE_MAPPING: + raise ValueError("unsupported v2607 language_to_tokenizer_mapping") + expected_vocab, expected_stacking = TOKENIZER_PROFILE_DIMENSIONS[profile] + if (text_vocab, frame_stacking) != (expected_vocab, expected_stacking): + raise ValueError( + f"Magpie tokenizer profile {profile} requires text_vocab_size={expected_vocab} " + f"and frame_stacking_factor={expected_stacking}, got " + f"text_vocab_size={text_vocab} and frame_stacking_factor={frame_stacking}" + ) + return profile diff --git a/docs/tts/models.md b/docs/tts/models.md index 07a17ad..67df1e7 100644 --- a/docs/tts/models.md +++ b/docs/tts/models.md @@ -46,6 +46,15 @@ they are not part of the GGUF. Extract the `.nemo` and pass that directory to the server as `--tts.tokenizer-model-dir` (here `models/magpie-tts/extracted`). The model-specific IPA/text tokenizer assets are loaded from this directory. +The GGUF and extracted directory must come from the same model revision. The +runtime recognizes the exact v2602 and v2607 tokenizer layouts from +`model_config.yaml` and rejects unknown layouts or a profile mismatch at +startup. Newly converted GGUFs record the profile explicitly; older v2602 and +v2607 GGUFs are identified from their text-vocabulary and frame-stacking +dimensions. In particular, v2602 uses the Hindi character tokenizer, while +v2607 uses the bundled Hindi IPA dictionary and also changes tokenizer order, +Japanese ASCII casing, Italian/Vietnamese tokenizer names, and the supported +language set. Japanese tokenization requires a build with `NEMO_SPEECH_TTS_WITH_JA=ON` (disabled by default), which builds Open JTalk, MeCab, and the NAIST dictionary. Mandarin requires `NEMO_SPEECH_TTS_WITH_ZH=ON` (disabled by default) and diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index d1b0176..8e56893 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -345,6 +345,17 @@ MagpieStreamingRuntime::speakerNames() const { return out; } +const std::string& +MagpieStreamingRuntime::tokenizerProfile() const { + static const std::string empty; + return impl_ ? impl_->magpie.tokenizer_profile : empty; +} + +int +MagpieStreamingRuntime::textVocabSize() const { + return impl_ ? impl_->magpie.hparams.text_vocab_size : 0; +} + void stream_latency_metrics::begin(int64_t now_us) { *this = {}; diff --git a/src/tts/magpietts/magpietts.h b/src/tts/magpietts/magpietts.h index bbdd825..ed7dd99 100644 --- a/src/tts/magpietts/magpietts.h +++ b/src/tts/magpietts/magpietts.h @@ -143,6 +143,8 @@ class MagpieStreamingRuntime { int sampleRate() const; int speakerCount() const; std::vector speakerNames() const; + const std::string& tokenizerProfile() const; + int textVocabSize() const; bool synthesize( magpie_stream_params& params, const std::vector& tokens, const magpie_pcm_callback& pcm_callback, stream_run_metrics& metrics); diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index 4ea6432..f6bc126 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -617,6 +617,8 @@ MagpieModel::operator=(MagpieModel&& other) noexcept { if (this != &other) { reset(); hparams = other.hparams; + tokenizer_profile = std::move(other.tokenizer_profile); + nemo_version = std::move(other.nemo_version); gguf = other.gguf; ctx = other.ctx; backend = other.backend; @@ -678,6 +680,8 @@ MagpieModel::reset() { ctx = nullptr; } hparams = {}; + tokenizer_profile.clear(); + nemo_version.clear(); cuda_unified_memory = false; text_embedding = nullptr; audio_embeddings.clear(); @@ -721,6 +725,8 @@ magpietts_model_load_impl( h.mask_token_id = gguf_i32(model.gguf, "magpietts.mask_token_id", h.mask_token_id); h.frame_stacking_factor = gguf_i32(model.gguf, "magpietts.frame_stacking_factor", h.frame_stacking_factor); + model.tokenizer_profile = gguf_string(model.gguf, "magpietts.tokenizer_profile"); + model.nemo_version = gguf_string(model.gguf, "magpietts.nemo_version"); const int32_t stored_stacked_codebooks = gguf_i32(model.gguf, "magpietts.stacked_audio_codebooks", h.stacked_audio_codebooks()); h.n_embd = gguf_i32(model.gguf, "magpietts.embedding_dim", h.n_embd); @@ -805,6 +811,17 @@ magpietts_model_load_impl( fprintf(stderr, "invalid frame_stacking_factor=%d\n", h.frame_stacking_factor); return false; } + if (model.tokenizer_profile.empty()) { + model.tokenizer_profile = magpietts_infer_tokenizer_profile(h); + } + if (!magpietts_tokenizer_profile_matches(model.tokenizer_profile, h)) { + fprintf( + stderr, + "unsupported or inconsistent Magpie tokenizer profile '%s': " + "text_vocab_size=%d frame_stacking_factor=%d\n", + model.tokenizer_profile.c_str(), h.text_vocab_size, h.frame_stacking_factor); + return false; + } if (h.audio_codebooks < 1 || stored_stacked_codebooks != h.stacked_audio_codebooks()) { fprintf( stderr, diff --git a/src/tts/magpietts/model.h b/src/tts/magpietts/model.h index 4edf464..93bbd66 100644 --- a/src/tts/magpietts/model.h +++ b/src/tts/magpietts/model.h @@ -92,6 +92,23 @@ struct magpietts_hparams { std::vector apply_prior_to_layers; }; +inline std::string +magpietts_infer_tokenizer_profile(const magpietts_hparams& h) { + if (h.text_vocab_size == 2362 && h.frame_stacking_factor == 1) { + return "v2602"; + } + if (h.text_vocab_size == 3359 && h.frame_stacking_factor == 2) { + return "v2607"; + } + return {}; +} + +inline bool +magpietts_tokenizer_profile_matches( + const std::string& profile, const magpietts_hparams& h) { + return profile == magpietts_infer_tokenizer_profile(h) && !profile.empty(); +} + inline bool magpietts_unstack_codes( const std::vector& stacked_codes, const magpietts_hparams& h, @@ -244,6 +261,8 @@ class MagpieModel { bool loaded() const { return gguf != nullptr && ctx != nullptr && backend != nullptr; } magpietts_hparams hparams; + std::string tokenizer_profile; + std::string nemo_version; gguf_context* gguf = nullptr; ggml_context* ctx = nullptr; diff --git a/src/tts/magpietts/runtime.cpp b/src/tts/magpietts/runtime.cpp index 5d086cd..cf18abf 100644 --- a/src/tts/magpietts/runtime.cpp +++ b/src/tts/magpietts/runtime.cpp @@ -111,6 +111,8 @@ class MagpieTtsRuntime::Impl { config_.magpie_cpu, config_.codec_cpu)) { throw std::runtime_error("failed to load MagpieTTS/NanoCodec GGUFs"); } + tokenizer_profile_ = stream_->tokenizerProfile(); + text_vocab_size_ = stream_->textVocabSize(); speaker_names_ = stream_->speakerNames(); if (speaker_names_.empty()) { @@ -125,6 +127,8 @@ class MagpieTtsRuntime::Impl { int speaker_count() const { return stream_->speakerCount(); } const std::vector& speaker_names() const { return speaker_names_; } const std::string& model_name() const { return model_name_; } + const std::string& tokenizer_profile() const { return tokenizer_profile_; } + int text_vocab_size() const { return text_vocab_size_; } MagpieSynthesisStats synthesize( const std::vector& tokens, const MagpieSynthesisOptions& options, @@ -255,6 +259,8 @@ class MagpieTtsRuntime::Impl { std::unique_ptr stream_; std::vector speaker_names_; std::string model_name_; + std::string tokenizer_profile_; + int text_vocab_size_ = 0; std::mutex mutex_; }; @@ -283,6 +289,16 @@ MagpieTtsRuntime::model_name() const { return impl_->model_name(); } +const std::string& +MagpieTtsRuntime::tokenizer_profile() const { + return impl_->tokenizer_profile(); +} + +int +MagpieTtsRuntime::text_vocab_size() const { + return impl_->text_vocab_size(); +} + MagpieSynthesisStats MagpieTtsRuntime::synthesize( const std::vector& tokens, const MagpieSynthesisOptions& options, diff --git a/src/tts/magpietts/runtime.h b/src/tts/magpietts/runtime.h index 2b71287..aafea86 100644 --- a/src/tts/magpietts/runtime.h +++ b/src/tts/magpietts/runtime.h @@ -129,6 +129,8 @@ class MagpieTtsRuntime { int speaker_count() const; const std::vector& speaker_names() const; const std::string& model_name() const; + const std::string& tokenizer_profile() const; + int text_vocab_size() const; MagpieSynthesisStats synthesize( const std::vector& tokens, const MagpieSynthesisOptions& options, diff --git a/src/tts/synthesizer.cpp b/src/tts/synthesizer.cpp index 5f33e7b..01accf9 100644 --- a/src/tts/synthesizer.cpp +++ b/src/tts/synthesizer.cpp @@ -61,6 +61,15 @@ struct Synthesizer::Impl { if (!config.tokenizer_model_dir.empty()) { tokenizer = std::make_unique( std::move(config.tokenizer_model_dir), config.tokenizer); + if (tokenizer->profile_id() != runtime.tokenizer_profile() || + tokenizer->text_vocab_size() != runtime.text_vocab_size()) { + throw std::invalid_argument( + "Magpie model/tokenizer mismatch: GGUF requires tokenizer profile '" + + runtime.tokenizer_profile() + "' with text vocabulary " + + std::to_string(runtime.text_vocab_size()) + ", but tokenizer directory is '" + + tokenizer->profile_id() + "' with text vocabulary " + + std::to_string(tokenizer->text_vocab_size())); + } } } diff --git a/src/tts/tokenizer/mandarin_tokenizer.cpp b/src/tts/tokenizer/mandarin_tokenizer.cpp index b193f66..befc39f 100644 --- a/src/tts/tokenizer/mandarin_tokenizer.cpp +++ b/src/tts/tokenizer/mandarin_tokenizer.cpp @@ -26,8 +26,6 @@ namespace fs = std::filesystem; namespace { -constexpr int kMandarinOffset = 349; - struct utf8_char { uint32_t codepoint = 0; std::string text; @@ -155,14 +153,16 @@ ascii_upper(std::string text) { class mandarin_tokenizer::impl { public: - explicit impl(const fs::path& model_dir) { + explicit impl(const fs::path& model_dir, int offset, const std::string& phoneme_dict) + : offset_(offset) { data_dir_ = find_data_dir(model_dir); if (data_dir_.empty()) { throw std::runtime_error( "failed to find Mandarin G2P data; set MAGPIE_MANDARIN_G2P_DIR"); } - const fs::path phoneme_path = find_phoneme_dict(model_dir); - if (phoneme_path.empty()) { + const fs::path phoneme_path = + phoneme_dict.empty() ? find_phoneme_dict(model_dir) : model_dir / phoneme_dict; + if (!fs::is_regular_file(phoneme_path)) { throw std::runtime_error("failed to find Mandarin pinyin-to-phoneme dictionary"); } @@ -326,10 +326,12 @@ class mandarin_tokenizer::impl { "Mandarin vocabulary does not match the Magpie model (expected 109 tokens)"); } for (size_t index = 0; index < tokens.size(); ++index) { - token_to_id_[tokens[index]] = kMandarinOffset + static_cast(index); + token_to_id_[tokens[index]] = offset_ + static_cast(index); } } + int offset_ = 349; + void append_word_pinyin(const std::string& word, std::vector& output) const { const auto chars = decode_utf8(word); for (size_t begin = 0; begin < chars.size();) { @@ -378,8 +380,9 @@ class mandarin_tokenizer::impl { } }; -mandarin_tokenizer::mandarin_tokenizer(const fs::path& model_dir) - : impl_(std::make_unique(model_dir)) {} +mandarin_tokenizer::mandarin_tokenizer( + const fs::path& model_dir, int offset, std::string phoneme_dict) + : impl_(std::make_unique(model_dir, offset, phoneme_dict)) {} mandarin_tokenizer::~mandarin_tokenizer() = default; diff --git a/src/tts/tokenizer/mandarin_tokenizer.h b/src/tts/tokenizer/mandarin_tokenizer.h index 12c3746..70f8801 100644 --- a/src/tts/tokenizer/mandarin_tokenizer.h +++ b/src/tts/tokenizer/mandarin_tokenizer.h @@ -9,7 +9,9 @@ class mandarin_tokenizer { public: - explicit mandarin_tokenizer(const std::filesystem::path& model_dir); + explicit mandarin_tokenizer( + const std::filesystem::path& model_dir, int offset = 349, + std::string phoneme_dict = {}); ~mandarin_tokenizer(); mandarin_tokenizer(const mandarin_tokenizer&) = delete; diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index 22d3740..cfc9ff7 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -17,6 +17,482 @@ namespace nemo_speech::tts { namespace { +enum class tokenizer_kind { ipa, byt5, hindi_chars, mandarin, japanese, arabic }; + +struct tokenizer_entry { + std::string name; + tokenizer_kind kind = tokenizer_kind::byt5; + int offset = 0; + int size = 0; + std::string locale; + std::string grapheme_case; + std::string grapheme_prefix; + std::string ascii_letter_case; + std::string phoneme_dict; + std::string heteronyms; + bool apostrophe = true; + bool pad_with_space = false; +}; + +struct tokenizer_profile { + std::string id; + std::string nemo_version; + int text_vocab_size = 0; + int eos_id = 0; + std::vector entries; + std::map language_mapping; + + const tokenizer_entry& entry_for_language(const std::string& language) const { + const auto mapping = language_mapping.find(language); + if (mapping == language_mapping.end()) { + throw std::invalid_argument("unsupported tokenizer language '" + language + "'"); + } + const auto entry = std::find_if( + entries.begin(), entries.end(), [&](const tokenizer_entry& value) { + return value.name == mapping->second; + }); + if (entry == entries.end()) { + throw std::runtime_error( + "tokenizer profile '" + id + "' maps language '" + language + + "' to missing tokenizer '" + mapping->second + "'"); + } + return *entry; + } +}; + +std::string +trim(std::string value) { + value.erase(value.begin(), std::find_if(value.begin(), value.end(), [](unsigned char c) { + return !std::isspace(c); + })); + value.erase( + std::find_if(value.rbegin(), value.rend(), [](unsigned char c) { return !std::isspace(c); }) + .base(), + value.end()); + return value; +} + +std::string +unquote(std::string value) { + value = trim(std::move(value)); + if (value.size() >= 2 && + ((value.front() == '\'' && value.back() == '\'') || + (value.front() == '"' && value.back() == '"'))) { + value = value.substr(1, value.size() - 2); + } + return value; +} + +int +yaml_indent(const std::string& line) { + int indent = 0; + while (indent < static_cast(line.size()) && line[static_cast(indent)] == ' ') { + ++indent; + } + return indent; +} + +std::vector +yaml_lines(const std::string& contents) { + std::vector lines; + std::istringstream input(contents); + std::string line; + while (std::getline(input, line)) { + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + lines.push_back(std::move(line)); + } + return lines; +} + +std::string +top_level_scalar(const std::vector& lines, const std::string& key) { + const std::string prefix = key + ":"; + for (const auto& line : lines) { + if (yaml_indent(line) != 0 || line.rfind(prefix, 0) != 0) { + continue; + } + return unquote(line.substr(prefix.size())); + } + return {}; +} + +std::map +yaml_child_blocks( + const std::vector& lines, const std::string& section, + std::vector& order) { + std::map blocks; + bool active = false; + std::string current; + for (const auto& line : lines) { + if (!active) { + if (line == section + ":") { + active = true; + } + continue; + } + if (!line.empty() && yaml_indent(line) == 0) { + break; + } + if (yaml_indent(line) == 2) { + const std::string value = trim(line); + if (!value.empty() && value.back() == ':') { + current = value.substr(0, value.size() - 1); + order.push_back(current); + blocks[current] = {}; + continue; + } + } + if (!current.empty()) { + blocks[current] += line; + blocks[current].push_back('\n'); + } + } + return blocks; +} + +std::string +block_scalar(const std::string& block, const std::string& key) { + std::istringstream input(block); + std::string line; + const std::string prefix = key + ":"; + while (std::getline(input, line)) { + const std::string value = trim(line); + if (value.rfind(prefix, 0) == 0) { + return unquote(value.substr(prefix.size())); + } + } + return {}; +} + +std::string +artifact_name(const std::string& value) { + static const std::string prefix = "nemo:"; + if (value.rfind(prefix, 0) == 0) { + return value.substr(prefix.size()); + } + return value; +} + +void +require_block_value( + const std::map& blocks, const std::string& tokenizer, + const std::string& key, const std::string& expected) { + const auto it = blocks.find(tokenizer); + if (it == blocks.end()) { + throw std::runtime_error("missing tokenizer config block '" + tokenizer + "'"); + } + const std::string actual = block_scalar(it->second, key); + if (actual != expected) { + throw std::runtime_error( + "unsupported tokenizer config: '" + tokenizer + "." + key + "' is '" + actual + + "', expected '" + expected + "'"); + } +} + +void +require_common_tokenizer_values( + const std::map& blocks, const std::string& tokenizer, + const std::string& target, const std::string& apostrophe, + const std::string& pad_with_space) { + require_block_value(blocks, tokenizer, "_target_", target); + require_block_value(blocks, tokenizer, "punct", "true"); + require_block_value(blocks, tokenizer, "apostrophe", apostrophe); + require_block_value(blocks, tokenizer, "pad_with_space", pad_with_space); +} + +void +require_ipa_values( + const std::map& blocks, const tokenizer_entry& item, + const std::string& configured_locale = {}) { + require_common_tokenizer_values( + blocks, item.name, + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer", + item.apostrophe ? "true" : "false", item.pad_with_space ? "true" : "false"); + require_block_value(blocks, item.name, "locale", configured_locale); + const std::string grapheme_case = block_scalar(blocks.at(item.name), "grapheme_case"); + const std::string effective_case = grapheme_case.empty() ? "upper" : grapheme_case; + if (effective_case != item.grapheme_case) { + throw std::runtime_error( + "unsupported tokenizer config: '" + item.name + ".grapheme_case' is '" + + effective_case + "', expected '" + item.grapheme_case + "'"); + } + require_block_value(blocks, item.name, "grapheme_prefix", item.grapheme_prefix); +} + +void +require_byt5_values( + const std::map& blocks, const std::string& tokenizer) { + require_block_value(blocks, tokenizer, "_target_", "AutoTokenizer"); + require_block_value(blocks, tokenizer, "pretrained_model", "google/byt5-small"); +} + +std::map +parse_language_mapping(const std::vector& lines) { + std::map result; + bool active = false; + for (const auto& line : lines) { + if (!active) { + if (line == "language_to_tokenizer_mapping:") { + active = true; + } + continue; + } + if (!line.empty() && yaml_indent(line) == 0) { + break; + } + if (yaml_indent(line) != 2) { + continue; + } + const std::string value = trim(line); + const size_t colon = value.find(':'); + if (colon == std::string::npos) { + continue; + } + std::string tokenizer = trim(value.substr(colon + 1)); + if (tokenizer.size() < 2 || tokenizer.front() != '[' || tokenizer.back() != ']') { + throw std::runtime_error( + "unsupported language_to_tokenizer_mapping value '" + tokenizer + "'"); + } + tokenizer = trim(tokenizer.substr(1, tokenizer.size() - 2)); + if (tokenizer.find(',') != std::string::npos) { + throw std::runtime_error("multiple tokenizer candidates are not supported"); + } + result.emplace(value.substr(0, colon), unquote(std::move(tokenizer))); + } + return result; +} + +tokenizer_entry +entry( + std::string name, tokenizer_kind kind, int offset, int size, std::string locale = {}, + std::string grapheme_case = {}, std::string grapheme_prefix = {}, + bool apostrophe = true, bool pad_with_space = false) { + tokenizer_entry value; + value.name = std::move(name); + value.kind = kind; + value.offset = offset; + value.size = size; + value.locale = std::move(locale); + value.grapheme_case = std::move(grapheme_case); + value.grapheme_prefix = std::move(grapheme_prefix); + value.apostrophe = apostrophe; + value.pad_with_space = pad_with_space; + return value; +} + +void +attach_artifacts( + tokenizer_profile& profile, const std::map& blocks, + const fs::path& root) { + for (auto& item : profile.entries) { + const auto block = blocks.find(item.name); + if (block == blocks.end()) { + continue; + } + item.phoneme_dict = artifact_name(block_scalar(block->second, "phoneme_dict")); + item.heteronyms = artifact_name(block_scalar(block->second, "heteronyms")); + for (const std::string* asset : {&item.phoneme_dict, &item.heteronyms}) { + if (!asset->empty() && !fs::is_regular_file(root / *asset)) { + throw std::runtime_error( + "tokenizer profile '" + profile.id + "' references missing asset '" + + (root / *asset).string() + "'"); + } + } + if ((item.kind == tokenizer_kind::ipa || item.kind == tokenizer_kind::mandarin) && + item.phoneme_dict.empty()) { + throw std::runtime_error( + "tokenizer profile '" + profile.id + "' has no phoneme_dict for '" + + item.name + "'"); + } + } +} + +void +validate_profile_layout(const tokenizer_profile& profile) { + int offset = 0; + for (const auto& item : profile.entries) { + if (item.offset != offset || item.size <= 0) { + throw std::runtime_error( + "invalid tokenizer offset table for profile '" + profile.id + "'"); + } + offset += item.size; + } + if (offset + 2 != profile.text_vocab_size || profile.eos_id != offset + 1) { + throw std::runtime_error( + "invalid tokenizer vocabulary size for profile '" + profile.id + "'"); + } +} + +tokenizer_profile +load_tokenizer_profile(const fs::path& root) { + const fs::path config_path = root / "model_config.yaml"; + const std::string contents = read_file(config_path); + const auto lines = yaml_lines(contents); + std::vector order; + const auto blocks = yaml_child_blocks(lines, "text_tokenizers", order); + const std::vector v2602_order = { + "english_phoneme", "spanish_phoneme", "german_phoneme", + "mandarin_phoneme", "japanese_phoneme", "french_chartokenizer", + "hindi_chartokenizer", "italian_phoneme", "vietnamese_phoneme", + "text_ce_tokenizer", + }; + const std::vector v2607_order = { + "english_phoneme", "text_ce_tokenizer", + "spanish_phoneme", "german_phoneme", + "mandarin_phoneme", "japanese_phoneme", + "portuguese_Brazilian_phoneme", "hindi_phoneme", + "arabic_AE_chartokenizer", "arabic_SA_chartokenizer", + "arabic_MSA_chartokenizer", "french_chartokenizer", + "italian_chartokenizer", "vietnamese_chartokenizer", + "korean_chartokenizer", + }; + + tokenizer_profile profile; + profile.nemo_version = top_level_scalar(lines, "nemo_version"); + if (order == v2602_order) { + profile.id = "v2602"; + profile.text_vocab_size = 2362; + profile.eos_id = 2361; + profile.entries = { + entry("english_phoneme", tokenizer_kind::ipa, 0, 96, "en-US", "upper", "", true, false), + entry("spanish_phoneme", tokenizer_kind::ipa, 96, 103, "es-ES", "upper", "", true, true), + entry("german_phoneme", tokenizer_kind::ipa, 199, 150, "de-DE", "mixed", "#", true, true), + entry("mandarin_phoneme", tokenizer_kind::mandarin, 349, 109), + entry("japanese_phoneme", tokenizer_kind::japanese, 458, 175), + entry("french_chartokenizer", tokenizer_kind::byt5, 633, 384), + entry("hindi_chartokenizer", tokenizer_kind::hindi_chars, 1017, 191), + entry("italian_phoneme", tokenizer_kind::byt5, 1208, 384), + entry("vietnamese_phoneme", tokenizer_kind::byt5, 1592, 384), + entry("text_ce_tokenizer", tokenizer_kind::byt5, 1976, 384), + }; + profile.language_mapping = { + {"en", "english_phoneme"}, {"es", "spanish_phoneme"}, + {"de", "german_phoneme"}, {"fr", "french_chartokenizer"}, + {"it", "italian_phoneme"}, {"vi", "vietnamese_phoneme"}, + {"zh", "mandarin_phoneme"}, {"hi", "hindi_chartokenizer"}, + {"ja", "japanese_phoneme"}, + }; + if (profile.nemo_version != "2.6.0rc0") { + throw std::runtime_error( + "unsupported v2602 nemo_version '" + profile.nemo_version + "'"); + } + } else if (order == v2607_order) { + profile.id = "v2607"; + profile.text_vocab_size = 3359; + profile.eos_id = 3358; + profile.entries = { + entry("english_phoneme", tokenizer_kind::ipa, 0, 96, "en-US", "upper", "", true, false), + entry("text_ce_tokenizer", tokenizer_kind::byt5, 96, 384), + entry("spanish_phoneme", tokenizer_kind::ipa, 480, 103, "es-ES", "upper", "", true, true), + entry("german_phoneme", tokenizer_kind::ipa, 583, 150, "de-DE", "mixed", "#", true, true), + entry("mandarin_phoneme", tokenizer_kind::mandarin, 733, 109), + entry("japanese_phoneme", tokenizer_kind::japanese, 842, 175), + entry("portuguese_Brazilian_phoneme", tokenizer_kind::ipa, 1017, 111, "pt-BR", "upper", "#", true, true), + entry("hindi_phoneme", tokenizer_kind::ipa, 1128, 201, "hi-IN", "upper", "", true, true), + entry("arabic_AE_chartokenizer", tokenizer_kind::arabic, 1329, 164), + entry("arabic_SA_chartokenizer", tokenizer_kind::arabic, 1493, 164), + entry("arabic_MSA_chartokenizer", tokenizer_kind::arabic, 1657, 164), + entry("french_chartokenizer", tokenizer_kind::byt5, 1821, 384), + entry("italian_chartokenizer", tokenizer_kind::byt5, 2205, 384), + entry("vietnamese_chartokenizer", tokenizer_kind::byt5, 2589, 384), + entry("korean_chartokenizer", tokenizer_kind::byt5, 2973, 384), + }; + profile.language_mapping = parse_language_mapping(lines); + const std::map expected_mapping = { + {"en", "english_phoneme"}, {"de", "german_phoneme"}, + {"es", "spanish_phoneme"}, {"fr", "french_chartokenizer"}, + {"it", "italian_chartokenizer"}, {"vi", "vietnamese_chartokenizer"}, + {"zh", "mandarin_phoneme"}, {"hi", "hindi_phoneme"}, + {"ja", "japanese_phoneme"}, {"pt-BR", "portuguese_Brazilian_phoneme"}, + {"ko", "korean_chartokenizer"}, {"ar-AE", "arabic_AE_chartokenizer"}, + {"ar-SA", "arabic_SA_chartokenizer"}, {"ar-MSA", "arabic_MSA_chartokenizer"}, + }; + if (profile.language_mapping != expected_mapping) { + throw std::runtime_error("unsupported v2607 language_to_tokenizer_mapping"); + } + profile.language_mapping = { + {"en", "english_phoneme"}, {"de", "german_phoneme"}, + {"es", "spanish_phoneme"}, {"fr", "french_chartokenizer"}, + {"it", "italian_chartokenizer"}, {"vi", "vietnamese_chartokenizer"}, + {"zh", "mandarin_phoneme"}, {"hi", "hindi_phoneme"}, + {"ja", "japanese_phoneme"}, {"pt-br", "portuguese_Brazilian_phoneme"}, + {"ko", "korean_chartokenizer"}, {"ar-ae", "arabic_AE_chartokenizer"}, + {"ar-sa", "arabic_SA_chartokenizer"}, {"ar-msa", "arabic_MSA_chartokenizer"}, + }; + if (profile.nemo_version != "2.8.0rc0") { + throw std::runtime_error( + "unsupported v2607 nemo_version '" + profile.nemo_version + "'"); + } + } else { + std::ostringstream found; + for (size_t i = 0; i < order.size(); ++i) { + if (i != 0) found << ", "; + found << order[i]; + } + throw std::runtime_error( + "unsupported Magpie tokenizer layout in " + config_path.string() + ": [" + + found.str() + "]"); + } + + for (auto& item : profile.entries) { + if (item.kind == tokenizer_kind::japanese) { + item.ascii_letter_case = block_scalar(blocks.at(item.name), "ascii_letter_case"); + } + } + + for (const auto& item : profile.entries) { + switch (item.kind) { + case tokenizer_kind::ipa: + require_ipa_values( + blocks, item, item.locale == "en-US" ? std::string{} : item.locale); + if (item.name == "portuguese_Brazilian_phoneme") { + require_block_value(blocks, item.name, "locale_specific_punct", "false"); + } + break; + case tokenizer_kind::byt5: + require_byt5_values(blocks, item.name); + break; + case tokenizer_kind::hindi_chars: + require_common_tokenizer_values( + blocks, item.name, + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "HindiCharsTokenizer", + "true", "true"); + break; + case tokenizer_kind::mandarin: + require_common_tokenizer_values( + blocks, item.name, + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ChinesePhonemesTokenizer", + "true", "true"); + require_block_value(blocks, item.name, "ascii_letter_case", "upper"); + break; + case tokenizer_kind::japanese: + require_common_tokenizer_values( + blocks, item.name, + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "JapanesePhonemeTokenizer", + "false", "true"); + require_block_value( + blocks, item.name, "ascii_letter_case", item.ascii_letter_case); + break; + case tokenizer_kind::arabic: + require_common_tokenizer_values( + blocks, item.name, + "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." + "ArabicCharsTokenizer", + "true", "true"); + require_block_value(blocks, item.name, "charset_version", "1"); + break; + } + } + attach_artifacts(profile, blocks, root); + validate_profile_layout(profile); + return profile; +} + std::string trim_language(std::string value) { value.erase(value.begin(), std::find_if(value.begin(), value.end(), [](unsigned char c) { @@ -150,9 +626,7 @@ class MagpieNativeTokenizer::Impl { if (model_dir_.empty()) { throw std::invalid_argument("tokenizer model directory is required"); } - std::ifstream config_file(fs::path(model_dir_) / "model_config.yaml"); - std::string contents((std::istreambuf_iterator(config_file)), {}); - v2607_ = contents.find("portuguese_Brazilian_phoneme") != std::string::npos; + profile_ = load_tokenizer_profile(model_dir_); } MagpieTokenizationResult tokenize( @@ -166,40 +640,35 @@ class MagpieNativeTokenizer::Impl { p.model = model_dir_; p.text = text; p.language = MagpieNativeTokenizer::normalize_language_code(language_code); - p.v2607 = v2607_; p.chunk_text_transform = chunk_text_transform; - const std::string tokenizer_name = tokenizer_for_language(p.language); - if (tokenizer_name.empty()) { + const tokenizer_entry* selected = nullptr; + try { + selected = &profile_.entry_for_language(p.language); + } + catch (const std::invalid_argument&) { throw std::invalid_argument("unsupported language_code '" + language_code + "'"); } - if (!supports_native(p)) { + if (!supports(*selected)) { throw std::invalid_argument( "native Magpie tokenizer is not available for language_code '" + language_code + "'"); } + p.tokenizer_name = selected->name; + p.locale = selected->locale; + p.grapheme_case = selected->grapheme_case; + p.grapheme_prefix = selected->grapheme_prefix; + p.ascii_letter_case = selected->ascii_letter_case; + p.phoneme_dict = selected->phoneme_dict; + p.heteronyms = selected->heteronyms; + p.offset = selected->offset; + p.eos_id = profile_.eos_id; + p.expected_vocab_size = selected->size; + p.apostrophe = selected->apostrophe; + p.pad_with_space = selected->pad_with_space; p.sentence_chunking = should_tokenize_by_sentence(text, p.language, config_.sentence_limit); - tokenizer_result native = tokenize_native(p); - if (v2607_) { - int offset_delta = 0; - if (p.language == "es" || p.language == "de" || p.language == "zh" || - p.language == "ja") - offset_delta = 384; - else if (p.language == "fr") - offset_delta = 1188; - else if (p.language == "it" || p.language == "vi") - offset_delta = 997; - else if (p.language == "hi") - offset_delta = 111; - for (auto& chunk : native.chunks) - for (int& token : chunk.tokens) { - if (token == native.eos_id) - token = 3358; - else - token += offset_delta; - } - } + tokenizer_result native = tokenize_native(p, *selected); MagpieTokenizationResult out; out.language = native.language; out.tokenizer_name = native.tokenizer_name; @@ -219,31 +688,77 @@ class MagpieNativeTokenizer::Impl { } std::vector supported_language_codes() const { - std::vector languages = {"en-US", "es-ES", "de-DE", "fr-FR", - "it-IT", "vi-VN", "hi-IN"}; - if (v2607_) { - languages.insert(languages.end(), {"ar-AE", "ar-SA", "ar-MSA", "ko-KR", "pt-BR"}); + static const std::vector> canonical = { + {"en", "en-US"}, {"es", "es-ES"}, {"de", "de-DE"}, + {"fr", "fr-FR"}, {"it", "it-IT"}, {"vi", "vi-VN"}, + {"zh", "zh-CN"}, {"hi", "hi-IN"}, {"ja", "ja-JP"}, + {"ar-ae", "ar-AE"}, {"ar-sa", "ar-SA"}, {"ar-msa", "ar-MSA"}, + {"ko", "ko-KR"}, {"pt-br", "pt-BR"}, + }; + std::vector languages; + for (const auto& [normalized, code] : canonical) { + const auto mapping = profile_.language_mapping.find(normalized); + if (mapping == profile_.language_mapping.end()) { + continue; + } + const auto item = std::find_if( + profile_.entries.begin(), profile_.entries.end(), [&](const tokenizer_entry& entry) { + return entry.name == mapping->second; + }); + if (item != profile_.entries.end() && supports(*item)) { + languages.push_back(code); + } } -#ifdef NEMO_SPEECH_TTS_WITH_ZH - languages.emplace_back("zh-CN"); -#endif -#ifdef NEMO_SPEECH_TTS_WITH_JA - languages.emplace_back("ja-JP"); -#endif return languages; } + const std::string& profile_id() const { return profile_.id; } + int text_vocab_size() const { return profile_.text_vocab_size; } + private: - tokenizer_result tokenize_native(const params& p) const { - if (p.language == "en" || p.language == "es" || p.language == "de") { - return tokenize_ipa(p); + bool supports(const tokenizer_entry& entry) const { + if (entry.kind == tokenizer_kind::japanese) { +#ifdef NEMO_SPEECH_TTS_WITH_JA + return !find_openjtalk_dictionary_dir(model_dir_).empty(); +#else + return false; +#endif } + if (entry.kind == tokenizer_kind::mandarin) { #ifdef NEMO_SPEECH_TTS_WITH_ZH - if (p.language == "zh") { - return tokenize_mandarin(p); + return mandarin_tokenizer_available(model_dir_); +#else + return false; +#endif } + return true; + } + + tokenizer_result tokenize_native(const params& p, const tokenizer_entry& entry) const { + switch (entry.kind) { + case tokenizer_kind::ipa: + return tokenize_ipa(p); + case tokenizer_kind::byt5: + return run_byt5_native(p); + case tokenizer_kind::hindi_chars: + return run_hindi_native(p); + case tokenizer_kind::arabic: + return run_arabic_native(p); + case tokenizer_kind::mandarin: +#ifdef NEMO_SPEECH_TTS_WITH_ZH + return tokenize_mandarin(p); +#else + break; +#endif + case tokenizer_kind::japanese: +#ifdef NEMO_SPEECH_TTS_WITH_JA + return run_japanese_native(p); +#else + break; #endif - return run_native(p); + } + throw std::invalid_argument( + "native Magpie tokenizer is not available for language '" + p.language + "'"); } #ifdef NEMO_SPEECH_TTS_WITH_ZH @@ -255,7 +770,9 @@ class MagpieNativeTokenizer::Impl { const mandarin_tokenizer& mandarin_tokenizer_for_model() const { std::lock_guard lock(cache_mutex_); if (!mandarin_cache_) { - mandarin_cache_ = std::make_unique(model_dir_); + const auto& entry = profile_.entry_for_language("zh"); + mandarin_cache_ = std::make_unique( + model_dir_, entry.offset, entry.phoneme_dict); } return *mandarin_cache_; } @@ -266,13 +783,20 @@ class MagpieNativeTokenizer::Impl { throw std::runtime_error( "native IPA tokenization requires an extracted Magpie .nemo directory"); } - const ipa_config cfg = ipa_config_for_language(p.language); + const ipa_config cfg = ipa_config_for_params(p); const ipa_tokenizer& tok = ipa_tokenizer_for(p.language, cfg); + if (tok.vocab_size() != p.expected_vocab_size) { + throw std::runtime_error( + "tokenizer '" + p.tokenizer_name + "' vocabulary has " + + std::to_string(tok.vocab_size()) + " entries; expected " + + std::to_string(p.expected_vocab_size)); + } tokenizer_result result; result.language = p.language; result.tokenizer_name = cfg.tokenizer_name; - const int pad_id = ipa_pad_id_for_config(cfg); + result.eos_id = p.eos_id; + const int pad_id = tok.pad_id(); for (const std::string& sentence : tokenizer_input_units(p, p.text)) { chunk ch; ch.text = sentence; @@ -297,7 +821,7 @@ class MagpieNativeTokenizer::Impl { std::string model_dir_; MagpieTokenizerConfig config_; - bool v2607_ = false; + tokenizer_profile profile_; mutable std::mutex cache_mutex_; mutable std::map> ipa_cache_; #ifdef NEMO_SPEECH_TTS_WITH_ZH @@ -361,6 +885,16 @@ MagpieNativeTokenizer::supported_language_codes() const { return impl_->supported_language_codes(); } +const std::string& +MagpieNativeTokenizer::profile_id() const { + return impl_->profile_id(); +} + +int +MagpieNativeTokenizer::text_vocab_size() const { + return impl_->text_vocab_size(); +} + std::string ensure_terminal_punctuation(const std::string& text, const std::string& language_code) { const size_t last = text.find_last_not_of(" \t\r\n"); diff --git a/src/tts/tokenizer/tokenizer.h b/src/tts/tokenizer/tokenizer.h index 6b2b366..9853752 100644 --- a/src/tts/tokenizer/tokenizer.h +++ b/src/tts/tokenizer/tokenizer.h @@ -83,6 +83,8 @@ class MagpieNativeTokenizer { static std::string normalize_language_code(const std::string& language_code); std::vector supported_language_codes() const; + const std::string& profile_id() const; + int text_vocab_size() const; private: class Impl; diff --git a/src/tts/tokenizer/tokenizer_impl.cpp b/src/tts/tokenizer/tokenizer_impl.cpp index 2989f0e..8d26fd3 100644 --- a/src/tts/tokenizer/tokenizer_impl.cpp +++ b/src/tts/tokenizer/tokenizer_impl.cpp @@ -53,7 +53,18 @@ struct params { fs::path model; std::string text; std::string language = "en"; - bool v2607 = false; + std::string tokenizer_name; + std::string locale; + std::string grapheme_case; + std::string grapheme_prefix; + std::string ascii_letter_case; + std::string phoneme_dict; + std::string heteronyms; + int offset = 0; + int eos_id = 0; + int expected_vocab_size = 0; + bool apostrophe = true; + bool pad_with_space = false; bool sentence_chunking = true; std::function chunk_text_transform; }; @@ -66,7 +77,7 @@ struct chunk { struct tokenizer_result { std::string language; std::string tokenizer_name; - int eos_id = 2361; + int eos_id = 0; std::vector chunks; }; @@ -103,15 +114,6 @@ split_utf8(const std::string& s) { return out; } -static std::string -join_utf8(const std::vector& chars) { - std::string out; - for (const auto& c : chars) { - out += c; - } - return out; -} - static bool is_ascii_alnum(char c) { return (c >= '0' && c <= '9') || (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z'); @@ -332,8 +334,9 @@ hindi_tokens() { #ifdef NEMO_SPEECH_TTS_WITH_JA static std::vector -japanese_tokens() { - return {" ", "0", "1", "ァ", "ア", "ィ", "イ", "ゥ", "ウ", "ェ", "エ", "ォ", "オ", +japanese_tokens(bool lowercase_ascii = false) { + std::vector tokens = { + " ", "0", "1", "ァ", "ア", "ィ", "イ", "ゥ", "ウ", "ェ", "エ", "ォ", "オ", "カ", "ガ", "キ", "ギ", "ク", "グ", "ケ", "ゲ", "コ", "ゴ", "サ", "ザ", "シ", "ジ", "ス", "ズ", "セ", "ゼ", "ソ", "ゾ", "タ", "ダ", "チ", "ヂ", "ッ", "ツ", "ヅ", "テ", "デ", "ト", "ド", "ナ", "ニ", "ヌ", "ネ", "ノ", "ハ", "バ", "パ", @@ -347,6 +350,14 @@ japanese_tokens() { "《", "》", "「", "」", "『", "』", "【", "】", "〒", "〓", "〔", "〕", "〖", "〗", "〘", "〙", "〚", "〛", "〜", "〽", "・", "・・・", "ー", "﹅", "﹆", "！", "＊", "？", "｟", "｠", "", ""}; + if (lowercase_ascii) { + for (auto& token : tokens) { + if (token.size() == 1 && token[0] >= 'A' && token[0] <= 'Z') { + token[0] = static_cast(token[0] - 'A' + 'a'); + } + } + } + return tokens; } #endif @@ -372,6 +383,9 @@ ipa_punct(const std::string& locale) { } else if (locale == "es-ES") { p.insert("¿"); p.insert("¡"); + } else if (locale == "hi-IN") { + p.insert("।"); + p.insert("॥"); } return std::vector(p.begin(), p.end()); } @@ -851,60 +865,6 @@ pad_short_text_chunk_before_eos(chunk& ch, int eos_id, int pad_id) { ch.tokens.insert(ch.tokens.begin() + 1, pad_count, pad_id); } -static std::string -tokenizer_for_language(const std::string& lang) { - if (lang == "en") - return "english_phoneme"; - if (lang == "de") - return "german_phoneme"; - if (lang == "es") - return "spanish_phoneme"; - if (lang == "fr") - return "french_chartokenizer"; - if (lang == "it") - return "italian_phoneme"; - if (lang == "vi") - return "vietnamese_phoneme"; - if (lang == "pt-br") - return "portuguese_Brazilian_phoneme"; - if (lang == "ko") - return "korean_chartokenizer"; - if (lang == "ar-ae") - return "arabic_AE_chartokenizer"; - if (lang == "ar-sa") - return "arabic_SA_chartokenizer"; - if (lang == "ar-msa") - return "arabic_MSA_chartokenizer"; -#ifdef NEMO_SPEECH_TTS_WITH_ZH - if (lang == "zh") - return "mandarin_phoneme"; -#endif - if (lang == "hi") - return "hindi_chartokenizer"; -#ifdef NEMO_SPEECH_TTS_WITH_JA - if (lang == "ja") - return "japanese_phoneme"; -#endif - return ""; -} - -static fs::path -find_file_containing(const fs::path& root, const std::string& needle) { - if (!fs::is_directory(root)) { - return {}; - } - for (const auto& entry : fs::directory_iterator(root)) { - if (!entry.is_regular_file()) { - continue; - } - const std::string name = entry.path().filename().string(); - if (name.find(needle) != std::string::npos) { - return entry.path(); - } - } - return {}; -} - struct ipa_config { std::string tokenizer_name; int offset = 0; @@ -918,45 +878,21 @@ struct ipa_config { }; static ipa_config -ipa_config_for_language(const std::string& lang) { - if (lang == "en") { - return {"english_phoneme", - 0, - "ipa_cmudict", - "heteronyms-052722", - "en-US", - "upper", - "", - true, - false}; - } - if (lang == "es") { - return {"spanish_phoneme", 96, "es_ES", "", "es-ES", "upper", "", true, true}; - } - if (lang == "de") { - return { - "german_phoneme", - 199, - "de_nv230119.dict", - "de_nv230119.heteronym", - "de-DE", - "mixed", - "#", - true, - true}; - } - if (lang == "pt-br") { - return {"portuguese_Brazilian_phoneme", - 1017, - "pt_br_prondict", - "", - "pt-BR", - "upper", - "#", - true, - true}; - } - throw std::runtime_error("no native IPA tokenizer for language " + lang); +ipa_config_for_params(const params& p) { + if (p.tokenizer_name.empty() || p.phoneme_dict.empty()) { + throw std::runtime_error("incomplete IPA tokenizer profile for language " + p.language); + } + return { + p.tokenizer_name, + p.offset, + p.phoneme_dict, + p.heteronyms, + p.locale, + p.grapheme_case, + p.grapheme_prefix, + p.apostrophe, + p.pad_with_space, + }; } static std::vector @@ -1009,18 +945,35 @@ exact_ipa_tokens(const std::string& tokenizer_name) { "ɪ", "ɲ", "ɹ", "ɾ", "ʁ", "ʃ", "ʊ", "ʌ", "ʎ", "ʒ", "ʲ", "ˈ", "ˌ", "ː", "̃", "θ", "ẽ", " ", "", ""}; } + if (tokenizer_name == "hindi_phoneme") { + return { + "!", "\"", "'", "(", ")", ",", "-", ".", "/", "0", "1", "2", "3", "4", + "5", "6", "7", "8", "9", ":", ";", "?", "A", "B", "C", "D", "E", "F", + "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", + "U", "V", "W", "X", "Y", "Z", "[", "]", "a", "b", "c", "d", "e", "f", + "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", + "v", "w", "x", "z", "{", "}", "À", "É", "ã", "æ", "ð", "õ", "ĩ", "ŋ", + "ũ", "ɑ", "ɔ", "ɖ", "ə", "ɚ", "ɛ", "ɝ", "ɟ", "ɡ", "ɣ", "ɪ", "ɭ", "ɲ", + "ɳ", "ɹ", "ɾ", "ʂ", "ʃ", "ʈ", "ʊ", "ʋ", "ʌ", "ʒ", "ʰ", "ˈ", "ˌ", "ː", + "̃", "̩", "θ", "χ", "ँ", "ं", "ः", "अ", "आ", "इ", "ई", "उ", "ऊ", "ऋ", + "ऌ", "ऍ", "ऎ", "ए", "ऐ", "ऑ", "ओ", "औ", "क", "ख", "ग", "घ", "ङ", "च", + "छ", "ज", "झ", "ञ", "ट", "ठ", "ड", "ढ", "ण", "त", "थ", "द", "ध", "न", + "ऩ", "प", "फ", "ब", "भ", "म", "य", "र", "ऱ", "ल", "ळ", "ऴ", "व", "श", + "ष", "स", "ह", "ऺ", "़", "ऽ", "ा", "ि", "ी", "ु", "ू", "ृ", "ॅ", "ॆ", + "े", "ै", "ॉ", "ॊ", "ो", "ौ", "्", "ॐ", "॓", "ॠ", "ॡ", "ॢ", "।", "॥", + "॰", "ẽ", " ", "", "", + }; + } return {}; } -static int -ipa_pad_id_for_config(const ipa_config& cfg) { - return token_id_for_symbol(exact_ipa_tokens(cfg.tokenizer_name), cfg.offset, ""); -} - class ipa_tokenizer { public: ipa_tokenizer(const fs::path& root, ipa_config cfg) : cfg_(std::move(cfg)) { load(root); } + int vocab_size() const { return static_cast(tokens_.size()); } + int pad_id() const { return cfg_.offset + token_to_id_.at(""); } + std::vector encode(const std::string& raw_text) const { std::string text = raw_text; if (cfg_.locale == "en-US") { @@ -1087,17 +1040,16 @@ class ipa_tokenizer { std::unordered_set punct_; void load(const fs::path& root) { - const fs::path dict_path = find_file_containing(root, cfg_.dict_hint); - if (dict_path.empty()) { - throw std::runtime_error( - "failed to find tokenizer dictionary containing '" + cfg_.dict_hint + "'"); + const fs::path dict_path = root / cfg_.dict_hint; + if (!fs::is_regular_file(dict_path)) { + throw std::runtime_error("failed to find tokenizer dictionary '" + dict_path.string() + "'"); } for (const auto& p : ipa_punct(cfg_.locale)) { punct_.insert(p); } if (!cfg_.heteronym_hint.empty()) { - const fs::path heteronym_path = find_file_containing(root, cfg_.heteronym_hint); - if (!heteronym_path.empty()) { + const fs::path heteronym_path = root / cfg_.heteronym_hint; + if (fs::is_regular_file(heteronym_path)) { std::istringstream in(read_file(heteronym_path)); std::string line; while (std::getline(in, line)) { @@ -1173,12 +1125,12 @@ class ipa_tokenizer { tokens_ = exact_ipa_tokens(cfg_.tokenizer_name); if (tokens_.empty()) { tokens_ = std::vector(symbols.begin(), symbols.end()); + if (std::find(tokens_.begin(), tokens_.end(), " ") == tokens_.end()) { + tokens_.push_back(" "); + } + tokens_.push_back(""); + tokens_.push_back(""); } - if (std::find(tokens_.begin(), tokens_.end(), " ") == tokens_.end()) { - tokens_.push_back(" "); - } - tokens_.push_back(""); - tokens_.push_back(""); for (size_t i = 0; i < tokens_.size(); ++i) { token_to_id_[tokens_[i]] = (int)i; @@ -1367,42 +1319,20 @@ class ipa_tokenizer { }; static tokenizer_result -run_byt5_native(const params& p, int offset, const std::string& tokenizer_name) { +run_byt5_native(const params& p) { tokenizer_result result; result.language = p.language; - result.tokenizer_name = tokenizer_name; + result.tokenizer_name = p.tokenizer_name; + result.eos_id = p.eos_id; for (const std::string& sentence : tokenizer_input_units(p, p.text)) { chunk ch; ch.text = sentence; for (unsigned char b : sentence) { - ch.tokens.push_back(offset + (int)b + 3); + ch.tokens.push_back(p.offset + (int)b + 3); } - ch.tokens.push_back(offset + 1); + ch.tokens.push_back(p.offset + 1); ch.tokens.push_back(result.eos_id); - pad_short_text_chunk_before_eos(ch, result.eos_id, offset); - result.chunks.push_back(std::move(ch)); - } - return result; -} - -static tokenizer_result -run_ipa_native(const params& p) { - if (!fs::is_directory(p.model)) { - throw std::runtime_error( - "native IPA tokenization requires an extracted Magpie .nemo directory"); - } - const ipa_config cfg = ipa_config_for_language(p.language); - ipa_tokenizer tok(p.model, cfg); - tokenizer_result result; - result.language = p.language; - result.tokenizer_name = cfg.tokenizer_name; - const int pad_id = ipa_pad_id_for_config(cfg); - for (const std::string& sentence : tokenizer_input_units(p, p.text)) { - chunk ch; - ch.text = sentence; - ch.tokens = tok.encode(sentence); - ch.tokens.push_back(result.eos_id); - pad_short_text_chunk_before_eos(ch, result.eos_id, pad_id); + pad_short_text_chunk_before_eos(ch, result.eos_id, p.offset); result.chunks.push_back(std::move(ch)); } return result; @@ -1410,17 +1340,20 @@ run_ipa_native(const params& p) { static tokenizer_result run_hindi_native(const params& p) { - static const int offset = 1017; const std::vector tokens = hindi_tokens(); + if (static_cast(tokens.size()) != p.expected_vocab_size) { + throw std::runtime_error("Hindi character vocabulary does not match tokenizer profile"); + } std::unordered_map token_to_id; for (size_t i = 0; i < tokens.size(); ++i) { token_to_id[tokens[i]] = (int)i; } - const int pad_id = token_id_for_symbol(tokens, offset, ""); + const int pad_id = token_id_for_symbol(tokens, p.offset, ""); tokenizer_result result; result.language = p.language; - result.tokenizer_name = "hindi_chartokenizer"; + result.tokenizer_name = p.tokenizer_name; + result.eos_id = p.eos_id; for (const std::string& sentence : tokenizer_input_units(p, replace_all(p.text, "’", "'"))) { chunk ch; ch.text = sentence; @@ -1444,7 +1377,7 @@ run_hindi_native(const params& p) { for (const auto& c : chars) { const auto it = token_to_id.find(c); if (it != token_to_id.end()) { - ch.tokens.push_back(offset + it->second); + ch.tokens.push_back(p.offset + it->second); } } ch.tokens.push_back(result.eos_id); @@ -1474,14 +1407,18 @@ arabic_tokens() { } static tokenizer_result -run_arabic_native(const params& p, int offset, const std::string& tokenizer_name) { +run_arabic_native(const params& p) { const auto tokens = arabic_tokens(); + if (static_cast(tokens.size()) != p.expected_vocab_size) { + throw std::runtime_error("Arabic vocabulary does not match tokenizer profile"); + } std::unordered_map ids; for (size_t i = 0; i < tokens.size(); ++i) ids[tokens[i]] = (int)i; tokenizer_result result; result.language = p.language; - result.tokenizer_name = tokenizer_name; - const int pad_id = token_id_for_symbol(tokens, offset, ""); + result.tokenizer_name = p.tokenizer_name; + result.eos_id = p.eos_id; + const int pad_id = token_id_for_symbol(tokens, p.offset, ""); for (const std::string& sentence : tokenizer_input_units(p, p.text)) { chunk ch; ch.text = sentence; @@ -1497,7 +1434,7 @@ run_arabic_native(const params& p, int offset, const std::string& tokenizer_name while (!symbols.empty() && symbols.back() == " ") symbols.pop_back(); symbols.insert(symbols.begin(), " "); symbols.push_back(" "); - for (const auto& symbol : symbols) ch.tokens.push_back(offset + ids[symbol]); + for (const auto& symbol : symbols) ch.tokens.push_back(p.offset + ids[symbol]); ch.tokens.push_back(result.eos_id); pad_short_text_chunk_before_eos(ch, result.eos_id, pad_id); result.chunks.push_back(std::move(ch)); @@ -1510,7 +1447,8 @@ static tokenizer_result run_mandarin_native(const params& p, const mandarin_tokenizer& tok) { tokenizer_result result; result.language = p.language; - result.tokenizer_name = "mandarin_phoneme"; + result.tokenizer_name = p.tokenizer_name; + result.eos_id = p.eos_id; for (const std::string& sentence : tokenizer_input_units(p, p.text)) { chunk ch; ch.text = sentence; @@ -1525,12 +1463,14 @@ run_mandarin_native(const params& p, const mandarin_tokenizer& tok) { #ifdef NEMO_SPEECH_TTS_WITH_JA static bool -is_ascii_upper_word(const std::string& text) { +is_ascii_letter_word(const std::string& text, bool lowercase) { if (text.empty()) { return false; } for (const unsigned char c : text) { - if (c < 'A' || c > 'Z') { + const unsigned char first = lowercase ? 'a' : 'A'; + const unsigned char last = lowercase ? 'z' : 'Z'; + if (c < first || c > last) { return false; } } @@ -1568,13 +1508,13 @@ process_japanese_chain( } static std::vector -japanese_g2p(openjtalk_frontend& frontend, const std::string& text) { +japanese_g2p(openjtalk_frontend& frontend, const std::string& text, bool lowercase_ascii) { const std::vector words = frontend.run(text); std::vector result; std::vector current_chain; - const std::unordered_map token_to_id = [] { + const std::unordered_map token_to_id = [lowercase_ascii] { std::unordered_map ids; - const auto tokens = japanese_tokens(); + const auto tokens = japanese_tokens(lowercase_ascii); for (size_t i = 0; i < tokens.size(); ++i) { ids[tokens[i]] = (int)i; } @@ -1583,7 +1523,7 @@ japanese_g2p(openjtalk_frontend& frontend, const std::string& text) { for (size_t idx = 0; idx < words.size(); ++idx) { const auto& word = words[idx]; - if (is_ascii_upper_word(word.text)) { + if (is_ascii_letter_word(word.text, lowercase_ascii)) { process_japanese_chain(current_chain, result); current_chain.clear(); const auto chars = split_utf8(word.text); @@ -1620,7 +1560,7 @@ japanese_g2p(openjtalk_frontend& frontend, const std::string& text) { static tokenizer_result run_japanese_native(const params& p) { - static const int offset = 458; + const bool lowercase_ascii = p.ascii_letter_case == "lower"; const fs::path dictionary_dir = find_openjtalk_dictionary_dir(p.model); if (dictionary_dir.empty()) { throw std::runtime_error( @@ -1628,22 +1568,27 @@ run_japanese_native(const params& p) { "open_jtalk_dic"); } const auto frontend = openjtalk_frontend_for_dictionary(dictionary_dir); - const std::vector tokens = japanese_tokens(); + const std::vector tokens = japanese_tokens(lowercase_ascii); + if (static_cast(tokens.size()) != p.expected_vocab_size) { + throw std::runtime_error("Japanese vocabulary does not match tokenizer profile"); + } std::unordered_map token_to_id; for (size_t i = 0; i < tokens.size(); ++i) { token_to_id[tokens[i]] = (int)i; } - const int pad_id = token_id_for_symbol(tokens, offset, ""); + const int pad_id = token_id_for_symbol(tokens, p.offset, ""); tokenizer_result result; result.language = p.language; - result.tokenizer_name = "japanese_phoneme"; - for (const std::string& sentence : tokenizer_input_units(p, ascii_upper(p.text))) { + result.tokenizer_name = p.tokenizer_name; + result.eos_id = p.eos_id; + const std::string cased_text = lowercase_ascii ? ascii_lower(p.text) : ascii_upper(p.text); + for (const std::string& sentence : tokenizer_input_units(p, cased_text)) { chunk ch; ch.text = sentence; std::vector symbols; const std::string space = " "; - for (const auto& symbol : japanese_g2p(*frontend, sentence)) { + for (const auto& symbol : japanese_g2p(*frontend, sentence, lowercase_ascii)) { if (symbol == space) { if (!symbols.empty() && symbols.back() != space) { symbols.push_back(symbol); @@ -1661,7 +1606,7 @@ run_japanese_native(const params& p) { for (const auto& symbol : symbols) { const auto it = token_to_id.find(symbol); if (it != token_to_id.end()) { - ch.tokens.push_back(offset + it->second); + ch.tokens.push_back(p.offset + it->second); } } ch.tokens.push_back(result.eos_id); @@ -1671,79 +1616,3 @@ run_japanese_native(const params& p) { return result; } #endif - -static bool -supports_native(const params& p) { - if (p.language == "fr" || p.language == "it" || p.language == "vi") { - return true; - } - if (p.language == "hi") { - return true; - } - if ((p.language == "pt-br" || p.language == "ko" || p.language == "ar-ae" || - p.language == "ar-sa" || p.language == "ar-msa") && - p.v2607) { - return p.language != "pt-br" || fs::is_directory(p.model); - } -#ifdef NEMO_SPEECH_TTS_WITH_JA - if (p.language == "ja") { - return !find_openjtalk_dictionary_dir(p.model).empty(); - } -#endif -#ifdef NEMO_SPEECH_TTS_WITH_ZH - if (p.language == "zh") { - return mandarin_tokenizer_available(p.model); - } -#endif - if ((p.language == "en" || p.language == "es" || p.language == "de") && - fs::is_directory(p.model)) { - return true; - } - return false; -} - -static tokenizer_result -run_native(const params& p) { - if (p.language == "fr") { - return run_byt5_native(p, 633, "french_chartokenizer"); - } - if (p.language == "it") { - return run_byt5_native(p, 1208, "italian_phoneme"); - } - if (p.language == "vi") { - return run_byt5_native(p, 1592, "vietnamese_phoneme"); - } - if (p.language == "hi") { - return run_hindi_native(p); - } - if (p.language == "pt-br") { - return run_ipa_native(p); - } - if (p.language == "ko") { - return run_byt5_native(p, 2973, "korean_chartokenizer"); - } - if (p.language == "ar-ae") { - return run_arabic_native(p, 1329, "arabic_AE_chartokenizer"); - } - if (p.language == "ar-sa") { - return run_arabic_native(p, 1493, "arabic_SA_chartokenizer"); - } - if (p.language == "ar-msa") { - return run_arabic_native(p, 1657, "arabic_MSA_chartokenizer"); - } -#ifdef NEMO_SPEECH_TTS_WITH_JA - if (p.language == "ja") { - return run_japanese_native(p); - } -#endif -#ifdef NEMO_SPEECH_TTS_WITH_ZH - if (p.language == "zh") { - const mandarin_tokenizer tok(p.model); - return run_mandarin_native(p, tok); - } -#endif - if (p.language == "en" || p.language == "es" || p.language == "de") { - return run_ipa_native(p); - } - throw std::runtime_error("native tokenizer is not available for language '" + p.language + "'"); -} diff --git a/tests/conversion/tts_tokenizer_profiles_test.py b/tests/conversion/tts_tokenizer_profiles_test.py new file mode 100644 index 0000000..475872c --- /dev/null +++ b/tests/conversion/tts_tokenizer_profiles_test.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) + +from conversion.tts_tokenizer_profiles import ( # noqa: E402 + TOKENIZER_ORDERS, + TOKENIZER_PROFILE_NEMO_VERSIONS, + TOKENIZER_TARGETS, + V2607_LANGUAGE_MAPPING, + tokenizer_profile, +) + + +class TtsTokenizerProfilesTest(unittest.TestCase): + @staticmethod + def _config(profile: str) -> dict: + tokenizers = { + name: {"_target_": TOKENIZER_TARGETS[profile][name]} + for name in TOKENIZER_ORDERS[profile] + } + tokenizers["japanese_phoneme"]["g2p"] = { + "ascii_letter_case": "upper" if profile == "v2602" else "lower" + } + config = { + "nemo_version": TOKENIZER_PROFILE_NEMO_VERSIONS[profile], + "text_tokenizers": tokenizers, + } + if profile == "v2607": + tokenizers["hindi_phoneme"]["locale"] = "hi-IN" + tokenizers["hindi_phoneme"]["g2p"] = {"locale": "hi-IN"} + tokenizers["portuguese_Brazilian_phoneme"]["locale_specific_punct"] = False + config["language_to_tokenizer_mapping"] = V2607_LANGUAGE_MAPPING.copy() + return config + + def test_known_profiles_require_matching_dimensions(self) -> None: + self.assertEqual(tokenizer_profile(self._config("v2602"), 2362, 1), "v2602") + self.assertEqual(tokenizer_profile(self._config("v2607"), 3359, 2), "v2607") + with self.assertRaisesRegex(ValueError, "requires text_vocab_size"): + tokenizer_profile(self._config("v2602"), 3359, 2) + with self.assertRaisesRegex(ValueError, "requires text_vocab_size"): + tokenizer_profile(self._config("v2607"), 2362, 1) + + def test_unknown_or_reordered_layout_is_rejected(self) -> None: + config = self._config("v2607") + config["text_tokenizers"]["unexpected_tokenizer"] = {} + with self.assertRaisesRegex(ValueError, "unsupported Magpie tokenizer layout"): + tokenizer_profile(config, 3359, 2) + + reordered = list(TOKENIZER_ORDERS["v2607"]) + reordered[0], reordered[1] = reordered[1], reordered[0] + with self.assertRaisesRegex(ValueError, "unsupported Magpie tokenizer layout"): + tokenizer_profile( + {"text_tokenizers": {name: {} for name in reordered}}, 3359, 2 + ) + + def test_critical_tokenizer_config_changes_are_rejected(self) -> None: + config = self._config("v2607") + config["text_tokenizers"]["hindi_phoneme"]["_target_"] = "HindiCharsTokenizer" + with self.assertRaisesRegex(ValueError, "tokenizer target for hindi_phoneme"): + tokenizer_profile(config, 3359, 2) + + config = self._config("v2607") + config["text_tokenizers"]["japanese_phoneme"]["g2p"][ + "ascii_letter_case" + ] = "upper" + with self.assertRaisesRegex(ValueError, "Japanese ascii_letter_case"): + tokenizer_profile(config, 3359, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/cpp/tts/CMakeLists.txt b/tests/cpp/tts/CMakeLists.txt index d78656b..b2e53fe 100644 --- a/tests/cpp/tts/CMakeLists.txt +++ b/tests/cpp/tts/CMakeLists.txt @@ -46,6 +46,16 @@ endif() if(NEMO_SPEECH_TTS_WITH_ZH) target_compile_definitions(test_tokenizer_single_chars PRIVATE NEMO_SPEECH_TTS_WITH_ZH=1) endif() +if(DEFINED ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2602}) + add_test( + NAME tokenizer_v2602 + COMMAND test_tokenizer_single_chars $ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2602}) +endif() +if(DEFINED ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2607}) + add_test( + NAME tokenizer_v2607 + COMMAND test_tokenizer_single_chars $ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2607}) +endif() if(NEMO_SPEECH_TTS_WITH_ZH) add_executable(test_tokenizer_mandarin test_tokenizer_mandarin.cpp) @@ -53,6 +63,16 @@ if(NEMO_SPEECH_TTS_WITH_ZH) test_tokenizer_mandarin PRIVATE nemo_speech_tts_tokenizer Threads::Threads) target_compile_definitions(test_tokenizer_mandarin PRIVATE MANDARIN_TEST_DATA_DIR="${CMAKE_CURRENT_SOURCE_DIR}/data") + if(DEFINED ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2602}) + add_test( + NAME tokenizer_mandarin_v2602 + COMMAND test_tokenizer_mandarin $ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2602}) + endif() + if(DEFINED ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2607}) + add_test( + NAME tokenizer_mandarin_v2607 + COMMAND test_tokenizer_mandarin $ENV{NEMO_SPEECH_TEST_TOKENIZER_DIR_V2607}) + endif() endif() add_executable(test_tts_terminal_punctuation test_tts_terminal_punctuation.cpp) diff --git a/tests/cpp/tts/test_grpc_tts_config.cpp b/tests/cpp/tts/test_grpc_tts_config.cpp index 58430e1..9f7448c 100644 --- a/tests/cpp/tts/test_grpc_tts_config.cpp +++ b/tests/cpp/tts/test_grpc_tts_config.cpp @@ -111,6 +111,8 @@ main() { synthesizer_config.tokenizer_model_dir = tokenizer_dir; synthesizer_config.default_language_code = "en-US"; auto synthesizer = std::make_shared(std::move(synthesizer_config)); + const std::vector supported_languages = + synthesizer->supported_language_codes(); nemo_speech::GrpcTtsService service(std::move(synthesizer)); nr_tts::RivaSynthesisConfigRequest req; @@ -120,7 +122,6 @@ main() { bool ok = true; ok &= expect(status.ok(), "GetRivaSynthesisConfig succeeds"); - const std::vector supported_languages = tts::supported_language_codes(); ok &= expect( resp.model_config_size() == static_cast(supported_languages.size()), "one TTS model config is returned per supported language"); diff --git a/tests/cpp/tts/test_magpietts_frame_stacking.cpp b/tests/cpp/tts/test_magpietts_frame_stacking.cpp index d094651..225849c 100644 --- a/tests/cpp/tts/test_magpietts_frame_stacking.cpp +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -9,6 +9,27 @@ namespace tts = nemo_speech::tts; int main() { + tts::magpietts_hparams v2602; + v2602.text_vocab_size = 2362; + v2602.frame_stacking_factor = 1; + tts::magpietts_hparams v2607; + v2607.text_vocab_size = 3359; + v2607.frame_stacking_factor = 2; + if (tts::magpietts_infer_tokenizer_profile(v2602) != "v2602" || + tts::magpietts_infer_tokenizer_profile(v2607) != "v2607" || + !tts::magpietts_tokenizer_profile_matches("v2602", v2602) || + !tts::magpietts_tokenizer_profile_matches("v2607", v2607) || + tts::magpietts_tokenizer_profile_matches("v2602", v2607) || + tts::magpietts_tokenizer_profile_matches("v2607", v2602)) { + std::fprintf(stderr, "tokenizer profile compatibility check failed\n"); + return 1; + } + v2607.text_vocab_size = 2362; + if (!tts::magpietts_infer_tokenizer_profile(v2607).empty()) { + std::fprintf(stderr, "unknown tokenizer dimensions were accepted\n"); + return 1; + } + tts::magpietts_hparams h; h.audio_codebooks = 2; h.frame_stacking_factor = 2; diff --git a/tests/cpp/tts/test_tokenizer_mandarin.cpp b/tests/cpp/tts/test_tokenizer_mandarin.cpp index a8f367d..3085c6d 100644 --- a/tests/cpp/tts/test_tokenizer_mandarin.cpp +++ b/tests/cpp/tts/test_tokenizer_mandarin.cpp @@ -83,10 +83,16 @@ main(int argc, char** argv) { } nemo_speech::tts::MagpieNativeTokenizer tokenizer(argv[1]); + const bool v2607 = tokenizer.profile_id() == "v2607"; std::vector> expected; bool ok = true; for (size_t index = 0; index < utterances.size(); ++index) { expected.push_back(parse_tokens(golden_lines[index])); + if (v2607) { + for (auto& token : expected.back()) { + token = token == 2361 ? 3358 : token + 384; + } + } const auto result = tokenizer.tokenize(utterances[index], "zh-CN"); ok &= result.language == "zh"; ok &= result.tokenizer_name == "mandarin_phoneme"; diff --git a/tests/cpp/tts/test_tokenizer_single_chars.cpp b/tests/cpp/tts/test_tokenizer_single_chars.cpp index f8e029f..3cf0326 100644 --- a/tests/cpp/tts/test_tokenizer_single_chars.cpp +++ b/tests/cpp/tts/test_tokenizer_single_chars.cpp @@ -2,7 +2,6 @@ // SPDX-License-Identifier: Apache-2.0 #include #include -#include #include #include #include @@ -73,6 +72,21 @@ check_chunk_count( return false; } +bool +check_supported_languages( + const tts::MagpieNativeTokenizer& tokenizer, const std::vector& expected) { + const auto actual = tokenizer.supported_language_codes(); + if (actual == expected) { + return true; + } + std::fprintf(stderr, "supported language mismatch\nexpected:"); + for (const auto& language : expected) std::fprintf(stderr, " %s", language.c_str()); + std::fprintf(stderr, "\nactual:"); + for (const auto& language : actual) std::fprintf(stderr, " %s", language.c_str()); + std::fprintf(stderr, "\n"); + return false; +} + int count_words_ascii_space(const std::string& text) { int count = 0; @@ -131,13 +145,35 @@ main(int argc, char** argv) { } tts::MagpieNativeTokenizer tokenizer(argv[1]); - std::ifstream tokenizer_config(std::string(argv[1]) + "/model_config.yaml"); - const std::string config_text( - (std::istreambuf_iterator(tokenizer_config)), std::istreambuf_iterator()); - const bool v2607 = config_text.find("portuguese_Brazilian_phoneme") != std::string::npos; + const bool v2607 = tokenizer.profile_id() == "v2607"; if (v2607) { bool v2607_ok = true; + if (tokenizer.text_vocab_size() != 3359) { + std::fprintf(stderr, "v2607 tokenizer reported the wrong text vocabulary size\n"); + v2607_ok = false; + } + std::vector expected_languages = { + "en-US", "es-ES", "de-DE", "fr-FR", "it-IT", "vi-VN", +#ifdef NEMO_SPEECH_TTS_WITH_ZH + "zh-CN", +#endif + "hi-IN", +#ifdef NEMO_SPEECH_TTS_WITH_JA + "ja-JP", +#endif + "ar-AE", "ar-SA", "ar-MSA", "ko-KR", "pt-BR", + }; + v2607_ok &= check_supported_languages(tokenizer, expected_languages); v2607_ok &= check_tokens(tokenizer, "A", "en-US", {90, 94, 94, 94, 88, 3358}); + v2607_ok &= check_tokens(tokenizer, "T", "es-ES", {580, 581, 581, 511, 580, 3358}); + v2607_ok &= check_tokens( + tokenizer, "er ist sehr gut", "de-DE", + {730, 615, 628, 730, 619, 629, 630, 730, 691, 716, 672, 718, 711, 730, 617, 631, + 630, 730, 3358}); + v2607_ok &= check_tokens( + tokenizer, "नमस्ते दुनिया।", "hi-IN", + {1326, 1190, 1216, 1189, 1234, 1195, 1196, 1182, 1239, 1326, 1279, 1305, 1281, + 1303, 1288, 1302, 1322, 1326, 3358}); v2607_ok &= check_tokens( tokenizer, "مرحبا", "ar-AE", {1329, 1405, 1391, 1387, 1382, 1381, 1329, 3358}); v2607_ok &= check_tokens( @@ -150,9 +186,35 @@ main(int argc, char** argv) { 3124, 3022, 2974, 3358}); v2607_ok &= check_tokens( tokenizer, "Olá!", "pt-BR", {1125, 1082, 1079, 1119, 1070, 1017, 1125, 3358}); + v2607_ok &= check_tokens(tokenizer, "A", "fr-FR", {1889, 1821, 1821, 1821, 1822, 3358}); + v2607_ok &= check_tokens(tokenizer, "A", "it-IT", {2273, 2205, 2205, 2205, 2206, 3358}); + v2607_ok &= check_tokens(tokenizer, "A", "vi-VN", {2657, 2589, 2589, 2589, 2590, 3358}); +#ifdef NEMO_SPEECH_TTS_WITH_JA + v2607_ok &= check_tokens(tokenizer, "あ", "ja-JP", {842, 1015, 843, 846, 842, 3358}); +#else + v2607_ok &= check_unsupported(tokenizer, "こんにちは世界。", "ja-JP"); +#endif +#ifndef NEMO_SPEECH_TTS_WITH_ZH + v2607_ok &= check_unsupported(tokenizer, "你好。", "zh-CN"); +#endif return v2607_ok ? 0 : 1; } + if (tokenizer.profile_id() != "v2602" || tokenizer.text_vocab_size() != 2362) { + std::fprintf(stderr, "v2602 tokenizer reported the wrong profile metadata\n"); + return 1; + } bool ok = true; + std::vector expected_languages = { + "en-US", "es-ES", "de-DE", "fr-FR", "it-IT", "vi-VN", +#ifdef NEMO_SPEECH_TTS_WITH_ZH + "zh-CN", +#endif + "hi-IN", +#ifdef NEMO_SPEECH_TTS_WITH_JA + "ja-JP", +#endif + }; + ok &= check_supported_languages(tokenizer, expected_languages); ok &= check_tokens(tokenizer, "A", "en-US", {90, 94, 94, 53, 84, 2361}); ok &= check_tokens(tokenizer, "T", "en-US", {90, 94, 94, 65, 56, 2361}); From 3c48ea2c127cab758713d959ebda33854ba68c5c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:53:46 +0000 Subject: [PATCH 09/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- conversion/tts_tokenizer_profiles.py | 26 +--- src/runtime/ggml/backend.cpp | 8 +- src/tts/magpietts/decoder.cpp | 92 +++++------ src/tts/magpietts/magpietts_cuda_sampling.cu | 104 +++++++------ src/tts/magpietts/magpietts_cuda_sampling.h | 4 +- src/tts/magpietts/model.h | 3 +- src/tts/tokenizer/mandarin_tokenizer.h | 3 +- src/tts/tokenizer/tokenizer.cpp | 144 ++++++++++-------- src/tts/tokenizer/tokenizer_impl.cpp | 69 ++++----- .../conversion/tts_tokenizer_profiles_test.py | 8 +- tests/cpp/tts/test_grpc_tts_config.cpp | 3 +- tests/cpp/tts/test_tokenizer_single_chars.cpp | 10 +- 12 files changed, 236 insertions(+), 238 deletions(-) diff --git a/conversion/tts_tokenizer_profiles.py b/conversion/tts_tokenizer_profiles.py index dd9b1d8..5176514 100644 --- a/conversion/tts_tokenizer_profiles.py +++ b/conversion/tts_tokenizer_profiles.py @@ -49,9 +49,7 @@ "v2607": "2.8.0rc0", } -IPA_TARGET = ( - "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer" -) +IPA_TARGET = "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers.IPATokenizer" BYT5_TARGET = "AutoTokenizer" TOKENIZER_TARGETS = { "v2602": { @@ -137,9 +135,7 @@ def tokenizer_profile(cfg: dict[str, Any], text_vocab: int, frame_stacking: int) raise ValueError(f"unsupported Magpie tokenizer layout: {list(order)}") profile = matches[0] if str(cfg.get("nemo_version", "")) != TOKENIZER_PROFILE_NEMO_VERSIONS[profile]: - raise ValueError( - f"unsupported {profile} nemo_version: {cfg.get('nemo_version')!r}" - ) + raise ValueError(f"unsupported {profile} nemo_version: {cfg.get('nemo_version')!r}") for name, target in TOKENIZER_TARGETS[profile].items(): if tokenizers[name].get("_target_") != target: raise ValueError( @@ -147,24 +143,14 @@ def tokenizer_profile(cfg: dict[str, Any], text_vocab: int, frame_stacking: int) f"{tokenizers[name].get('_target_')!r}" ) expected_japanese_case = "upper" if profile == "v2602" else "lower" - japanese_case = ( - tokenizers["japanese_phoneme"].get("g2p", {}).get("ascii_letter_case") - ) + japanese_case = tokenizers["japanese_phoneme"].get("g2p", {}).get("ascii_letter_case") if japanese_case != expected_japanese_case: - raise ValueError( - f"unsupported {profile} Japanese ascii_letter_case: {japanese_case!r}" - ) + raise ValueError(f"unsupported {profile} Japanese ascii_letter_case: {japanese_case!r}") if profile == "v2607": hindi = tokenizers["hindi_phoneme"] - if ( - hindi.get("locale") != "hi-IN" - or hindi.get("g2p", {}).get("locale") != "hi-IN" - ): + if hindi.get("locale") != "hi-IN" or hindi.get("g2p", {}).get("locale") != "hi-IN": raise ValueError("unsupported v2607 Hindi IPA locale") - if ( - tokenizers["portuguese_Brazilian_phoneme"].get("locale_specific_punct") - is not False - ): + if tokenizers["portuguese_Brazilian_phoneme"].get("locale_specific_punct") is not False: raise ValueError("unsupported v2607 Portuguese punctuation mode") if cfg.get("language_to_tokenizer_mapping") != V2607_LANGUAGE_MAPPING: raise ValueError("unsupported v2607 language_to_tokenizer_mapping") diff --git a/src/runtime/ggml/backend.cpp b/src/runtime/ggml/backend.cpp index 4fca3b7..ac60330 100644 --- a/src/runtime/ggml/backend.cpp +++ b/src/runtime/ggml/backend.cpp @@ -53,9 +53,8 @@ BackendManager::init_backends() { if (borrowed_gpu_backend_ != nullptr) { dev = ggml_backend_get_device(borrowed_gpu_backend_); const auto dev_type = dev ? ggml_backend_dev_type(dev) : GGML_BACKEND_DEVICE_TYPE_CPU; - if (dev == nullptr || - (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && - dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU)) { + if (dev == nullptr || (dev_type != GGML_BACKEND_DEVICE_TYPE_GPU && + dev_type != GGML_BACKEND_DEVICE_TYPE_IGPU)) { throw std::runtime_error("borrowed GPU backend is not a GPU device"); } auto* buft = ggml_backend_dev_buffer_type(dev); @@ -147,7 +146,8 @@ std::vector BackendManager::get_backends() { std::vector handles; handles.reserve(backends.size() + (borrowed_gpu_backend_ ? 1 : 0)); - if (borrowed_gpu_backend_) handles.push_back(borrowed_gpu_backend_); + if (borrowed_gpu_backend_) + handles.push_back(borrowed_gpu_backend_); for (const auto& b : backends) handles.push_back(b.get()); return handles; } diff --git a/src/tts/magpietts/decoder.cpp b/src/tts/magpietts/decoder.cpp index 56ec3c5..faa6536 100644 --- a/src/tts/magpietts/decoder.cpp +++ b/src/tts/magpietts/decoder.cpp @@ -502,8 +502,7 @@ class PersistentDecoderModule final : public ggml_runtime::Module { "magpietts.decoder.runtime.slot_ids"); std::vector alignment_outputs; - for (int layer_index = 0; layer_index < static_cast(tr.layers.size()); - ++layer_index) { + for (int layer_index = 0; layer_index < static_cast(tr.layers.size()); ++layer_index) { const magpietts_layer& layer = tr.layers[layer_index]; ggml_tensor* residual = x; ggml_tensor* cur = layer_norm(ctx, x, layer.norm_self); @@ -518,27 +517,26 @@ class PersistentDecoderModule final : public ggml_runtime::Module { ggml_tensor* q = split_heads(0); ggml_tensor* k = split_heads(static_cast(tr.n_embd) * element); ggml_tensor* v = split_heads(static_cast(2 * tr.n_embd) * element); - auto kv = session->model_tensor_container->get_tensor_by_name( - runtime_kv_name(layer_index)); + auto kv = + session->model_tensor_container->get_tensor_by_name(runtime_kv_name(layer_index)); ggml_tensor* heads = ggml_fused_attn_cached( ctx, q, k, v, nullptr, kv.tensor, slots.tensor, cache_meta.tensor, cache_len_, 1.0f / std::sqrt(static_cast(d_head)), true); ggml_tensor* merged = ggml_reshape_2d( - ctx, ggml_permute(ctx, heads, 0, 2, 1, 3), tr.n_embd, - kMagpieCfgLanes); + ctx, ggml_permute(ctx, heads, 0, 2, 1, 3), tr.n_embd, kMagpieCfgLanes); x = ggml_add(ctx, residual, linear(ctx, layer.self_o, merged)); // Text cross-attention applies only to the conditional lane. if (tr.has_cross && layer.has_cross) { ggml_tensor* cond = ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], 0); - ggml_tensor* uncond = - ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], x->nb[1]); + ggml_tensor* uncond = ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], x->nb[1]); ggml_tensor* cross_in = layer_norm(ctx, cond, layer.norm_xattn_query); ggml_tensor* last_attn = nullptr; - const bool apply_prior = tr.apply_attention_prior && - runtime_layer_selected(tr.apply_prior_to_layers, layer_index); - const bool collect = runtime_layer_selected( - tr.estimate_alignment_from_layers, layer_index); + const bool apply_prior = + tr.apply_attention_prior && + runtime_layer_selected(tr.apply_prior_to_layers, layer_index); + const bool collect = + runtime_layer_selected(tr.estimate_alignment_from_layers, layer_index); ggml_tensor* cross = cross_attention_cached( ctx, tr, layer, cross_kv_, layer_index, cross_in, apply_prior ? prior.tensor : nullptr, collect ? &last_attn : nullptr, true); @@ -556,12 +554,13 @@ class PersistentDecoderModule final : public ggml_runtime::Module { cur = causal_conv1d(ctx, cur, layer.ff_out); x = ggml_add(ctx, residual, cur); } - if (tr.norm_out) x = layer_norm(ctx, x, tr.norm_out); + if (tr.norm_out) + x = layer_norm(ctx, x, tr.norm_out); - ggml_tensor* cond = ggml_cont_2d( - ctx, ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], 0), tr.n_embd, 1); - ggml_tensor* uncond = ggml_cont_2d( - ctx, ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], x->nb[1]), tr.n_embd, 1); + ggml_tensor* cond = + ggml_cont_2d(ctx, ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], 0), tr.n_embd, 1); + ggml_tensor* uncond = + ggml_cont_2d(ctx, ggml_view_2d(ctx, x, tr.n_embd, 1, x->nb[1], x->nb[1]), tr.n_embd, 1); ggml_set_name(cond, "magpietts_decoder_runtime_hidden_cond"); ggml_set_name(uncond, "magpietts_decoder_runtime_hidden_uncond"); ggml_runtime::TensorBag outputs; @@ -584,9 +583,7 @@ class PersistentDecoderModule final : public ggml_runtime::Module { sum = sum ? ggml_add(ctx, sum, per_layer) : per_layer; } ggml_tensor* mean = ggml_scale( - ctx, sum, - 1.0f / - static_cast(alignment_outputs.size() * tr.n_cross_head)); + ctx, sum, 1.0f / static_cast(alignment_outputs.size() * tr.n_cross_head)); ggml_set_name(mean, "magpietts_decoder_runtime_alignment_mean"); outputs.add_tensor({mean, bf_ctx.buft}); } @@ -648,12 +645,12 @@ class MagpieDecoder::PersistentDecoderRuntime { } ggml_context* ctx = new_graph_context(); const size_t element = sizeof(float); - const size_t source_layer_bytes = - static_cast(cond.n_ctx) * cond.n_embd * element; + const size_t source_layer_bytes = static_cast(cond.n_ctx) * cond.n_embd * element; const size_t copy_elements = static_cast(cond.n_tokens) * cond.n_embd; const size_t destination_token = static_cast(cache_len_ - cond.n_tokens); for (int layer = 0; layer < cond.n_layers; ++layer) { - auto arena = session_.model_tensor_container->get_tensor_by_name(runtime_kv_name(layer)); + auto arena = + session_.model_tensor_container->get_tensor_by_name(runtime_kv_name(layer)); for (int plane = 0; plane < 2; ++plane) { ggml_tensor* dst_base = arena.tensor; ggml_tensor* cond_src_base = plane == 0 ? cond.memory_k : cond.memory_v; @@ -663,10 +660,9 @@ class MagpieDecoder::PersistentDecoderRuntime { ggml_tensor* src = ggml_view_1d( ctx, const_cast(sources[lane]), copy_elements, static_cast(layer) * source_layer_bytes); - const size_t dst_offset = - static_cast(plane) * dst_base->nb[2] + - static_cast(lane) * dst_base->nb[1] + - destination_token * cond.n_embd * element; + const size_t dst_offset = static_cast(plane) * dst_base->nb[2] + + static_cast(lane) * dst_base->nb[1] + + destination_token * cond.n_embd * element; ggml_tensor* dst = ggml_view_1d(ctx, dst_base, copy_elements, dst_offset); ggml_backend_tensor_copy_async(model_.backend, model_.backend, src, dst); } @@ -682,24 +678,26 @@ class MagpieDecoder::PersistentDecoderRuntime { bool eval( const std::vector>& audio_codes, DecoderKvCache& cond_kv, DecoderKvCache& uncond_kv, decoder_result& cond_result, decoder_result& uncond_result, - magpietts_backend_tensor* cond_hidden_out, - magpietts_backend_tensor* uncond_hidden_out, + magpietts_backend_tensor* cond_hidden_out, magpietts_backend_tensor* uncond_hidden_out, const magpietts_decoder_attention* attention) { const ggml_nvtx::range nvtx_range("magpietts_persistent_decoder_eval"); const magpietts_hparams& h = model_.hparams; if (!cond_hidden_out || !uncond_hidden_out || !cond_hidden_out->tensor || - !uncond_hidden_out->tensor || static_cast(audio_codes.size()) != h.audio_codebooks || - audio_codes.empty()) { + !uncond_hidden_out->tensor || + static_cast(audio_codes.size()) != h.audio_codebooks || audio_codes.empty()) { return false; } const size_t raw_len = audio_codes[0].size(); - if (raw_len == 0 || raw_len % h.frame_stacking_factor != 0) return false; + if (raw_len == 0 || raw_len % h.frame_stacking_factor != 0) + return false; for (const auto& codes : audio_codes) { - if (codes.size() != raw_len) return false; + if (codes.size() != raw_len) + return false; } - const int total_len = h.baked_context_length + - static_cast(raw_len / h.frame_stacking_factor); - if (total_len != n_tokens_ + 1 || n_tokens_ >= cache_len_) return false; + const int total_len = + h.baked_context_length + static_cast(raw_len / h.frame_stacking_factor); + if (total_len != n_tokens_ + 1 || n_tokens_ >= cache_len_) + return false; std::vector tokens(static_cast(h.stacked_audio_codebooks())); const size_t frame_start = raw_len - static_cast(h.frame_stacking_factor); @@ -716,7 +714,8 @@ class MagpieDecoder::PersistentDecoderRuntime { ring_head_, ring_head_, valid_tokens_, valid_tokens_}; std::vector log_prior(static_cast(text_len_), 0.0f); if (attention && attention->prior) { - if (static_cast(attention->prior->size()) != text_len_) return false; + if (static_cast(attention->prior->size()) != text_len_) + return false; for (int i = 0; i < text_len_; ++i) { log_prior[static_cast(i)] = std::log(std::max((*attention->prior)[static_cast(i)], 1.0e-20f)); @@ -724,10 +723,14 @@ class MagpieDecoder::PersistentDecoderRuntime { } std::vector inputs = { - {"magpietts.decoder.runtime.tokens", GGML_TYPE_I32, tokens.data(), + {"magpietts.decoder.runtime.tokens", + GGML_TYPE_I32, + tokens.data(), {h.stacked_audio_codebooks()}}, {"magpietts.decoder.runtime.position", GGML_TYPE_I32, &position, {1}}, - {"magpietts.decoder.runtime.cache_meta", GGML_TYPE_I32, cache_meta, + {"magpietts.decoder.runtime.cache_meta", + GGML_TYPE_I32, + cache_meta, {kMagpieCfgLanes, 2}}, {"magpietts.decoder.runtime.prior", GGML_TYPE_F32, log_prior.data(), {text_len_}}}; @@ -852,15 +855,14 @@ MagpieDecoder::evalCachedPair( "MagpieTTS decoder runtime: fixed-shape CUDA graph, CFG batch=2, " "device K/V arena enabled\n"); } - if (persistent_runtime_ && - !persistent_runtime_->sequence_matches(cond_kv.n_tokens)) { + if (persistent_runtime_ && !persistent_runtime_->sequence_matches(cond_kv.n_tokens)) { // Reuse the graph and reseed the cache suffix. persistent_runtime_->seed(cond_kv, uncond_kv); } - if (persistent_runtime_ && persistent_runtime_->eval( - audio_codes, cond_kv, uncond_kv, cond_result, - uncond_result, cond_hidden_out, uncond_hidden_out, - attention)) { + if (persistent_runtime_ && + persistent_runtime_->eval( + audio_codes, cond_kv, uncond_kv, cond_result, uncond_result, cond_hidden_out, + uncond_hidden_out, attention)) { return true; } persistent_runtime_.reset(); diff --git a/src/tts/magpietts/magpietts_cuda_sampling.cu b/src/tts/magpietts/magpietts_cuda_sampling.cu index 3172eb0..5fd454b 100644 --- a/src/tts/magpietts/magpietts_cuda_sampling.cu +++ b/src/tts/magpietts/magpietts_cuda_sampling.cu @@ -2,11 +2,10 @@ // SPDX-License-Identifier: Apache-2.0 #include -#include - #include #include #include +#include #include "magpietts_cuda_sampling.h" @@ -139,8 +138,7 @@ magpietts_sample_codebooks_kernel( const magpietts_cuda_sampling_config config = *config_ptr; - using block_sort = cub::BlockRadixSort< - float, MAGPIETTS_CUDA_BLOCK_SIZE, items_per_thread, int>; + using block_sort = cub::BlockRadixSort; __shared__ typename block_sort::TempStorage sort_storage; __shared__ double s_sums[MAGPIETTS_CUDA_BLOCK_SIZE]; @@ -159,12 +157,12 @@ magpietts_sample_codebooks_kernel( #pragma unroll for (int item = 0; item < items_per_thread; ++item) { const int id = (int)threadIdx.x * items_per_thread + item; - thread_vals[item] = id < vocab_size - ? sampled_logit( - logits_cond, logits_uncond, off, id, audio_codebook_size, - audio_eos_id, config.use_cfg != 0, config.cfg_scale, - config.forbid_audio_eos != 0) - : -INFINITY; + thread_vals[item] = + id < vocab_size + ? sampled_logit( + logits_cond, logits_uncond, off, id, audio_codebook_size, audio_eos_id, + config.use_cfg != 0, config.cfg_scale, config.forbid_audio_eos != 0) + : -INFINITY; thread_ids[item] = id; } block_sort(sort_storage).SortDescendingBlockedToStriped(thread_vals, thread_ids); @@ -192,8 +190,7 @@ magpietts_sample_codebooks_kernel( double local_sum = 0.0; for (int i = threadIdx.x; i < k; i += blockDim.x) { if (isfinite(top_vals[i])) { - local_sum += - exp((double)(top_vals[i] - max_logit) / (double)config.temperature); + local_sum += exp((double)(top_vals[i] - max_logit) / (double)config.temperature); } } s_sums[threadIdx.x] = local_sum; @@ -210,8 +207,7 @@ magpietts_sample_codebooks_kernel( double acc = 0.0; for (int i = 0; i < k; ++i) { if (isfinite(top_vals[i])) { - acc += - exp((double)(top_vals[i] - max_logit) / (double)config.temperature); + acc += exp((double)(top_vals[i] - max_logit) / (double)config.temperature); } if (target <= acc) { sampled = top_ids[i]; @@ -350,7 +346,8 @@ magpietts_cuda_sampler_configure( sampler->h_config->seed = seed; sampler->h_config->use_cfg = use_cfg ? 1 : 0; sampler->h_config->forbid_audio_eos = forbid_audio_eos ? 1 : 0; - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -369,7 +366,8 @@ magpietts_cuda_sampler_upload_config( err = cudaGraphAddMemcpyNode1D( &node, sampler->sequence_graph, deps, dependency_count, sampler->d_config, sampler->h_config, sizeof(*sampler->d_config), cudaMemcpyHostToDevice); - if (err == cudaSuccess) sampler->sequence_tail = node; + if (err == cudaSuccess) + sampler->sequence_tail = node; } else { err = cudaMemcpyAsync( sampler->d_config, sampler->h_config, sizeof(*sampler->d_config), @@ -379,7 +377,8 @@ magpietts_cuda_sampler_upload_config( set_error(error, error_size, "failed to upload CUDA sampler config", err); return false; } - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -405,7 +404,8 @@ magpietts_cuda_sampler_sequence_build_active(const magpietts_cuda_sampler* sampl void magpietts_cuda_sampler_sequence_mark_warm(magpietts_cuda_sampler* sampler) { - if (sampler) sampler->sequence_warm = true; + if (sampler) + sampler->sequence_warm = true; } bool @@ -425,7 +425,8 @@ magpietts_cuda_sampler_sequence_begin_build( sampler->sequence_graph = graph; sampler->sequence_tail = nullptr; sampler->sequence_build_active = true; - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -461,14 +462,17 @@ magpietts_cuda_sampler_sequence_finish_build_and_launch( return false; } sampler->sequence_exec = exec; - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } void magpietts_cuda_sampler_sequence_abort_build(magpietts_cuda_sampler* sampler) { - if (!sampler || !sampler->sequence_build_active) return; - if (sampler->sequence_graph) cudaGraphDestroy(sampler->sequence_graph); + if (!sampler || !sampler->sequence_build_active) + return; + if (sampler->sequence_graph) + cudaGraphDestroy(sampler->sequence_graph); sampler->sequence_graph = nullptr; sampler->sequence_tail = nullptr; sampler->sequence_build_active = false; @@ -476,7 +480,8 @@ magpietts_cuda_sampler_sequence_abort_build(magpietts_cuda_sampler* sampler) { void magpietts_cuda_sampler_sequence_disable(magpietts_cuda_sampler* sampler) { - if (!sampler) return; + if (!sampler) + return; magpietts_cuda_sampler_sequence_abort_build(sampler); if (sampler->sequence_exec) { cudaGraphExecDestroy(sampler->sequence_exec); @@ -502,7 +507,8 @@ magpietts_cuda_sampler_sequence_launch( set_error(error, error_size, "failed to launch CUDA local sequence graph", err); return false; } - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -525,7 +531,8 @@ magpietts_cuda_sampler_sequence_add_ggml_graph( return false; } sampler->sequence_tail = node; - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -549,7 +556,8 @@ magpietts_cuda_sampler_sequence_add_device_copy( return false; } sampler->sequence_tail = node; - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -594,8 +602,8 @@ magpietts_cuda_sample_codebooks_device( bool magpietts_cuda_sample_codebooks_device_configured( magpietts_cuda_sampler* sampler, const float* logits_cond, const float* logits_uncond, - int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, - int codebook_offset, int output_offset, char* error, size_t error_size) { + int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, int codebook_offset, + int output_offset, char* error, size_t error_size) { if (!sampler || !logits_cond) { set_error(error, error_size, "invalid CUDA sampler device arguments"); return false; @@ -620,18 +628,16 @@ magpietts_cuda_sample_codebooks_device_configured( int32_t* codes = sampler->d_codes; int32_t* argmax = sampler->d_argmax; void* kernel_args[] = { - &logits_cond, &logits_uncond, &codebooks, &vocab_size, - &audio_codebook_size, &audio_eos_id, &config, &codebook_offset, - &output_offset, &top_ids, &top_vals, &codes, - &argmax}; + &logits_cond, &logits_uncond, &codebooks, &vocab_size, &audio_codebook_size, + &audio_eos_id, &config, &codebook_offset, &output_offset, &top_ids, + &top_vals, &codes, &argmax}; cudaKernelNodeParams params{}; - params.func = vocab_size <= MAGPIETTS_CUDA_SMALL_VOCAB - ? reinterpret_cast( - magpietts_sample_codebooks_kernel< - MAGPIETTS_CUDA_SMALL_ITEMS_PER_THREAD>) - : reinterpret_cast( - magpietts_sample_codebooks_kernel< - MAGPIETTS_CUDA_MAX_ITEMS_PER_THREAD>); + params.func = + vocab_size <= MAGPIETTS_CUDA_SMALL_VOCAB + ? reinterpret_cast( + magpietts_sample_codebooks_kernel) + : reinterpret_cast( + magpietts_sample_codebooks_kernel); params.gridDim = dim3((unsigned int)codebooks, 1, 1); params.blockDim = dim3(MAGPIETTS_CUDA_BLOCK_SIZE, 1, 1); params.sharedMemBytes = 0; @@ -640,14 +646,15 @@ magpietts_cuda_sample_codebooks_device_configured( cudaGraphNode_t node = nullptr; const cudaGraphNode_t* deps = sampler->sequence_tail ? &sampler->sequence_tail : nullptr; const size_t dependency_count = sampler->sequence_tail ? 1 : 0; - const cudaError_t err = cudaGraphAddKernelNode( - &node, sampler->sequence_graph, deps, dependency_count, ¶ms); + const cudaError_t err = + cudaGraphAddKernelNode(&node, sampler->sequence_graph, deps, dependency_count, ¶ms); if (err != cudaSuccess) { set_error(error, error_size, "failed to add sampler to CUDA local sequence", err); return false; } sampler->sequence_tail = node; - if (error && error_size > 0) error[0] = '\0'; + if (error && error_size > 0) + error[0] = '\0'; return true; } @@ -655,16 +662,14 @@ magpietts_cuda_sample_codebooks_device_configured( magpietts_sample_codebooks_kernel <<stream>>>( logits_cond, logits_uncond, codebooks, vocab_size, audio_codebook_size, - audio_eos_id, sampler->d_config, codebook_offset, output_offset, - sampler->d_top_ids, sampler->d_top_vals, sampler->d_codes, - sampler->d_argmax); + audio_eos_id, sampler->d_config, codebook_offset, output_offset, sampler->d_top_ids, + sampler->d_top_vals, sampler->d_codes, sampler->d_argmax); } else { magpietts_sample_codebooks_kernel <<stream>>>( logits_cond, logits_uncond, codebooks, vocab_size, audio_codebook_size, - audio_eos_id, sampler->d_config, codebook_offset, output_offset, - sampler->d_top_ids, sampler->d_top_vals, sampler->d_codes, - sampler->d_argmax); + audio_eos_id, sampler->d_config, codebook_offset, output_offset, sampler->d_top_ids, + sampler->d_top_vals, sampler->d_codes, sampler->d_argmax); } cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { @@ -693,8 +698,7 @@ magpietts_cuda_copy_sampled_code_to_device( if (sampler->sequence_build_active) { return magpietts_cuda_sampler_sequence_add_device_copy( - sampler, sampler->d_codes + codebook, dst_device, sizeof(int32_t), error, - error_size); + sampler, sampler->d_codes + codebook, dst_device, sizeof(int32_t), error, error_size); } cudaError_t err = cudaSuccess; diff --git a/src/tts/magpietts/magpietts_cuda_sampling.h b/src/tts/magpietts/magpietts_cuda_sampling.h index 1660276..5e7fa75 100644 --- a/src/tts/magpietts/magpietts_cuda_sampling.h +++ b/src/tts/magpietts/magpietts_cuda_sampling.h @@ -58,8 +58,8 @@ bool magpietts_cuda_sample_codebooks_device( // Launch using the most recently uploaded configuration. bool magpietts_cuda_sample_codebooks_device_configured( magpietts_cuda_sampler* sampler, const float* logits_cond, const float* logits_uncond, - int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, - int codebook_offset, int output_offset, char* error, size_t error_size); + int codebooks, int vocab_size, int audio_codebook_size, int audio_eos_id, int codebook_offset, + int output_offset, char* error, size_t error_size); bool magpietts_cuda_copy_sampled_code_to_device( magpietts_cuda_sampler* sampler, int codebook, void* dst_device, char* error, diff --git a/src/tts/magpietts/model.h b/src/tts/magpietts/model.h index 8a668be..0dcf138 100644 --- a/src/tts/magpietts/model.h +++ b/src/tts/magpietts/model.h @@ -104,8 +104,7 @@ magpietts_infer_tokenizer_profile(const magpietts_hparams& h) { } inline bool -magpietts_tokenizer_profile_matches( - const std::string& profile, const magpietts_hparams& h) { +magpietts_tokenizer_profile_matches(const std::string& profile, const magpietts_hparams& h) { return profile == magpietts_infer_tokenizer_profile(h) && !profile.empty(); } diff --git a/src/tts/tokenizer/mandarin_tokenizer.h b/src/tts/tokenizer/mandarin_tokenizer.h index 70f8801..262a10f 100644 --- a/src/tts/tokenizer/mandarin_tokenizer.h +++ b/src/tts/tokenizer/mandarin_tokenizer.h @@ -10,8 +10,7 @@ class mandarin_tokenizer { public: explicit mandarin_tokenizer( - const std::filesystem::path& model_dir, int offset = 349, - std::string phoneme_dict = {}); + const std::filesystem::path& model_dir, int offset = 349, std::string phoneme_dict = {}); ~mandarin_tokenizer(); mandarin_tokenizer(const mandarin_tokenizer&) = delete; diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index cfc9ff7..551bf15 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -48,9 +48,8 @@ struct tokenizer_profile { throw std::invalid_argument("unsupported tokenizer language '" + language + "'"); } const auto entry = std::find_if( - entries.begin(), entries.end(), [&](const tokenizer_entry& value) { - return value.name == mapping->second; - }); + entries.begin(), entries.end(), + [&](const tokenizer_entry& value) { return value.name == mapping->second; }); if (entry == entries.end()) { throw std::runtime_error( "tokenizer profile '" + id + "' maps language '" + language + @@ -75,9 +74,8 @@ trim(std::string value) { std::string unquote(std::string value) { value = trim(std::move(value)); - if (value.size() >= 2 && - ((value.front() == '\'' && value.back() == '\'') || - (value.front() == '"' && value.back() == '"'))) { + if (value.size() >= 2 && ((value.front() == '\'' && value.back() == '\'') || + (value.front() == '"' && value.back() == '"'))) { value = value.substr(1, value.size() - 2); } return value; @@ -194,8 +192,7 @@ require_block_value( void require_common_tokenizer_values( const std::map& blocks, const std::string& tokenizer, - const std::string& target, const std::string& apostrophe, - const std::string& pad_with_space) { + const std::string& target, const std::string& apostrophe, const std::string& pad_with_space) { require_block_value(blocks, tokenizer, "_target_", target); require_block_value(blocks, tokenizer, "punct", "true"); require_block_value(blocks, tokenizer, "apostrophe", apostrophe); @@ -267,8 +264,8 @@ parse_language_mapping(const std::vector& lines) { tokenizer_entry entry( std::string name, tokenizer_kind kind, int offset, int size, std::string locale = {}, - std::string grapheme_case = {}, std::string grapheme_prefix = {}, - bool apostrophe = true, bool pad_with_space = false) { + std::string grapheme_case = {}, std::string grapheme_prefix = {}, bool apostrophe = true, + bool pad_with_space = false) { tokenizer_entry value; value.name = std::move(name); value.kind = kind; @@ -303,8 +300,8 @@ attach_artifacts( if ((item.kind == tokenizer_kind::ipa || item.kind == tokenizer_kind::mandarin) && item.phoneme_dict.empty()) { throw std::runtime_error( - "tokenizer profile '" + profile.id + "' has no phoneme_dict for '" + - item.name + "'"); + "tokenizer profile '" + profile.id + "' has no phoneme_dict for '" + item.name + + "'"); } } } @@ -333,19 +330,25 @@ load_tokenizer_profile(const fs::path& root) { std::vector order; const auto blocks = yaml_child_blocks(lines, "text_tokenizers", order); const std::vector v2602_order = { - "english_phoneme", "spanish_phoneme", "german_phoneme", - "mandarin_phoneme", "japanese_phoneme", "french_chartokenizer", - "hindi_chartokenizer", "italian_phoneme", "vietnamese_phoneme", - "text_ce_tokenizer", + "english_phoneme", "spanish_phoneme", "german_phoneme", "mandarin_phoneme", + "japanese_phoneme", "french_chartokenizer", "hindi_chartokenizer", "italian_phoneme", + "vietnamese_phoneme", "text_ce_tokenizer", }; const std::vector v2607_order = { - "english_phoneme", "text_ce_tokenizer", - "spanish_phoneme", "german_phoneme", - "mandarin_phoneme", "japanese_phoneme", - "portuguese_Brazilian_phoneme", "hindi_phoneme", - "arabic_AE_chartokenizer", "arabic_SA_chartokenizer", - "arabic_MSA_chartokenizer", "french_chartokenizer", - "italian_chartokenizer", "vietnamese_chartokenizer", + "english_phoneme", + "text_ce_tokenizer", + "spanish_phoneme", + "german_phoneme", + "mandarin_phoneme", + "japanese_phoneme", + "portuguese_Brazilian_phoneme", + "hindi_phoneme", + "arabic_AE_chartokenizer", + "arabic_SA_chartokenizer", + "arabic_MSA_chartokenizer", + "french_chartokenizer", + "italian_chartokenizer", + "vietnamese_chartokenizer", "korean_chartokenizer", }; @@ -357,8 +360,10 @@ load_tokenizer_profile(const fs::path& root) { profile.eos_id = 2361; profile.entries = { entry("english_phoneme", tokenizer_kind::ipa, 0, 96, "en-US", "upper", "", true, false), - entry("spanish_phoneme", tokenizer_kind::ipa, 96, 103, "es-ES", "upper", "", true, true), - entry("german_phoneme", tokenizer_kind::ipa, 199, 150, "de-DE", "mixed", "#", true, true), + entry( + "spanish_phoneme", tokenizer_kind::ipa, 96, 103, "es-ES", "upper", "", true, true), + entry( + "german_phoneme", tokenizer_kind::ipa, 199, 150, "de-DE", "mixed", "#", true, true), entry("mandarin_phoneme", tokenizer_kind::mandarin, 349, 109), entry("japanese_phoneme", tokenizer_kind::japanese, 458, 175), entry("french_chartokenizer", tokenizer_kind::byt5, 633, 384), @@ -368,9 +373,9 @@ load_tokenizer_profile(const fs::path& root) { entry("text_ce_tokenizer", tokenizer_kind::byt5, 1976, 384), }; profile.language_mapping = { - {"en", "english_phoneme"}, {"es", "spanish_phoneme"}, - {"de", "german_phoneme"}, {"fr", "french_chartokenizer"}, - {"it", "italian_phoneme"}, {"vi", "vietnamese_phoneme"}, + {"en", "english_phoneme"}, {"es", "spanish_phoneme"}, + {"de", "german_phoneme"}, {"fr", "french_chartokenizer"}, + {"it", "italian_phoneme"}, {"vi", "vietnamese_phoneme"}, {"zh", "mandarin_phoneme"}, {"hi", "hindi_chartokenizer"}, {"ja", "japanese_phoneme"}, }; @@ -385,12 +390,17 @@ load_tokenizer_profile(const fs::path& root) { profile.entries = { entry("english_phoneme", tokenizer_kind::ipa, 0, 96, "en-US", "upper", "", true, false), entry("text_ce_tokenizer", tokenizer_kind::byt5, 96, 384), - entry("spanish_phoneme", tokenizer_kind::ipa, 480, 103, "es-ES", "upper", "", true, true), - entry("german_phoneme", tokenizer_kind::ipa, 583, 150, "de-DE", "mixed", "#", true, true), + entry( + "spanish_phoneme", tokenizer_kind::ipa, 480, 103, "es-ES", "upper", "", true, true), + entry( + "german_phoneme", tokenizer_kind::ipa, 583, 150, "de-DE", "mixed", "#", true, true), entry("mandarin_phoneme", tokenizer_kind::mandarin, 733, 109), entry("japanese_phoneme", tokenizer_kind::japanese, 842, 175), - entry("portuguese_Brazilian_phoneme", tokenizer_kind::ipa, 1017, 111, "pt-BR", "upper", "#", true, true), - entry("hindi_phoneme", tokenizer_kind::ipa, 1128, 201, "hi-IN", "upper", "", true, true), + entry( + "portuguese_Brazilian_phoneme", tokenizer_kind::ipa, 1017, 111, "pt-BR", "upper", + "#", true, true), + entry( + "hindi_phoneme", tokenizer_kind::ipa, 1128, 201, "hi-IN", "upper", "", true, true), entry("arabic_AE_chartokenizer", tokenizer_kind::arabic, 1329, 164), entry("arabic_SA_chartokenizer", tokenizer_kind::arabic, 1493, 164), entry("arabic_MSA_chartokenizer", tokenizer_kind::arabic, 1657, 164), @@ -401,25 +411,39 @@ load_tokenizer_profile(const fs::path& root) { }; profile.language_mapping = parse_language_mapping(lines); const std::map expected_mapping = { - {"en", "english_phoneme"}, {"de", "german_phoneme"}, - {"es", "spanish_phoneme"}, {"fr", "french_chartokenizer"}, - {"it", "italian_chartokenizer"}, {"vi", "vietnamese_chartokenizer"}, - {"zh", "mandarin_phoneme"}, {"hi", "hindi_phoneme"}, - {"ja", "japanese_phoneme"}, {"pt-BR", "portuguese_Brazilian_phoneme"}, - {"ko", "korean_chartokenizer"}, {"ar-AE", "arabic_AE_chartokenizer"}, - {"ar-SA", "arabic_SA_chartokenizer"}, {"ar-MSA", "arabic_MSA_chartokenizer"}, + {"en", "english_phoneme"}, + {"de", "german_phoneme"}, + {"es", "spanish_phoneme"}, + {"fr", "french_chartokenizer"}, + {"it", "italian_chartokenizer"}, + {"vi", "vietnamese_chartokenizer"}, + {"zh", "mandarin_phoneme"}, + {"hi", "hindi_phoneme"}, + {"ja", "japanese_phoneme"}, + {"pt-BR", "portuguese_Brazilian_phoneme"}, + {"ko", "korean_chartokenizer"}, + {"ar-AE", "arabic_AE_chartokenizer"}, + {"ar-SA", "arabic_SA_chartokenizer"}, + {"ar-MSA", "arabic_MSA_chartokenizer"}, }; if (profile.language_mapping != expected_mapping) { throw std::runtime_error("unsupported v2607 language_to_tokenizer_mapping"); } profile.language_mapping = { - {"en", "english_phoneme"}, {"de", "german_phoneme"}, - {"es", "spanish_phoneme"}, {"fr", "french_chartokenizer"}, - {"it", "italian_chartokenizer"}, {"vi", "vietnamese_chartokenizer"}, - {"zh", "mandarin_phoneme"}, {"hi", "hindi_phoneme"}, - {"ja", "japanese_phoneme"}, {"pt-br", "portuguese_Brazilian_phoneme"}, - {"ko", "korean_chartokenizer"}, {"ar-ae", "arabic_AE_chartokenizer"}, - {"ar-sa", "arabic_SA_chartokenizer"}, {"ar-msa", "arabic_MSA_chartokenizer"}, + {"en", "english_phoneme"}, + {"de", "german_phoneme"}, + {"es", "spanish_phoneme"}, + {"fr", "french_chartokenizer"}, + {"it", "italian_chartokenizer"}, + {"vi", "vietnamese_chartokenizer"}, + {"zh", "mandarin_phoneme"}, + {"hi", "hindi_phoneme"}, + {"ja", "japanese_phoneme"}, + {"pt-br", "portuguese_Brazilian_phoneme"}, + {"ko", "korean_chartokenizer"}, + {"ar-ae", "arabic_AE_chartokenizer"}, + {"ar-sa", "arabic_SA_chartokenizer"}, + {"ar-msa", "arabic_MSA_chartokenizer"}, }; if (profile.nemo_version != "2.8.0rc0") { throw std::runtime_error( @@ -428,12 +452,13 @@ load_tokenizer_profile(const fs::path& root) { } else { std::ostringstream found; for (size_t i = 0; i < order.size(); ++i) { - if (i != 0) found << ", "; + if (i != 0) + found << ", "; found << order[i]; } throw std::runtime_error( - "unsupported Magpie tokenizer layout in " + config_path.string() + ": [" + - found.str() + "]"); + "unsupported Magpie tokenizer layout in " + config_path.string() + ": [" + found.str() + + "]"); } for (auto& item : profile.entries) { @@ -475,8 +500,7 @@ load_tokenizer_profile(const fs::path& root) { "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." "JapanesePhonemeTokenizer", "false", "true"); - require_block_value( - blocks, item.name, "ascii_letter_case", item.ascii_letter_case); + require_block_value(blocks, item.name, "ascii_letter_case", item.ascii_letter_case); break; case tokenizer_kind::arabic: require_common_tokenizer_values( @@ -689,11 +713,10 @@ class MagpieNativeTokenizer::Impl { std::vector supported_language_codes() const { static const std::vector> canonical = { - {"en", "en-US"}, {"es", "es-ES"}, {"de", "de-DE"}, - {"fr", "fr-FR"}, {"it", "it-IT"}, {"vi", "vi-VN"}, - {"zh", "zh-CN"}, {"hi", "hi-IN"}, {"ja", "ja-JP"}, - {"ar-ae", "ar-AE"}, {"ar-sa", "ar-SA"}, {"ar-msa", "ar-MSA"}, - {"ko", "ko-KR"}, {"pt-br", "pt-BR"}, + {"en", "en-US"}, {"es", "es-ES"}, {"de", "de-DE"}, {"fr", "fr-FR"}, + {"it", "it-IT"}, {"vi", "vi-VN"}, {"zh", "zh-CN"}, {"hi", "hi-IN"}, + {"ja", "ja-JP"}, {"ar-ae", "ar-AE"}, {"ar-sa", "ar-SA"}, {"ar-msa", "ar-MSA"}, + {"ko", "ko-KR"}, {"pt-br", "pt-BR"}, }; std::vector languages; for (const auto& [normalized, code] : canonical) { @@ -702,9 +725,8 @@ class MagpieNativeTokenizer::Impl { continue; } const auto item = std::find_if( - profile_.entries.begin(), profile_.entries.end(), [&](const tokenizer_entry& entry) { - return entry.name == mapping->second; - }); + profile_.entries.begin(), profile_.entries.end(), + [&](const tokenizer_entry& entry) { return entry.name == mapping->second; }); if (item != profile_.entries.end() && supports(*item)) { languages.push_back(code); } @@ -771,8 +793,8 @@ class MagpieNativeTokenizer::Impl { std::lock_guard lock(cache_mutex_); if (!mandarin_cache_) { const auto& entry = profile_.entry_for_language("zh"); - mandarin_cache_ = std::make_unique( - model_dir_, entry.offset, entry.phoneme_dict); + mandarin_cache_ = + std::make_unique(model_dir_, entry.offset, entry.phoneme_dict); } return *mandarin_cache_; } diff --git a/src/tts/tokenizer/tokenizer_impl.cpp b/src/tts/tokenizer/tokenizer_impl.cpp index 8d26fd3..13009cc 100644 --- a/src/tts/tokenizer/tokenizer_impl.cpp +++ b/src/tts/tokenizer/tokenizer_impl.cpp @@ -336,20 +336,19 @@ hindi_tokens() { static std::vector japanese_tokens(bool lowercase_ascii = false) { std::vector tokens = { - " ", "0", "1", "ァ", "ア", "ィ", "イ", "ゥ", "ウ", "ェ", "エ", "ォ", "オ", - "カ", "ガ", "キ", "ギ", "ク", "グ", "ケ", "ゲ", "コ", "ゴ", "サ", "ザ", "シ", - "ジ", "ス", "ズ", "セ", "ゼ", "ソ", "ゾ", "タ", "ダ", "チ", "ヂ", "ッ", "ツ", - "ヅ", "テ", "デ", "ト", "ド", "ナ", "ニ", "ヌ", "ネ", "ノ", "ハ", "バ", "パ", - "ヒ", "ビ", "ピ", "フ", "ブ", "プ", "ヘ", "ベ", "ペ", "ホ", "ボ", "ポ", "マ", - "ミ", "ム", "メ", "モ", "ャ", "ヤ", "ュ", "ユ", "ョ", "ヨ", "ラ", "リ", "ル", - "レ", "ロ", "ヮ", "ワ", "ヲ", "ン", "ヴ", "ヵ", "ヶ", "ー", "A", "B", "C", - "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", - "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "!", "\"", "(", - ")", ",", "-", ".", "/", ":", ";", "?", "[", "]", "{", "}", "«", - "»", "•", "‥", "…", "‹", "›", "※", "◦", "、", "。", "〃", "〈", "〉", - "《", "》", "「", "」", "『", "』", "【", "】", "〒", "〓", "〔", "〕", "〖", - "〗", "〘", "〙", "〚", "〛", "〜", "〽", "・", "・・・", "ー", "﹅", "﹆", "！", - "＊", "？", "｟", "｠", "", ""}; + " ", "0", "1", "ァ", "ア", "ィ", "イ", "ゥ", "ウ", "ェ", "エ", "ォ", "オ", "カ", + "ガ", "キ", "ギ", "ク", "グ", "ケ", "ゲ", "コ", "ゴ", "サ", "ザ", "シ", "ジ", "ス", + "ズ", "セ", "ゼ", "ソ", "ゾ", "タ", "ダ", "チ", "ヂ", "ッ", "ツ", "ヅ", "テ", "デ", + "ト", "ド", "ナ", "ニ", "ヌ", "ネ", "ノ", "ハ", "バ", "パ", "ヒ", "ビ", "ピ", "フ", + "ブ", "プ", "ヘ", "ベ", "ペ", "ホ", "ボ", "ポ", "マ", "ミ", "ム", "メ", "モ", "ャ", + "ヤ", "ュ", "ユ", "ョ", "ヨ", "ラ", "リ", "ル", "レ", "ロ", "ヮ", "ワ", "ヲ", "ン", + "ヴ", "ヵ", "ヶ", "ー", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", + "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", + "Y", "Z", "!", "\"", "(", ")", ",", "-", ".", "/", ":", ";", "?", "[", + "]", "{", "}", "«", "»", "•", "‥", "…", "‹", "›", "※", "◦", "、", "。", + "〃", "〈", "〉", "《", "》", "「", "」", "『", "』", "【", "】", "〒", "〓", "〔", + "〕", "〖", "〗", "〘", "〙", "〚", "〛", "〜", "〽", "・", "・・・", "ー", "﹅", "﹆", + "！", "＊", "？", "｟", "｠", "", ""}; if (lowercase_ascii) { for (auto& token : tokens) { if (token.size() == 1 && token[0] >= 'A' && token[0] <= 'Z') { @@ -883,15 +882,8 @@ ipa_config_for_params(const params& p) { throw std::runtime_error("incomplete IPA tokenizer profile for language " + p.language); } return { - p.tokenizer_name, - p.offset, - p.phoneme_dict, - p.heteronyms, - p.locale, - p.grapheme_case, - p.grapheme_prefix, - p.apostrophe, - p.pad_with_space, + p.tokenizer_name, p.offset, p.phoneme_dict, p.heteronyms, p.locale, + p.grapheme_case, p.grapheme_prefix, p.apostrophe, p.pad_with_space, }; } @@ -947,21 +939,19 @@ exact_ipa_tokens(const std::string& tokenizer_name) { } if (tokenizer_name == "hindi_phoneme") { return { - "!", "\"", "'", "(", ")", ",", "-", ".", "/", "0", "1", "2", "3", "4", - "5", "6", "7", "8", "9", ":", ";", "?", "A", "B", "C", "D", "E", "F", - "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", - "U", "V", "W", "X", "Y", "Z", "[", "]", "a", "b", "c", "d", "e", "f", - "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", - "v", "w", "x", "z", "{", "}", "À", "É", "ã", "æ", "ð", "õ", "ĩ", "ŋ", - "ũ", "ɑ", "ɔ", "ɖ", "ə", "ɚ", "ɛ", "ɝ", "ɟ", "ɡ", "ɣ", "ɪ", "ɭ", "ɲ", - "ɳ", "ɹ", "ɾ", "ʂ", "ʃ", "ʈ", "ʊ", "ʋ", "ʌ", "ʒ", "ʰ", "ˈ", "ˌ", "ː", - "̃", "̩", "θ", "χ", "ँ", "ं", "ः", "अ", "आ", "इ", "ई", "उ", "ऊ", "ऋ", - "ऌ", "ऍ", "ऎ", "ए", "ऐ", "ऑ", "ओ", "औ", "क", "ख", "ग", "घ", "ङ", "च", - "छ", "ज", "झ", "ञ", "ट", "ठ", "ड", "ढ", "ण", "त", "थ", "द", "ध", "न", - "ऩ", "प", "फ", "ब", "भ", "म", "य", "र", "ऱ", "ल", "ळ", "ऴ", "व", "श", - "ष", "स", "ह", "ऺ", "़", "ऽ", "ा", "ि", "ी", "ु", "ू", "ृ", "ॅ", "ॆ", - "े", "ै", "ॉ", "ॊ", "ो", "ौ", "्", "ॐ", "॓", "ॠ", "ॡ", "ॢ", "।", "॥", - "॰", "ẽ", " ", "", "", + "!", "\"", "'", "(", ")", ",", "-", ".", "/", "0", "1", "2", "3", "4", "5", "6", + "7", "8", "9", ":", ";", "?", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", + "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", + "[", "]", "a", "b", "c", "d", "e", "f", "h", "i", "j", "k", "l", "m", "n", "o", + "p", "q", "r", "s", "t", "u", "v", "w", "x", "z", "{", "}", "À", "É", "ã", "æ", + "ð", "õ", "ĩ", "ŋ", "ũ", "ɑ", "ɔ", "ɖ", "ə", "ɚ", "ɛ", "ɝ", "ɟ", "ɡ", "ɣ", "ɪ", + "ɭ", "ɲ", "ɳ", "ɹ", "ɾ", "ʂ", "ʃ", "ʈ", "ʊ", "ʋ", "ʌ", "ʒ", "ʰ", "ˈ", "ˌ", "ː", + "̃", "̩", "θ", "χ", "ँ", "ं", "ः", "अ", "आ", "इ", "ई", "उ", "ऊ", "ऋ", "ऌ", "ऍ", + "ऎ", "ए", "ऐ", "ऑ", "ओ", "औ", "क", "ख", "ग", "घ", "ङ", "च", "छ", "ज", "झ", "ञ", + "ट", "ठ", "ड", "ढ", "ण", "त", "थ", "द", "ध", "न", "ऩ", "प", "फ", "ब", "भ", "म", + "य", "र", "ऱ", "ल", "ळ", "ऴ", "व", "श", "ष", "स", "ह", "ऺ", "़", "ऽ", "ा", "ि", + "ी", "ु", "ू", "ृ", "ॅ", "ॆ", "े", "ै", "ॉ", "ॊ", "ो", "ौ", "्", "ॐ", "॓", "ॠ", + "ॡ", "ॢ", "।", "॥", "॰", "ẽ", " ", "", "", }; } return {}; @@ -1042,7 +1032,8 @@ class ipa_tokenizer { void load(const fs::path& root) { const fs::path dict_path = root / cfg_.dict_hint; if (!fs::is_regular_file(dict_path)) { - throw std::runtime_error("failed to find tokenizer dictionary '" + dict_path.string() + "'"); + throw std::runtime_error( + "failed to find tokenizer dictionary '" + dict_path.string() + "'"); } for (const auto& p : ipa_punct(cfg_.locale)) { punct_.insert(p); diff --git a/tests/conversion/tts_tokenizer_profiles_test.py b/tests/conversion/tts_tokenizer_profiles_test.py index 475872c..d6c908e 100644 --- a/tests/conversion/tts_tokenizer_profiles_test.py +++ b/tests/conversion/tts_tokenizer_profiles_test.py @@ -58,9 +58,7 @@ def test_unknown_or_reordered_layout_is_rejected(self) -> None: reordered = list(TOKENIZER_ORDERS["v2607"]) reordered[0], reordered[1] = reordered[1], reordered[0] with self.assertRaisesRegex(ValueError, "unsupported Magpie tokenizer layout"): - tokenizer_profile( - {"text_tokenizers": {name: {} for name in reordered}}, 3359, 2 - ) + tokenizer_profile({"text_tokenizers": {name: {} for name in reordered}}, 3359, 2) def test_critical_tokenizer_config_changes_are_rejected(self) -> None: config = self._config("v2607") @@ -69,9 +67,7 @@ def test_critical_tokenizer_config_changes_are_rejected(self) -> None: tokenizer_profile(config, 3359, 2) config = self._config("v2607") - config["text_tokenizers"]["japanese_phoneme"]["g2p"][ - "ascii_letter_case" - ] = "upper" + config["text_tokenizers"]["japanese_phoneme"]["g2p"]["ascii_letter_case"] = "upper" with self.assertRaisesRegex(ValueError, "Japanese ascii_letter_case"): tokenizer_profile(config, 3359, 2) diff --git a/tests/cpp/tts/test_grpc_tts_config.cpp b/tests/cpp/tts/test_grpc_tts_config.cpp index 9f7448c..d6cf3ff 100644 --- a/tests/cpp/tts/test_grpc_tts_config.cpp +++ b/tests/cpp/tts/test_grpc_tts_config.cpp @@ -111,8 +111,7 @@ main() { synthesizer_config.tokenizer_model_dir = tokenizer_dir; synthesizer_config.default_language_code = "en-US"; auto synthesizer = std::make_shared(std::move(synthesizer_config)); - const std::vector supported_languages = - synthesizer->supported_language_codes(); + const std::vector supported_languages = synthesizer->supported_language_codes(); nemo_speech::GrpcTtsService service(std::move(synthesizer)); nr_tts::RivaSynthesisConfigRequest req; diff --git a/tests/cpp/tts/test_tokenizer_single_chars.cpp b/tests/cpp/tts/test_tokenizer_single_chars.cpp index 3cf0326..1ed87e2 100644 --- a/tests/cpp/tts/test_tokenizer_single_chars.cpp +++ b/tests/cpp/tts/test_tokenizer_single_chars.cpp @@ -153,7 +153,7 @@ main(int argc, char** argv) { v2607_ok = false; } std::vector expected_languages = { - "en-US", "es-ES", "de-DE", "fr-FR", "it-IT", "vi-VN", + "en-US", "es-ES", "de-DE", "fr-FR", "it-IT", "vi-VN", #ifdef NEMO_SPEECH_TTS_WITH_ZH "zh-CN", #endif @@ -168,12 +168,12 @@ main(int argc, char** argv) { v2607_ok &= check_tokens(tokenizer, "T", "es-ES", {580, 581, 581, 511, 580, 3358}); v2607_ok &= check_tokens( tokenizer, "er ist sehr gut", "de-DE", - {730, 615, 628, 730, 619, 629, 630, 730, 691, 716, 672, 718, 711, 730, 617, 631, - 630, 730, 3358}); + {730, 615, 628, 730, 619, 629, 630, 730, 691, 716, 672, 718, 711, 730, 617, 631, 630, + 730, 3358}); v2607_ok &= check_tokens( tokenizer, "नमस्ते दुनिया।", "hi-IN", - {1326, 1190, 1216, 1189, 1234, 1195, 1196, 1182, 1239, 1326, 1279, 1305, 1281, - 1303, 1288, 1302, 1322, 1326, 3358}); + {1326, 1190, 1216, 1189, 1234, 1195, 1196, 1182, 1239, 1326, 1279, 1305, 1281, 1303, + 1288, 1302, 1322, 1326, 3358}); v2607_ok &= check_tokens( tokenizer, "مرحبا", "ar-AE", {1329, 1405, 1391, 1387, 1382, 1381, 1329, 3358}); v2607_ok &= check_tokens( From a898c78a4bdea44a1551029856c5a863f4c751d2 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 17:58:31 +0000 Subject: [PATCH 10/19] Fix Arabic terminal punctuation Signed-off-by: Anand Joseph --- src/tts/tokenizer/tokenizer.cpp | 3 ++- tests/cpp/tts/test_tts_terminal_punctuation.cpp | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index cfc9ff7..dc3d48b 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -928,7 +928,8 @@ ensure_terminal_punctuation(const std::string& text, const std::string& language return result; }; - static const std::array terminals = {".", "?", "!", "？", "！", "。", "।"}; + static const std::array terminals = { + ".", "?", "!", "？", "！", "。", "।", "؟"}; for (const std::string& existing : terminals) { if (last + 1 >= existing.size()) { const size_t marker_start = last + 1 - existing.size(); diff --git a/tests/cpp/tts/test_tts_terminal_punctuation.cpp b/tests/cpp/tts/test_tts_terminal_punctuation.cpp index 599d2d1..f655621 100644 --- a/tests/cpp/tts/test_tts_terminal_punctuation.cpp +++ b/tests/cpp/tts/test_tts_terminal_punctuation.cpp @@ -47,6 +47,9 @@ main() { check(ensure_terminal_punctuation("नमस्ते।", "hi"), "नमस्ते।", "existing Hindi danda"); check(ensure_terminal_punctuation("你好。", "zh"), "你好。", "existing Chinese full stop"); check(ensure_terminal_punctuation("你好？", "zh"), "你好？", "existing Chinese question mark"); + check( + ensure_terminal_punctuation("كيف حالك؟", "ar-MSA"), "كيف حالك؟", + "existing Arabic question mark"); check( ensure_terminal_punctuation("你好！", "zh"), "你好！", "existing Chinese exclamation mark"); check( From cedbeb2fe94fe1952f43d957fb5b0be3d3276cc8 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 18:00:03 +0000 Subject: [PATCH 11/19] Update documentation for language config Signed-off-by: Anand Joseph --- docs/tts/configuration.md | 32 ++++++++++++++++++-------------- 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/docs/tts/configuration.md b/docs/tts/configuration.md index b7ec8f0..2a834fc 100644 --- a/docs/tts/configuration.md +++ b/docs/tts/configuration.md @@ -41,20 +41,24 @@ when present. The older split layout (`classify/tokenize_and_classify.far`, The optional `riva_server` adapter implements `RivaSpeechSynthesis.Synthesize`, `SynthesizeOnline`, and `GetRivaSynthesisConfig`. It takes plain text in -`SynthesizeSpeechRequest.text`, supports native Magpie tokenizers for `en`, -`es`, `de`, `fr`, `it`, `vi`, `zh`, `hi`, `ja`, `ar-AE`, `ar-SA`, `ar-MSA`, -`ko`, and `pt-BR`, and returns `LINEAR_PCM` -s16le at the NanoCodec sample rate. Japanese and Mandarin are included when -`NEMO_SPEECH_TTS_WITH_JA` and `NEMO_SPEECH_TTS_WITH_ZH`, respectively, are -enabled at build time; both default to `OFF`. Native tokenizers are cached by -language. Mandarin uses bundled Jieba and pypinyin-compatible data together -with the model's pinyin-to-phoneme dictionary. Set `MAGPIE_MANDARIN_G2P_DIR` -only to override the bundled Mandarin data directory. - -`GetRivaSynthesisConfig` advertises every compiled-in TTS language in -`language_code` and exposes the per-language dotted voice names in -`voices_by_language`. The legacy `voice_name`, `subvoices`, and `voices` -parameters remain available for clients that assemble voice names themselves. +`SynthesizeSpeechRequest.text` and returns `LINEAR_PCM` s16le at the NanoCodec +sample rate. The loaded Magpie tokenizer/model determines the supported +languages. Both v2602 and v2607 support `en`, `es`, `de`, `fr`, `it`, `vi`, +`zh`, `hi`, and `ja`. Arabic (`ar-AE`, `ar-SA`, and `ar-MSA`), Korean (`ko`), +and Brazilian Portuguese (`pt-BR`) require a matching v2607 tokenizer and +model; v2602 does not support or advertise those codes. Japanese and Mandarin +also require the respective `NEMO_SPEECH_TTS_WITH_JA` and +`NEMO_SPEECH_TTS_WITH_ZH` build options, both of which default to `OFF`. Native +tokenizers are cached by language. Mandarin uses bundled Jieba and +pypinyin-compatible data together with the model's pinyin-to-phoneme +dictionary. Set `MAGPIE_MANDARIN_G2P_DIR` only to override the bundled Mandarin +data directory. + +`GetRivaSynthesisConfig` advertises the TTS languages supported by the loaded +Magpie tokenizer/model in `language_code` and exposes the per-language dotted +voice names in `voices_by_language`. The legacy `voice_name`, `subvoices`, and +`voices` parameters remain available for clients that assemble voice names +themselves. TTS auto-enables when `tts.magpie-model`, `tts.codec-model`, and `tts.tokenizer-model-dir` are all set; force with `tts.enabled`. From 28b7d0c4cd7ef499e019a684c6a3afe693d825e5 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 18:08:23 +0000 Subject: [PATCH 12/19] Fixes for frame stacking Signed-off-by: Anand Joseph --- server/http/http_server.cpp | 5 ++ src/tts/magpietts/lt.cpp | 84 +++++++++++++++++-- src/tts/magpietts/lt.h | 4 + src/tts/magpietts/magpietts.cpp | 18 +++- src/tts/magpietts/model.cpp | 27 ++++-- tests/cpp/CMakeLists.txt | 7 ++ tests/cpp/common/test_http_server_config.cpp | 44 ++++++++++ tests/cpp/tts/CMakeLists.txt | 1 + .../cpp/tts/test_magpietts_frame_stacking.cpp | 24 ++++++ 9 files changed, 198 insertions(+), 16 deletions(-) create mode 100644 tests/cpp/common/test_http_server_config.cpp diff --git a/server/http/http_server.cpp b/server/http/http_server.cpp index 175644e..fee4fc2 100644 --- a/server/http/http_server.cpp +++ b/server/http/http_server.cpp @@ -329,6 +329,7 @@ class TtsPreemptionCoordinator { uint64_t claim() { std::unique_lock lock(mutex_); const uint64_t generation = ++newest_generation_; + ready_.notify_all(); ready_.wait(lock, [&] { return !active_ || generation != newest_generation_; }); if (generation != newest_generation_) return 0; @@ -379,6 +380,10 @@ struct Server::Impl { Impl(EngineRegistry& engines, ServerConfig config) : models(engines), config(std::move(config)) { + if (this->config.preempt_tts && this->config.threads < 2) { + throw std::invalid_argument( + "tts.preempt requires at least two HTTP workers (http.threads >= 2)"); + } const bool has_cert = !this->config.tls_certificate.empty(); const bool has_key = !this->config.tls_private_key.empty(); if (has_cert != has_key) diff --git a/src/tts/magpietts/lt.cpp b/src/tts/magpietts/lt.cpp index dca2032..6a39c84 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include #include @@ -1172,6 +1173,68 @@ local_transformer_graph_bank_eval_cuda( } #endif +bool +magpietts_stack_forced_code_frames( + const std::vector>& raw_frames, size_t first_frame, + const magpietts_hparams& h, std::vector& stacked_codes) { + stacked_codes.clear(); + if (h.audio_codebooks < 1 || h.frame_stacking_factor < 1) { + fprintf( + stderr, "invalid forced-code layout: codebooks=%d frame_stacking_factor=%d\n", + h.audio_codebooks, h.frame_stacking_factor); + return false; + } + + const int64_t stacked_count = + static_cast(h.audio_codebooks) * h.frame_stacking_factor; + if (stacked_count > std::numeric_limits::max()) { + fprintf( + stderr, "forced-code layout has too many stacked codebooks: %lld\n", + static_cast(stacked_count)); + return false; + } + + const size_t lane_count = static_cast(h.frame_stacking_factor); + if (raw_frames.size() % lane_count != 0) { + fprintf( + stderr, + "incomplete forced-code input: %zu raw frame(s) cannot be grouped into %zu-frame " + "stacks\n", + raw_frames.size(), lane_count); + return false; + } + if (first_frame % lane_count != 0) { + fprintf( + stderr, "forced-code stack starts at unaligned raw frame %zu (stacking factor %zu)\n", + first_frame, lane_count); + return false; + } + if (first_frame > raw_frames.size() || raw_frames.size() - first_frame < lane_count) { + const size_t available = + first_frame < raw_frames.size() ? raw_frames.size() - first_frame : 0; + fprintf( + stderr, + "incomplete forced-code input at frame %zu: found %zu consecutive frame(s), " + "expected %zu\n", + first_frame, available, lane_count); + return false; + } + + stacked_codes.reserve(static_cast(stacked_count)); + for (size_t lane = 0; lane < lane_count; ++lane) { + const auto& frame = raw_frames[first_frame + lane]; + if (frame.size() != static_cast(h.audio_codebooks)) { + fprintf( + stderr, "forced-code frame %zu has %zu codebooks, expected %d\n", + first_frame + lane, frame.size(), h.audio_codebooks); + stacked_codes.clear(); + return false; + } + stacked_codes.insert(stacked_codes.end(), frame.begin(), frame.end()); + } + return true; +} + static bool sample_local_codebooks_impl( const magpietts_model& model, const std::vector& cond_hidden, @@ -1182,13 +1245,24 @@ sample_local_codebooks_impl( const std::vector* forced_codes) { const ggml_nvtx::range nvtx_range("magpietts_sample_local_codebooks"); const auto& h = model.hparams; + codes.clear(); + argmax_codes.clear(); + const int32_t stacked_codebooks = h.stacked_audio_codebooks(); + if (stacked_codebooks < 1) { + fprintf(stderr, "invalid stacked codebook count: %d\n", stacked_codebooks); + return false; + } + if (forced_codes && forced_codes->size() != static_cast(stacked_codebooks)) { + fprintf( + stderr, "forced-code input has %zu stacked codebooks, expected %d\n", + forced_codes->size(), stacked_codebooks); + return false; + } if (!local_graphs.beginFrame(model, use_cfg)) { return false; } - codes.clear(); - argmax_codes.clear(); std::vector prev; - for (int c = 0; c < h.stacked_audio_codebooks(); ++c) { + for (int c = 0; c < stacked_codebooks; ++c) { std::vector logits; if (use_cfg) { std::vector uncond; @@ -1209,9 +1283,7 @@ sample_local_codebooks_impl( logits, h, temperature, top_k, rng, forbid_audio_eos); const int greedy = MagpieCodebookSampler::argmaxFromLogits(logits, h, forbid_audio_eos); dump_local_codebook_logits(logit_dump, c, sampled, greedy, logits); - const int emitted = forced_codes && (int)forced_codes->size() == h.stacked_audio_codebooks() - ? (*forced_codes)[c] - : sampled; + const int emitted = forced_codes ? (*forced_codes)[c] : sampled; codes.push_back(emitted); argmax_codes.push_back(greedy); prev.push_back(emitted); diff --git a/src/tts/magpietts/lt.h b/src/tts/magpietts/lt.h index 7120268..a473723 100644 --- a/src/tts/magpietts/lt.h +++ b/src/tts/magpietts/lt.h @@ -23,6 +23,10 @@ struct LocalCodebookLogitDump { int frame_index = 0; }; +bool magpietts_stack_forced_code_frames( + const std::vector>& raw_frames, size_t first_frame, + const magpietts_hparams& h, std::vector& stacked_codes); + class LocalCodebookSampler { public: LocalCodebookSampler(const magpietts_model& model, int threads); diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index 8e56893..1c245d9 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1485,13 +1485,23 @@ stream_magpie_to_audio( logit_dump.step = step; logit_dump.frame_index = sample_frame_index; } + std::vector stacked_forced_codes; + const std::vector* forced_codes = nullptr; + const size_t first_forced_frame = + static_cast(step) * + static_cast(h.frame_stacking_factor); + if (chunk_index == 0 && first_forced_frame < forced_code_frames.size()) { + if (!magpietts_stack_forced_code_frames( + forced_code_frames, first_forced_frame, h, + stacked_forced_codes)) { + return cancel_worker(); + } + forced_codes = &stacked_forced_codes; + } if (!local_sampler->sample( cond.hidden_last, uncond.hidden_last, params.use_cfg, h.cfg_scale, h.temperature, h.top_k, forbid_eos, rng, next_codes, argmax_codes, - dump_logits ? &logit_dump : nullptr, - chunk_index == 0 && step < (int)forced_code_frames.size() - ? &forced_code_frames[(size_t)step] - : nullptr)) { + dump_logits ? &logit_dump : nullptr, forced_codes)) { return cancel_worker(); } } else { diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index f6bc126..fc6d661 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -725,10 +725,29 @@ magpietts_model_load_impl( h.mask_token_id = gguf_i32(model.gguf, "magpietts.mask_token_id", h.mask_token_id); h.frame_stacking_factor = gguf_i32(model.gguf, "magpietts.frame_stacking_factor", h.frame_stacking_factor); + if (h.audio_codebooks < 1 || h.frame_stacking_factor < 1) { + fprintf( + stderr, "invalid stacked audio layout: codebooks=%d frame_stacking_factor=%d\n", + h.audio_codebooks, h.frame_stacking_factor); + return false; + } + const int64_t stacked_audio_codebooks_64 = + static_cast(h.audio_codebooks) * h.frame_stacking_factor; + if (stacked_audio_codebooks_64 > std::numeric_limits::max()) { + fprintf( + stderr, + "stacked audio layout overflows int32: codebooks=%d frame_stacking_factor=%d " + "slots=%lld\n", + h.audio_codebooks, h.frame_stacking_factor, + static_cast(stacked_audio_codebooks_64)); + return false; + } + const int32_t expected_stacked_codebooks = + static_cast(stacked_audio_codebooks_64); model.tokenizer_profile = gguf_string(model.gguf, "magpietts.tokenizer_profile"); model.nemo_version = gguf_string(model.gguf, "magpietts.nemo_version"); const int32_t stored_stacked_codebooks = - gguf_i32(model.gguf, "magpietts.stacked_audio_codebooks", h.stacked_audio_codebooks()); + gguf_i32(model.gguf, "magpietts.stacked_audio_codebooks", expected_stacked_codebooks); h.n_embd = gguf_i32(model.gguf, "magpietts.embedding_dim", h.n_embd); h.n_ffn = gguf_i32(model.gguf, "magpietts.ffn_dim", h.n_ffn); h.n_ctx = gguf_i32(model.gguf, "magpietts.context_length", h.n_ctx); @@ -807,10 +826,6 @@ magpietts_model_load_impl( return false; } - if (h.frame_stacking_factor < 1) { - fprintf(stderr, "invalid frame_stacking_factor=%d\n", h.frame_stacking_factor); - return false; - } if (model.tokenizer_profile.empty()) { model.tokenizer_profile = magpietts_infer_tokenizer_profile(h); } @@ -822,7 +837,7 @@ magpietts_model_load_impl( model.tokenizer_profile.c_str(), h.text_vocab_size, h.frame_stacking_factor); return false; } - if (h.audio_codebooks < 1 || stored_stacked_codebooks != h.stacked_audio_codebooks()) { + if (stored_stacked_codebooks != expected_stacked_codebooks) { fprintf( stderr, "invalid stacked audio layout: codebooks=%d frame_stacking_factor=%d stored_slots=%d\n", diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index b0b362a..5e8e559 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -20,6 +20,13 @@ add_executable(test_subtitles target_link_libraries(test_subtitles PRIVATE nemo_speech_common) add_test(NAME subtitles COMMAND test_subtitles) +if(TARGET nemo_speech_http_server) + add_executable(test_http_server_config + ${CMAKE_SOURCE_DIR}/tests/cpp/common/test_http_server_config.cpp) + target_link_libraries(test_http_server_config PRIVATE nemo_speech_http_server) + add_test(NAME http_server_config COMMAND test_http_server_config) +endif() + find_package(Python3 COMPONENTS Interpreter QUIET) if(Python3_Interpreter_FOUND AND (NEMO_SPEECH_BUILD_ASR OR NEMO_SPEECH_BUILD_DIAR OR diff --git a/tests/cpp/common/test_http_server_config.cpp b/tests/cpp/common/test_http_server_config.cpp new file mode 100644 index 0000000..3ce7d99 --- /dev/null +++ b/tests/cpp/common/test_http_server_config.cpp @@ -0,0 +1,44 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#include +#include +#include + +#include "engine_registry.h" +#include "http_server.h" + +int +main() { + nemo_speech::EngineRegistry engines; + nemo_speech::http::ServerConfig config; + config.threads = 1; + config.preempt_tts = true; + + bool rejected = false; + try { + nemo_speech::http::Server server(engines, config); + } + catch (const std::invalid_argument& error) { + if (std::string(error.what()).find("http.threads >= 2") != std::string::npos) { + rejected = true; + } else { + std::cerr << "FAIL: unexpected validation error: " << error.what() << '\n'; + return 1; + } + } + if (!rejected) { + std::cerr << "FAIL: single-worker TTS preemption was accepted\n"; + return 1; + } + + config.threads = 2; + try { + nemo_speech::http::Server server(engines, config); + } + catch (const std::exception& error) { + std::cerr << "FAIL: two-worker TTS preemption was rejected: " << error.what() << '\n'; + return 1; + } + + return 0; +} diff --git a/tests/cpp/tts/CMakeLists.txt b/tests/cpp/tts/CMakeLists.txt index b2e53fe..0d558fe 100644 --- a/tests/cpp/tts/CMakeLists.txt +++ b/tests/cpp/tts/CMakeLists.txt @@ -10,6 +10,7 @@ target_link_libraries(test_magpietts_attention_prior PRIVATE nemo_speech_tts_mag add_executable(test_magpietts_frame_stacking test_magpietts_frame_stacking.cpp) target_link_libraries(test_magpietts_frame_stacking PRIVATE nemo_speech_tts_magpietts) +add_test(NAME magpietts_frame_stacking COMMAND test_magpietts_frame_stacking) if(GGML_CUDA AND NEMO_SPEECH_GGML_PATCHED) add_executable(test_magpietts_cached_attention test_magpietts_cached_attention.cpp) diff --git a/tests/cpp/tts/test_magpietts_frame_stacking.cpp b/tests/cpp/tts/test_magpietts_frame_stacking.cpp index 225849c..c304c86 100644 --- a/tests/cpp/tts/test_magpietts_frame_stacking.cpp +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -3,6 +3,7 @@ #include #include +#include "tts/magpietts/lt.h" #include "tts/magpietts/model.h" namespace tts = nemo_speech::tts; @@ -44,6 +45,29 @@ main() { return 1; } + const std::vector> forced_frames = { + {10, 11}, {20, 21}, {30, 31}, {40, 41}}; + std::vector stacked_forced; + if (!tts::magpietts_stack_forced_code_frames(forced_frames, 0, h, stacked_forced) || + stacked_forced != stacked) { + std::fprintf(stderr, "forced-code frame stacking failed\n"); + return 1; + } + const std::vector> incomplete_forced_frames = { + {10, 11}, {20, 21}, {30, 31}}; + if (tts::magpietts_stack_forced_code_frames( + incomplete_forced_frames, 0, h, stacked_forced) || + !stacked_forced.empty()) { + std::fprintf(stderr, "incomplete forced-code frame group was accepted\n"); + return 1; + } + const std::vector> malformed_forced_frames = {{10}, {20, 21}}; + if (tts::magpietts_stack_forced_code_frames( + malformed_forced_frames, 0, h, stacked_forced)) { + std::fprintf(stderr, "malformed forced-code frame was accepted\n"); + return 1; + } + std::vector greedy = stacked; greedy[3] = h.audio_eos_id; if (tts::magpietts_first_eos_lane(stacked, greedy, h) != 1) { From a8b23662b2b3f8726c4b9349fdc91e7e633f3e87 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 18:27:40 +0000 Subject: [PATCH 13/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/tts/magpietts/lt.cpp | 3 +-- src/tts/magpietts/model.cpp | 3 +-- src/tts/tokenizer/tokenizer.cpp | 3 +-- tests/cpp/tts/test_magpietts_frame_stacking.cpp | 6 ++---- 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/src/tts/magpietts/lt.cpp b/src/tts/magpietts/lt.cpp index 6a39c84..dd4f068 100644 --- a/src/tts/magpietts/lt.cpp +++ b/src/tts/magpietts/lt.cpp @@ -1185,8 +1185,7 @@ magpietts_stack_forced_code_frames( return false; } - const int64_t stacked_count = - static_cast(h.audio_codebooks) * h.frame_stacking_factor; + const int64_t stacked_count = static_cast(h.audio_codebooks) * h.frame_stacking_factor; if (stacked_count > std::numeric_limits::max()) { fprintf( stderr, "forced-code layout has too many stacked codebooks: %lld\n", diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index f3623cb..1d0014e 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -736,8 +736,7 @@ magpietts_model_load_impl( static_cast(stacked_audio_codebooks_64)); return false; } - const int32_t expected_stacked_codebooks = - static_cast(stacked_audio_codebooks_64); + const int32_t expected_stacked_codebooks = static_cast(stacked_audio_codebooks_64); model.tokenizer_profile = gguf_string(model.gguf, "magpietts.tokenizer_profile"); model.nemo_version = gguf_string(model.gguf, "magpietts.nemo_version"); const int32_t stored_stacked_codebooks = diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index 5efc2f6..6394b99 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -950,8 +950,7 @@ ensure_terminal_punctuation(const std::string& text, const std::string& language return result; }; - static const std::array terminals = { - ".", "?", "!", "？", "！", "。", "।", "؟"}; + static const std::array terminals = {".", "?", "!", "？", "！", "。", "।", "؟"}; for (const std::string& existing : terminals) { if (last + 1 >= existing.size()) { const size_t marker_start = last + 1 - existing.size(); diff --git a/tests/cpp/tts/test_magpietts_frame_stacking.cpp b/tests/cpp/tts/test_magpietts_frame_stacking.cpp index c304c86..847e3aa 100644 --- a/tests/cpp/tts/test_magpietts_frame_stacking.cpp +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -55,15 +55,13 @@ main() { } const std::vector> incomplete_forced_frames = { {10, 11}, {20, 21}, {30, 31}}; - if (tts::magpietts_stack_forced_code_frames( - incomplete_forced_frames, 0, h, stacked_forced) || + if (tts::magpietts_stack_forced_code_frames(incomplete_forced_frames, 0, h, stacked_forced) || !stacked_forced.empty()) { std::fprintf(stderr, "incomplete forced-code frame group was accepted\n"); return 1; } const std::vector> malformed_forced_frames = {{10}, {20, 21}}; - if (tts::magpietts_stack_forced_code_frames( - malformed_forced_frames, 0, h, stacked_forced)) { + if (tts::magpietts_stack_forced_code_frames(malformed_forced_frames, 0, h, stacked_forced)) { std::fprintf(stderr, "malformed forced-code frame was accepted\n"); return 1; } From 1e83b0e19f6cb1b27ce79e5ef52d81947783347f Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 18:35:12 +0000 Subject: [PATCH 14/19] Gate preempt checks on TTS support. Signed-off-by: Anand Joseph --- tests/cli/cli_contract_test.py | 6 +- tests/conversion/tts_index_layout_test.py | 90 +++++++++++++++++++++++ 2 files changed, 94 insertions(+), 2 deletions(-) create mode 100644 tests/conversion/tts_index_layout_test.py diff --git a/tests/cli/cli_contract_test.py b/tests/cli/cli_contract_test.py index 6bad8ee..8c8f1de 100644 --- a/tests/cli/cli_contract_test.py +++ b/tests/cli/cli_contract_test.py @@ -74,8 +74,10 @@ def main() -> None: assert error["exit_code"] in (3, 4), error serve_help = run(binary, "serve", "--help") + has_tts_options = "--tts-model" in serve_help.stdout if "--cors-origin" in serve_help.stdout: - assert "--tts.preempt" in serve_help.stdout + if has_tts_options: + assert "--tts.preempt" in serve_help.stdout with socket.socket() as stalled_server: stalled_server.bind(("127.0.0.1", 0)) stalled_server.listen() @@ -147,7 +149,7 @@ def stall_response() -> None: binary, "--json", "serve", - "--tts.preempt", + *(["--tts.preempt"] if has_tts_options else []), *missing_arguments, "--no-warmup", ), diff --git a/tests/conversion/tts_index_layout_test.py b/tests/conversion/tts_index_layout_test.py new file mode 100644 index 0000000..b37b9ec --- /dev/null +++ b/tests/conversion/tts_index_layout_test.py @@ -0,0 +1,90 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import sys +import unittest +from pathlib import Path +from unittest import mock + +import torch + +ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(ROOT)) + +from conversion.tts import ( # noqa: E402 + _indexed_weight_indices, + _require_contiguous_indices, + add_metadata, +) + + +class TtsIndexLayoutTest(unittest.TestCase): + def test_contiguous_audio_embedding_indexes(self) -> None: + state = { + "audio_embeddings.0.weight": object(), + "audio_embeddings.1.weight": object(), + "audio_embeddings.2.weight": object(), + } + indices = _indexed_weight_indices(state, "audio_embeddings.") + _require_contiguous_indices("audio embedding", indices, 3) + self.assertEqual(indices, [0, 1, 2]) + + def test_sparse_audio_embedding_indexes_are_rejected(self) -> None: + state = { + "audio_embeddings.0.weight": object(), + "audio_embeddings.2.weight": object(), + } + indices = _indexed_weight_indices(state, "audio_embeddings.") + with self.assertRaisesRegex(ValueError, "must be contiguous"): + _require_contiguous_indices("audio embedding", indices, 2) + + def test_duplicate_lt_head_indexes_are_rejected(self) -> None: + state = { + "local_transformer_out_projections.0.weight": object(), + "local_transformer_out_projections.00.weight": object(), + } + indices = _indexed_weight_indices(state, "local_transformer_out_projections.") + with self.assertRaisesRegex(ValueError, "must be contiguous"): + _require_contiguous_indices("local transformer output projection", indices, 2) + + def test_non_numeric_index_is_rejected(self) -> None: + state = {"audio_embeddings.first.weight": object()} + with self.assertRaisesRegex(ValueError, "non-numeric weight index"): + _indexed_weight_indices(state, "audio_embeddings.") + + def test_sparse_audio_embeddings_fail_before_metadata_is_written(self) -> None: + writer = mock.Mock() + config = {"encoder": {}, "decoder": {}, "frame_stacking_factor": 1} + state = { + "audio_embeddings.0.weight": object(), + "audio_embeddings.2.weight": object(), + } + with self.assertRaisesRegex(ValueError, "must be contiguous"): + add_metadata(writer, config, state) # type: ignore[arg-type] + self.assertFalse(writer.method_calls) + + def test_duplicate_lt_heads_fail_before_metadata_is_written(self) -> None: + writer = mock.Mock() + config = {"encoder": {}, "decoder": {}, "frame_stacking_factor": 1} + state = { + "audio_embeddings.0.weight": torch.zeros((16, 1)), + "audio_embeddings.1.weight": torch.zeros((16, 1)), + "text_embedding.weight": torch.zeros((1, 1)), + "_baked_embedding_T": torch.tensor(1), + "_baked_embedding_D": torch.tensor(1), + "baked_context_embedding_len": torch.tensor([1]), + "final_proj.weight": torch.zeros((32, 1)), + "local_transformer_out_projections.0.weight": torch.zeros((1, 1)), + "local_transformer_out_projections.00.weight": torch.zeros((1, 1)), + } + with mock.patch("conversion.tts.tokenizer_profile", return_value="test"): + with self.assertRaisesRegex(ValueError, "must be contiguous"): + add_metadata(writer, config, state) + self.assertFalse(writer.method_calls) + + +if __name__ == "__main__": + unittest.main() From 98f97f54da69cee87f0a21ebbf7007106944fc70 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 18:36:01 +0000 Subject: [PATCH 15/19] Validate sequential tensor indexes. Signed-off-by: Anand Joseph --- conversion/tts.py | 57 ++++++++++++++++++++++++++++++++--------------- 1 file changed, 39 insertions(+), 18 deletions(-) diff --git a/conversion/tts.py b/conversion/tts.py index adc7f75..b58f859 100644 --- a/conversion/tts.py +++ b/conversion/tts.py @@ -80,29 +80,52 @@ def add_i32_array(writer: gguf.GGUFWriter, key: str, value: Any) -> None: writer.add_array(key, items) +def _indexed_weight_indices(sd: dict[str, torch.Tensor], prefix: str) -> list[int]: + suffix = ".weight" + indices: list[int] = [] + for name in sd: + if not name.startswith(prefix) or not name.endswith(suffix): + continue + raw_index = name[len(prefix) : -len(suffix)] + if not raw_index.isascii() or not raw_index.isdigit(): + raise ValueError(f"{prefix} contains a non-numeric weight index: {name}") + indices.append(int(raw_index)) + return indices + + +def _require_contiguous_indices(label: str, indices: list[int], expected_count: int) -> None: + expected = list(range(expected_count)) + actual = sorted(indices) + if actual != expected: + raise ValueError( + f"{label} indexes must be contiguous from 0 through {expected_count - 1}: " + f"found={actual}" + ) + + def add_metadata( writer: gguf.GGUFWriter, cfg: dict[str, Any], sd: dict[str, torch.Tensor] ) -> dict[str, Any]: - audio_vocab = int(sd["audio_embeddings.0.weight"].shape[0]) - codebook_size = audio_vocab - SPECIAL_AUDIO_TOKENS - text_vocab = int(sd["text_embedding.weight"].shape[0]) - baked_t = int(sd["_baked_embedding_T"].item()) - baked_d = int(sd["_baked_embedding_D"].item()) - baked_lens = [int(x) for x in sd["baked_context_embedding_len"].tolist()] - encoder = cfg["encoder"] decoder = cfg["decoder"] lt_hidden = int(cfg.get("local_transformer_hidden_dim", 256)) frame_stacking = int(cfg.get("frame_stacking_factor", 1)) - profile = tokenizer_profile(cfg, text_vocab, frame_stacking) - n_stacked_codebooks = int( - len([k for k in sd if k.startswith("audio_embeddings.") and k.endswith(".weight")]) - ) + audio_embedding_indices = _indexed_weight_indices(sd, "audio_embeddings.") + n_stacked_codebooks = len(audio_embedding_indices) if frame_stacking < 1 or n_stacked_codebooks == 0 or n_stacked_codebooks % frame_stacking: raise ValueError( "audio embedding count must be a positive multiple of frame_stacking_factor: " f"embeddings={n_stacked_codebooks} frame_stacking_factor={frame_stacking}" ) + _require_contiguous_indices("audio embedding", audio_embedding_indices, n_stacked_codebooks) + + audio_vocab = int(sd["audio_embeddings.0.weight"].shape[0]) + codebook_size = audio_vocab - SPECIAL_AUDIO_TOKENS + text_vocab = int(sd["text_embedding.weight"].shape[0]) + baked_t = int(sd["_baked_embedding_T"].item()) + baked_d = int(sd["_baked_embedding_D"].item()) + baked_lens = [int(x) for x in sd["baked_context_embedding_len"].tolist()] + profile = tokenizer_profile(cfg, text_vocab, frame_stacking) n_codebooks = n_stacked_codebooks // frame_stacking expected_logits = n_stacked_codebooks * audio_vocab if int(sd["final_proj.weight"].shape[0]) != expected_logits: @@ -110,18 +133,16 @@ def add_metadata( "final_proj rows do not match stacked audio layout: " f"rows={sd['final_proj.weight'].shape[0]} expected={expected_logits}" ) - n_lt_heads = len( - [ - k - for k in sd - if k.startswith("local_transformer_out_projections.") and k.endswith(".weight") - ] - ) + lt_head_indices = _indexed_weight_indices(sd, "local_transformer_out_projections.") + n_lt_heads = len(lt_head_indices) if n_lt_heads != n_stacked_codebooks: raise ValueError( "local transformer output heads do not match audio embeddings: " f"heads={n_lt_heads} embeddings={n_stacked_codebooks}" ) + _require_contiguous_indices( + "local transformer output projection", lt_head_indices, n_stacked_codebooks + ) inf = cfg.get("inference_parameters", {}) From ca4769c0ce3c15e6d2e3d5429cad4108836707b6 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 18:55:46 +0000 Subject: [PATCH 16/19] Decoder fixes Signed-off-by: Anand Joseph --- src/tts/magpietts/decoder.cpp | 33 ++++++++++++++++++++------------- src/tts/magpietts/decoder.h | 3 ++- src/tts/magpietts/magpietts.cpp | 13 +++++++------ src/tts/magpietts/model.cpp | 14 ++++++++------ 4 files changed, 37 insertions(+), 26 deletions(-) diff --git a/src/tts/magpietts/decoder.cpp b/src/tts/magpietts/decoder.cpp index faa6536..109af5a 100644 --- a/src/tts/magpietts/decoder.cpp +++ b/src/tts/magpietts/decoder.cpp @@ -616,13 +616,11 @@ class PersistentDecoderModule final : public ggml_runtime::Module { class MagpieDecoder::PersistentDecoderRuntime { public: PersistentDecoderRuntime( - const magpietts_model& model, const DecoderCrossKvCache& cross_kv, int text_len) + const magpietts_model& model, const DecoderCrossKvCache& cross_kv, int text_len, + int stacked_position_budget) : model_(model), cross_kv_(&cross_kv), text_len_(text_len), - cache_len_( - model.hparams.baked_context_length + - (model.hparams.max_decoder_steps + model.hparams.frame_stacking_factor - 1) / - model.hparams.frame_stacking_factor - - 1), + stacked_position_budget_(stacked_position_budget), + cache_len_(model.hparams.baked_context_length + stacked_position_budget - 1), backend_manager_(ggml_runtime::Params{true, 0, nullptr}, model.backend), module_(model, cross_kv, text_len, cache_len_), session_(backend_manager_, &module_, nullptr) { @@ -633,8 +631,10 @@ class MagpieDecoder::PersistentDecoderRuntime { session_.setup(); } - bool matches(const DecoderCrossKvCache* cross_kv, int text_len) const { - return cross_kv == cross_kv_ && text_len == text_len_; + bool matches( + const DecoderCrossKvCache* cross_kv, int text_len, int stacked_position_budget) const { + return cross_kv == cross_kv_ && text_len == text_len_ && + stacked_position_budget == stacked_position_budget_; } bool sequence_matches(int n_tokens) const { return n_tokens == n_tokens_; } @@ -772,6 +772,7 @@ class MagpieDecoder::PersistentDecoderRuntime { const magpietts_model& model_; const DecoderCrossKvCache* cross_kv_ = nullptr; int text_len_ = 0; + int stacked_position_budget_ = 0; int cache_len_ = 0; int n_tokens_ = 0; int valid_tokens_ = 0; @@ -829,7 +830,8 @@ MagpieDecoder::evalCachedPair( const std::vector& text_cond, int text_len, const std::vector>& audio_codes, int speaker, int threads, DecoderKvCache& cond_kv, DecoderKvCache& uncond_kv, decoder_result& cond_result, - decoder_result& uncond_result, magpietts_cuda_sample_request* cuda_sample, + decoder_result& uncond_result, int stacked_position_budget, + magpietts_cuda_sample_request* cuda_sample, const magpietts_backend_tensor* text_cond_device, magpietts_backend_tensor* cond_hidden_out, magpietts_backend_tensor* uncond_hidden_out, DecoderCrossKvCache* cond_cross_kv, const magpietts_decoder_attention* attention) const { @@ -840,15 +842,18 @@ MagpieDecoder::evalCachedPair( cond_kv.n_tokens > 0 && cond_kv.n_tokens == uncond_kv.n_tokens; if (persistent_candidate) { try { - if (persistent_runtime_ && !persistent_runtime_->matches(cond_cross_kv, text_len)) { - // Cross-cache address or shape changes require a new graph. + if (persistent_runtime_ && + !persistent_runtime_->matches( + cond_cross_kv, text_len, stacked_position_budget)) { + // Cross-cache address, shape, or request budget changes require a new graph. persistent_runtime_.reset(); cond_kv.clear(); uncond_kv.clear(); } if (cond_kv.n_tokens > 0 && !persistent_runtime_) { persistent_runtime_ = - std::make_unique(model_, *cond_cross_kv, text_len); + std::make_unique( + model_, *cond_cross_kv, text_len, stacked_position_budget); persistent_runtime_->seed(cond_kv, uncond_kv); fprintf( stderr, @@ -871,7 +876,9 @@ MagpieDecoder::evalCachedPair( } catch (const std::exception& e) { fprintf(stderr, "MagpieTTS persistent decoder failed: %s\n", e.what()); - return false; + persistent_runtime_.reset(); + cond_kv.clear(); + uncond_kv.clear(); } } return decoder_eval_cached_pair_impl( diff --git a/src/tts/magpietts/decoder.h b/src/tts/magpietts/decoder.h index a298053..84569f2 100644 --- a/src/tts/magpietts/decoder.h +++ b/src/tts/magpietts/decoder.h @@ -122,7 +122,8 @@ class MagpieDecoder { const std::vector& text_cond, int text_len, const std::vector>& audio_codes, int speaker, int threads, DecoderKvCache& cond_kv, DecoderKvCache& uncond_kv, decoder_result& cond_result, - decoder_result& uncond_result, magpietts_cuda_sample_request* cuda_sample = nullptr, + decoder_result& uncond_result, int stacked_position_budget, + magpietts_cuda_sample_request* cuda_sample = nullptr, const magpietts_backend_tensor* text_cond_device = nullptr, magpietts_backend_tensor* cond_hidden_out = nullptr, magpietts_backend_tensor* uncond_hidden_out = nullptr, diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index 7bf9a0a..b72a66c 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1366,7 +1366,8 @@ stream_magpie_to_audio( ? decoder.evalCachedPair( text_cond, text_len, audio_codes, params.speaker, params.threads, cond_kv, uncond_kv, cond, uncond, - nullptr, &text_cond_device, &cond_hidden_device, + max_decoder_positions, nullptr, &text_cond_device, + &cond_hidden_device, &uncond_hidden_device, &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( @@ -1419,8 +1420,8 @@ stream_magpie_to_audio( ? decoder.evalCachedPair( text_cond, text_len, audio_codes, params.speaker, params.threads, cond_kv, uncond_kv, cond, uncond, - &cuda_sample, &text_cond_device, nullptr, nullptr, - &cond_cross_kv, decoder_attention_arg) + max_decoder_positions, &cuda_sample, &text_cond_device, + nullptr, nullptr, &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( text_cond, text_len, audio_codes, params.speaker, params.threads, cond, uncond, &cuda_sample, @@ -1454,9 +1455,9 @@ stream_magpie_to_audio( params.use_kv_cache ? decoder.evalCachedPair( text_cond, text_len, audio_codes, params.speaker, - params.threads, cond_kv, uncond_kv, cond, uncond, nullptr, - nullptr, nullptr, nullptr, &cond_cross_kv, - decoder_attention_arg) + params.threads, cond_kv, uncond_kv, cond, uncond, + max_decoder_positions, nullptr, nullptr, nullptr, nullptr, + &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( text_cond, text_len, audio_codes, params.speaker, params.threads, cond, uncond, nullptr, nullptr, nullptr, diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index f3623cb..0d910d5 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -2057,8 +2057,9 @@ MagpieCodeGenerator::generate( ? (params.use_kv_cache ? decoder.evalCachedPair( text_cond, (int)tokens.size(), audio_codes, params.speaker, - params.threads, cond_kv, uncond_kv, cond, uncond, nullptr, - &text_cond_device, &cond_hidden_device, + params.threads, cond_kv, uncond_kv, cond, uncond, + max_decoder_positions, nullptr, &text_cond_device, + &cond_hidden_device, &uncond_hidden_device, &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( @@ -2109,8 +2110,8 @@ MagpieCodeGenerator::generate( ? decoder.evalCachedPair( text_cond, (int)tokens.size(), audio_codes, params.speaker, params.threads, cond_kv, uncond_kv, cond, uncond, - &cuda_sample, &text_cond_device, nullptr, nullptr, - &cond_cross_kv, decoder_attention_arg) + max_decoder_positions, &cuda_sample, &text_cond_device, + nullptr, nullptr, &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( text_cond, (int)tokens.size(), audio_codes, params.speaker, params.threads, cond, uncond, &cuda_sample, @@ -2143,8 +2144,9 @@ MagpieCodeGenerator::generate( params.use_kv_cache ? decoder.evalCachedPair( text_cond, (int)tokens.size(), audio_codes, params.speaker, - params.threads, cond_kv, uncond_kv, cond, uncond, nullptr, - nullptr, nullptr, nullptr, &cond_cross_kv, decoder_attention_arg) + params.threads, cond_kv, uncond_kv, cond, uncond, + max_decoder_positions, nullptr, nullptr, nullptr, nullptr, + &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( text_cond, (int)tokens.size(), audio_codes, params.speaker, params.threads, cond, uncond, nullptr, nullptr, nullptr, nullptr, From fc4a2ecccba6e201262894fc9e2ee7670d4dc95d Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 19:17:11 +0000 Subject: [PATCH 17/19] Magpie fixes Signed-off-by: Anand Joseph --- src/tts/magpietts/magpietts.cpp | 7 +++++-- src/tts/magpietts/model.cpp | 7 +++++-- src/tts/tokenizer/tokenizer.cpp | 4 +++- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index b72a66c..f8b83e3 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1329,6 +1329,8 @@ stream_magpie_to_audio( (h.max_decoder_steps + h.frame_stacking_factor - 1) / h.frame_stacking_factor; for (int step = 0; step < max_decoder_positions; ++step) { const ggml_nvtx::range nvtx_step("magpietts_stream_generation_step"); + const int frames_remaining = + h.max_decoder_steps - step * h.frame_stacking_factor; if (codec_worker.is_failed()) { codec_worker.join(); return false; @@ -1563,7 +1565,7 @@ stream_magpie_to_audio( } const int eos_lane = forbid_eos ? -1 : magpietts_first_eos_lane(next_codes, argmax_codes, h); - const bool has_eos = eos_lane >= 0; + const bool has_eos = eos_lane >= 0 && eos_lane < frames_remaining; bool terminate_after_frame = false; if (has_eos) { ggml_nvtx::mark("magpietts_stream_eos"); @@ -1599,7 +1601,8 @@ stream_magpie_to_audio( } decoder_frames_generated += h.frame_stacking_factor; - const int frames_to_emit = has_eos ? eos_lane : h.frame_stacking_factor; + const int frames_to_emit = std::min( + frames_remaining, has_eos ? eos_lane : h.frame_stacking_factor); if (!suppress_nonfinal_codec_output) { for (int lane = 0; lane < frames_to_emit; ++lane) { const auto& frame = codec_frames[(size_t)lane]; diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index 0d910d5..f938376 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -2029,6 +2029,8 @@ MagpieCodeGenerator::generate( for (int step = 0; step < max_decoder_positions; ++step) { const ggml_nvtx::range nvtx_step("magpietts_generate_step"); const int64_t frame_start_us = ggml_time_us(); + const int frames_remaining = + h.max_decoder_steps - step * h.frame_stacking_factor; if (step % 10 == 0) { fprintf(stderr, "%s decoding frame %d/%d\n", label, step, max_decoder_positions); } @@ -2199,8 +2201,9 @@ MagpieCodeGenerator::generate( audio_codes[c].push_back(next_codes[c + lane * h.audio_codebooks]); } } - for (int lane = 0; lane < (eos_lane >= 0 ? eos_lane : h.frame_stacking_factor); - ++lane) { + const int frames_to_emit = std::min( + frames_remaining, eos_lane >= 0 ? eos_lane : h.frame_stacking_factor); + for (int lane = 0; lane < frames_to_emit; ++lane) { generated_frames.push_back(codec_frames[(size_t)lane]); } if (eos_lane >= 0) { diff --git a/src/tts/tokenizer/tokenizer.cpp b/src/tts/tokenizer/tokenizer.cpp index 5efc2f6..9e18381 100644 --- a/src/tts/tokenizer/tokenizer.cpp +++ b/src/tts/tokenizer/tokenizer.cpp @@ -500,7 +500,9 @@ load_tokenizer_profile(const fs::path& root) { "nemo.collections.common.tokenizers.text_to_speech.tts_tokenizers." "JapanesePhonemeTokenizer", "false", "true"); - require_block_value(blocks, item.name, "ascii_letter_case", item.ascii_letter_case); + require_block_value( + blocks, item.name, "ascii_letter_case", + profile.id == "v2602" ? "upper" : "lower"); break; case tokenizer_kind::arabic: require_common_tokenizer_values( From 9cffd5bed04315da299380309bf150e8f58f73a8 Mon Sep 17 00:00:00 2001 From: Anand Joseph Date: Thu, 27 Aug 2026 19:18:13 +0000 Subject: [PATCH 18/19] More Magpie fixes Signed-off-by: Anand Joseph --- src/tts/magpietts/magpietts.cpp | 7 ++- src/tts/magpietts/model.cpp | 49 +++++++++++++------ src/tts/magpietts/model.h | 10 ++++ .../cpp/tts/test_magpietts_frame_stacking.cpp | 22 +++++++++ 4 files changed, 70 insertions(+), 18 deletions(-) diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index f8b83e3..7c9f478 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1331,6 +1331,9 @@ stream_magpie_to_audio( const ggml_nvtx::range nvtx_step("magpietts_stream_generation_step"); const int frames_remaining = h.max_decoder_steps - step * h.frame_stacking_factor; + if (frames_remaining <= 0) { + break; + } if (codec_worker.is_failed()) { codec_worker.join(); return false; @@ -1601,8 +1604,8 @@ stream_magpie_to_audio( } decoder_frames_generated += h.frame_stacking_factor; - const int frames_to_emit = std::min( - frames_remaining, has_eos ? eos_lane : h.frame_stacking_factor); + const int frames_to_emit = magpietts_frames_to_emit( + frames_remaining, h.frame_stacking_factor, has_eos ? eos_lane : -1); if (!suppress_nonfinal_codec_output) { for (int lane = 0; lane < frames_to_emit; ++lane) { const auto& frame = codec_frames[(size_t)lane]; diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index f938376..a2efbd1 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -2031,6 +2031,9 @@ MagpieCodeGenerator::generate( const int64_t frame_start_us = ggml_time_us(); const int frames_remaining = h.max_decoder_steps - step * h.frame_stacking_factor; + if (frames_remaining <= 0) { + break; + } if (step % 10 == 0) { fprintf(stderr, "%s decoding frame %d/%d\n", label, step, max_decoder_positions); } @@ -2201,31 +2204,45 @@ MagpieCodeGenerator::generate( audio_codes[c].push_back(next_codes[c + lane * h.audio_codebooks]); } } - const int frames_to_emit = std::min( - frames_remaining, eos_lane >= 0 ? eos_lane : h.frame_stacking_factor); + const int frames_to_emit = magpietts_frames_to_emit( + frames_remaining, h.frame_stacking_factor, eos_lane); for (int lane = 0; lane < frames_to_emit; ++lane) { generated_frames.push_back(codec_frames[(size_t)lane]); } if (eos_lane >= 0) { ggml_nvtx::mark("magpietts_eos"); fprintf(stderr, "%s EOS detected at frame %d\n", label, step); - break; } - bool first_frame = false; - const int64_t frame_done_us = ggml_time_us(); - const double inter_ms = metrics.record_frame(frame_done_us, first_frame); - const double frame_latency_ms = (double)(frame_done_us - frame_start_us) / 1000.0; - if (step < 4 || step % 10 == 0) { - if (first_frame) { - fprintf( - stderr, "%s frame %d latency=%.2f ms ttff=%.2f ms\n", label, step, - frame_latency_ms, metrics.ttff_ms); - } else { - fprintf( - stderr, "%s frame %d latency=%.2f ms inter=%.2f ms\n", label, step, - frame_latency_ms, inter_ms); + if (frames_to_emit > 0) { + const int64_t frame_done_us = ggml_time_us(); + bool first_frame = false; + double inter_ms = 0.0; + for (int lane = 0; lane < frames_to_emit; ++lane) { + bool lane_is_first = false; + const double lane_inter_ms = + metrics.record_frame(frame_done_us, lane_is_first); + if (lane == 0) { + first_frame = lane_is_first; + inter_ms = lane_inter_ms; + } } + const double frame_latency_ms = + (double)(frame_done_us - frame_start_us) / 1000.0; + if (step < 4 || step % 10 == 0) { + if (first_frame) { + fprintf( + stderr, "%s frame %d latency=%.2f ms ttff=%.2f ms\n", label, step, + frame_latency_ms, metrics.ttff_ms); + } else { + fprintf( + stderr, "%s frame %d latency=%.2f ms inter=%.2f ms\n", label, step, + frame_latency_ms, inter_ms); + } + } + } + if (eos_lane >= 0) { + break; } } } diff --git a/src/tts/magpietts/model.h b/src/tts/magpietts/model.h index 0dcf138..da7c3eb 100644 --- a/src/tts/magpietts/model.h +++ b/src/tts/magpietts/model.h @@ -12,6 +12,7 @@ #include "magpietts_cuda_sampling.h" #endif +#include #include #include #include @@ -144,6 +145,15 @@ magpietts_first_eos_lane( return -1; } +inline int +magpietts_frames_to_emit(int frames_remaining, int frame_stacking_factor, int eos_lane) { + if (frames_remaining <= 0 || frame_stacking_factor <= 0) { + return 0; + } + const int available_frames = eos_lane >= 0 ? eos_lane : frame_stacking_factor; + return std::min(frames_remaining, std::max(0, available_frames)); +} + struct magpietts_layer { ggml_tensor* norm_self = nullptr; ggml_tensor* self_qkv = nullptr; diff --git a/tests/cpp/tts/test_magpietts_frame_stacking.cpp b/tests/cpp/tts/test_magpietts_frame_stacking.cpp index c304c86..9ee816b 100644 --- a/tests/cpp/tts/test_magpietts_frame_stacking.cpp +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -78,5 +78,27 @@ main() { std::fprintf(stderr, "unexpected EOS lane\n"); return 1; } + + h.max_decoder_steps = 5; + int emitted_frames = 0; + int final_position_frames = 0; + const int decoder_positions = + (h.max_decoder_steps + h.frame_stacking_factor - 1) / h.frame_stacking_factor; + for (int step = 0; step < decoder_positions; ++step) { + const int frames_remaining = + h.max_decoder_steps - step * h.frame_stacking_factor; + final_position_frames = tts::magpietts_frames_to_emit( + frames_remaining, h.frame_stacking_factor, -1); + emitted_frames += final_position_frames; + } + if (emitted_frames != h.max_decoder_steps || final_position_frames != 1) { + std::fprintf(stderr, "non-divisible decoder frame budget was exceeded\n"); + return 1; + } + if (tts::magpietts_frames_to_emit(1, h.frame_stacking_factor, 0) != 0 || + tts::magpietts_frames_to_emit(1, h.frame_stacking_factor, 1) != 1) { + std::fprintf(stderr, "partial final stacked frame did not preserve EOS handling\n"); + return 1; + } return 0; } From ab1d9339dc0034380516ff735d9228f335fd55b1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 19:22:50 +0000 Subject: [PATCH 19/19] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/tts/magpietts/decoder.cpp | 15 ++++++--------- src/tts/magpietts/magpietts.cpp | 11 +++++------ src/tts/magpietts/model.cpp | 16 ++++++---------- tests/cpp/tts/test_magpietts_frame_stacking.cpp | 7 +++---- 4 files changed, 20 insertions(+), 29 deletions(-) diff --git a/src/tts/magpietts/decoder.cpp b/src/tts/magpietts/decoder.cpp index 109af5a..fa6881d 100644 --- a/src/tts/magpietts/decoder.cpp +++ b/src/tts/magpietts/decoder.cpp @@ -831,10 +831,9 @@ MagpieDecoder::evalCachedPair( const std::vector>& audio_codes, int speaker, int threads, DecoderKvCache& cond_kv, DecoderKvCache& uncond_kv, decoder_result& cond_result, decoder_result& uncond_result, int stacked_position_budget, - magpietts_cuda_sample_request* cuda_sample, - const magpietts_backend_tensor* text_cond_device, magpietts_backend_tensor* cond_hidden_out, - magpietts_backend_tensor* uncond_hidden_out, DecoderCrossKvCache* cond_cross_kv, - const magpietts_decoder_attention* attention) const { + magpietts_cuda_sample_request* cuda_sample, const magpietts_backend_tensor* text_cond_device, + magpietts_backend_tensor* cond_hidden_out, magpietts_backend_tensor* uncond_hidden_out, + DecoderCrossKvCache* cond_cross_kv, const magpietts_decoder_attention* attention) const { const bool persistent_candidate = cuda_sample == nullptr && cond_hidden_out != nullptr && uncond_hidden_out != nullptr && cond_cross_kv != nullptr && cond_cross_kv->validFor(model_, text_len) && @@ -843,17 +842,15 @@ MagpieDecoder::evalCachedPair( if (persistent_candidate) { try { if (persistent_runtime_ && - !persistent_runtime_->matches( - cond_cross_kv, text_len, stacked_position_budget)) { + !persistent_runtime_->matches(cond_cross_kv, text_len, stacked_position_budget)) { // Cross-cache address, shape, or request budget changes require a new graph. persistent_runtime_.reset(); cond_kv.clear(); uncond_kv.clear(); } if (cond_kv.n_tokens > 0 && !persistent_runtime_) { - persistent_runtime_ = - std::make_unique( - model_, *cond_cross_kv, text_len, stacked_position_budget); + persistent_runtime_ = std::make_unique( + model_, *cond_cross_kv, text_len, stacked_position_budget); persistent_runtime_->seed(cond_kv, uncond_kv); fprintf( stderr, diff --git a/src/tts/magpietts/magpietts.cpp b/src/tts/magpietts/magpietts.cpp index 7c9f478..5ff1528 100644 --- a/src/tts/magpietts/magpietts.cpp +++ b/src/tts/magpietts/magpietts.cpp @@ -1329,8 +1329,7 @@ stream_magpie_to_audio( (h.max_decoder_steps + h.frame_stacking_factor - 1) / h.frame_stacking_factor; for (int step = 0; step < max_decoder_positions; ++step) { const ggml_nvtx::range nvtx_step("magpietts_stream_generation_step"); - const int frames_remaining = - h.max_decoder_steps - step * h.frame_stacking_factor; + const int frames_remaining = h.max_decoder_steps - step * h.frame_stacking_factor; if (frames_remaining <= 0) { break; } @@ -1372,9 +1371,8 @@ stream_magpie_to_audio( text_cond, text_len, audio_codes, params.speaker, params.threads, cond_kv, uncond_kv, cond, uncond, max_decoder_positions, nullptr, &text_cond_device, - &cond_hidden_device, - &uncond_hidden_device, &cond_cross_kv, - decoder_attention_arg) + &cond_hidden_device, &uncond_hidden_device, + &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( text_cond, text_len, audio_codes, params.speaker, params.threads, cond, uncond, nullptr, @@ -1426,7 +1424,8 @@ stream_magpie_to_audio( text_cond, text_len, audio_codes, params.speaker, params.threads, cond_kv, uncond_kv, cond, uncond, max_decoder_positions, &cuda_sample, &text_cond_device, - nullptr, nullptr, &cond_cross_kv, decoder_attention_arg) + nullptr, nullptr, &cond_cross_kv, + decoder_attention_arg) : decoder.evalPair( text_cond, text_len, audio_codes, params.speaker, params.threads, cond, uncond, &cuda_sample, diff --git a/src/tts/magpietts/model.cpp b/src/tts/magpietts/model.cpp index 0e7f53f..baf64c6 100644 --- a/src/tts/magpietts/model.cpp +++ b/src/tts/magpietts/model.cpp @@ -2028,8 +2028,7 @@ MagpieCodeGenerator::generate( for (int step = 0; step < max_decoder_positions; ++step) { const ggml_nvtx::range nvtx_step("magpietts_generate_step"); const int64_t frame_start_us = ggml_time_us(); - const int frames_remaining = - h.max_decoder_steps - step * h.frame_stacking_factor; + const int frames_remaining = h.max_decoder_steps - step * h.frame_stacking_factor; if (frames_remaining <= 0) { break; } @@ -2063,8 +2062,7 @@ MagpieCodeGenerator::generate( text_cond, (int)tokens.size(), audio_codes, params.speaker, params.threads, cond_kv, uncond_kv, cond, uncond, max_decoder_positions, nullptr, &text_cond_device, - &cond_hidden_device, - &uncond_hidden_device, &cond_cross_kv, + &cond_hidden_device, &uncond_hidden_device, &cond_cross_kv, decoder_attention_arg) : decoder.evalPair( text_cond, (int)tokens.size(), audio_codes, params.speaker, @@ -2203,8 +2201,8 @@ MagpieCodeGenerator::generate( audio_codes[c].push_back(next_codes[c + lane * h.audio_codebooks]); } } - const int frames_to_emit = magpietts_frames_to_emit( - frames_remaining, h.frame_stacking_factor, eos_lane); + const int frames_to_emit = + magpietts_frames_to_emit(frames_remaining, h.frame_stacking_factor, eos_lane); for (int lane = 0; lane < frames_to_emit; ++lane) { generated_frames.push_back(codec_frames[(size_t)lane]); } @@ -2219,15 +2217,13 @@ MagpieCodeGenerator::generate( double inter_ms = 0.0; for (int lane = 0; lane < frames_to_emit; ++lane) { bool lane_is_first = false; - const double lane_inter_ms = - metrics.record_frame(frame_done_us, lane_is_first); + const double lane_inter_ms = metrics.record_frame(frame_done_us, lane_is_first); if (lane == 0) { first_frame = lane_is_first; inter_ms = lane_inter_ms; } } - const double frame_latency_ms = - (double)(frame_done_us - frame_start_us) / 1000.0; + const double frame_latency_ms = (double)(frame_done_us - frame_start_us) / 1000.0; if (step < 4 || step % 10 == 0) { if (first_frame) { fprintf( diff --git a/tests/cpp/tts/test_magpietts_frame_stacking.cpp b/tests/cpp/tts/test_magpietts_frame_stacking.cpp index 8cb1346..7290e33 100644 --- a/tests/cpp/tts/test_magpietts_frame_stacking.cpp +++ b/tests/cpp/tts/test_magpietts_frame_stacking.cpp @@ -83,10 +83,9 @@ main() { const int decoder_positions = (h.max_decoder_steps + h.frame_stacking_factor - 1) / h.frame_stacking_factor; for (int step = 0; step < decoder_positions; ++step) { - const int frames_remaining = - h.max_decoder_steps - step * h.frame_stacking_factor; - final_position_frames = tts::magpietts_frames_to_emit( - frames_remaining, h.frame_stacking_factor, -1); + const int frames_remaining = h.max_decoder_steps - step * h.frame_stacking_factor; + final_position_frames = + tts::magpietts_frames_to_emit(frames_remaining, h.frame_stacking_factor, -1); emitted_frames += final_position_frames; } if (emitted_frames != h.max_decoder_steps || final_position_frames != 1) {